diff --git a/cpp/tensorrt_llm/kernels/kdaDecode/kdaDecodeLegacy.cu b/cpp/tensorrt_llm/kernels/kdaDecode/kdaDecodeLegacy.cu index ca77fc702e29..6d34478a8507 100644 --- a/cpp/tensorrt_llm/kernels/kdaDecode/kdaDecodeLegacy.cu +++ b/cpp/tensorrt_llm/kernels/kdaDecode/kdaDecodeLegacy.cu @@ -275,6 +275,8 @@ __device__ __forceinline__ Sum2 block_reduce_sum2_for(float x, float y, float* s block_y = lane < kReduceWarps ? scratch[kReduceWarps + lane] : 0.0f; block_x = warp_reduce_sum(block_x); block_y = warp_reduce_sum(block_y); + // Lane 0 overwrites partials that other lanes of this warp read above. + __syncwarp(); if (lane == 0) { scratch[0] = block_x; @@ -349,6 +351,8 @@ __device__ __forceinline__ Sum2 block_reduce_sum2_active_for(float x, float y, f block_y = lane < kReduceWarps ? scratch[kReduceWarps + lane] : 0.0f; block_x = warp_reduce_sum(block_x); block_y = warp_reduce_sum(block_y); + // Lane 0 overwrites partials that other lanes of this warp read above. + __syncwarp(); if (lane == 0) { scratch[0] = block_x; diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu index e714466dbb3a..60baff165e2b 100644 --- a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu +++ b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu @@ -411,9 +411,17 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(bf16_t c float2 f[2] = {__bfloat1622float2(v2[0]), __bfloat1622float2(v2[1])}; if constexpr (FULL_N12) { - if (n == AN - 1 && lane == 0) + if (n == AN - 1) { - cute::arrive_barrier(plan.bar_consumed[chunk_slot]); + // The producer refills the slots through the async proxy (TMA): a + // cross-proxy fence orders this lane's generic-proxy reads before it. + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + // Every lane's reads of the chunk's slots before lane 0 releases them. + __syncwarp(); + if (lane == 0) + { + cute::arrive_barrier(plan.bar_consumed[chunk_slot]); + } } } tmem_st_32dp32bNx<4>( @@ -492,6 +500,11 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(bf16_t c } if constexpr (!FULL_N12) { + // The producer refills the slots through the async proxy (TMA): a cross-proxy fence orders this + // lane's generic-proxy reads before it. + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + // Every lane's reads of the chunk's slots before lane 0 releases them to the producer. + __syncwarp(); if (lane == 0) { cute::arrive_barrier(plan.bar_consumed[chunk_slot]); @@ -551,6 +564,10 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(bf16_t c { cross_warp_tail(lane); } + // ws_stats is one buffer reused every chunk. The barrier above orders this chunk's writes before the + // reads; this one orders the reads before the next chunk's writes, so a warp that has finished + // reading cannot overwrite its row while a slower warp still reads it. + cutlass::arch::NamedBarrier::sync(CONSUMER_THREADS, 0); float logit_n[N_CHUNK]; #pragma unroll for (int n = 0; n < N_CHUNK; n++) @@ -972,6 +989,9 @@ __global__ void __launch_bounds__(BLK, 1) } cutlass::arch::NamedBarrier::sync(CONSUMER_THREADS, 1); } + // The producer refills the tile through the async proxy (TMA): a cross-proxy fence orders this thread's + // generic-proxy reads before it. + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); cute::arrive_barrier(plan.bar_consumed[slot]); } } @@ -1027,12 +1047,18 @@ __global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_single_cta_kernel(bf16 bf16_t const* __restrict__ layer_res, bf16_t const* __restrict__ layer_res_add, bf16_t const* __restrict__ res_w, bf16_t const* __restrict__ rms_w, bf16_t const* __restrict__ output_rms_w, bf16_t* __restrict__ updated_layer_res, bf16_t* __restrict__ output, float* __restrict__ rsigma_out, float* __restrict__ probs_out, - float* __restrict__ logits_out, float rms_eps, float output_rms_eps, int num_tokens) + float* __restrict__ logits_out, float rms_eps, float output_rms_eps, int num_tokens, bool early_trigger) { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 if constexpr (ENABLE_PDL) { cudaGridDependencySynchronize(); + // The dependent grid still waits for this one to complete before reading its + // output; triggering here lets it launch and stream its weights while this kernel runs. + if (early_trigger) + { + cudaTriggerProgrammaticLaunchCompletion(); + } } constexpr int H = 7168; @@ -1302,7 +1328,10 @@ __global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_single_cta_kernel(bf16 if constexpr (ENABLE_PDL) { - cudaTriggerProgrammaticLaunchCompletion(); + if (!early_trigger) + { + cudaTriggerProgrammaticLaunchCompletion(); + } } #else if (cute::thread0()) @@ -1312,11 +1341,17 @@ __global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_single_cta_kernel(bf16 #endif } +// Decode handoff options for the s1 kernels (see AttnResFwdParams). +struct S1Handoff +{ + bool early_trigger = false; +}; + template static void launch_s1_single_cta(bf16_t const* block_residual, bf16_t const* layer_residual, bf16_t const* layer_residual_add, bf16_t const* res_weight, bf16_t const* rms_weight, bf16_t const* output_rms_weight, bf16_t* updated_layer_residual, bf16_t* output, float* rsigma, float* probs, - float* logits, float rms_eps, float output_rms_eps, int num_tokens, cudaStream_t stream) + float* logits, float rms_eps, float output_rms_eps, int num_tokens, cudaStream_t stream, S1Handoff handoff = {}) { if (tensorrt_llm::common::getEnvEnablePDL()) { @@ -1332,14 +1367,14 @@ static void launch_s1_single_cta(bf16_t const* block_residual, bf16_t const* lay config.numAttrs = 1; cudaLaunchKernelEx(&config, kernel, block_residual, layer_residual, layer_residual_add, res_weight, rms_weight, output_rms_weight, updated_layer_residual, output, rsigma, probs, logits, rms_eps, output_rms_eps, - num_tokens); + num_tokens, handoff.early_trigger); } else { attn_res_fwd_s1_single_cta_kernel <<>>(block_residual, layer_residual, layer_residual_add, res_weight, rms_weight, output_rms_weight, updated_layer_residual, output, rsigma, probs, logits, rms_eps, output_rms_eps, - num_tokens); + num_tokens, false); } } @@ -1361,13 +1396,22 @@ __global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_splitk_kernel(bf16_t c bf16_t const* __restrict__ layer_res, bf16_t const* __restrict__ layer_res_add, bf16_t const* __restrict__ res_w, bf16_t const* __restrict__ rms_w, bf16_t const* __restrict__ output_rms_w, bf16_t* __restrict__ updated_layer_res, bf16_t* __restrict__ output, float* __restrict__ rsigma_out, float* __restrict__ probs_out, - float* __restrict__ logits_out, float rms_eps, float output_rms_eps, int num_tokens) + float* __restrict__ logits_out, float rms_eps, float output_rms_eps, int num_tokens, bool early_trigger) { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 - if constexpr (ENABLE_PDL) + auto const wait_for_previous_grid = [&] { - cudaGridDependencySynchronize(); - } + if constexpr (ENABLE_PDL) + { + cudaGridDependencySynchronize(); + // See attn_res_fwd_s1_single_cta_kernel. + if (early_trigger) + { + cudaTriggerProgrammaticLaunchCompletion(); + } + } + }; + wait_for_previous_grid(); namespace cg = cooperative_groups; constexpr int H = 7168; @@ -1420,6 +1464,7 @@ __global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_splitk_kernel(bf16_t c float sq[N] = {}; float dot[N] = {}; + // Every candidate in one pass, so all loads are in flight together. #pragma unroll for (int ki = tid; ki < K_PER_CTA; ki += THREADS) { @@ -1619,7 +1664,10 @@ __global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_splitk_kernel(bf16_t c if constexpr (ENABLE_PDL) { - cudaTriggerProgrammaticLaunchCompletion(); + if (!early_trigger) + { + cudaTriggerProgrammaticLaunchCompletion(); + } } #else if (cute::thread0()) @@ -1633,13 +1681,14 @@ template : &attn_res_fwd_s1_splitk_kernel; static std::once_flag attrs_set[2][64]; @@ -1659,7 +1708,7 @@ static void launch_s1_splitk(bf16_t const* block_residual, bf16_t const* layer_r void* args[] = {const_cast(&block_residual), const_cast(&layer_residual), const_cast(&layer_residual_add), const_cast(&res_weight), const_cast(&rms_weight), const_cast(&output_rms_weight), &updated_layer_residual, &output, &rsigma, &probs, &logits, &rms_eps, - &output_rms_eps, &num_tokens}; + &output_rms_eps, &num_tokens, &early_trigger}; cudaLaunchConfig_t config{}; // One cluster per token. clusterDim stays GROUPS, so clusters are contiguous // spans of the grid and each token's DSM exchange stays within its own. @@ -1868,18 +1917,22 @@ static void launchAttnResDecodeRmsNorm(AttnResFwdParams const& params, cudaStrea auto const* layer_residual_add = FUSE_LAYER_ADD ? params.layerResidualAdd : nullptr; auto* updated_layer_residual = FUSE_LAYER_ADD ? params.updatedLayerResidual : nullptr; + // The dependent kernel launches (and streams its weights) while this one runs: it waits for + // this grid before reading the output. + constexpr bool early_trigger = true; + S1Handoff const handoff{early_trigger}; if constexpr (N <= 4) { launch_s1_single_cta(params.blockResidual, params.layerResidual, layer_residual_add, params.resWeight, params.rmsWeight, params.outputRmsWeight, updated_layer_residual, params.output, nullptr, - nullptr, nullptr, params.rmsEps, params.outputRmsEps, params.seqLen, stream); + nullptr, nullptr, params.rmsEps, params.outputRmsEps, params.seqLen, stream, handoff); } else { launch_s1_splitk(params.blockResidual, params.layerResidual, layer_residual_add, params.resWeight, params.rmsWeight, params.outputRmsWeight, updated_layer_residual, params.output, nullptr, - nullptr, nullptr, params.rmsEps, params.outputRmsEps, params.seqLen, stream); + nullptr, nullptr, params.rmsEps, params.outputRmsEps, params.seqLen, stream, handoff); } } diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.cu b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.cu index 7245f724134a..ee734add0661 100644 --- a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.cu +++ b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.cu @@ -650,6 +650,11 @@ __global__ void __launch_bounds__(BLK, 1) default: __builtin_unreachable(); } } + // The producer refills the slots through the async proxy (TMA): a cross-proxy fence orders this lane's + // generic-proxy reads before it. + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + // Every lane's reads of the chunk's slots before lane 0 releases them to the producer. + __syncwarp(); if (lane == 0) { mbarrier_arrive(plan.bar_consumed[chunk_slot]); diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md index 61be49f5be0d..87069886e5be 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md @@ -181,8 +181,10 @@ Perf is measured, never gated. ## Status of every record in this tree -**The catalog is fully certified on sm_103. The targets construct but have -never executed.** +**The catalog is certified: 19 entries on sm_103, and 18 on sm_100 (B200 / +GB200): the 16 Kimi K3 entries, where their first caller runs, and +`cublas_mm` and `flashinfer_rmsnorm`. The targets construct but have never +executed.** Two things voided every receipt in the move: each catalog test file was rewritten, and the targets moved from sm_100 (B200) to sm_103 (GB300), where diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index d1b2a7840899..16c6b9769eb6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -59,6 +59,7 @@ MODELING_V2_ROUTERS = { "GptOssForCausalLM": "models.gpt_oss.routing", "DeepseekV3ForCausalLM": "models.deepseek_v3.routing", + "KimiK3ForConditionalGeneration": "models.kimi_k3_vl.routing", } #: Backends whose model construction reaches ``modeling_v2_resolve``. The diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/k3_situ_mul.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/k3_situ_mul.md new file mode 100644 index 000000000000..14f89ee0710f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/k3_situ_mul.md @@ -0,0 +1,105 @@ +--- +receipts: + sm_100: {status: passed, tests: 4} +--- + +# k3_situ_mul + +**Wraps** `torch.ops.trtllm.k3_situ_mul` (one call). + +## Semantics + +The SiTU-gated multiply of the Kimi K3 dense MLP (`SituAndMul` in `tensorrt_llm/_torch/modules/situ.py`) for at +most 8 tokens, as one CuTe DSL kernel that lets the next kernel launch at once under programmatic dependent launch. +The last dim of the input holds the gate half followed by the up half: + +``` +K = gu.shape[1] // 2; g = gu[m, k], u = gu[m, K + k] 0 <= m < M <= 8, 0 <= k < K, in fp32 +a = (beta * tanh(g / beta)) * sigmoid(g) sigmoid(g) = 1 / (1 + exp(-g)) +v = linear_beta * tanh(u / linear_beta) if linear_beta is not None, else v = u +out[m, k] = bf16(a * v) +``` + +The arithmetic is fp32 in this order, with IEEE division and the full-precision `exp` and `tanh`, and the result +is rounded to bf16 once, round-to-nearest-even. That is the order of `SituAndMul`'s eager path. The `exp` and +`tanh` implementations are not torch's or Triton's, so the result is not bit-identical to `SituAndMul` (eager or +fused); they agree within the tolerance below. Rows are independent: a token's output row has the same bits +whatever `M` is. + +Fusion boundary: the single call computes the activation only. The caller owns the gate_up projection that +produces `gu` and the down projection that consumes the output. There is no quantization of the output. + +## Signature + +```python +def k3_situ_mul(gu: torch.Tensor, beta: float = 1.0, linear_beta: float | None = None) -> torch.Tensor +``` + +| Argument | Shape | Dtype | Layout | Device | +|---|---|---|---|---| +| `gu` | `[M, 2 * K]`, gate half first, `1 <= M <= 8`, `K % 8 == 0` | bf16 | contiguous, 16-byte-aligned | CUDA | +| `beta` | scalar | Python float | — | — | +| `linear_beta` | scalar or `None` | Python float, not 0 (the wrapper asserts) | — | — | +| returns | `[M, K]` | bf16 | newly allocated, contiguous | `gu`'s device | + +Every element of the returned tensor is written; `gu` is not mutated. + +### Certified arguments + +The dense MLP of Kimi K3 at TP16 (gate_up output `[M, 4224]`, activation `[M, 2112]`) at every `M` in 1..8, with +`(beta, linear_beta)` = `(4.0, 25.0)` (the Kimi K3 checkpoint's) and `(1.0, None)` (the defaults). Inputs are +`gu ~ 2 * N(0, 1)`, rounded to bf16. The gate per cell: `max|y - ref| / max|ref| <= 8e-3` against the formula +above evaluated in fp64 on the same bf16 input, identical bits on a repeated call, and each `M`'s rows +bit-identical to the same rows of the 8-row call. Also certified, with both settings at `M` 1 and 8: calls +captured in a CUDA graph after an eager call per key, replayed with `gu` rewritten in place, return the bits of +eager calls on the new `gu`. Other widths and `beta` / `linear_beta` values are accepted but not certified. + +## Metadata consumed + +None. The op reads no attention metadata, KV cache or module state; every input is an argument, and the kernel +runs on the current CUDA stream of `gu`'s device. + +One process-global cache sits behind it, and it is result-neutral: the compiled kernel per key +`(K, linear_beta is not None, PDL on/off)`. The first call with a new key compiles the kernel with the CuTe DSL, +which costs host time; `M`, `beta` and the value of `linear_beta` are runtime arguments and never recompile. That +first call refuses to run under CUDA-graph capture: it raises `RuntimeError` ("run once per shape outside +CUDA-graph capture first") before launching anything. Call every key once eagerly, then capture; captured and +later calls reuse the compiled kernel. The PDL setting is read from `TRTLLM_ENABLE_PDL` on every call, so +changing it within a process adds a key. + +## Preconditions + +The op checks these and raises `ValueError` ("k3_situ_mul: unsupported call ...") when one fails: + +- `gu` is a 2-D contiguous CUDA bf16 tensor with `1 <= M <= 8` rows. +- `gu.shape[1] % 16 == 0`, so each half is a whole number of 16-byte vectors. +- `gu` starts at a 16-byte-aligned address. + +Certified refusals: 0 and 9 rows, a width of 4216, a `gu` starting 2 bytes past a 16-byte boundary, fp16, a +row-strided `gu` and a 1-D `gu`. + +The wrapper adds one check, because the op does not fail on it: `linear_beta` must not be `0.0`. The op hands the +kernel `linear_beta or 1.0`, so `linear_beta=0.0` would run as `1.0` (`v = tanh(u)`) with no error, where the +formula gives `0 * tanh(u / 0)`. The wrapper raises `AssertionError` before calling the op (certified). + +Not checked by the op: + +- `K >= 8`; an empty `gu` (width 0) passes the op's check. +- The architecture. The kernel needs no SM 10.x feature (no tcgen05, no clusters; `griddepcontrol` needs SM 9.0 + or newer), but only sm_100 has been measured. +- The CuTe DSL (`cutlass`) and `cuda-python` (`cuda.bindings`) are importable. The op imports the kernel module + after its check, so without them a call raises `ImportError`. + +## Notes + +- Programmatic dependent launch (PDL). With `TRTLLM_ENABLE_PDL` unset or `1` the kernel is launched with PDL; any + other value launches it without. There is no `trigger_early` argument: every CTA executes + `griddepcontrol.launch_dependents` first, then `griddepcontrol.wait`, and only then reads `gu`. So the next + kernel on the stream, if launched with PDL, launches at once and can stream its own weights while this kernel + and its predecessor run (the CTM GEMVs do); it must execute `griddepcontrol.wait` + (`cudaGridDependencySynchronize`) before it reads the output, as those GEMVs do before they read their + activation. Kernels launched without PDL, torch's included, start after this one completes as usual. None of + this changes results. +- Two identical calls return identical bits. +- Grid: `ceil(K / 1024)` x 8 CTAs of 128 threads, 8 columns per thread; rows past `M` exit at once. +- Only sm_100 has been measured. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/k3_situ_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/k3_situ_mul.py new file mode 100644 index 000000000000..6118c1fbc2ef --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/k3_situ_mul.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SiTU-gated multiply (SituAndMul) of a gate_up output for M <= 8 tokens via the CTM (CuTe DSL) kernel.""" + +import torch + +import tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv.op # noqa: F401 (registers torch.ops.trtllm.k3_situ_mul) + + +def k3_situ_mul( + gu: torch.Tensor, beta: float = 1.0, linear_beta: float | None = None +) -> torch.Tensor: + """Return `SituAndMul(beta, linear_beta)(gu)`, gate half first, in one k3_situ_mul call.""" + # The op hands the kernel `linear_beta or 1.0`, so 0.0 would run as linear_beta=1.0 (up half + # tanh(u)) with no error, where the formula gives 0 * tanh(u / 0). + assert linear_beta is None or linear_beta != 0.0, ( + "linear_beta=0.0 would silently run as 1.0; pass None to leave the up half unscaled" + ) + return torch.ops.trtllm.k3_situ_mul(gu, beta=beta, linear_beta=linear_beta) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/k3_mla_attn_vb_out.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/k3_mla_attn_vb_out.md new file mode 100644 index 000000000000..2b10e316f27b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/k3_mla_attn_vb_out.md @@ -0,0 +1,178 @@ +--- +receipts: + sm_100: {status: passed, tests: 9} +--- + +# k3_mla_attn_vb_out + +**Wraps** `torch.ops.trtllm.k3_mla_attn_vb_out` (one call). The same contract covers two sibling wrappers, one call +each: `k3_mla_attn_out` (`torch.ops.trtllm.k3_mla_attn_out`, the attention output itself) and `k3_mla_attn` +(`torch.ops.trtllm.k3_mla_attn`, the same into a new tensor, page offset 0). All three take a caller-owned +`K3MlaAttnWorkspace` (*State*). + +## Semantics + +Kimi K3's MLA decode attention for `R <= 8` requests of `T <= 8` tokens over the paged latent cache, with v_b and +the output gate in the same launch. `q` is `fused_q` (from `attention/k3_mla_qkv`), request-major: rows +`i T .. i T + T - 1` are request `i`'s tokens. Request `i`'s cache rows `kv_i[k]`, `k < L_i = seq_len[i]`, are row +`(page_table[i][k // 64] + page_offset) * 64 + k % 64` of `pool` (512 latent columns, then 64 rope columns). Per +request `i`, token `t` and head `h`: + +``` +o[t, h] = softmax_k( softmax_scale * q[t, h] . kv_i[k] ) @ kv_i[k, :512] k <= L_i - T + t (bottom-right causal) +y[t, h] = bf16( bf16(o[t, h]) @ w_vb[h]^T ) # v_b, 128 per head +out[t, 128 h : 128 h + 128] = y[t, h] # without gate + = bf16(y[t, h] * s[t, h]) # with gate: s = gate[t, gate_col0 + 128 h .. + 128] +``` + +`gate` holds the output gate's sigmoid values (the op multiplies, it does not apply the sigmoid); the gated output is +bit-identical to `bf16(plain * s)` (certified). `k3_mla_attn_out` writes `o` itself (`[M, heads * 512]`, bf16) into +`out`; `k3_mla_attn` returns it in a new tensor and is bit-identical to `k3_mla_attn_out` at page offset 0 +(certified). + +How it computes: one cluster of 16 CTAs per (request, group of 6 heads), split-KV over 128-row tiles; fp32 softmax; +each CTA's partial `O / l` (a convex combination of V rows) is kept in fp16, and the 16 partials are merged in fp32 +in a fixed order. More than `CLUSTER_WAVE` = 7 clusters (`R x heads / 6`) would not co-reside on GB200; when all +their CTAs fit on the SMs at once (one per SM: at most 9 clusters on 148 SMs) the launch takes the **no-cluster +mode**, which exchanges the merge statistics through the workspace instead of cluster shared memory. The merge sums +the same values in the same order in both modes, so the outputs do not depend on the mode. Request `i`'s rows are +computed as the one-request call on its own rows, pages and length would compute them (the op test, +`test_k3_mla_attn.py`). + +Accuracy (certified bound): every request's output within `1e-2` (max relative, per request) of a float64 +reference: the attention in float64 with the causal mask above, rounded to bf16; v_b in float64; the gate applied +to the bf16-rounded product. + +Fusion boundary. Inside: the attention over the cache, v_b, the output gate. Outside: `fused_q` and the cache append +(`attention/k3_mla_qkv`, which must store the step's rows before this call reads them), and the output projection +that consumes `out`. + +## Signature + +```python +def k3_mla_attn_vb_out( + q: torch.Tensor, + pool: torch.Tensor, + row_stride: int, + page_table: torch.Tensor, + page_offset: int, + seq_len: torch.Tensor, + softmax_scale: float, + w_vb: torch.Tensor, + out: torch.Tensor, + workspace: K3MlaAttnWorkspace, + gate: Optional[torch.Tensor] = None, + gate_col0: int = 0, +) -> None + +def k3_mla_attn_out(q, pool, row_stride, page_table, page_offset, seq_len, softmax_scale, out, + workspace: K3MlaAttnWorkspace) -> None + +def k3_mla_attn(q, pool, row_stride, page_table, seq_len, softmax_scale, + workspace: K3MlaAttnWorkspace) -> torch.Tensor +``` + +### Certified arguments + +| Argument | Shape | Dtype | Layout | Device | +|---|---|---|---|---| +| `q` | `[M, heads * 576]`, `M = R T`, `R <= 8`, `T <= 8`, `heads` a multiple of 6 (6 at TP16, 24 at TP4) | bf16 | contiguous | CUDA | +| `pool` | the layer's paged latent pool, flat | bf16 | contiguous | CUDA | +| `row_stride` | elements per cache row, `>= 576`, a multiple of 8 | Python int | — | — | +| `page_table` | `[R, W]` with unit column stride and any row stride `>= W` (rows of `kv_cache_block_offsets`), or `[W]` when `R = 1` | int32 | see shape | CUDA | +| `page_offset` | the layer's slot in a layer-interleaved pool | Python int | — | — | +| `seq_len` | `[R]`, `L_i >= T`: each request's length including its `T` new rows | int32 | contiguous | CUDA | +| `softmax_scale` | `1 / (sqrt(qk_nope + qk_rope) * q_scaling)` (the view's) | Python float | — | — | +| `w_vb` | `[heads, 128, 512]` (`v_b_proj`) | bf16 | contiguous | CUDA | +| `out` | `[M, heads * 128]` (`k3_mla_attn_out`: `[M, heads * 512]`) | bf16 | contiguous | CUDA | +| `workspace` | a `K3MlaAttnWorkspace` of `q`'s device for `heads / 6` head groups | — | — | — | +| `gate` | `None`, or `[M, C]` with unit column stride and `gate_col0 + heads * 128 <= C` | bf16 | rows of any stride | CUDA | +| `gate_col0` | `>= 0` | Python int | — | — | +| returns | `None` (`k3_mla_attn`: `[M, heads * 512]`, newly allocated) | bf16 | contiguous | = `q.device` | + +Certified at `(R, T, heads)` = `(3, 8, 6)` (3 clusters), `(8, 1, 6)` and `(2, 8, 24)` (8 clusters: the no-cluster +mode), with 64 to 4100 cached rows per request before a step (pages crossed, several tiles per CTA), over a real +`KVCacheManager` (below). + +## State + +**Object.** `K3MlaAttnWorkspace` (`catalog/attention/k3_mla_attn_workspace.py`), one per device and head-group +count (`heads / 6`: 1 at TP16, 4 at TP4), owned by the caller. The pool is the KV cache manager's (as +`attention/k3_mla_qkv`); this call only reads it. + +**Contents and size.** One fp16 buffer, `attn_workspace_elems(groups)` elements: the per-CTA partials of 8 requests +x `groups` x 16 CTAs (64 KiB each: 8 MiB per head group), then the no-cluster mode's fp32 `(m, l)` exchange (48 KiB +per group) and two int32 arrival counters per (request, head group), each on its own 128-byte line (2 KiB per +group): 8,439,808 bytes per head group, 33,759,232 at TP4. The size does not depend on `R` or `T`: every call fits. + +**Who creates it, and when.** The target, in `post_load_weights`, with `K3MlaAttnWorkspace.create(device, groups)`: + +- eager: it allocates, so it refuses to run under CUDA-graph capture (`RuntimeError`, certified); +- it zeroes the counters and leaves the partials as allocated (a call reads only words it wrote, below), then + synchronizes the device, so the workspace is ready on any stream. + +**Which ops may share one object.** All three forms, for every MLA layer of the device with the same head-group +count: one workspace serves a model's layers. Two workspaces are independent (certified: the layers' calls +alternating between two workspaces over three no-cluster steps; every output bit-identical to the same call on a new +workspace, each workspace's counters counting only its own launches). + +**Call-order invariant.** Calls on one workspace run one at a time. Every access a call makes to the workspace +(partials, exchange, counters) follows its grid-dependency wait, so a call may directly follow another on the same +stream; the earlier call must have completed by the time the later one passes its wait. That holds for calls on one +stream when every kernel between them waits on its predecessor or is launched without PDL, as the kernels of a decode +step do. Calls that overlap on one workspace (two streams, or a graph replay beside an eager call) overwrite each +other's partials and counters; they are outside the contract (not measured). + +**What a later launch reads.** Only the counters, and only in the no-cluster mode: a launch reads each of its +(request, head group) counters after its grid-dependency wait and before it arrives, then waits for 16 arrivals past +`count & ~15`. The partials and the exchange are written and read within one call: the op test refills them with NaN +before every call and the outputs keep their bits (`test_k3_mla_attn.py::test_attn_workspace_poison`). Before its +grid-dependency wait a launch reads only the page table, the lengths and the cache pages before the one holding row +`L_i - T` (earlier steps wrote them); `q` and the pages from that one on after it. + +**How it is re-armed.** Never. The counters only grow: each no-cluster launch adds 16 to both counters of each of its +requests' head groups and leaves the others; cluster-mode launches leave them all (certified: counters equal to 16 x +the launches after a sequence of eager calls and graph replays). They are compared by signed difference, so they run +through the int32 wrap (op test: counters started at `2^31 - 16` and at `-16` give the bits of zeroed ones). A +launch must start with its counters at a multiple of 16, which every complete launch leaves; `create()` starts them +at 0. + +**Why the test drives call sequences.** One workspace serves every layer and step, eagerly and inside captured +graphs, and its counters carry over from launch to launch. The test runs 3 layers x 2 decode steps on one shared +workspace eagerly, then one step's 3 calls captured as a CUDA graph on the same workspace and replayed for 2 more +steps with rewritten `q` and the metadata prepared for each step. Every output is +compared bit for bit with the same call on a new workspace and within `1e-2` with the reference; in the no-cluster +mode the shared workspace's counters are checked at the end. + +**What a wrong workspace does.** Measured (the test's negative control): a one-head-group call given a workspace +laid out for four head groups, an fp32 tensor, or a buffer 8 elements short raises `ValueError` naming the expected +workspace, before any launch; neither `out` nor the workspace changes. + +## Metadata consumed + +Through `k3_mla_decode_view` (see `attention/k3_mla_qkv`): the page table and lengths are views of the prepared +metadata's `kv_cache_block_offsets` and `kv_lens_cuda_runtime`, `page_offset` the layer's slot, `row_stride` the +pool's token stride, `softmax_scale` from the attention module's `qk_nope_head_dim`, `qk_rope_head_dim` and +`q_scaling`. + +## Preconditions + +- `q`, the page table and the lengths describe `R <= 8` requests of the same `T <= 8` tokens, `heads` a multiple of 6, + `row_stride >= 576` and a multiple of 8; `w_vb`, `out` and `gate` as above. Anything else raises `ValueError` + before any launch (the op test, `test_attn_rejects`), as does a workspace of another layout, dtype or device + (certified, above). +- `seq_len[i] >= T`: the step's `T` new rows are in the cache (stored by `k3_mla_qkv` first). +- The first call of each configuration compiles the kernel and must run outside CUDA-graph capture (it raises + `RuntimeError` inside one). +- The no-cluster mode needs its 16 x `R x groups` CTAs co-resident (one per SM); the op takes it only when they fit, + and otherwise launches clusters in more than one wave. +- SM 100 / SM 103 only (CuTe DSL `tcgen05` kernel). + +## Notes + +- The op-level test is `tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_mla_attn.py`: every split of up to 8 + tokens and the R x 8 verify steps against the float64 reference and the stock CuTe DSL MLA decode + (`trtllm::cute_dsl_mla_decode_fp16_blackwell`), the TP4 head count, interleaved pools of rows of 640, the + workspace poison and counter-wrap checks above, refusals, and (with `K3_BASE_TRTLLM`) batch-1 identity against + the unmodified single-request kernel. +- `k3_mla_attn` takes no page offset: it fits a pool with one layer per block, or layer 0 of an interleaved one. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/k3_mla_attn_vb_out.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/k3_mla_attn_vb_out.py new file mode 100644 index 000000000000..8f52772f1565 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/k3_mla_attn_vb_out.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3's MLA decode attention over the paged latent cache, with v_b and the output gate in the same launch, over +a caller-owned :class:`K3MlaAttnWorkspace`; also the plain forms (the attention output itself).""" + +from typing import Optional + +import torch + +from tensorrt_llm._torch.cute_dsl_kernels.k3_mla import ( # noqa: F401 — registers the ops + op as _k3_mla_op, +) + +from .k3_mla_attn_workspace import K3MlaAttnWorkspace + + +def k3_mla_attn_vb_out( + q: torch.Tensor, + pool: torch.Tensor, + row_stride: int, + page_table: torch.Tensor, + page_offset: int, + seq_len: torch.Tensor, + softmax_scale: float, + w_vb: torch.Tensor, + out: torch.Tensor, + workspace: K3MlaAttnWorkspace, + gate: Optional[torch.Tensor] = None, + gate_col0: int = 0, +) -> None: + """Write ``out`` [M, heads * 128] = per head ``bf16(bf16(o) @ w_vb[h]^T)`` (times the gate's sigmoid with + ``gate``) for the decode attention o of R <= 8 requests of T <= 8 tokens over their pages of ``pool``. Uses + ``workspace``'s partial slots and, past 7 clusters, advances its counters.""" + torch.ops.trtllm.k3_mla_attn_vb_out( + q, + pool, + row_stride, + page_table, + page_offset, + seq_len, + softmax_scale, + w_vb, + out, + workspace.buffer, + gate, + gate_col0, + ) + + +def k3_mla_attn_out( + q: torch.Tensor, + pool: torch.Tensor, + row_stride: int, + page_table: torch.Tensor, + page_offset: int, + seq_len: torch.Tensor, + softmax_scale: float, + out: torch.Tensor, + workspace: K3MlaAttnWorkspace, +) -> None: + """Write the attention output itself, ``out`` [M, heads * 512] bf16 (no v_b), over ``workspace``.""" + torch.ops.trtllm.k3_mla_attn_out( + q, pool, row_stride, page_table, page_offset, seq_len, softmax_scale, out, workspace.buffer + ) + + +def k3_mla_attn( + q: torch.Tensor, + pool: torch.Tensor, + row_stride: int, + page_table: torch.Tensor, + seq_len: torch.Tensor, + softmax_scale: float, + workspace: K3MlaAttnWorkspace, +) -> torch.Tensor: + """Return the attention output [M, heads * 512] bf16 in a new tensor (page offset 0), over ``workspace``.""" + return torch.ops.trtllm.k3_mla_attn( + q, pool, row_stride, page_table, seq_len, softmax_scale, workspace.buffer + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/k3_mla_attn_workspace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/k3_mla_attn_workspace.py new file mode 100644 index 000000000000..3514dcad6265 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/k3_mla_attn_workspace.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Caller-owned state of Kimi K3's MLA decode attention: the per-CTA partial slots and the no-cluster mode's arrival +counters of every ``attention/k3_mla_attn_vb_out`` call (and its ``k3_mla_attn`` / ``k3_mla_attn_out`` forms). + +A state type, not an entry: it launches nothing per call. Its constructor is eager; the target builds one per device +and head-group count in ``post_load_weights`` (before any CUDA-graph capture) and passes it to every MLA attention +call of that shape. The contract is the ``## State`` section of ``k3_mla_attn_vb_out.md``. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass(eq=False) +class K3MlaAttnWorkspace: + """One device's attention workspace for calls of ``groups`` head groups (the rank's heads / 6): fp16 partial + slots for 8 requests x ``groups`` x 16 CTAs, then the no-cluster mode's fp32 (m, l) exchange and its int32 + arrival counters, which only grow. Calls on one workspace must run one at a time (see the contract's + ``## State``).""" + + buffer: torch.Tensor + """fp16 [attn_workspace_elems(groups)]: the words the op's ``workspace`` argument names.""" + groups: int + + @classmethod + def create(cls, device: torch.device, groups: int) -> "K3MlaAttnWorkspace": + """Allocate and arm a workspace for ``groups`` head groups on ``device``: the counters zeroed, the partial + slots left as allocated (a call reads only the words it wrote). Eager: it allocates, so it refuses to run + under CUDA-graph capture; it returns once the arming is complete, so the workspace is ready on any stream.""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_mla import op + + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("K3MlaAttnWorkspace.create allocates: call it before capture") + device = torch.device(device) + buffer = op.make_attn_workspace(device, groups) + torch.cuda.synchronize(device) + return cls(buffer=buffer, groups=groups) + + @property + def device(self) -> torch.device: + return self.buffer.device diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/k3_mla_qkv.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/k3_mla_qkv.md new file mode 100644 index 000000000000..fb766b8d41e1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/k3_mla_qkv.md @@ -0,0 +1,155 @@ +--- +receipts: + sm_100: {status: passed, tests: 8} +--- + +# k3_mla_qkv + +**Wraps** `torch.ops.trtllm.k3_mla_qkv` (one call). The same contract covers two sibling wrappers, one call each: +`k3_mla_q` (`torch.ops.trtllm.k3_mla_q`, the query path alone) and `k3_mla_qkv_out` +(`torch.ops.trtllm.k3_mla_qkv_out`, the cache rows stored densely instead of into the pool). + +## Semantics + +Kimi K3's MLA decode query path for `M <= 64` tokens in one launch, plus the step's latent KV rows stored into the +paged latent cache. Per token `t` and head `h`, from the fused projection's rows `ag` (fp32 statistics, one bf16 +rounding wherever the unfused model rounds): + +``` +q_n[t] = bf16(rmsnorm(ag[t, :1536]) * w_qa) # q_a_layernorm +q[t, h] = bf16(q_n[t] @ w_qb[192 h : 192 h + 192]^T) # q_b_proj: 128 nope | 64 pe +fused_q[t, h] = [ bf16(q[t, h, :128] @ w_kb[h]^T) | q[t, h, 128:] ] # k_b absorption (512) | q_pe (64) +row[t] = [ bf16(rmsnorm(ag[t, 1536:2048]) * w_kv) | ag[t, 2048:2112] ] # kv_a_layernorm (512) | rope (64) +``` + +Kimi K3 is NoPE: `q_pe` and the rope columns are copied, not rotated. `fused_q` is returned in a new tensor +`[M, heads * 576]`; `row[t]` is stored into `pool`. The `M` tokens are `R` requests of `T = M / R` tokens each, +request-major: token `t = i T + u` is token `u` of request `i`, at position `p = seq_len[i] - T + u`, stored as row +`(page_table[i][p // 64] + page_offset) * 64 + p % 64` of `row_stride` elements (64-bit element index). A token +with `p < 0` is not stored. Nothing else in `pool` is written (certified: every layer's whole pool compared after +every step). + +- `k3_mla_q` returns the same `fused_q` bits and stores nothing (certified). +- `k3_mla_qkv_out` returns the same `fused_q` bits and stores `row[t]` densely into `kv_out[t]` (`[M, 576]`), the + same bits `k3_mla_qkv` stores into the pool (certified). +- `fused_q` rows are computed in chunks of 8, 16 or 32 tokens by call size; every 8-token chunk of rows is + bit-identical to the call on those 8 rows alone (the op test, `test_k3_mla_q.py::test_q`). + +Accuracy (certified bounds): `fused_q` within `2e-2` (max relative, per half) of a float64 reference that keeps the +model's bf16 roundings; the latent columns of `row` within `1e-2` of an fp32 reference (`(x * rrms) * w`, one bf16 +rounding); the rope columns bit-exact. + +Fusion boundary. Inside: both RMSNorms, the q_b GEMM, the k_b absorption, the cache append. Outside: the fused +projection that produced `ag` (`[q_a 1536 | kv_a latent 512 | rope 64 | gate heads * 128]`), the attention that reads +`fused_q` and the cache (`attention/k3_mla_attn_vb_out`), and the output projection. + +## Signature + +```python +def k3_mla_qkv( + ag: torch.Tensor, + w_qa: torch.Tensor, + eps: float, + w_qb: torch.Tensor, + w_kb: torch.Tensor, + w_kv: torch.Tensor, + kv_eps: float, + pool: torch.Tensor, + row_stride: int, + page_table: torch.Tensor, + page_offset: int, + seq_len: torch.Tensor, + trigger_early: bool = True, +) -> torch.Tensor + +def k3_mla_q(ag, w_qa, eps, w_qb, w_kb, trigger_early=True) -> torch.Tensor + +def k3_mla_qkv_out(ag, w_qa, eps, w_qb, w_kb, w_kv, kv_eps, kv_out, trigger_early=True) -> torch.Tensor +``` + +### Certified arguments + +| Argument | Shape | Dtype | Layout | Device | +|---|---|---|---|---| +| `ag` | `[M, C]`, `M` = `R T` up to 64; `C >= 2112`, a multiple of 8 (Kimi K3: `C = 2112 + heads * 128`) | bf16 | contiguous | CUDA | +| `w_qa`, `w_kv` | `[1536]`, `[512]` | bf16 | contiguous | CUDA | +| `eps`, `kv_eps` | scalar | Python float | — | — | +| `w_qb` | `[heads * 192, 1536]`, `heads` 6 (TP16) or 24 (TP4); the op takes 2 to 24 | bf16 | contiguous | CUDA | +| `w_kb` | `[heads, 512, 128]` (`k_b_proj_trans`) | bf16 | contiguous | CUDA | +| `pool` | the layer's paged latent pool, flat; at least 64 rows | bf16 | contiguous | CUDA | +| `row_stride` | elements per cache row, `>= 576`, a multiple of 8 (576 for the manager's MLA pool) | Python int | — | — | +| `page_table` | `[R, W]` int32 with unit column stride and any row stride `>= W` (rows of `kv_cache_block_offsets`), or `[W]` when `R = 1` | int32 | see shape | CUDA | +| `page_offset` | the layer's slot in a layer-interleaved pool, added to every page-table entry | Python int | — | — | +| `seq_len` | `[R]`, each request's length including its `T` new tokens | int32 | contiguous | CUDA | +| `trigger_early` | default `True`: dependents may launch early (they wait for the whole grid before reading) | Python bool | — | — | +| `kv_out` (`k3_mla_qkv_out`) | `[M, 576]` | bf16 | contiguous | CUDA | +| returns | `fused_q` `[M, heads * 576]`, newly allocated | bf16 | contiguous | = `ag.device` | + +Certified at `R` up to 8 and `T` 1 and 8 (no speculation; a DSpark verify step), 6 and 24 heads, over a real +`KVCacheManager` (below). + +## State + +**Object.** None of its own (kind P: it writes the caller's cache). `pool` is the KV cache manager's MLA latent +pool: `KVCacheManager.get_buffers(layer)` of a `SELFKONLY` manager (one latent head of 576, `kv_factor` 1, 64 tokens +per block), owned and sized by the manager. The addressing (`pool`, `row_stride`, `page_table`, `page_offset`, +`seq_len`) is the attention metadata's view of one generation step: `k3_mla_decode_view(attn, metadata, M)` +(`attention/backends/fmha/cute_dsl_mla.py`), computed per layer after `metadata.prepare()`. + +**What a call writes.** One 576-column row per token, at that token's position in its request (positions below 0 +skipped); nothing else (certified, above). The manager keeps every layer in one pool, interleaved by block: a layer's +rows sit in its slot of each block, and `page_offset` names the slot. + +**Which calls may share one object.** Every layer's call writes the same pool, kept apart by `page_offset` (each +layer's slot) and by position (each step's tokens). Two managers, such as two models' caches, are independent +(certified: calls alternating between two managers step by step and layer by layer; each pool holds exactly its own +rows). + +**Call order.** Calls of one step write disjoint rows, so their order does not matter; a step's rows must be stored +before the attention of the same layer reads them, which it does after its grid-dependency wait, so the attention +may directly follow this call. + +**What a launch reads before its grid-dependency wait.** The page-table rows and the lengths: host-filled metadata +buffers (`kv_cache_block_offsets`, `kv_lens_cuda_runtime`, written by `TrtllmAttentionMetadata.prepare()` before the +step runs or its graph replays). They must not be produced by the kernel directly before this one in the stream. +`ag` is read after the wait. + +**Re-arm.** Nothing to re-arm. + +**Why the test drives call sequences.** A misplaced row is silent: it lands in another layer's slot or another +position, and only a later read sees it. The test therefore keeps each layer's expected image from the manager's own +block ids (`get_batch_cache_indices`), not from the op's addressing, and compares the whole pool of every layer after +every step: 3 layers x 2 decode steps eagerly, then one step captured as a CUDA graph and replayed for 2 more steps +with rewritten inputs and the metadata prepared for each step. + +**What a wrong slot does.** Measured (the test's negative control): layer 1's call given layer 0's `page_offset`. +Nothing raises; layer 0's rows at the step's positions now hold layer 1's values, and layer 1's slot is not written. +The op trusts `page_offset`. + +## Metadata consumed + +Through `k3_mla_decode_view`: request `i`'s pages are `kv_cache_block_offsets[pool, num_contexts + i, 0, :]` (a +strided view, no copy), its length `kv_lens_cuda_runtime[num_contexts + i]`, `page_offset` the layer's index in the +pool, `row_stride` the pool's token stride. The view applies to generation-only steps (`num_contexts == 0`) of `R +<= 8` requests of the same `T <= 8` tokens, beam width 1, no tree speculation, 64-token pages, a bf16 pool, and +`heads` a multiple of 6 with `kv_lora_rank` 512 and `qk_rope_head_dim` 64; otherwise it returns a reason string and +the target takes its generic path. + +## Preconditions + +- `M <= 64`, and `page_table` / `seq_len` describe `R` requests of `T = M / R` tokens (`M % R == 0`; a `[W]` table + only when `R = 1`); `heads * 6 <= 148` (at most 24 heads) and, with the cache half, `heads >= 2`. Out of contract + calls raise `ValueError` before any launch, the pool unchanged (certified). +- `seq_len[i] >= T` for a request whose tokens should all be stored; a smaller length stores only the positions + `>= 0` (the op test, `test_qkv_short_length`). +- The first call of each configuration compiles the kernel and must run outside CUDA-graph capture (it raises + `RuntimeError` inside one). +- SM 100 / SM 103 only (CuTe DSL `tcgen05` kernel). + +## Notes + +- The op-level test is `tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_mla_q.py`: `fused_q` at `M` 1 to 64 + against the reference and the unfused chain (the model's RMSNorm, the q_b GEMM, the k_b bmm), the cache rows of R + x T steps in a sentinel-filled pool, an interleaved pool of rows of 640, the short-length case, refusals, and (with + `K3_BASE_TRTLLM`) batch-1 identity against the unmodified single-request kernel. +- The weights are TMA'd before the grid-dependency wait (EVICT_FIRST); the q_a and kv_a rows after it. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/k3_mla_qkv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/k3_mla_qkv.py new file mode 100644 index 000000000000..0e1345ff6f3a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/k3_mla_qkv.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3's MLA decode query path (q_a RMSNorm, q_b projection, k_b absorption) in one launch, which also stores the +step's latent KV rows into the paged latent cache; also the query path alone and the form that stores the rows +densely.""" + +import torch + +from tensorrt_llm._torch.cute_dsl_kernels.k3_mla import ( # noqa: F401 — registers the ops + op as _k3_mla_op, +) + + +def k3_mla_qkv( + ag: torch.Tensor, + w_qa: torch.Tensor, + eps: float, + w_qb: torch.Tensor, + w_kb: torch.Tensor, + w_kv: torch.Tensor, + kv_eps: float, + pool: torch.Tensor, + row_stride: int, + page_table: torch.Tensor, + page_offset: int, + seq_len: torch.Tensor, + trigger_early: bool = True, +) -> torch.Tensor: + """Return ``fused_q`` [M, heads * 576] bf16 of the M <= 64 rows of ``ag``, and store each token's cache row + ``[bf16(rmsnorm(ag[t, 1536:2048]) * w_kv) | ag[t, 2048:2112]]`` into ``pool`` at its request's position.""" + return torch.ops.trtllm.k3_mla_qkv( + ag, + w_qa, + eps, + w_qb, + w_kb, + w_kv, + kv_eps, + pool, + row_stride, + page_table, + page_offset, + seq_len, + trigger_early, + ) + + +def k3_mla_q( + ag: torch.Tensor, + w_qa: torch.Tensor, + eps: float, + w_qb: torch.Tensor, + w_kb: torch.Tensor, + trigger_early: bool = True, +) -> torch.Tensor: + """Return ``fused_q`` alone (no cache row stored): the same bits as ``k3_mla_qkv``'s.""" + return torch.ops.trtllm.k3_mla_q(ag, w_qa, eps, w_qb, w_kb, trigger_early) + + +def k3_mla_qkv_out( + ag: torch.Tensor, + w_qa: torch.Tensor, + eps: float, + w_qb: torch.Tensor, + w_kb: torch.Tensor, + w_kv: torch.Tensor, + kv_eps: float, + kv_out: torch.Tensor, + trigger_early: bool = True, +) -> torch.Tensor: + """Return ``fused_q`` and store the cache rows densely into ``kv_out`` [M, 576] instead of the pool.""" + return torch.ops.trtllm.k3_mla_qkv_out( + ag, w_qa, eps, w_qb, w_kb, w_kv, kv_eps, kv_out, trigger_early + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.md index 7b32f301dc29..7044e9449ced 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.md +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.md @@ -1,6 +1,8 @@ --- receipts: + # sm_103 was certified before the sm_100 key was added, on the same test file. sm_103: {status: passed, tests: 5} + sm_100: {status: passed, tests: 5} --- # cublas_mm diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv.md new file mode 100644 index 000000000000..13497756a6b3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv.md @@ -0,0 +1,139 @@ +--- +receipts: + sm_100: {status: passed, tests: 4} +--- + +# k3_ctm_gemv + +**Wraps** `torch.ops.trtllm.k3_ctm_gemv` (one call). + +## Semantics + +The decode GEMV of a Kimi K3 projection whose reduction dimension fits in shared memory: at most 8 bf16 tokens +times a bf16 weight in `nn.Linear` layout (the bias-free `F.linear(x, weight)`), as one CuTe DSL kernel on the +tensor cores. + +``` +y[m, n] = bf16( sum_k x[m, k] * weight[n, k] ) 0 <= m < M <= 8, 0 <= n < N, 0 <= k < K +``` + +The bf16 products are accumulated in fp32 by the tensor cores (tcgen05 MMAs of 128 weight rows x 8 token columns +x 16 k, into tensor memory) and the sum is rounded to bf16 once, round-to-nearest-even. The summation order is +fixed by `K` and `split`: + +- `K` is cut into 128-column k-tiles. Each block of 128 output rows is one CTA (`split=1`) or one cluster of + `split` CTAs. +- With `split=1` the CTA accumulates all k-tiles in ascending order. +- With `split=2` or `4`, rank `r` of the cluster accumulates k-tiles `r, r + split, r + 2 * split, ...` in + ascending order into its own fp32 partial (when `split` does not divide the k-tiles, the first + `k_tiles % split` ranks take one more). The rank that owns an output row adds the `split` partials in rank order + `0, 1, ...` in fp32 and rounds that sum once. + +`push` only chooses how partials travel to the owning rank: DSMEM stores plus a cluster-scope release arrive +(`False`), or 16-byte `st.async` stores that complete the owner's barrier by bytes (`True`). The sums and their +order are the same either way, and `push` has no effect at `split=1`. Token rows are independent: rows past `M` +enter the MMA as zeros and are not stored, so a token's output row has the same bits whatever `M` is. + +Fusion boundary: the single call computes the GEMV only. There is no bias, activation, residual, all-reduce +(under tensor parallelism the per-rank partial is returned as is) or quantization; the caller owns those. + +## Signature + +```python +def k3_ctm_gemv( + x: torch.Tensor, + weight: torch.Tensor, + trigger_early: bool = True, + split: int = 1, + push: bool = False, +) -> torch.Tensor +``` + +| Argument | Shape | Dtype | Layout | Device | +|---|---|---|---|---| +| `x` | `[M, K]`, `1 <= M <= 8` | bf16 | contiguous, 16-byte-aligned start | CUDA | +| `weight` | `[N, K]` (`nn.Linear` layout) | bf16 | contiguous, 16-byte-aligned start | CUDA, `x`'s device | +| `trigger_early` | scalar | bool: let PDL dependents launch early | — | — | +| `split` | scalar | int: 1, 2 or 4 CTAs per 128 output rows | — | — | +| `push` | scalar | bool: split-K transport | — | — | +| returns | `[M, N]` | bf16 | newly allocated, contiguous | `x`'s device | + +Every element of the returned tensor is written; the inputs are not mutated. + +### Certified arguments + +The Kimi K3 TP16 per-rank shapes with the `split` and `push` their call sites pass, `trigger_early=True`, at +every `M` in 1..8: + +| Cell | `N` | `K` | `split` | `push` | +|---|---|---|---|---| +| MLA o_proj | 7168 | 768 | 1 | False | +| MLA o_proj | 7168 | 768 | 2 | False | +| drafter o_proj | 7168 | 384 | 1 | True | +| drafter o_proj, synthetic drafter | 7168 | 256 | 1 | True | + +Inputs are `x ~ N(0, 1)` and `weight ~ 0.03 * N(0, 1)`, rounded to bf16. The gate per cell: +`max|y - ref| / max|ref| <= 8e-3` against the fp64 product of the same bf16 inputs, identical bits on a repeated +call, and each `M`'s rows bit-identical to the same rows of the 8-row call. Also certified, at the MLA o_proj: + +- with `split=2`, `push=True` and `trigger_early=False` return the bits of the `push=False`, + `trigger_early=True` call, at every `M`; +- with `split` 1 and 2, `M` 1 and 8: calls captured in a CUDA graph after an eager call per key, replayed with + `x` rewritten in place, return the bits of eager calls on the new `x`. + +`split=4` and other shapes the preconditions admit are accepted but not certified. + +## Metadata consumed + +None. The op reads no attention metadata, KV cache or module state; every input is an argument, and the kernel +runs on the current CUDA stream of `x`'s device. + +One process-global cache sits behind it, and it is result-neutral: the compiled kernel per key +`(N, K, split, trigger_early, push, PDL on/off)`. The first call with a new key compiles the kernel with the CuTe +DSL, which costs host time; `M` is a runtime argument and never recompiles. That first call refuses to run under +CUDA-graph capture: it raises `RuntimeError` ("run once per shape outside CUDA-graph capture first") before +launching anything. Call every key once eagerly, then capture; captured and later calls reuse the compiled +kernel. The PDL setting is read from `TRTLLM_ENABLE_PDL` on every call, so changing it within a process adds a +key. + +## Preconditions + +The op checks these and raises `ValueError` ("k3_ctm_gemv: unsupported call ...") when one fails: + +- `x` is a 2-D contiguous CUDA bf16 tensor with `1 <= M <= 8` rows. +- `weight` is a 2-D contiguous bf16 tensor with `weight.shape[1] == x.shape[1]`. +- `N % 128 == 0` and `K % 128 == 0`. +- `split` is 1, 2 or 4, `split <= K / 128`, and no rank holds more than 6 k-tiles: + `ceil(K / 128 / split) <= 6`, i.e. `K <= 768` at `split=1`, `K <= 1536` at 2, `K <= 3072` at 4. + +Certified refusals: 0 and 9 rows, fp16 `x`, a row-strided `x`, `N = 7104`, `K = 704`, `split=3`, and `K = 896` +(7 k-tiles) at `split=1`. + +Not checked by the op: + +- `x` and `weight` start at 16-byte-aligned addresses. The op passes both to the kernel with a declared 16-byte + alignment and loads them by TMA; a view whose start is not a multiple of 8 elements past an aligned allocation + is outside the contract. Row slices `t[a:b]` of these shapes are aligned. +- `weight` is on `x`'s device; the op checks only `x.is_cuda`. +- The GPU has compute capability 10.x: the kernel uses tcgen05 MMA, tensor memory and, at `split > 1`, + thread-block clusters. +- The CuTe DSL (`cutlass`) and `cuda-python` (`cuda.bindings`) are importable. The op's check imports the kernel + module, so without them a call raises `ImportError`, not `ValueError`. + +## Notes + +- Programmatic dependent launch (PDL). With `TRTLLM_ENABLE_PDL` unset or `1` the kernel is launched with PDL; any + other value launches it without. Each CTA issues its whole weight read (TMA, L2 evict-first) without waiting + for the grid dependency; only the read of `x` follows `griddepcontrol.wait`. So `weight` must not be written by + work still running ahead of this call on the stream, while `x` may be. With `trigger_early=True` each CTA executes + `griddepcontrol.launch_dependents` right after issuing its weight loads: the next kernel on the stream, if + launched with PDL, can start while this one runs, and it must execute `griddepcontrol.wait` + (`cudaGridDependencySynchronize`) before it reads `y`. Kernels launched without PDL, torch's included, start + after this one completes as usual. With `trigger_early=False` dependents launch when this grid completes. + None of this changes results. +- Two identical calls return identical bits. The bits depend on `split`, whose summation orders can differ in the + last bf16 place, but not on `push`, `trigger_early`, the PDL setting or `M`. +- The result is not bit-identical to cuBLAS: `F.linear` sums in another order. The two agree within the + tolerance above. +- Grid: `(N / 128) * split` CTAs of 256 threads, in clusters of `split` when `split > 1`. +- Only sm_100 has been measured. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv.py new file mode 100644 index 000000000000..d8d1de17eab3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 decode GEMV x @ weight^T for M <= 8 bf16 tokens via the CTM (CuTe DSL) kernel.""" + +import torch + +import tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv.op # noqa: F401 (registers torch.ops.trtllm.k3_ctm_gemv*) + + +def k3_ctm_gemv( + x: torch.Tensor, + weight: torch.Tensor, + trigger_early: bool = True, + split: int = 1, + push: bool = False, +) -> torch.Tensor: + """Return `bf16(x @ weight.T)`, fp32-accumulated, in one k3_ctm_gemv call.""" + return torch.ops.trtllm.k3_ctm_gemv( + x, weight, trigger_early=trigger_early, split=split, push=push + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv_long.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv_long.md new file mode 100644 index 000000000000..2af2cbcde62c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv_long.md @@ -0,0 +1,156 @@ +--- +receipts: + sm_100: {status: passed, tests: 4} +--- + +# k3_ctm_gemv_long + +**Wraps** `torch.ops.trtllm.k3_ctm_gemv_long` (one call). + +## Semantics + +The decode GEMV of a Kimi K3 projection with a long reduction dimension (the hidden size 7168, or any `K` too long +to hold in shared memory): at most 8 bf16 tokens times a bf16 weight in `nn.Linear` layout, as one CuTe DSL kernel +that splits each 128-row block of the weight over a cluster of `split` CTAs and streams it through a `ring` of +shared-memory stages. Output columns from `sig_col0` on can hold the sigmoid of the product instead. + +``` +acc[m, n] = sum_k x[m, k] * weight[n, k] fp32 accumulation, 0 <= m < M <= 8, 0 <= n < N +y[m, n] = bf16(acc[m, n]) n < sig_col0, or every n when sig_col0 < 0 +y[m, n] = bf16(sigmoid(bf16(acc[m, n]))) n >= sig_col0 >= 0; sigmoid(v) = 1 / (1 + exp(-v)) in fp32 +``` + +The bf16 products are accumulated in fp32 by the tensor cores (tcgen05 MMAs of 128 weight rows x 8 token columns +x 16 k, into tensor memory). Summation order: `K` is cut into 128-column k-tiles; the last one may be a 64-column +half, for which the TMA reads zeros past `K` in both operands. Rank `r` of a block's cluster accumulates k-tiles +`r, r + split, r + 2 * split, ...` in ascending order into its own fp32 partial (when `split` does not divide the +k-tiles, the first `k_tiles % split` ranks take one more). Each 32-row quarter `q` of the block has an owning rank +(`q`, or `q // 2` at `split=2`), which adds the `split` partials of its rows in rank order `0, 1, ...` in fp32 and +rounds the sum once, round-to-nearest-even. A sigmoid column rounds that sum to bf16, takes the sigmoid of the +bf16 value in fp32 (IEEE division, full-precision `exp`) and rounds again: its bits are exactly `torch.sigmoid` of +the bf16 value the same call returns with `sig_col0=-1`. That is the form of the MLA output gate +`bf16(sigmoid(g))` in the fused `[W_a; W_g]` projection. + +`ring` (the weight stages each CTA keeps in shared memory) and `push` (how partials travel to the owning rank: +DSMEM stores plus a cluster-scope release arrive, or 16-byte `st.async` stores completing the owner's barrier by +bytes) change neither the sums nor their order. Rows past `M` enter the MMA as zeros and are not stored, and rows +past `N` are not stored, so a token's output row has the same bits whatever `M` is. + +Fusion boundary: the single call computes the GEMV and, from `sig_col0` on, the sigmoid. There is no bias, other +activation, residual, all-reduce (under tensor parallelism the per-rank partial is returned as is) or +quantization; the caller owns those. + +## Signature + +```python +def k3_ctm_gemv_long( + x: torch.Tensor, + weight: torch.Tensor, + sig_col0: int = -1, + split: int = 6, + ring: int = 5, + trigger_early: bool = True, + push: bool = False, +) -> torch.Tensor +``` + +| Argument | Shape | Dtype | Layout | Device | +|---|---|---|---|---| +| `x` | `[M, K]`, `1 <= M <= 8` | bf16 | contiguous, 16-byte-aligned start | CUDA | +| `weight` | `[N, K]` (`nn.Linear` layout) | bf16 | contiguous, 16-byte-aligned start | CUDA, `x`'s device | +| `sig_col0` | scalar | int: first sigmoid column; `< 0` or `>= N`: none | — | — | +| `split` | scalar | int: 2, 4, 5, 6, 7 or 8 CTAs per 128 output rows | — | — | +| `ring` | scalar | int: weight k-tiles staged per CTA (see Preconditions) | — | — | +| `trigger_early` | scalar | bool: let PDL dependents launch early | — | — | +| `push` | scalar | bool: split-K transport | — | — | +| returns | `[M, N]` | bf16 | newly allocated, contiguous | `x`'s device | + +Every element of the returned tensor is written; the inputs are not mutated. + +### Certified arguments + +The Kimi K3 TP16 per-rank shapes with the `sig_col0`, `split`, `ring` and `push` their call sites pass, +`trigger_early=True`, at every `M` in 1..8: + +| Cell | `N` | `K` | `sig_col0` | `split` | `ring` | `push` | +|---|---|---|---|---|---|---| +| MLA `[W_a; W_g]`, gate columns as sigmoid | 2880 | 7168 | 2112 | 6 | 6 | True | +| dense MLP gate_up | 4224 | 7168 | -1 | 4 | 5 | False | +| dense MLP down (last k-tile a half) | 7168 | 2112 | -1 | 2 | 6 | False | +| drafter qkv | 512 | 7168 | -1 | 8 | 6 | True | +| drafter gate_up | 1792 | 7168 | -1 | 8 | 6 | True | +| drafter gate_up, synthetic drafter | 1536 | 7168 | -1 | 8 | 6 | True | +| KDA q/k/v/g/f_a/b (last block 8 rows) | 3208 | 7168 | -1 | 5 | 6 | True | + +Inputs are `x ~ N(0, 1)` and `weight ~ 0.02 * N(0, 1)`, rounded to bf16. The gate per cell: the `sig_col0=-1` +output within `max|y - ref| / max|ref| <= 8e-3` of the fp64 product of the same bf16 inputs; with `sig_col0 >= 0`, +the columns before it bit-identical to that output and the columns from it bit-identical to `torch.sigmoid` of it; +identical bits on a repeated call; and each `M`'s rows bit-identical to the same rows of the 8-row call. Also +certified: + +- one flag changed from the call site's values returns the call site's bits, at every `M`: `push=True` and + `ring=3` at the dense MLP down, `push=True` and `trigger_early=False` at the dense MLP gate_up; +- at the MLA `[W_a; W_g]` and dense MLP down cells, `M` 1 and 8: calls captured in a CUDA graph after an eager + call per key, replayed with `x` rewritten in place, return the bits of eager calls on the new `x`. + +Other values the preconditions admit are accepted but not certified. + +## Metadata consumed + +None. The op reads no attention metadata, KV cache or module state; every input is an argument, and the kernel +runs on the current CUDA stream of `x`'s device. + +One process-global cache sits behind it, and it is result-neutral: the compiled kernel per key +`(N, K, split, ring, trigger_early, push, PDL on/off)`. The first call with a new key compiles the kernel with the +CuTe DSL, which costs host time; `M` and `sig_col0` are runtime arguments and never recompile. That first call +refuses to run under CUDA-graph capture: it raises `RuntimeError` ("run once per shape outside CUDA-graph capture +first") before launching anything. Call every key once eagerly, then capture; captured and later calls reuse the +compiled kernel. The PDL setting is read from `TRTLLM_ENABLE_PDL` on every call, so changing it within a process +adds a key. + +## Preconditions + +The op checks these and raises `ValueError` ("k3_ctm_gemv_long: unsupported call ...") when one fails: + +- `x` is a 2-D contiguous CUDA bf16 tensor with `1 <= M <= 8` rows. +- `weight` is a 2-D contiguous bf16 tensor with `weight.shape[1] == x.shape[1]` and `N >= 1`. `N` needs no + divisibility: the last 128-row block may be partial. +- `K % 64 == 0`. +- `split` is 2, 4, 5, 6, 7 or 8. +- `1 <= ring <= ceil(K / 128) // split` (so `ceil(K / 128) >= split`), and the ring and the resident activation + fit 216 KiB of shared memory: `ring * 32 KiB + (ceil(K / 128) // split + 1) * 2 KiB <= 216 KiB`. That caps + `ring` at 6, and at 5 for `K = 7168` with `split=4`. + +Certified refusals, at `N = 7168`, `K = 2112` (17 k-tiles): 0 and 9 rows, fp16 `x`, a row-strided `x`, +`K = 2080`, `split` 3 and 9, `ring=0`, `ring=3` at `split=8` (2 k-tiles per rank) and `ring=7` at `split=2` +(past shared memory). + +Not checked by the op: + +- `x` and `weight` start at 16-byte-aligned addresses. The op passes both to the kernel with a declared 16-byte + alignment and loads them by TMA; a view whose start is not a multiple of 8 elements past an aligned allocation + is outside the contract. Row slices `t[a:b]` are aligned, since `K % 64 == 0`. +- `weight` is on `x`'s device; the op checks only `x.is_cuda`. +- The GPU has compute capability 10.x: the kernel uses tcgen05 MMA, tensor memory and thread-block clusters. +- The CuTe DSL (`cutlass`) and `cuda-python` (`cuda.bindings`) are importable. The op's check imports the kernel + module, so without them a call raises `ImportError`, not `ValueError`. + +## Notes + +- Programmatic dependent launch (PDL). With `TRTLLM_ENABLE_PDL` unset or `1` the kernel is launched with PDL; any + other value launches it without. Each CTA fills its weight ring (TMA, L2 evict-first) and prefetches the rest of + its weight k-tiles into L2 without waiting for the grid dependency; only the read of `x` follows + `griddepcontrol.wait`. So `weight` must not be written by work still running ahead of this call on the stream, + while `x` may be. With `trigger_early=True` each CTA executes `griddepcontrol.launch_dependents` right after + issuing that first weight traffic: the next kernel on the stream, if launched with PDL, can start while this one + runs, and it must execute `griddepcontrol.wait` (`cudaGridDependencySynchronize`) before it reads `y`. Kernels + launched without PDL, torch's included, start after this one completes as usual. With `trigger_early=False` + dependents launch when this grid completes. None of this changes results. +- Two identical calls return identical bits. The bits depend on `split`, whose summation orders can differ in the + last bf16 place, but not on `ring`, `push`, `trigger_early`, the PDL setting or `M`. `k3_ctm_gemv_wide` at the + same split returns the same bits for each token. +- Grid: `ceil(N / 128) * split` CTAs of 256 threads in clusters of `split`; the op does not limit it to one + wave. +- The result is not bit-identical to cuBLAS: `F.linear` sums in another order. The two agree within the + tolerance above. +- Only sm_100 has been measured. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv_long.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv_long.py new file mode 100644 index 000000000000..a55c8bc21b7c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv_long.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 long-K decode GEMV x @ weight^T for M <= 8 tokens via the split-K CTM (CuTe DSL) kernel.""" + +import torch + +import tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv.op # noqa: F401 (registers torch.ops.trtllm.k3_ctm_gemv*) + + +def k3_ctm_gemv_long( + x: torch.Tensor, + weight: torch.Tensor, + sig_col0: int = -1, + split: int = 6, + ring: int = 5, + trigger_early: bool = True, + push: bool = False, +) -> torch.Tensor: + """Return `bf16(x @ weight.T)`, columns >= `sig_col0` (if >= 0) as `bf16(sigmoid(.))`, in one call.""" + return torch.ops.trtllm.k3_ctm_gemv_long( + x, + weight, + sig_col0=sig_col0, + split=split, + ring=ring, + trigger_early=trigger_early, + push=push, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv_swiglu.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv_swiglu.md new file mode 100644 index 000000000000..560710673676 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv_swiglu.md @@ -0,0 +1,138 @@ +--- +receipts: + sm_100: {status: passed, tests: 4} +--- + +# k3_ctm_gemv_swiglu + +**Wraps** `torch.ops.trtllm.k3_ctm_gemv_swiglu` (one call). + +## Semantics + +The down projection of a SwiGLU MLP with the activation folded into it: the `k3_ctm_gemv` GEMV whose input is +`silu(gate) * up` of a gate_up output, computed inside the kernel and never written to memory. + +``` +K = weight.shape[1]; g = gu[m, k], u = gu[m, K + k] (gate columns first) +act[m, k] = bf16( (g * sigmoid(g)) * u ) in fp32; sigmoid(g) = 1 / (1 + exp(-g)) +y[m, n] = bf16( sum_k act[m, k] * weight[n, k] ) 0 <= m < M <= 8 +``` + +The activation is evaluated in fp32 in that order, with IEEE division and the full-precision `exp`, and rounded +to bf16 once, as `silu_and_mul` rounds its output. The GEMV is `k3_ctm_gemv`'s: the products are accumulated in +fp32 by the tensor cores (tcgen05 MMA into tensor memory) and the sum is rounded to bf16 once, round-to-nearest-even. +Summation order: `K` is cut into 128-column k-tiles; each block of 128 output rows runs on `split` CTAs (a +cluster when `split > 1`), rank `r` accumulating k-tiles `r, r + split, ...` in ascending order into its own fp32 +partial (when `split` does not divide the k-tiles, the first `k_tiles % split` ranks take one more); the rank that +owns an output row adds the `split` partials in rank order in fp32 and rounds once. `push` only chooses how +partials travel to the owner (DSMEM stores plus a release arrive, or `st.async` stores completing the owner's +barrier by bytes); the sums and their order are the same. Rows past `M` are zeros and are not stored, so a token's +output row has the same bits whatever `M` is. + +The activation is the formula above, not a call of another kernel. torch's +`(F.silu(gate.float()) * up.float()).bfloat16()` computes the same function, but not necessarily with the same fp32 +operations, so individual activation values can differ from it in the last bf16 place. + +Fusion boundary: the single call computes the activation and the GEMV. The caller owns the gate_up projection +that produces `gu`. There is no bias, all-reduce (under tensor parallelism the per-rank partial is returned as is) +or quantization. + +## Signature + +```python +def k3_ctm_gemv_swiglu( + gu: torch.Tensor, + weight: torch.Tensor, + trigger_early: bool = True, + split: int = 2, + push: bool = False, +) -> torch.Tensor +``` + +| Argument | Shape | Dtype | Layout | Device | +|---|---|---|---|---| +| `gu` | `[M, 2 * K]`, gate columns first, `1 <= M <= 8` | bf16 | contiguous, 16-byte-aligned start | CUDA | +| `weight` | `[N, K]` (`nn.Linear` layout) | bf16 | contiguous, 16-byte-aligned start | CUDA, `gu`'s device | +| `trigger_early` | scalar | bool: let PDL dependents launch early | — | — | +| `split` | scalar | int: 1, 2 or 4 CTAs per 128 output rows | — | — | +| `push` | scalar | bool: split-K transport | — | — | +| returns | `[M, N]` | bf16 | newly allocated, contiguous | `gu`'s device | + +Every element of the returned tensor is written; the inputs are not mutated. + +### Certified arguments + +The Kimi K3 TP16 per-rank drafter down projection with the `split` and `push` its call site passes, +`trigger_early=True`, at every `M` in 1..8: + +| Cell | `N` | `K` | `split` | `push` | +|---|---|---|---|---| +| drafter down | 7168 | 896 | 2 | True | +| drafter down, synthetic drafter | 7168 | 768 | 2 | True | + +`K = 896` is 7 k-tiles, so the two ranks hold 4 and 3. Inputs are `gu ~ 2 * N(0, 1)` and +`weight ~ 0.02 * N(0, 1)`, rounded to bf16. The gate per cell: `max|y - ref| / max|ref| <= 8e-3` against the fp64 +product of `weight` and torch's bf16 activation `(F.silu(gate.float()) * up.float()).bfloat16()`, identical bits on +a repeated call, and each `M`'s rows bit-identical to the same rows of the 8-row call. Also certified, at the +drafter down: + +- `push=False` and `trigger_early=False` return the call site's bits, at every `M`; +- at `M` 1 and 8, calls captured in a CUDA graph after an eager call per key, replayed with `gu` rewritten in + place, return the bits of eager calls on the new `gu`. + +`split` 1 and 4 and other shapes the preconditions admit are accepted but not certified. + +## Metadata consumed + +None. The op reads no attention metadata, KV cache or module state; every input is an argument, and the kernel +runs on the current CUDA stream of `gu`'s device. + +One process-global cache sits behind it, and it is result-neutral: the compiled kernel per key +`(N, K, split, trigger_early, push, PDL on/off)`. The first call with a new key compiles the kernel with the CuTe +DSL, which costs host time; `M` is a runtime argument and never recompiles. That first call refuses to run under +CUDA-graph capture: it raises `RuntimeError` ("run once per shape outside CUDA-graph capture first") before +launching anything. Call every key once eagerly, then capture; captured and later calls reuse the compiled +kernel. The PDL setting is read from `TRTLLM_ENABLE_PDL` on every call, so changing it within a process adds a +key. + +## Preconditions + +The op checks these and raises `ValueError` ("k3_ctm_gemv_swiglu: unsupported call ...") when one fails: + +- `gu` is a 2-D contiguous CUDA bf16 tensor with `1 <= M <= 8` rows and `gu.shape[1] == 2 * weight.shape[1]`. +- `weight` is a 2-D contiguous bf16 tensor. +- `N % 128 == 0` and `K % 128 == 0`. +- `split` is 1, 2 or 4, `split <= K / 128`, and no rank holds more than 6 k-tiles: + `ceil(K / 128 / split) <= 6`, i.e. `K <= 768` at `split=1`, `K <= 1536` at 2, `K <= 3072` at 4. + +Certified refusals: 0 and 9 rows, fp16 `gu`, a row-strided `gu`, a `gu` of width `1664 != 2 * 896`, `N = 7104`, +`split=3`, and `K = 896` (7 k-tiles) at `split=1`. + +Not checked by the op: + +- `gu` and `weight` start at 16-byte-aligned addresses. The op passes both to the kernel with a declared 16-byte + alignment and loads them by TMA; a view whose start is not a multiple of 8 elements past an aligned allocation + is outside the contract. Row slices `t[a:b]` of these shapes are aligned. +- `weight` is on `gu`'s device; the op checks only `gu.is_cuda`. +- The GPU has compute capability 10.x: the kernel uses tcgen05 MMA, tensor memory and, at `split > 1`, + thread-block clusters. +- The CuTe DSL (`cutlass`) and `cuda-python` (`cuda.bindings`) are importable. The op's check imports the kernel + module, so without them a call raises `ImportError`, not `ValueError`. + +## Notes + +- Programmatic dependent launch (PDL). With `TRTLLM_ENABLE_PDL` unset or `1` the kernel is launched with PDL; any + other value launches it without. Each CTA issues its whole weight read (TMA, L2 evict-first) without waiting + for the grid dependency; only the read of `gu` follows `griddepcontrol.wait`. So `weight` must not be written by + work still running ahead of this call on the stream, while `gu` may be. With `trigger_early=True` each CTA + executes `griddepcontrol.launch_dependents` right after issuing its weight loads: the next kernel on the stream, + if launched with PDL, can start while this one runs, and it must execute `griddepcontrol.wait` + (`cudaGridDependencySynchronize`) before it reads `y`. Kernels launched without PDL, torch's included, start + after this one completes as usual. With `trigger_early=False` dependents launch when this grid completes. None + of this changes results. +- Two identical calls return identical bits. The bits depend on `split`, whose summation orders can differ in the + last bf16 place, but not on `push`, `trigger_early`, the PDL setting or `M`. +- The result is not bit-identical to `silu_and_mul` followed by cuBLAS: the activation can differ in the last bf16 + place (see Semantics) and `F.linear` sums in another order. They agree within the tolerance above. +- Grid: `(N / 128) * split` CTAs of 256 threads, in clusters of `split` when `split > 1`. +- Only sm_100 has been measured. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv_swiglu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv_swiglu.py new file mode 100644 index 000000000000..ffb558ebe793 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv_swiglu.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SwiGLU down projection silu_and_mul(gu) @ weight^T for M <= 8 tokens via the CTM (CuTe DSL) kernel.""" + +import torch + +import tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv.op # noqa: F401 (registers torch.ops.trtllm.k3_ctm_gemv*) + + +def k3_ctm_gemv_swiglu( + gu: torch.Tensor, + weight: torch.Tensor, + trigger_early: bool = True, + split: int = 2, + push: bool = False, +) -> torch.Tensor: + """Return `bf16(silu_and_mul(gu) @ weight.T)`, gate columns first, in one k3_ctm_gemv_swiglu call.""" + return torch.ops.trtllm.k3_ctm_gemv_swiglu( + gu, weight, trigger_early=trigger_early, split=split, push=push + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv_wide.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv_wide.md new file mode 100644 index 000000000000..d2e27a39d893 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv_wide.md @@ -0,0 +1,156 @@ +--- +receipts: + sm_100: {status: passed, tests: 4} +--- + +# k3_ctm_gemv_wide + +**Wraps** `torch.ops.trtllm.k3_ctm_gemv_wide` (one call). + +## Semantics + +The `k3_ctm_gemv_long` GEMV for up to 64 bf16 tokens (a wide decode step, e.g. speculative verification of several +requests): all `M` token columns go through one tensor-core MMA of 16, 32 or 64 columns, the smallest that holds +`M`, on the long kernel's split-K clusters and weight ring. Split and ring are not arguments; the op picks them +from the shape and the device. Output columns from `sig_col0` on can hold the sigmoid of the product, or the whole +output can be fp32. + +``` +acc[m, n] = sum_k x[m, k] * weight[n, k] fp32 accumulation, 0 <= m < M <= 64, 0 <= n < N +out_fp32=False: y[m, n] = bf16(acc[m, n]) n < sig_col0, or every n when sig_col0 < 0 + y[m, n] = bf16(sigmoid(bf16(acc[m, n]))) n >= sig_col0 >= 0; sigmoid in fp32 +out_fp32=True: y[m, n] = acc[m, n] fp32, not rounded +``` + +The bf16 products are accumulated in fp32 by the tensor cores (tcgen05 MMA into tensor memory). `split` and `ring` +come from `wide_config(N, K, token tile, SM count)` in the op module: the largest split in (8, 7, 6, 5, 4, 2) whose +`ceil(N / 128) * split` CTAs fit one wave of the device's SMs (with `ceil(K / 128) >= split`), then the deepest +weight ring that fits shared memory beside an activation ring of up to 3 stages. The summation order is +`k3_ctm_gemv_long`'s at that split: `K` in 128-column k-tiles (the last may be a 64-column half, read as zeros past +`K`), rank `r` of a 128-row block's cluster accumulating k-tiles `r, r + split, ...` in ascending order into its +own fp32 partial (the first `k_tiles % split` ranks one more when `split` does not divide them), and the `split` +partials of a token added in rank order in fp32 by the rank that reduces that token, then rounded once, +round-to-nearest-even. So a token's row is bit-identical to the row `k3_ctm_gemv_long` returns for that token at +the same split, with any ring and either transport, and it does not depend on `M`. A sigmoid column rounds the sum +to bf16, takes the sigmoid in fp32 (IEEE division, full-precision `exp`) and rounds again: its bits are exactly +`torch.sigmoid` of the bf16 value the same call returns with `sig_col0=-1`. + +The split follows the device's SM count, so the bits can differ between devices with different SM counts. On +148- and 152-SM devices the certified shapes run at the splits listed under *Certified arguments*. + +Fusion boundary: the single call computes the GEMV and either the sigmoid from `sig_col0` on or an fp32 output. +There is no bias, other activation, residual, all-reduce (under tensor parallelism the per-rank partial is +returned as is) or quantization; the caller owns those. + +## Signature + +```python +def k3_ctm_gemv_wide( + x: torch.Tensor, + weight: torch.Tensor, + sig_col0: int = -1, + out_fp32: bool = False, +) -> torch.Tensor +``` + +| Argument | Shape | Dtype | Layout | Device | +|---|---|---|---|---| +| `x` | `[M, K]`, `1 <= M <= 64` | bf16 | contiguous, 16-byte-aligned start | CUDA | +| `weight` | `[N, K]` (`nn.Linear` layout) | bf16 | contiguous, 16-byte-aligned start | CUDA, `x`'s device | +| `sig_col0` | scalar | int: first sigmoid column, `< N`; `< 0`: none | — | — | +| `out_fp32` | scalar | bool: fp32 output, no rounding (needs `sig_col0 < 0`) | — | — | +| returns | `[M, N]` | bf16, or fp32 with `out_fp32` | newly allocated, contiguous | `x`'s device | + +Every element of the returned tensor is written; the inputs are not mutated. + +### Certified arguments + +The Kimi K3 TP16 per-rank shapes of the projections of a wide decode step, at every `M` in 1..64 (all three token +tiles): + +| Cell | `N` | `K` | `sig_col0` | `out_fp32` | split at 148 / 152 SMs | +|---|---|---|---|---|---| +| KDA q/k/v/g/f_a/b (last block 8 rows) | 3208 | 7168 | -1 | False | 5 | +| MLA `[W_a; W_g]`, gate columns as sigmoid | 2880 | 7168 | 2112 | False | 6 | +| KDA / MLA o_proj | 7168 | 768 | -1 | False | 2 | +| MoE head (latent down slice; router rows) | 280 | 7168 | -1 | True | 8 | +| MoE tail (5 k-tiles) | 7168 | 640 | -1 | False | 2 | +| shared expert gate_up | 768 | 7168 | -1 | False | 8 | +| dense MLP gate_up | 4224 | 7168 | -1 | False | 4 | +| dense MLP down (last k-tile a half) | 7168 | 2112 | -1 | False | 2 | +| drafter qkv | 512 | 7168 | -1 | False | 8 | +| drafter gate_up | 1792 | 7168 | -1 | False | 8 | +| drafter o_proj | 7168 | 384 | -1 | False | 2 | + +Inputs are `x ~ N(0, 1)` and `weight ~ 0.02 * N(0, 1)`, rounded to bf16. The gate per cell: the `sig_col0=-1` +output within `max|y - ref| / max|ref| <= 8e-3` (bf16 output) or `<= 1e-4` (fp32 output) of the fp64 product of +the same bf16 inputs; with `sig_col0 >= 0`, the columns before it bit-identical to that output and the columns +from it bit-identical to `torch.sigmoid` of it; identical bits on a repeated call; and each `M`'s rows +bit-identical to the same rows of the 64-row call. Also certified: + +- at the five shapes that `k3_ctm_gemv_long` runs at up to 8 tokens (MLA `[W_a; W_g]` with split 6, ring 6, + `push=True`; dense MLP gate_up 4, 5, False; dense MLP down 2, 6, False; drafter qkv and drafter gate_up 8, 6, + True), at `M` 1..8 and 16, 24, ..., 64: every token's row is bit-identical to `k3_ctm_gemv_long`'s for that + token, in calls of up to 8 tokens, wherever the device's split equals that call's (all five at 148 or 152 SMs); +- at the MLA `[W_a; W_g]` and MoE head cells, `M` 1 and 64: calls captured in a CUDA graph after an eager call per + key, replayed with `x` rewritten in place, return the bits of eager calls on the new `x`. + +Other shapes the preconditions admit are accepted but not certified. + +## Metadata consumed + +None. The op reads no attention metadata, KV cache or module state; every input is an argument, and the kernel +runs on the current CUDA stream of `x`'s device. The device's SM count enters through `wide_config`. + +One process-global cache sits behind it, and it is result-neutral: the compiled kernel per key +`(N, K, split, ring, activation ring, token tile, out_fp32, PDL on/off)`, where split and the rings follow from +`(N, K, token tile)` and the SM count. On one device that is one compile per `(N, K, out_fp32)` and token tile +(16 for `M <= 16`, 32 for `M <= 32`, 64 above). The first call with a new key compiles the kernel with the CuTe +DSL, which costs host time; `M` within a token tile and `sig_col0` are runtime arguments and never recompile. That +first call refuses to run under CUDA-graph capture: it raises `RuntimeError` ("run once per shape outside +CUDA-graph capture first") before launching anything. Call every key once eagerly, then capture; captured and +later calls reuse the compiled kernel. The PDL setting is read from `TRTLLM_ENABLE_PDL` on every call, so changing +it within a process adds a key. + +## Preconditions + +The op checks these and raises `ValueError` ("k3_ctm_gemv_wide: unsupported call ...") when one fails: + +- `x` is a 2-D contiguous CUDA bf16 tensor with `1 <= M <= 64` rows whose start address is 16-byte aligned. +- `weight` is a 2-D contiguous bf16 tensor with `weight.shape[1] == x.shape[1]` whose start address is 16-byte + aligned. +- `sig_col0 < N`, and `sig_col0 < 0` when `out_fp32`. +- `K % 64 == 0` and `N % 8 == 0`. +- A split exists: `2 * ceil(N / 128) <= SM count` and `ceil(K / 128) >= 2`. On a 148-SM device that is + `N <= 9472`. + +Certified refusals, at `N = 3208`, `K = 7168`: 0 and 65 rows, an `x` starting 2 bytes past a 16-byte boundary, +`sig_col0 = N`, `sig_col0 = 100` with `out_fp32`, `N = 3204`, `K = 7136`, and fp16 `x`. + +Not checked by the op: + +- `weight` is on `x`'s device; the op checks only `x.is_cuda`. +- The GPU has compute capability 10.x: the kernel uses tcgen05 MMA, tensor memory and thread-block clusters. +- The CuTe DSL (`cutlass`) and `cuda-python` (`cuda.bindings`) are importable. The op's check imports the kernel + module, so without them a call raises `ImportError`, not `ValueError`. + +## Notes + +- Programmatic dependent launch (PDL). With `TRTLLM_ENABLE_PDL` unset or `1` the kernel is launched with PDL; any + other value launches it without. Each CTA fills its weight ring (TMA, L2 evict-first) and prefetches the rest of + its weight k-tiles into L2 without waiting for the grid dependency; only the reads of `x` follow + `griddepcontrol.wait`. So `weight` must not be written by work still running ahead of this call on the stream, + while `x` may be. There is no `trigger_early` argument: each CTA always executes + `griddepcontrol.launch_dependents` right after issuing that first weight traffic, as `k3_ctm_gemv_long` does + with `trigger_early=True`. The next kernel on the stream, if launched with PDL, can start while this one runs, + and it must execute `griddepcontrol.wait` (`cudaGridDependencySynchronize`) before it reads `y`. Kernels + launched without PDL, torch's included, start after this one completes as usual. None of this changes results. +- There is no `push` argument either: the split-K partials always travel by 16-byte `st.async` stores, and each + rank reduces the partials of its own share of the tokens. +- Two identical calls return identical bits. The bits depend on the split, so on the device's SM count, but not + on `M` or the PDL setting. +- The result is not bit-identical to cuBLAS: `F.linear` sums in another order. The two agree within the + tolerances above. +- Grid: `ceil(N / 128) * split` CTAs of 256 threads in clusters of `split`; the choice of split keeps that at + most the SM count. +- Only sm_100 has been measured. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv_wide.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv_wide.py new file mode 100644 index 000000000000..4788c47145d9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_ctm_gemv_wide.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 wide decode GEMV x @ weight^T for M <= 64 tokens via the split-K CTM (CuTe DSL) kernel.""" + +import torch + +import tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv.op # noqa: F401 (registers torch.ops.trtllm.k3_ctm_gemv*) + + +def k3_ctm_gemv_wide( + x: torch.Tensor, + weight: torch.Tensor, + sig_col0: int = -1, + out_fp32: bool = False, +) -> torch.Tensor: + """Return `x @ weight.T` in bf16 (columns >= `sig_col0` as sigmoid) or fp32, in one call.""" + return torch.ops.trtllm.k3_ctm_gemv_wide(x, weight, sig_col0=sig_col0, out_fp32=out_fp32) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_decode_gemv.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_decode_gemv.md new file mode 100644 index 000000000000..2ee95e8c7c17 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_decode_gemv.md @@ -0,0 +1,113 @@ +--- +receipts: + sm_100: {status: passed, tests: 4} +--- + +# k3_decode_gemv + +**Wraps** `torch.ops.trtllm.k3_decode_gemv` (one call). + +## Semantics + +The product of a decode step's activation rows with a projection weight, for at most 8 rows, in bf16: + +``` +y[m, n] = bf16(acc[m, n]) acc[m, n] = sum over k of x[m, k] * weight[n, k], accumulated in fp32 + m < M <= 8, n < N, k < K +``` + +The products are accumulated in fp32 by tcgen05 MMAs (M 128 x N 8 x K 16; the weight is the 128-row operand, the +activation the 8 token columns) into a TMEM accumulator, and each sum is rounded once, to nearest, to bf16. The op +picks one of two kernels from the weight's shape: + +- **Short K** (`N % 128 == 0`, `K % 64 == 0`, `K <= 768`): one CTA per 128-row weight tile accumulates the whole of + `K`, in k order, into one fp32 accumulator. Its whole weight slice is resident in shared memory, one stage per + 128-column k-tile. +- **Split K** (`K % 512 == 0`, `K > 768`, any `N`): a cluster of 4 CTAs per 128-row weight tile. Rank `r` + accumulates k-tiles `r, r + 4, r + 8, ...` in fp32; rank `r` also owns rows `32 r .. 32 r + 31` of the tile, the + other ranks push their fp32 partials of those rows into its shared memory, and it adds the four partials in rank + order, `((p0 + p1) + p2) + p3`, before the one bf16 rounding. A last tile that runs past `N` computes zero-filled + rows there and does not store them. + +The kernel always computes 8 token columns (rows of `x` past `M` arrive as zeros from the TMA), so row `m` of the +result depends on row `m` of `x` only: the result of an `M`-row call is bit-identical to the same rows of an 8-row +call (certified). The summation order is fixed, so the result is bit-identical from run to run (certified). It is not +bit-identical to cuBLAS (`F.linear`), whose accumulation order differs. + +Fusion boundary. Inside: the product and its bf16 rounding. Outside: any bias, activation, scaling or quantization, +the tensor-parallel reduction of a row-parallel projection's partial outputs, and flattening leading dims into `M`. +`x` and `weight` are read only; the result is a new tensor. + +## Signature + +```python +def k3_decode_gemv(x: torch.Tensor, weight: torch.Tensor, trigger_early: bool = True) -> torch.Tensor +``` + +| Argument | Shape | Dtype | Layout | Device | +|---|---|---|---|---| +| `x` | `[M, K]`, `1 <= M <= 8` | bf16 | contiguous, 16-byte aligned | CUDA | +| `weight` | `[N, K]`, short K or split K (*Semantics*) | bf16 | contiguous, 16-byte aligned | CUDA, `x`'s device | +| `trigger_early` | scalar | Python bool | — | — | +| returns | `[M, N]` | bf16 | contiguous, newly allocated | `x`'s device | + +`trigger_early`: under programmatic dependent launch (PDL), release the next kernel on the stream as soon as every CTA +has issued its weight loads (short K: its whole slice; split K: the first 6 of its k-tiles, or all if it has fewer), +rather than when this grid ends. Scheduling only (*Notes*). + +### Certified arguments + +| Argument | Shape | Dtype | Layout | Device | +|---|---|---|---|---| +| `x` | `[M, K]`, `M` = 1, 2, ..., 8 | bf16 | contiguous | CUDA | +| `weight` | `[7168, 768]` (short K) and `[3208, 7168]` (split K) | bf16 | contiguous | CUDA | +| `trigger_early` | `True` and `False` | bool | — | — | +| returns | `[M, N]` | bf16 | contiguous | CUDA | + +The weights are Kimi K3's TP16 per-rank KDA o_proj slice (`[7168, 768]`) and KDA input projection slice +(`[3208, 7168]`, whose 26th tile has 8 rows). Every cell: the error is at most 8e-3 of `max |ref|` against an fp64 +torch product; a rerun is bit-identical; the rows are bit-identical to the same rows of the 8-row call; +`trigger_early=False` is bit-identical to `True`. Also certified: a CUDA graph of `M` 1 and `M` 8 calls on each +weight, captured after eager calls and replayed with `x` rewritten in place, bit-identical to eager calls on the same +rows; and the refusals under *Preconditions*. + +## Metadata consumed + +None besides the arguments and the current CUDA stream (the kernel is launched there), plus two pieces of process +state: + +- A per-process compile cache keyed by `(N, K, trigger_early, PDL)`. `M` is a runtime argument, so one compiled + kernel serves every `M`. The first call for a key compiles the kernel (seconds) and must be made eagerly: under + CUDA-graph capture it raises `RuntimeError` ("must run once per shape outside CUDA-graph capture first"). The cache + is result-neutral. +- `TRTLLM_ENABLE_PDL` (default `"1"`), read on every call: whether the kernel is launched with PDL. It is part of the + cache key and changes scheduling, not results. The test runs with the default. + +## Preconditions + +- `x` and `weight` bf16, 2-D and contiguous; `x.shape[1] == weight.shape[1]`; `1 <= M <= 8`; `x` on a CUDA device. +- `(N, K)` taken by one of the kernels: short K (`N % 128 == 0`, `K % 64 == 0`, `ceil(K / 128) <= 6`) or split K + (`K % 128 == 0`, `K / 128 > 6` and a multiple of 4, `N >= 1`). +- A call outside these raises `ValueError` ("k3_decode_gemv: unsupported call ...") before launching anything. + Certified: `M` = 0 and 9, `[128, 1152]` (9 k-tiles: too many for short K, not a multiple of 4 for split K) and + `[200, 768]` (short K needs whole 128-row tiles, split K more than 6 k-tiles). +- Not checked by the op: `weight` on `x`'s device; 16-byte-aligned data pointers (the op declares that alignment to + the kernel; row slices `x[a:b]` of these shapes are aligned); an SM 10.x GPU (the kernels use tcgen05); the CuTe DSL + (`cutlass`) and `cuda-python` (`cuda.bindings`), which the call imports, so without them it raises `ImportError`. +- Under PDL each CTA loads its weight before it waits for the kernel it follows on the stream, so `weight` must not + be written by that kernel; a projection weight is constant during inference. `x` is read, and `y` written, after + the wait. + +## Notes + +- PDL: with `TRTLLM_ENABLE_PDL` = `"1"` the kernel is launched with programmatic dependent launch. Each short-K CTA + issues the loads of its whole weight slice at launch; each split-K CTA loads up to 6 of its k-tiles into its + shared-memory ring and prefetches the rest into L2. Only the activation load waits for the preceding kernel + (`griddepcontrol.wait`), so behind a long predecessor the call pays for the activation load, the MMAs and the store. + A dependent released early by `trigger_early` must itself wait for this grid before reading `y`, as every PDL kernel + waits for its predecessor before reading its output. +- The weight's shared-memory loads carry an L2 EVICT_FIRST hint: a projection's weight streams once per step without + evicting the rest of L2. +- Grid: short K, `N / 128` CTAs of 256 threads; split K, `4 x ceil(N / 128)` CTAs in clusters of 4. +- `trtllm::k3_decode_gemv_tail` (the row-parallel MoE tail on the same kernel, with an RMS-scaled latent part) is a + separate op and not this entry. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_decode_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_decode_gemv.py new file mode 100644 index 000000000000..17b9d2b2f005 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_decode_gemv.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 decode GEMV: ``x @ weight^T`` for at most 8 bf16 rows on a CuTe DSL kernel (short-K or split-K, chosen +by the weight's shape).""" + +import torch + +import tensorrt_llm._torch.cute_dsl_kernels.k3_decode_gemv.op # noqa: F401 — registers the op + + +def k3_decode_gemv( + x: torch.Tensor, weight: torch.Tensor, trigger_early: bool = True +) -> torch.Tensor: + """Return ``x @ weight^T`` as a new bf16 ``[M, N]`` tensor (fp32 accumulation, one bf16 rounding) in one + k3_decode_gemv call.""" + return torch.ops.trtllm.k3_decode_gemv(x, weight, trigger_early) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_head_gemv.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_head_gemv.md new file mode 100644 index 000000000000..85ace6ac6cb2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_head_gemv.md @@ -0,0 +1,200 @@ +--- +receipts: + sm_100: {status: passed, tests: 5} +--- + +# k3_head_gemv + +**Wraps** `torch.ops.trtllm.k3_head_gemv` (one call). + +## Semantics + +The product of a decode step's activation rows with a large weight (an lm_head vocabulary shard), for at most 8 rows, +in bf16: + +``` +y[m, n] = bf16(acc[m, n]) acc[m, n] = sum over k of x[m, k] * weight[n, k], accumulated in fp32 + m < M <= 8, n < N, k < K +``` + +The weight is cut into 128-row tiles and 128-column k-tiles. The products are accumulated in fp32 by tcgen05 MMAs +(M 128 x N 8 x K 16; the weight is the 128-row operand, the activation the 8 token columns) into TMEM accumulators, +and each sum is rounded once, to nearest, to bf16. Under the stream-K schedule (certified): + +- The `T = (N / 128) x (K / 128)` (tile, k-tile) items, tile-major, are cut into `G = min(SMs, T)` equal contiguous + ranges, one persistent CTA per range: CTA `c` owns items `[c T / G, (c + 1) T / G)`. +- A tile whose items all lie in one range is accumulated by that CTA in k order and stored. +- A tile split over several CTAs has one piece per CTA, in k order. Each piece accumulates its k-tiles in k order in + fp32. Pieces 1, 2, ... store their fp32 partials in the workspace and raise their flag words; the CTA of piece 0 + (the finalizer) waits for those flags and adds the partials after its own accumulator in k order, + `((acc_0 + acc_1) + acc_2) + ...`, before the one bf16 rounding. + +The partition depends only on `(N, K, G)`, so every sum has a fixed order and the result is bit-identical from run to +run on one GPU (certified). A GPU with another SM count partitions the items differently, so its sums may round +differently. The kernel always computes 8 token columns (rows of `x` past `M` arrive as zeros from the TMA), so the +result of an `M`-row call is bit-identical to the same rows of an 8-row call (certified). It is not bit-identical to +cuBLAS (`F.linear`), whose accumulation order differs. + +The `dynamic` schedule (a workspace created with `schedule="dynamic"`; accepted, not certified) cuts each tile's K +into chunks of `chunk_tiles` k-tiles; one persistent CTA per SM takes one (tile, chunk) unit and claims the rest from +a counter, and the last unit of a tile to finish adds the tile's unit partials in chunk order before the one bf16 +rounding. Its order is fixed too. + +Fusion boundary. Inside: the product and its bf16 rounding. Outside: any bias or logit processing (soft-capping, +temperature, fp32 conversion) and the gather of the vocabulary shards across ranks. `x` and `weight` are read only; +the call writes the result and the workspace's buffers (*State*), nothing else. + +## Signature + +```python +def k3_head_gemv( + x: torch.Tensor, + weight: torch.Tensor, + workspace: K3HeadGemvWorkspace, + keep_tiles: int = 0, + ring: int = 6, + prefetch: int = 16, +) -> torch.Tensor +``` + +The wrapper passes `workspace.partials`, `workspace.flags` and `workspace.claim` as the op's state buffers and +`workspace.chunk_tiles` and `workspace.schedule` as its schedule, so a call always runs the schedule its workspace was +created for. `K3HeadGemvWorkspace` is importable from this entry's wrapper module. + +| Argument | Shape | Dtype | Layout | Device | +|---|---|---|---|---| +| `x` | `[M, K]`, `1 <= M <= 8` | bf16 | contiguous, 16-byte aligned | CUDA | +| `weight` | `[N, K]`, `N % 128 == 0`, `K % 128 == 0` | bf16 | contiguous, 16-byte aligned | CUDA, `x`'s device | +| `workspace` | a `K3HeadGemvWorkspace` for `(N, K)` on `x`'s device (*State*) | — | — | — | +| `keep_tiles` | scalar | Python int | — | — | +| `ring` | scalar | Python int, 1-6 | — | — | +| `prefetch` | scalar | Python int, `>= 0` | — | — | +| returns | `[M, N]` | bf16 | contiguous, the first `M` rows of a new `[8, N]` buffer | `x`'s device | + +`keep_tiles`: weight tiles below it load at normal L2 priority, the rest with an EVICT_FIRST hint (*Notes*). `ring`: +the shared-memory pipeline stages. `prefetch` (stream-K): the k-tiles past the ring that each CTA prefetches into L2 +before its grid-dependency wait. + +### Certified arguments + +| Argument | Shape | Dtype | Layout | Device | +|---|---|---|---|---| +| `x` | `[M, 7168]`, `M` = 1, 2, ..., 8 | bf16 | contiguous | CUDA | +| `weight` | `[10240, 7168]`, two different weights | bf16 | contiguous | CUDA | +| `workspace` | `K3HeadGemvWorkspace.create(10240, 7168, device)` (stream-K) | — | — | — | +| `keep_tiles` | 0 (every cell); 40 and 80 at `M` 1 and 8 | int | — | — | +| `ring`, `prefetch` | 6 and 16 (the defaults) | int | — | — | +| returns | `[M, 10240]` | bf16 | contiguous | CUDA | + +`[10240, 7168]` is Kimi K3's TP16 per-rank LM-head vocabulary shard (163840 / 16 rows). The `dynamic` schedule is +accepted but not certified. Every cell: the error is at most 8e-3 of `max |ref|` against an fp64 torch product; a +rerun is bit-identical; the rows are bit-identical to the same rows of the 8-row call; the workspace's words are all +zero after the call; `keep_tiles` 40 and 80 are bit-identical to 0. The call sequences and refusals certified on top +of the cells are listed under *State* and *Preconditions*. + +## State + +**Object.** `K3HeadGemvWorkspace` (`tensorrt_llm/_torch/cute_dsl_kernels/k3_head_gemv/op.py`, re-exported by this +entry's wrapper module), owned by the caller: a frozen dataclass of the weight shape (`n_out`, `k_in`), `schedule`, +`chunk_tiles` (0 for stream-K) and three tensors. + +**Contents and size.** With `tiles = N / 128`, `G = min(SMs, tiles x K / 128)` and `P` the most CTAs any tile is +split over (`k3_head_gemv_kernel.streamk_max_pieces(N, K, G)`), a stream-K workspace holds: + +- `partials`, fp32 `[tiles x P x 128 x 8]`: slot `(t, p)` holds the fp32 partial `[128 rows][8 tokens]` of piece + `p > 0` of tile `t` (slot `p = 0` is not used); +- `flags`, int32 `[tiles x P]`: word `(t, p)` is raised to 1 by piece `p > 0` of tile `t` once its partial is stored, + and lowered to 0 by the tile's finalizer once it has read the partial; +- `claim`, int32 `[1]`: not used by stream-K (the dynamic schedule's unit counter). + +At the certified shape on a GPU with 148 or 152 SMs: 80 tiles, 56 k-tiles, `P` = 3, so 240 flag words (960 B) and +245,760 partial floats (960 KiB). The sizes do not depend on `M`: one workspace serves every `M`. A dynamic workspace +holds the fp32 partial of every (tile, chunk) unit, one unit count per tile in `flags` and the unit counter in `claim`. + +**Who creates it, and when.** The target, once per weight shape and schedule, in `post_load_weights`, with +`K3HeadGemvWorkspace.create(n_out, k_in, device, schedule="streamk", chunk_tiles=0)`. It is single-GPU, not +collective, and eager: it allocates `partials` uninitialized and zeroes `flags` and `claim` on the current stream (the +first call must be ordered after that), and it raises `RuntimeError` under CUDA-graph capture (certified). An unknown +schedule raises `ValueError` (certified). The sizes follow the SM count of `device`, so the workspace is created for +the device its calls run on. No environment variable is read. + +**Which ops may share one object.** Only `k3_head_gemv` calls of the workspace's weight shape and schedule on its +device: any `M`, any weight of that shape, any `keep_tiles` (and any `ring` or `prefetch`, which do not change the +layout). At TP16 the target's LM head and DSpark's draft-logits head both have the `[10240, 7168]` shape. Certified: +`M` = 8, 8, 2, 7, 8, 1, 1, 8, 3, 8 back to back on one workspace, and two weights alternating on one workspace, every +call bit-identical to the same call on a fresh workspace. Two workspaces are independent: calls alternating between +two, each with its own weight, are all correct (certified). The op refuses, with `ValueError` before launching, a +workspace whose buffers do not fit the call: another dtype or device, another number of `flags` words, too few +`partials` (certified with a workspace created for `[5120, 7168]`). It checks nothing else, so the sharing rule is the +caller's. + +**Call-order invariant.** A launch on a workspace must start after the previous launch on it has completed: they use +the same flag words and partial slots. The calls that share a workspace are therefore ordered on one stream, which +holds under PDL too: a launch reads and writes the workspace only after its grid-dependency wait +(`griddepcontrol.wait`), that is, after the work before it on the stream has completed. Eager calls from a second +stream must not share the workspace unless the caller orders them (a synchronization or an event between, e.g., a +load-time call on one stream and the decode steps on another). Captured calls replay in capture order on the +replaying stream, so a graph holding calls on a workspace may be replayed on the stream that issues the eager calls +on it, between those calls. Certified: a graph of `M` 1, 8 and 3 calls on one workspace, replayed three times with `x` +rewritten in place and an eager `M` 8 call on the same workspace between replays, every result bit-identical to the +same call on a fresh workspace. Not exercised: two launches running at once on one workspace raise and lower each +other's words, so a finalizer can read the other launch's partial (a silently wrong result) or wait for a word the +other launch has already lowered (a hang). + +**What a later launch reads.** Every word of `flags`, which must be 0 when the launch starts: a finalizer waits until +the word of each later piece of its tile is non-zero and then reads that piece's slot of `partials`. `partials` are +written before they are read within one launch, so their contents between launches do not matter. A dynamic launch +also reads the per-tile unit counts and the unit counter, which must be 0 as well. + +**How it is re-armed.** By the launch itself. Each finalizer lowers its tile's words after reading the partials (in the +dynamic schedule the finalizing unit resets its tile's count, and the last claim rolls the counter back to 0), so a +launch that completes leaves every word at 0 for the next one, whatever its `M` (certified: all words zero after every +sequence above). The layout depends only on `(N, K, schedule, SMs)`; no call depends on an earlier call's `M`. A +launch that never completes can leave words raised; the workspace must then be recreated (`create()` zeroes them). + +## Metadata consumed + +Besides `workspace`, which is an explicit argument: + +- The current CUDA stream: the kernel is launched there. +- A per-process compile cache: stream-K kernels keyed by `(N, K, ring, G, P, prefetch, PDL)`, dynamic ones by + `(N, K, chunk_tiles, ring, G, PDL)`. `M` and `keep_tiles` are runtime arguments. The first call for a key compiles + the kernel (seconds) and must be made eagerly: under CUDA-graph capture it raises `RuntimeError` ("run once per + shape outside CUDA-graph capture first"). The cache is result-neutral. +- The device's SM count, read on every call: it sets `G`, hence the partition and the workspace sizes. +- `TRTLLM_ENABLE_PDL` (default `"1"`), read on every call: whether the kernel is launched with PDL. It is part of the + cache key and changes scheduling, not results. The test runs with the default. + +## Preconditions + +- `x` and `weight` bf16, 2-D and contiguous; `x.shape[1] == weight.shape[1]`; `1 <= M <= 8`; `x` on a CUDA device; + `N % 128 == 0` and `K % 128 == 0`. +- Stream-K: `1 <= ring <= 6`. Dynamic: the workspace's `chunk_tiles` divides `K / 128`. +- A call outside these raises `ValueError` ("k3_head_gemv: unsupported call ...") before launching anything + (certified: `M` = 0 and 9, `N` = 200). +- `workspace` created for this weight shape and schedule on `x`'s device: one whose buffers do not fit the call + raises `ValueError` ("... is not one for weight ...") before launching anything (certified, see *State*). +- `create()` runs outside CUDA-graph capture (else `RuntimeError`), and every key is compiled by an eager call before + a capture uses it (*Metadata consumed*). +- Not checked by the op: `weight` on `x`'s device; 16-byte-aligned data pointers (the op declares that alignment to + the kernel; row slices `x[a:b]` of these shapes are aligned); `prefetch >= 0`; an SM 10.x GPU (the kernel uses + tcgen05); the CuTe DSL (`cutlass`) and `cuda-python` (`cuda.bindings`), which the call and `create()` import, so + without them they raise `ImportError`. +- Under PDL each CTA loads weight tiles (its first `ring` k-tiles, and `prefetch` more into L2) before it waits for + the kernel it follows on the stream, so `weight` must not be written by that kernel; a head weight is constant + during inference. `x` is read, and the workspace and `y` written, after the wait. + +## Notes + +- PDL: with `TRTLLM_ENABLE_PDL` = `"1"` the kernel is launched with programmatic dependent launch. Before + `griddepcontrol.wait` a stream-K CTA only loads weight tiles; it releases its dependents once it has issued all of + its loads (a dynamic CTA, after its last claim). A dependent must itself wait for this grid before reading `y`, as + every PDL kernel waits for its predecessor before reading its output. +- L2: weight tiles `>= keep_tiles` are loaded with an EVICT_FIRST hint, so the streamed shard does not evict the rest + of L2; tiles `< keep_tiles` load at normal priority, so a later reader of the same weight can find them in L2. + Cache policy only: `keep_tiles` 0, 40 and 80 are bit-identical (certified). +- Grid: `G` persistent CTAs of 256 threads (dynamic: one per SM). Stream-K launches at most one CTA per (tile, + k-tile) item, so every CTA's range is non-empty and raises the flag its finalizer waits for, also for a weight with + fewer items than the GPU has SMs. Such weights are not certified here. +- The result is a view of the first `M` rows of a new `[8, N]` buffer: the kernel writes all 8 token rows, the rows + past `M` from the zero-filled activation rows. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_head_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_head_gemv.py new file mode 100644 index 000000000000..2158f38c825b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/k3_head_gemv.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 head GEMV: ``x @ weight^T`` for at most 8 bf16 rows and a large weight (an lm_head vocabulary shard) on a +persistent CuTe DSL kernel, over a caller-owned :class:`K3HeadGemvWorkspace`. + +``K3HeadGemvWorkspace`` is re-exported here so that a target creates it through the catalog: once per weight shape +and schedule, eagerly, before any CUDA-graph capture. The contract is the ``## State`` section of ``k3_head_gemv.md``. +""" + +import torch + +# Importing the op module also registers torch.ops.trtllm.k3_head_gemv. +from tensorrt_llm._torch.cute_dsl_kernels.k3_head_gemv.op import K3HeadGemvWorkspace + +__all__ = ["K3HeadGemvWorkspace", "k3_head_gemv"] + + +def k3_head_gemv( + x: torch.Tensor, + weight: torch.Tensor, + workspace: K3HeadGemvWorkspace, + keep_tiles: int = 0, + ring: int = 6, + prefetch: int = 16, +) -> torch.Tensor: + """Return ``x @ weight^T`` as bf16 ``[M, N]`` (fp32 accumulation, one bf16 rounding) in one k3_head_gemv call over + ``workspace``, whose schedule and ``chunk_tiles`` the call takes. The calls that share a workspace must be ordered + on one stream.""" + return torch.ops.trtllm.k3_head_gemv( + x, + weight, + workspace.partials, + workspace.flags, + workspace.claim, + keep_tiles, + workspace.chunk_tiles, + ring, + workspace.schedule, + prefetch, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml index 7382f31d99a1..f9103979de78 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml @@ -28,7 +28,7 @@ # limits, e.g. no float8 arithmetic, are noted in the entry docstring). # Entries wrapping trtllm ops carry all three. # -# ── RECEIPT STATUS: all 19 entries certified on sm_103 ──────────────────── +# ── RECEIPT STATUS: 19 entries certified on sm_103, 18 on sm_100 ────────── # # A receipt says this entry's test passed on a stated GPU architecture. The # architecture is the whole key: it is a real axis -- see the @@ -66,6 +66,11 @@ # exactly as written. They are true records of what was observed on another # device, and rewriting them would manufacture GB300 evidence that does not # exist. Read them as provenance; the frontmatter is the certification. +# +# The Kimi K3 entries (k3_*, attn_res_* and ssm/kda_decode) are certified on +# sm_100 (B200 / GB200), where their first caller runs, and their tests skip on +# any other architecture. cublas_mm and flashinfer_rmsnorm, which that caller +# also uses, keep their sm_103 receipts and add sm_100 ones. # ────────────────────────────────────────────────────────────────────────── # # 30 of the source catalog's 43 entries are here — the union of what the two @@ -133,11 +138,27 @@ entries: impl: torch.ops.trtllm.flashinfer_fused_add_rmsnorm summary: "In-place fused residual add + RMS norm: residual += x; x = rmsnorm(residual) * weight" + - path: norm/k3_embed_norm.py + impl: torch.ops.trtllm.k3_embed_norm + summary: "Kimi K3 decode-step embedding + layer 0's input RMSNorm in one CuTe DSL launch: raw[:] = table[ids] (int32 or int64 ids, N <= 64; zero rows for ids outside [0, V)) written bit for bit into the caller's bf16 [N,H] buffer (the model passes slot 0 of its attention-residual snapshot bank), and a fresh bf16 [N,H] raw * rsqrt(mean(raw^2) + eps) * weight with fp32 statistics, an approximate rsqrt and one rounding, H in (6144, 16384] a multiple of 1024; rows independent of N and bit-reproducible run to run; the op states bit-identity with the gather followed by flashinfer's CuTe DSL RMSNorm (the kernel behind norm/flashinfer_rmsnorm); compiles per (N, V, H, ids dtype, PDL) on its first, eager, call; under PDL it releases its dependents at entry, so a dependent waits on the grid dependency before reading raw or the output" + + - path: norm/attn_res_fwd.py + impl: torch.ops.trtllm.attn_res_fwd + summary: "Kimi K3 attention-residual selection: per token (B == 1), a softmax over N = K+1 candidates (the K bf16 [K,T,1,H] snapshots in order, then the bf16 [T,1,H] layer residual) of fp32 logits rsqrt(mean(v^2) + rms_eps) * , mixing the raw candidates into a fresh bf16 [T,1,H] output (one rounding) and returning the fp32 [N,T,1] rsigma / probs / logits too; N <= 12, T <= 16384, H a multiple of 1024 in [4096, 8192], sm_100 family only; kernel chosen by (T, N) -- a single-CTA or 8-CTA-cluster decode kernel at T == 1 for N in {1,2,4} / {8,12}, else a persistent TMEM online-softmax kernel -- each bit-reproducible run to run" + + - path: norm/attn_res_rmsnorm_fwd.py + impl: torch.ops.trtllm.attn_res_rmsnorm_fwd + summary: "Kimi K3 attention-residual selection fused with the RMSNorm that consumes it: attn_res_fwd's bf16 mixture (that rounding kept), normalized in fp32 with output_rms_eps, rounded to bf16, times the bf16 output_rms_weight and rounded again (KimiK3RMSNorm order), into a fresh bf16 [T,1,7168] tensor; H == 7168 only, N <= 12, T <= 16384, one CTA (N <= 4) or one 8-CTA cluster (N >= 5) per token so rows do not depend on T; with TRTLLM_ENABLE_PDL on, each CTA waits on its grid dependency and then releases its dependents before writing, so a PDL dependent must wait on the grid dependency before reading the output" + # ─── activation ──────────────────────────────────────────────── - path: activation/flashinfer_silu_and_mul.py impl: torch.ops.trtllm.flashinfer_silu_and_mul summary: "SwiGLU activation: silu(x[..., :d]) * x[..., d:] over a packed gate/up last dim" + - path: activation/k3_situ_mul.py + impl: torch.ops.trtllm.k3_situ_mul + summary: "SiTU-gated multiply (SituAndMul) of a packed [gate | up] bf16 gate_up output for M <= 8 tokens on a CuTe DSL kernel: bf16(beta*tanh(g/beta)*sigmoid(g) * v), v = linear_beta*tanh(u/linear_beta) or u, in fp32 with one rounding; launches its PDL dependents at once and reads gu after the grid wait; linear_beta=0.0 refused by the wrapper (the op would run it as 1.0)" + # ─── gemm ────────────────────────────────────────────────────── - path: gemm/cublas_mm.py impl: torch.ops.trtllm.cublas_mm @@ -151,6 +172,30 @@ entries: impl: torch.ops.trtllm.nvfp4_gemm summary: "NVFP4 x NVFP4 dense GEMM in nn.Linear layout with in-kernel block-scale dequantization: alpha * ([M,K/2] e2m1 + 128x4-swizzled e4m3 scales) @ ([N,K/2] e2m1 + swizzled scales)^T (+ per-column bias) -> fresh [M,N] bf16/fp16/fp32 buffer, fp32 accumulation, backend auto-selected (cutlass/cublaslt/cuda_core/cutedsl)" + - path: gemm/k3_ctm_gemv.py + impl: torch.ops.trtllm.k3_ctm_gemv + summary: "Kimi K3 decode GEMV for M <= 8 bf16 tokens on the CTM CuTe DSL kernel (SM 10.x): bf16(x @ W^T) for an [N, K] weight with N % 128 == 0 and at most 6 128-column k-tiles per CTA, fp32 tcgen05 accumulation with one rounding, one CTA or a 2- or 4-CTA split-K cluster per 128 rows (partials summed in rank order); the weight is read before the PDL grid wait" + + - path: gemm/k3_ctm_gemv_swiglu.py + impl: torch.ops.trtllm.k3_ctm_gemv_swiglu + summary: "SwiGLU down projection for M <= 8 tokens on the CTM CuTe DSL kernel (SM 10.x): bf16(bf16(silu(gate) * up) @ W^T) from a packed [gate | up] bf16 gate_up output, the activation computed in fp32 in the GEMV's shared memory and never stored, fp32 accumulation, 1/2/4-CTA split-K summed in rank order" + + - path: gemm/k3_ctm_gemv_long.py + impl: torch.ops.trtllm.k3_ctm_gemv_long + summary: "Long-K decode GEMV for M <= 8 bf16 tokens on the CTM CuTe DSL kernel (SM 10.x): bf16(x @ W^T) with each 128-row block split over a cluster of 2 or 4-8 CTAs, its weight streamed through a shared-memory ring filled before the PDL grid wait, partials summed in rank order; columns >= sig_col0 as bf16(sigmoid(bf16(.))) (torch.sigmoid of the plain output, bitwise); any N, K % 64 == 0" + + - path: gemm/k3_ctm_gemv_wide.py + impl: torch.ops.trtllm.k3_ctm_gemv_wide + summary: "Wide decode GEMV for M <= 64 bf16 tokens on the long CTM kernel (SM 10.x): one MMA of 16/32/64 token columns, split and ring picked from the shape and the SM count (one wave), each token's row bit-identical to k3_ctm_gemv_long's at the same split; bf16 output with optional sigmoid columns from sig_col0, or unrounded fp32; N % 8 == 0, K % 64 == 0" + + - path: gemm/k3_decode_gemv.py + impl: torch.ops.trtllm.k3_decode_gemv + summary: "Kimi K3 decode GEMV for at most 8 rows: bf16 x [M,K] @ weight [N,K]^T into a fresh bf16 [M,N], fp32 accumulation (tcgen05 into TMEM) and one rounding, on a CuTe DSL kernel the weight's shape selects -- short K (K <= 768, N a multiple of 128: one CTA per 128-row tile with the whole weight slice in shared memory) or split K (K a multiple of 512 above 768: 4-CTA clusters, the partials added in rank order); rows independent of M and bit-reproducible run to run, not bit-identical to cuBLAS; SM 10.x; compiles per (N, K, trigger_early, PDL) on its first, eager, call; under PDL the weight is loaded before the grid-dependency wait, so it must not be written by the preceding kernel" + + - path: gemm/k3_head_gemv.py + impl: torch.ops.trtllm.k3_head_gemv + summary: "Kimi K3 vocabulary-shard head GEMV for at most 8 rows: bf16 x [M,K] @ weight [N,K]^T (N, K multiples of 128) into bf16 [M,N] (the first M rows of a fresh [8,N]) on a persistent stream-K CuTe DSL kernel of min(SMs, tiles x k-tiles) CTAs, a split tile's fp32 piece partials added in k order before one rounding -- rows independent of M, bit-reproducible run to run on one GPU; over a caller-owned K3HeadGemvWorkspace (fp32 partials and one flag word per (tile, piece); created eagerly, before any capture, for one weight shape and schedule; one that does not fit the call is refused with ValueError) that every completed launch leaves at zero, so the calls sharing one are ordered on one stream, eager calls and graph replays alike; weight tiles past keep_tiles streamed EVICT_FIRST; the dynamic schedule accepted, not certified; SM 10.x; compiles per shape on its first, eager, call" + # ─── attention ───────────────────────────────────────────────── - path: attention/fused_qk_norm_rope.py impl: torch.ops.trtllm.fused_qk_norm_rope @@ -172,6 +217,35 @@ entries: impl: tensorrt_llm.bindings.internal.thop.attention summary: "Full attention core with fully explicit state (pybind binding, approved exception): paged KV-cache append + causal/padding-masked GQA FMHA over a caller-owned pool (bf16, or fp8-e4m3 with per-tensor kv scales) and explicit length/offset tensors, written into a caller buffer, on either context execution path (use_paged_context_fmha selects packed-QKV context FMHA, or paged-KV context FMHA so a context call may run over a cached prefix — KV-cache reuse and chunked prefill), with optional per-query-head attention sinks (one extra softmax-denominator logit, dropped from the output) and an optional per-call sliding window (attention_window_size keys ending at the query's absolute position; a pure mask — the append stays at absolute positions, so cyclic pool reuse is the caller's page mapping); MLA mode runs context prefill (in-kernel RoPE + latent append), no-append context over explicit K/V (latent_cache=None: cached-KV prefixes and chunked partial passes with softmax-stats output), and generation latent-MQA decode as separate calls over a paged latent pool (bf16 or fp8-e4m3)" + # Kimi K3 MLA decode (sm_100). The state type attention/k3_mla_attn_workspace.py is not an entry: it launches nothing + # per call. + - path: attention/k3_mla_qkv.py + impl: torch.ops.trtllm.k3_mla_qkv + summary: "Kimi K3's MLA decode query path for at most 64 tokens in one launch (q_a RMSNorm, q_b projection and k_b absorption into fused_q [M, heads x 576]) plus the step's latent KV rows (kv_a RMSNorm, rope columns) stored into the paged latent cache at each token's row; the page table is read before the grid-dependency wait, so under CUDA graphs it must be written before the graph runs; siblings: k3_mla_q (the query path alone) and k3_mla_qkv_out (the cache rows into a dense buffer)" + + - path: attention/k3_mla_attn_vb_out.py + impl: torch.ops.trtllm.k3_mla_attn_vb_out + summary: "Kimi K3's MLA decode attention of at most 8 requests of at most 8 tokens over the paged latent cache (heads a multiple of 6, causal bottom-right, a request's rows exactly those of its own R = 1 call) with v_b and the optional sigmoid output gate in the same launch, over a caller-owned K3MlaAttnWorkspace: per-CTA partials and the no-cluster mode's arrival counters, one workspace per device and head-group count shared by every layer, its calls one at a time in stream order; siblings: k3_mla_attn_out (the attention output itself) and k3_mla_attn" + + # ─── ssm ─────────────────────────────────────────────────────── + # KDA (Kimi Delta Attention) decode. The state type ssm/k3_kda_buffers.py is not an entry: it launches nothing per + # call. The pools every entry here updates belong to the cache manager (the V2 hybrid manager's per-layer views). + - path: ssm/kda_decode.py + impl: torch.ops.trtllm.kda_decode + summary: "One-token KDA decode of B requests: the causal depthwise conv over each request's conv window, the gated delta rule on its fp32 recurrent state and the gated RMS norm, the state (and the conv window) updated in place at each request's slot of the cache manager's pools; at Kimi K3's rank slice (6 heads, K = V = 128, conv width 4) B <= 5 runs the cluster kernel and 6 <= B <= 24 the legacy compact-heads kernel" + + - path: ssm/k3_kda_attn.py + impl: torch.ops.trtllm.k3_kda_attn + summary: "Kimi K3's KDA layer for one speculative-decoding request in one launch: the fused input projection of its golden token and 7 drafts and the KDA verify of ssm/k3_kda_verify on it, bit for bit the unfused pair, writing the slot's state, per-draft states and conv window in place, over a caller-owned K3KdaBuffers (three Lamport buffers and per-CTA indices, one set per device shared with ssm/k3_kda_decode_attn, its launches one at a time in stream order); the head CTAs read the slot's pools before the grid-dependency wait, so a launch must not directly follow another on the same pools; sibling: k3_kda_qkvg (the projection stream alone, a 104-CTA set of its own)" + + - path: ssm/k3_kda_decode_attn.py + impl: torch.ops.trtllm.k3_kda_decode_attn + summary: "Kimi K3's KDA layer for plain decode of R <= 8 requests in one launch: the fused input projection of one token per request and the one-token KDA decode of ssm/kda_decode on it, the state and conv window updated in place at each request's slot, over the same caller-owned K3KdaBuffers as ssm/k3_kda_attn (either may follow the other on one set); the head CTAs read the slots' pools before the grid-dependency wait" + + - path: ssm/k3_kda_verify.py + impl: torch.ops.trtllm.k3_kda_verify + summary: "Kimi K3's KDA speculative verify of N requests of 1 + num_spec tokens from the fused projection rows to the gated-norm core output, starting each request from the state after the drafts the sampler accepted last round (the per-token states of the V2 hybrid manager, selected by prev_num_accepted_tokens, read before the grid-dependency wait) and committing the state after every verify token, so the next round needs no replay" + # ─── moe ─────────────────────────────────────────────────────── - path: moe/noaux_tc_op.py impl: torch.ops.trtllm.noaux_tc_op diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/attn_res_fwd.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/attn_res_fwd.md new file mode 100644 index 000000000000..26a0f9da0ab0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/attn_res_fwd.md @@ -0,0 +1,153 @@ +--- +receipts: + sm_100: {status: passed, tests: 5} +--- + +# attn_res_fwd + +**Wraps** `torch.ops.trtllm.attn_res_fwd` (one call). + +## Semantics + +Attention-residual selection: Kimi K3's per-token mixing of residual-stream snapshots. For every token the op +scores `N = K + 1` candidate rows -- the `K` snapshots of `block_residual` in their stored order, then +`layer_residual` -- and returns the softmax-weighted sum of the raw candidates. + +For token `t` (`B == 1`) and candidate `n`, every step in fp32 (`H` is the hidden size): + +``` +V[n] = block_residual[n, t, 0, :] for n < K +V[K] = layer_residual[t, 0, :] +q = fp32(rms_weight) * fp32(res_weight) # elementwise; never rounded to bf16 +rsigma[n] = rsqrt(sum_h V[n, h]^2 / H + rms_eps) +logits[n] = rsigma[n] * sum_h V[n, h] * q[h] +probs = softmax(logits) # over the token's N candidates +output = bf16_rn(sum_n probs[n] * V[n]) # the only rounding +``` + +`logits[n]` is the dot product of `res_weight` with the RMSNorm of `V[n]` under weight `rms_weight`, the +normalized row kept in fp32, as in HF Kimi's `_apply_attn_res`. The RMSNorm only scores the candidates: `output` +mixes the raw rows, not their normalized form. `rsigma`, `probs` and `logits` are returned as the fp32 values the +formulas name. With `N == 1` the single probability is 1, and the formulas return `layer_residual` as `output`. + +The kernels are compiled with `--use_fast_math`: `rsqrt`, the softmax exponentials and the reciprocal of the +softmax denominator are hardware approximations, denormals flush to zero, and every sum runs in the kernel's own +order. The results are close to an fp32 torch evaluation of the formulas, not bit-identical to it. + +Fusion boundary: selection only. No residual add, no trailing RMSNorm (`attn_res_rmsnorm_fwd` fuses the norm that +follows), and no snapshot write: the caller owns the snapshot bank, including appending the running residual to +it. + +## Signature + +```python +def attn_res_fwd( + layer_residual: torch.Tensor, + block_residual: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + rms_eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] +``` + +| Argument | Shape | Dtype | Layout | Device | +|---|---|---|---|---| +| `layer_residual` | `[T, 1, H]` | bf16 | contiguous | CUDA | +| `block_residual` | `[K, T, 1, H]`, `K = N - 1 >= 0` | bf16 | contiguous | `layer_residual`'s (not checked) | +| `res_weight` | `H` elements (`[H]`, `[1, H]` or `[H, 1]`) | bf16 | contiguous | `layer_residual`'s (not checked) | +| `rms_weight` | `H` elements | bf16 | contiguous | `layer_residual`'s (not checked) | +| `rms_eps` | scalar | Python float, passed as fp32 | — | — | +| returns `[0]` (`output`) | `[T, 1, H]` | bf16 | contiguous, newly allocated | `layer_residual`'s | +| returns `[1]` (`rsigma`) | `[N, T, 1]` | fp32 | contiguous, newly allocated | `layer_residual`'s | +| returns `[2]` (`probs`) | `[N, T, 1]` | fp32 | contiguous, newly allocated | `layer_residual`'s | +| returns `[3]` (`logits`) | `[N, T, 1]` | fp32 | contiguous, newly allocated | `layer_residual`'s | + +On the candidate axis of `rsigma`, `probs` and `logits`, index `n < K` is snapshot `n` and index `K` is +`layer_residual`. `K = 0` is passed as a zero-size `[0, T, 1, H]` tensor. No input is mutated and no output +aliases an input. `res_weight` is the flattened `[1, H]` weight of the scoring `nn.Linear(H, 1)` projection; +`rms_weight` is the scoring RMSNorm's weight. + +### Certified arguments + +- `H = 7168` (Kimi K3), `B = 1`, bf16 inputs, `res_weight` and `rms_weight` of shape `[H]`. +- Every `N = 2 ... 9` -- the candidate counts Kimi K3 produces: its 93 layers append a snapshot every 12 layers, so + the bank holds at most `ceil(93 / 12) = 8` snapshots -- at `T = 1 ... 8` (decode, including multi-token + speculative steps) and at `T = 300` and `T = 2048` (prefill). +- The remaining dispatch branches and `N` extremes: `N = 1, 10, 11, 12` at `T = 1`; `N = 1, 12` at `T = 8`; + `N = 12` at `T = 1024`. +- `rms_eps = 1e-6` in every cell; `1e-5` and `1e-2` in four cells that cover the three kernels. `1e-2` exceeds the + test inputs' mean square of `2.5e-3`, so it pins where `rms_eps` enters. + +## Metadata consumed + +None. Stateless: no attention metadata, no workspace, no caller-visible state. The op caches per device the SM +count, the architecture check and the one-time raise of a kernel's dynamic shared-memory limit, and reads +`TRTLLM_ENABLE_PDL` once per process. These select launch parameters, never results. + +## Preconditions + +Each item is a `TORCH_CHECK` in `cpp/tensorrt_llm/thop/attnResOp.cpp`, evaluated in this order; a failure raises +`RuntimeError` with the quoted message. This entry's test does not exercise the rejections. + +1. `layer_residual.dim() == 3` and `block_residual.dim() == 4` (`attn_res_fwd: layer_residual must be [T, B, H]`, + `attn_res_fwd: block_residual must be [K, T, B, H]`). +2. All four tensors are CUDA tensors (`attn_res_fwd: all input tensors must be CUDA tensors`). +3. `layer_residual`'s device has compute capability `10.x` (`attn_res_fwd requires an sm_100-family (datacenter + Blackwell) GPU`). +4. `B == 1`; `1 <= N <= 12` with `N = block_residual.shape[0] + 1`; `1 <= T <= 16384`; `H` a multiple of 1024 in + `[4096, 8192]` (`attn_res_fwd: unsupported B=...`, `N=...`, `T=...`, `H=...`). +5. All four tensors are bf16 (`attn_res_fwd: must be bf16`) and contiguous (`attn_res_fwd: inputs must be + contiguous`). A strided view raises; it is never read as if it were dense. +6. `block_residual.shape == (N - 1, T, B, H)` (`attn_res_fwd: block_residual shape must match layer_residual`). +7. `res_weight.numel() == H` and `rms_weight.numel() == H` (`attn_res_fwd: must have H elements`). + +Not checked: that `block_residual`, `res_weight` and `rms_weight` are on `layer_residual`'s device. The kernels +dereference all four on that device, so the caller must keep them together (`attn_res_rmsnorm_fwd` checks it). + +## Notes + +- **Code paths.** The op picks a kernel from `(T, N)`. At `H = 7168`: + + | `T`, `N` | Kernel | Launch | + |---|---|---| + | `T == 1`, `N` in {1, 2, 4} | single-CTA decode kernel | 1 CTA of 256 threads | + | `T == 1`, `N` in {8, 12} | split-K decode kernel | 1 cluster of 8 CTAs of 256 threads | + | `T == 1024`, `N == 12` | online kernel, fixed-`N = 12` variant | persistent: (SM count - 1) CTAs of 288 threads | + | every other `(T, N)` | online kernel | persistent: one CTA of 288 threads per SM | + + The single-CTA kernel gives each of its 256 threads 28 of the 7168 elements and keeps every candidate's values + for them in registers. The split-K kernel gives each CTA of the cluster 896 of the 7168 elements, keeps them in + shared memory as fp32, and exchanges per-candidate sums through distributed shared memory. The online kernel is + warp-specialized: one warp streams candidate rows into shared memory with bulk asynchronous copies, eight warps + score them and keep the fp32 rows in Tensor Memory until the mixing pass, the softmax runs online over chunks of + four candidates, and each CTA loops over tokens. Other hidden sizes run the online kernel with other tilings, or + a row-tiled kernel at `N == 1` for `H` 4096 and 8192; none of them is certified here. +- **Architecture.** The op requires compute capability `10.x`, and the kernels are built only for the sm_100 + family target (`100f`). They use sm_100 instructions: paired fp32 arithmetic (`.f32x2`) in all three kernels, + Tensor Memory (`tcgen05`) and bulk asynchronous copies in the online kernel, which also takes 115,040 bytes of + dynamic shared memory per CTA at `H = 7168` (one CTA per SM), and an 8-CTA thread-block cluster with distributed + shared memory in the split-K kernel. +- **PDL.** `TRTLLM_ENABLE_PDL` (read once per process; unset or `1` enables it, `0` disables it) controls + programmatic dependent launch for the two decode kernels. When it is enabled they launch with programmatic + stream serialization, wait on the grid dependency (`cudaGridDependencySynchronize`) before reading any input, + and trigger their dependents (`cudaTriggerProgrammaticLaunchCompletion`) after their last store. The online + kernel launches without the attribute: it starts after its predecessor completes and releases its dependents + when it completes. A dependent kernel launched with PDL must wait on the grid dependency before it reads any + output: the trigger only lets it start, and only the wait makes this op's stores visible. PDL changes scheduling, + not arithmetic. +- **Determinism.** Each kernel reduces in a fixed order without atomics, so identical inputs give identical bits + run to run; the test asserts it in every cell. The kernels' summation orders differ from each other, so a token + that a decode kernel computes in a `T == 1` call can come out different in the last bits when the online kernel + computes it inside a larger call. The online kernel's per-token arithmetic depends on neither `T` nor the grid + size, except that its fixed-`N = 12` variant sums the softmax denominator in a different order. +- **Tolerance.** The test compares `output` and each fp32 output with an fp32 torch evaluation of *Semantics*, + using the metric of main's op test (`tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py`): + cosine similarity above 0.999 and relative L2 error below 3e-2. +- **CUDA graphs.** The call does not synchronize with the host, and its launch configuration depends on shapes + only. The test captures one call per kernel after an eager warm-up and checks that replay reproduces the eager + bits. +- **Kimi K3 call sites.** `modeling_kimi_linear.py` calls this op, followed by a separate RMSNorm, where the fused + `attn_res_rmsnorm_fwd` is not taken: above `KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS` tokens (1 by default) unless + `KIMI_K3_ATTN_RES_TOPOLOGY` routes the call to the persistent fused op, and on the layers whose pre-norm mixture + the speculative drafter captures. The `register_fake` in `custom_ops/cpp_custom_ops.py` reports the shapes and + dtypes above. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/attn_res_fwd.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/attn_res_fwd.py new file mode 100644 index 000000000000..f123f0ba7930 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/attn_res_fwd.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 attention-residual selection: a per-token softmax mixture of residual snapshots.""" + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* + + +def attn_res_fwd( + layer_residual: torch.Tensor, + block_residual: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + rms_eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Mix each token's N = K + 1 candidates: the K snapshots of block_residual, then layer_residual. + + A candidate's logit is its fp32 RMSNorm (weight `rms_weight`, `rms_eps`) dotted with `res_weight`; + the softmax of the logits weights the raw candidates. `layer_residual` is bf16 [T, 1, H] and + `block_residual` bf16 [K, T, 1, H], both contiguous. + + Returns `(output, rsigma, probs, logits)`: the bf16 [T, 1, H] mixture and the fp32 [N, T, 1] + per-candidate rsqrt(mean square + eps), softmax probabilities and logits, all newly allocated. + """ + return torch.ops.trtllm.attn_res_fwd( + layer_residual, block_residual, res_weight, rms_weight, rms_eps + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/attn_res_rmsnorm_fwd.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/attn_res_rmsnorm_fwd.md new file mode 100644 index 000000000000..b07775ef94e8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/attn_res_rmsnorm_fwd.md @@ -0,0 +1,162 @@ +--- +receipts: + sm_100: {status: passed, tests: 6} +--- + +# attn_res_rmsnorm_fwd + +**Wraps** `torch.ops.trtllm.attn_res_rmsnorm_fwd` (one call). + +## Semantics + +Attention-residual selection followed by the RMSNorm that consumes it, in one kernel. The selection is +`attn_res_fwd`'s: for every token, a softmax over `N = K + 1` candidate rows -- the `K` snapshots of +`block_residual` in their stored order, then `layer_residual` -- of fp32 logits, mixing the raw candidates. The +fused norm then normalizes that mixture. + +For token `t` (`B == 1`), every step in fp32 unless rounded explicitly (`H = 7168`): + +``` +V[n] = block_residual[n, t, 0, :] for n < K +V[K] = layer_residual[t, 0, :] +q = fp32(rms_weight) * fp32(res_weight) # elementwise; never rounded to bf16 +rsigma[n] = rsqrt(sum_h V[n, h]^2 / H + rms_eps) +logits[n] = rsigma[n] * sum_h V[n, h] * q[h] +probs = softmax(logits) # over the token's N candidates +mixed = bf16_rn(sum_n probs[n] * V[n]) # attn_res_fwd's output +r = rsqrt(sum_h mixed[h]^2 / H + output_rms_eps) +normed = bf16_rn(mixed * r) +output = bf16_rn(normed * fp32(output_rms_weight)) +``` + +The op keeps two bf16 rounding boundaries: + +- `mixed` is rounded to bf16 before the norm reads it, as when `attn_res_fwd` and a separate RMSNorm run in turn. +- The norm rounds twice, in `KimiK3RMSNorm`'s order: the normalized value is rounded to bf16, then multiplied by + the bf16 weight and rounded again. An RMSNorm that rounds once (normalize and scale in fp32, then cast, as the + `flashinfer_rmsnorm` entry describes) differs from this op by one unit in the last bf16 place in a large share of + the elements: about a fifth of them when both formulas are evaluated in torch on this entry's test inputs. + +As in `attn_res_fwd`, the kernels are compiled with `--use_fast_math`: `rsqrt`, the softmax exponentials and the +reciprocal of the softmax denominator are hardware approximations, denormals flush to zero, and every sum runs in +the kernel's own order. The output is close to an fp32 torch evaluation of the formulas, not bit-identical to it. + +Fusion boundary: the selection and the RMSNorm that follows it. The op returns only the normalized output -- not +the pre-norm mixture and not `attn_res_fwd`'s `rsigma`, `probs` or `logits`. No residual add (the separate op +`attn_res_add_rmsnorm_fwd` adds one first). + +## Signature + +```python +def attn_res_rmsnorm_fwd( + layer_residual: torch.Tensor, + block_residual: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + output_rms_weight: torch.Tensor, + rms_eps: float, + output_rms_eps: float, +) -> torch.Tensor +``` + +| Argument | Shape | Dtype | Layout | Device | +|---|---|---|---|---| +| `layer_residual` | `[T, 1, 7168]` | bf16 | contiguous | CUDA | +| `block_residual` | `[K, T, 1, 7168]`, `K = N - 1 >= 0` | bf16 | contiguous | `layer_residual`'s | +| `res_weight` | 7168 elements (`[H]`, `[1, H]` or `[H, 1]`) | bf16 | contiguous | `layer_residual`'s | +| `rms_weight` | 7168 elements | bf16 | contiguous | `layer_residual`'s | +| `output_rms_weight` | 7168 elements | bf16 | contiguous | `layer_residual`'s | +| `rms_eps` | scalar | Python float, passed as fp32 | — | — | +| `output_rms_eps` | scalar | Python float, passed as fp32 | — | — | +| returns | `[T, 1, 7168]` | bf16 | contiguous, newly allocated | `layer_residual`'s | + +`res_weight` and `rms_weight` score the candidates, with `rms_eps`; `output_rms_weight` and `output_rms_eps` +belong to the trailing norm. `K = 0` is passed as a zero-size `[0, T, 1, 7168]` tensor. No input is mutated and the +output aliases no input. + +### Certified arguments + +- `H = 7168`, `B = 1`, bf16 inputs, `res_weight`, `rms_weight` and `output_rms_weight` of shape `[H]`. +- Every `N = 2 ... 9` -- the candidate counts Kimi K3 produces: its 93 layers append a snapshot every 12 layers, so + the bank holds at most `ceil(93 / 12) = 8` snapshots -- at `T = 1 ... 8` (decode, including multi-token + speculative steps) and at `T = 32` and `T = 300`. +- `N = 1, 10, 11, 12` at `T = 1` and `T = 8`. +- `rms_eps = output_rms_eps = 1e-6` in every cell; the pairs `(1e-5, 1e-5)`, `(1e-2, 1e-6)` and `(1e-6, 1e-2)` in + three cells that cover both kernels. `1e-2` exceeds the mean square of both the test inputs and their mixture, + so the two unequal pairs pin which norm each eps feeds. + +## Metadata consumed + +None. Stateless: no attention metadata, no workspace, no caller-visible state. The op caches per device the +architecture check and the one-time raise of the split-K kernel's dynamic shared-memory limit, and reads +`TRTLLM_ENABLE_PDL` once per process. These select launch parameters, never results. + +## Preconditions + +Each item is a `TORCH_CHECK` in `cpp/tensorrt_llm/thop/attnResOp.cpp`, evaluated in this order; a failure raises +`RuntimeError` with the quoted message. This entry's test does not exercise the rejections. + +1. `layer_residual.dim() == 3` and `block_residual.dim() == 4` (`attn_res_rmsnorm_fwd: layer_residual must be + [T, B, H]`, `attn_res_rmsnorm_fwd: block_residual must be [K, T, B, H]`). +2. All five tensors are CUDA tensors on one device (`attn_res_rmsnorm_fwd: all input tensors must be CUDA + tensors`, `... must be on the same CUDA device`). +3. The device has compute capability `10.x`; `B == 1`; `1 <= N <= 12` with `N = block_residual.shape[0] + 1`; + `1 <= T <= 16384`; `H` a multiple of 1024 in `[4096, 8192]`. This is the contract check `attn_res_fwd` shares, + and its messages name that op: `attn_res_fwd requires an sm_100-family (datacenter Blackwell) GPU`, + `attn_res_fwd: unsupported B=...`, `N=...`, `T=...`, `H=...`. +4. `H == 7168` (`attn_res_rmsnorm_fwd: requires B=1 and H=7168, got T=... B=... H=...`): the kernels derive their + per-thread tiling from `H` at compile time. +5. All five tensors are bf16 (`attn_res_rmsnorm_fwd: must be bf16`) and contiguous + (`attn_res_rmsnorm_fwd: inputs must be contiguous`). A strided view raises; it is never read as if it were dense. +6. `block_residual.shape == (N - 1, T, B, H)` (`attn_res_rmsnorm_fwd: block_residual shape must match + layer_residual`). +7. `res_weight`, `rms_weight` and `output_rms_weight` each have `H` elements (`attn_res_rmsnorm_fwd: must + have H elements`). + +The kernels take any token count as a grid dimension, but the shared contract check still caps `T` at 16384. + +## Notes + +- **Code paths.** The kernel depends on `N` only, never on `T`: + + | `N` | Kernel | Launch | + |---|---|---| + | 1 ... 4 | single-CTA kernel | `T` CTAs of 256 threads, one per token | + | 5 ... 12 | split-K kernel | `T` clusters of 8 CTAs of 256 threads, one per token | + + They are the templates of `attn_res_fwd`'s two decode kernels, instantiated with the trailing norm for every `N` + (`attn_res_fwd` itself runs them only at `T == 1`, for five `N` values). The single-CTA kernel gives each of its + 256 threads 28 of the 7168 elements and keeps every candidate's values for them in registers as bf16. In the + split-K kernel each CTA of the cluster owns 896 of the 7168 elements, keeps them in shared memory as fp32, and + exchanges the per-candidate sums, then the mixture's sum of squares, through distributed shared memory. This op + has no persistent path; the persistent fused kernel is the separate op `attn_res_add_rmsnorm_persistent_fwd`. +- **Architecture.** The op requires compute capability `10.x`, and the kernels are built only for the sm_100 + family target (`100f`). Both use paired fp32 arithmetic (`.f32x2`); the split-K kernel also uses an 8-CTA + thread-block cluster with distributed shared memory and takes `3,652 * N` bytes of dynamic shared memory per CTA. +- **PDL: the output is written after the dependents are released.** `TRTLLM_ENABLE_PDL` (read once per process; + unset or `1` enables it, `0` disables it) controls programmatic dependent launch. When it is enabled both + kernels launch with programmatic stream serialization. Every CTA first waits on the grid dependency + (`cudaGridDependencySynchronize`), so the op reads its inputs only after the preceding kernel has completed, and + then immediately triggers its dependents (`cudaTriggerProgrammaticLaunchCompletion`), before it computes + anything. A dependent kernel launched with PDL can therefore start while this op is still writing `output`: it + must wait on the grid dependency before it reads `output` or overwrites any of this op's inputs. A kernel + launched without the PDL attribute is ordered by the stream as usual. With PDL disabled the kernels launch + normally and contain no grid-dependency instructions. PDL changes scheduling, not arithmetic. The test chains two + calls, the second reading the first's output as its `layer_residual`, and checks the chained result bit for bit + against the same call on a settled copy. +- **Determinism.** Every token is computed by its own CTA or cluster, with the same code at every `T`, reducing in + a fixed order without atomics. Identical inputs give identical bits run to run, which the test asserts in every + cell, and a token's result does not depend on `T` or on its position in the call. +- **Tolerance.** The test compares the output with an fp32 torch evaluation of *Semantics*, both rounding + boundaries included, using the metric of main's op test + (`tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py`): cosine similarity above 0.9999 + and relative L2 error below 5e-3. +- **CUDA graphs.** The call does not synchronize with the host, and its launch configuration depends on shapes + only. The test captures one call per kernel after an eager warm-up and checks that replay reproduces the eager + bits. +- **Kimi K3 call sites.** `modeling_kimi_linear.py` uses this op for the attention-residual + RMSNorm pairs before + attention, before the MLP on the layers that append a snapshot, and at the model output, for calls of at most + `KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS` tokens (1 by default, 32 when `KIMI_K3_ATTN_RES_TOPOLOGY` is on) that the + persistent topology does not take. Layers whose pre-norm mixture the speculative drafter captures run + `attn_res_fwd` and a separate norm instead. The `register_fake` in `custom_ops/cpp_custom_ops.py` returns + `empty_like(layer_residual)`. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/attn_res_rmsnorm_fwd.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/attn_res_rmsnorm_fwd.py new file mode 100644 index 000000000000..e5d9087477b3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/attn_res_rmsnorm_fwd.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 attention-residual selection fused with the RMSNorm that consumes it.""" + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* + + +def attn_res_rmsnorm_fwd( + layer_residual: torch.Tensor, + block_residual: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + output_rms_weight: torch.Tensor, + rms_eps: float, + output_rms_eps: float, +) -> torch.Tensor: + """Return rmsnorm(attn_res_fwd's bf16 mixture) * output_rms_weight as a new bf16 [T, 1, 7168] tensor. + + `res_weight`, `rms_weight` and `rms_eps` score the candidates as in `attn_res_fwd`; + `output_rms_weight` and `output_rms_eps` belong to the trailing norm, which rounds the normalized + value to bf16 before the weight multiply. With PDL enabled the kernel releases its dependents before + it writes the output, so a dependent PDL kernel must wait on the grid dependency before reading it. + """ + return torch.ops.trtllm.attn_res_rmsnorm_fwd( + layer_residual, + block_residual, + res_weight, + rms_weight, + output_rms_weight, + rms_eps, + output_rms_eps, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_rmsnorm.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_rmsnorm.md index 490f1d9f71ac..05373dea5da6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_rmsnorm.md +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/flashinfer_rmsnorm.md @@ -1,6 +1,8 @@ --- receipts: + # sm_103 was certified before the sm_100 key was added, on the same test file. sm_103: {status: passed, tests: 4} + sm_100: {status: passed, tests: 4} --- # flashinfer_rmsnorm diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/k3_embed_norm.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/k3_embed_norm.md new file mode 100644 index 000000000000..bda71f5dba06 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/k3_embed_norm.md @@ -0,0 +1,120 @@ +--- +receipts: + sm_100: {status: passed, tests: 5} +--- + +# k3_embed_norm + +**Wraps** `torch.ops.trtllm.k3_embed_norm` (one call). + +## Semantics + +A decode step's embedding rows and the first layer's input RMSNorm, in one launch. For `N <= 64` token ids, a bf16 +table `[V, H]` and a bf16 norm weight `[H]`: + +``` +row[t] = table[ids[t], :] if 0 <= ids[t] < V, else a zero row t < N +raw[t, :] = row[t] (written into the caller's raw) +ss[t] = sum over h of fp32(row[t, h])^2 fp32 +rstd[t] = rsqrt(ss[t] / H + eps) fp32, approximate (fast-math) rsqrt +out[t, h] = bf16((fp32(row[t, h]) * rstd[t]) * fp32(weight[h])) one bf16 rounding +``` + +`raw` receives the rows bit for bit. One 128-thread CTA computes one token's row: each thread squares and sums its +`H / 128` values, the sum is reduced across the warp's lanes, then across the 4 warps through shared memory, in a +fixed order; so the result is bit-identical from run to run (certified), and a row's result depends only on its id, +not on `N` or the other ids (certified: every row bit-identical to the same id's row of a 64-id call). An id outside +`[0, V)` reads nothing from the table: its `raw` row and its output row are zeros (certified). + +The op states that `out` is bit-identical to the gather (`trtllm::k3_embed`) followed by `flashinfer.norm.rmsnorm` on +flashinfer's CuTe DSL RMSNorm kernel, the kernel the catalog's `norm/flashinfer_rmsnorm` runs unless +`FLASHINFER_USE_CUDA_NORM=1`: it keeps that kernel's geometry for `6144 < H <= 16384` (one 128-thread CTA per row) +and its arithmetic order. The op's own test checks that bit for bit. This entry's test checks `out` against a native +fp64 torch RMSNorm, within 8e-3 of each row's largest magnitude, and `raw` bit for bit against a torch gather. + +Fusion boundary. Inside: the gather (zero rows for ids outside `[0, V)`), the copy of the rows into `raw`, the RMSNorm +and the weight scaling. Outside: producing the ids; allocating `raw` (the Kimi K3 model passes slot 0 of its +attention-residual snapshot bank `[S, N, H]`, so the rows become layer 0's first snapshot); any residual add, +`(1 + weight)` scaling or quantization. `ids`, `table` and `weight` are read only; every element of `raw` is written; +the result is a new tensor. + +## Signature + +```python +def k3_embed_norm( + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, +) -> torch.Tensor +``` + +| Argument | Shape | Dtype | Layout | Device | +|---|---|---|---|---| +| `ids` | `[N]`, `1 <= N <= 64` | int32 or int64 | 1-D (the op copies a strided one) | CUDA | +| `table` | `[V, H]`, `6144 < H <= 16384`, `H % 1024 == 0` | bf16 | contiguous, 16-byte aligned | CUDA, `ids`' device | +| `weight` | `[H]` | bf16 | contiguous, 16-byte aligned | CUDA, `ids`' device | +| `eps` | scalar | Python float | — | — | +| `raw` | `[N, H]`, overwritten | bf16 | contiguous, 16-byte aligned | CUDA, `ids`' device | +| returns | `[N, H]` | bf16 | contiguous, newly allocated | `ids`' device | + +### Certified arguments + +| Argument | Shape | Dtype | Layout | Device | +|---|---|---|---|---| +| `ids` | `[N]`, `N` = 1, 2, ..., 8 and 16, 24, ..., 64 | int32 (every `N`), int64 (`N` 1, 8, 64) | contiguous | CUDA | +| `table` | `[163840, 7168]` | bf16 | contiguous | CUDA | +| `weight` | `[7168]` | bf16 | contiguous | CUDA | +| `eps` | `1e-6` (every cell) and `1e-5` | float | — | — | +| `raw` | `[N, 7168]`: slot 0 (every cell) or slot 1 of a `[4, N, 7168]` bank | bf16 | contiguous | CUDA | +| returns | `[N, 7168]` | bf16 | contiguous | CUDA | + +The `N` are the token counts of Kimi K3's decode steps: up to 8 tokens, and DSpark verify steps of up to 8 requests of +8 tokens. The table has Kimi K3's shape (a replicated vocabulary of 163840), with rows at log-uniform scales from 1e-4 +to 1e2 and one all-zero row; the norm weight is around 1 with every 101st element negative. The ids are random in +`[0, V)` with 0, `V - 1`, the zero row, the smallest-scale row and repeats; and, at `N` = 8, ids outside `[0, V)` +(-1, `V`, `V + 7`, `-V`, `2^31 - 1`, `-2^31`, and for int64 also `2^32`, `2^32 + 5`, `3 - 2^32`, `2^40`). + +Every cell: `raw` bit-identical to the torch gather (zero rows outside `[0, V)`) and the bank's other slots +untouched; `out` within 8e-3 of each row's `max |ref|` against an fp64 torch RMSNorm, and exactly zero on zero rows; +a rerun bit-identical; every row bit-identical to the same id's row of the 64-id call; int64 ids bit-identical to +int32 ids. Also certified: a CUDA graph of an `N` = 8 call replayed with the ids rewritten in place, bit-identical to +eager calls; and the refusals under *Preconditions*. + +## Metadata consumed + +None besides the arguments and the current CUDA stream (the kernel is launched there), plus two pieces of process +state: + +- A per-process compile cache keyed by `(N, V, H, ids dtype, PDL)`: `N`, `V` and `H` are compiled into the kernel, + `eps` and the tensors' addresses are runtime arguments. The first call for a key compiles the kernel (seconds) and + must be made eagerly: under CUDA-graph capture it raises `RuntimeError` ("must run once per shape outside CUDA-graph + capture first"). The cache is result-neutral. +- `TRTLLM_ENABLE_PDL` (default `"1"`), read on every call: whether the kernel is launched with PDL. It is part of the + cache key and changes scheduling, not results. The test runs with the default. + +## Preconditions + +- `ids` 1-D, int32 or int64, `1 <= N <= 64`, on a CUDA device. +- `table` bf16, 2-D, contiguous, 16-byte-aligned data pointer, `6144 < H <= 16384` and `H % 1024 == 0`. +- `weight` bf16 of shape `(H,)`, contiguous, 16-byte aligned; `raw` bf16 of shape `(N, H)`, contiguous, 16-byte + aligned. +- A call outside these raises `ValueError` ("k3_embed_norm: unsupported call ...") before launching anything, leaving + `raw` untouched. Certified: `N` = 0 and 65, int16 ids, `raw` `[64, H]` for 8 ids, `raw` 2 bytes past a 16-byte + boundary, `weight` `[H / 2]`, and tables with `H` = 6144 and 7680. +- Not checked by the op: `table`, `weight` and `raw` on `ids`' device; `raw` not overlapping `ids`, `table` or + `weight`; a GPU with programmatic dependent launch (SM 9.0 or newer: the kernel uses `griddepcontrol`; the receipts + name the architectures it is certified on); the CuTe DSL (`cutlass`) and `cuda-python` (`cuda.bindings`), which the + call imports, so without them it raises `ImportError`. + +## Notes + +- PDL: with `TRTLLM_ENABLE_PDL` = `"1"` the kernel is launched with programmatic dependent launch. It releases its + dependents at entry and waits for the kernel it follows (`griddepcontrol.wait`) before reading anything, so the ids + may be the predecessor's output. A dependent must itself wait for this grid before reading `raw` or the result, as + every PDL kernel waits for its predecessor before reading its output. +- Grid: `N` CTAs of 128 threads; no CTA reads another's writes. +- The Kimi K3 table is replicated (no vocabulary shard), so the model never passes an id outside `[0, V)`; the zero + rows define the op on every input. Where `weight` is negative, a zero row's output holds `-0.0`. +- `trtllm::k3_embed` (the gather alone, into a new tensor) is a separate op and not this entry. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/k3_embed_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/k3_embed_norm.py new file mode 100644 index 000000000000..927ff24364cd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/norm/k3_embed_norm.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3's decode-step embedding and its first layer's input RMSNorm in one CuTe DSL launch: the rows ``table[ids]`` +written into a caller's buffer, and their RMSNorm.""" + +import torch + +import tensorrt_llm._torch.cute_dsl_kernels.k3_embed.op # noqa: F401 — registers the op + + +def k3_embed_norm( + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, +) -> torch.Tensor: + """Write ``table[ids]`` into ``raw`` (zero rows for ids outside ``[0, V)``) and return + ``raw * rsqrt(mean(raw^2) + eps) * weight`` as a new bf16 tensor, in one k3_embed_norm call.""" + return torch.ops.trtllm.k3_embed_norm(ids, table, weight, eps, raw) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/__init__.py new file mode 100644 index 000000000000..a0f916e7970b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""State-space (KDA) entries.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_attn.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_attn.md new file mode 100644 index 000000000000..67648d10ae0e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_attn.md @@ -0,0 +1,196 @@ +--- +receipts: + sm_100: {status: passed, tests: 8} +--- + +# k3_kda_attn + +**Wraps** `torch.ops.trtllm.k3_kda_attn` (one call), over a caller-owned `K3KdaBuffers` +(`catalog/ssm/k3_kda_buffers.py`). Sibling in this contract: `k3_kda_qkvg`, which wraps +`torch.ops.trtllm.k3_kda_qkvg` (the projection stream alone, one call). + +Kimi K3's KDA layer for one speculative-decoding request, at the TP16 rank slice: the fused input projection of the +request's golden token and its 7 drafts and the KDA verify of `ssm/k3_kda_verify` on it, in one launch. + +## Semantics + +`x` bf16 `[8, 7168]` is one request's golden token and 7 drafts. One launch computes: + +``` +# 1. projection, streamed by 26 clusters of 4 CTAs (w = [q | k | v | og | f_a | b | pad], 3208 rows) +y = x @ w^T +# q, k, f_a: bf16 of the four split-K partials of a cluster, summed in rank order +# v, og, b: bf16(half_0 + half_1), the fp32 sums of the two K halves +# 2. verify (exactly ssm/k3_kda_verify on y with the gate folded in, g_ext = None): +# starting state S = ssm[slot] if P == 0 else state_tok[slot, P - 1], P = pending[slot] +# per token t = 0..7: g = bf16(f_a @ w_fb^T); q, k, v = SiLU(conv4(window, new raw)); q, k L2-normalized, +# q *= scale; beta = sigmoid(b); decay = exp(lower_bound * sigmoid(exp(a_log) * (g + dt_bias))); +# S *= decay (per key); S += beta (v - S k) k^T; o = S q; +# out[t] = o * rsqrt(mean(o^2) + eps) * onorm_w * sigmoid(og) +``` + +and returns `out` bf16 `[8, 6, 128]`, the gated-norm core output. In place, at the request's slot: + +- `ssm[slot]` = the state after the golden token (t = 0); +- `state_tok[slot, t - 1]` = the state after draft t, t = 1..7; +- `cs_q` / `cs_k` / `cs_v[slot]` = the raw inputs at positions -2..7 around the golden token (the next round's + window starts at column P). + +The result is bit for bit that of `k3_kda_qkvg` followed by `ssm/k3_kda_verify` on the decoded rows, on copies of the +same pools (certified, every pool word). Repeated runs are bit-identical. + +**`k3_kda_qkvg`** runs phase 1 alone on `x` bf16 `[T <= 8, 7168]` and publishes the rows: buffer +`e = buffers.epoch[0]` (before the call) holds q, k and f_a as bf16 bits in `p1`, and v, og and b as the fp32 bits of +the two K-half partials in `part`. The projection is bf16 of their sum, the consumer's job. Measured against a float64 +`x @ w^T`: the rows are within 2.5e-3 of the projection's largest magnitude at T = 1, and 2.4e-3 at T = 3 and 8. + +## Signature + +```python +def k3_kda_attn( + x: torch.Tensor, + w: torch.Tensor, + w_fb: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + a_log: torch.Tensor, + dt_bias: torch.Tensor, + onorm_w: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + ssm: torch.Tensor, + state_tok: torch.Tensor, + slots: torch.Tensor, + pending: torch.Tensor, + buffers: K3KdaBuffers, + num_spec: int, + lower_bound: float, + scale: float, + eps: float, +) -> torch.Tensor + +def k3_kda_qkvg(x: torch.Tensor, w: torch.Tensor, buffers: K3KdaBuffers) -> None +``` + +| Argument | Shape | Dtype | Layout | Device | +|---|---|---|---|---| +| `x` | `[8, 7168]` (`k3_kda_qkvg`: `[T <= 8, 7168]`) | bf16 | contiguous | CUDA | +| `w` | `[3208, 7168]`: q, k, v, og (768 rows each), f_a (128), b (6), pad (2) | bf16 | contiguous | CUDA | +| `w_fb` | `[768, 128]` (f_b, out x in) | bf16 | contiguous | CUDA | +| `w_q`, `w_k`, `w_v` | `[768, 4]` conv taps, oldest input first | fp32 | dense | CUDA | +| `a_log` | `[6]` | fp32 | dense | CUDA | +| `dt_bias` | `[768]` | fp32 | dense | CUDA | +| `onorm_w` | `[128]` | fp32 | dense | CUDA | +| `cs_q`, `cs_k`, `cs_v` | `[pool, 768, 3 + num_spec]` | fp32 | channel stride 1 (dim-contiguous) | CUDA | +| `ssm` | `[pool, 6, 128, 128]` (V rows, K contiguous) | fp32 | each slot dense, slots at any stride | CUDA | +| `state_tok` | `[pool, num_spec, 6, 128, 128]` | fp32 | contiguous | CUDA | +| `slots` | `[1]`: the request's slot | int32 | any element offset | CUDA | +| `pending` | `[pool]`: drafts the sampler accepted last round, per slot | int32 | any element offset | CUDA | +| `buffers` | `K3KdaBuffers` made with `ctas=FUSED_CTAS` (128) | — | — | CUDA | +| `num_spec` | 7 | Python int | — | — | +| `lower_bound`, `scale`, `eps` | scalars (Kimi K3: -5.0, 128^-0.5, 1e-5) | Python float | — | — | +| returns | `[8, 6, 128]` | bf16 | contiguous, fresh | CUDA | + +`mutates_args`: `cs_q`, `cs_k`, `cs_v`, `ssm`, `state_tok`, and the set's `p1`, `part`, `epoch`. `k3_kda_qkvg` takes a +set made with `ctas=CTAS` (104) and mutates its `p1`, `part`, `epoch`; it returns None. + +## State + +**Object.** `K3KdaBuffers` (`catalog/ssm/k3_kda_buffers.py`), one per device, owned by the caller. + +**Contents and size.** The projection's three Lamport buffers and each CTA's buffer index: + +- `p1` int16 `[3 * 8 * 1664]`: per buffer, 8 token rows of the q, k and f_a columns (79,872 bytes); +- `part` int32 `[3 * 3 * 2 * 8 * 768]`: per buffer, the fp32 bits of the two K-half partials of v, og and b + (442,368 bytes); +- `epoch` int32 `[128]` (`[104]` for `k3_kda_qkvg`): each CTA's launch count mod 3. + +Every buffer word holds the sentinel (all ones, a NaN no finite GEMV produces; a computed all-ones word is stored as +the canonical NaN instead) until the launch's producer CTAs write it. The consumer CTAs poll the words they need until +none is the sentinel: the data words are the flags, with no fence or counter. 8 token rows are published per launch, +whatever the batch. + +**Who creates it, and when.** The target, in `post_load_weights`, with `K3KdaBuffers.create(device)`: + +- eager: it allocates, so it refuses to run under CUDA-graph capture; +- it arms every buffer word to the sentinel and every index to 0; +- `ctas` is `FUSED_CTAS` (128) for this op and `ssm/k3_kda_decode_attn`, or `CTAS` (104) for `k3_kda_qkvg`; any other + count raises. No environment variable is read. + +**Which ops may share one object.** `k3_kda_attn` and `ssm/k3_kda_decode_attn`, of every KDA layer on the device: +one set per device serves them all. Both run the same stream role, publish all 8 token rows, re-arm the same words +and move every CTA's index once per launch, so either may follow the other in any order. Certified: plain-decode and +verify launches interleaved on one set give the bits of the same launches on sets of their own; layers x steps on one +set, a set per layer, and two sets alternating between launches agree bit for bit. `k3_kda_qkvg` needs a set of its +own (104 CTAs); a 104-CTA set passed here raises `ValueError` before anything is written. + +**Call-order invariant.** Launches on one set run one at a time, in stream order: never on two streams at once. Each +launch reads `e = epoch[cta]` after its grid-dependency wait, writes buffer e, stores the sentinel into the same words +of buffer (e + 1) % 3 (the next launch's, which the launch before last read) and leaves `epoch[cta] = (e + 1) % 3`. +After any launch every CTA's index is equal. + +**What a later launch reads.** `epoch`, as the previous launch left it, and its own buffer's words, which the previous +launch re-armed to the sentinel. + +**How it is re-armed.** By the previous launch (above); the first launch on a new set writes buffer 0, which `create` +armed. The index stays in 0..2: a raw launch count would turn negative after 2^31 launches and its signed remainder +would index before the buffers. Certified with every index preset to 2^31 - 2: four launches give the bits of a fresh +set's and leave every index in 0..2 (the op's own test also checks with guard bands that nothing is written outside the +buffers). + +**The pools.** Not part of the object: they belong to the cache manager. With the V2 hybrid manager built with the +KDA replay caches and `kda_token_states`, layer `l`'s are `mamba_layer_cache(l).kda_conv_q` / `kda_conv_k` / +`kda_conv_v` (`cs_*`), `get_ssm_states(l)` (`ssm`, its slots strided by the manager's per-slot coalescing) and +`mamba_layer_cache(l).kda_state_tok`; `pending` is `prev_num_accepted_tokens`, one record shared by every layer and +written by the sampler's acceptance between steps. A launch on layer `l` writes only layer `l`'s views, at its slot. + +**The PDL rule.** The head CTAs read the slot's pools (state, records, conv window) before their grid-dependency wait. +A launch must therefore not directly follow another launch on the same pools in the stream: a kernel that waits, or a +non-PDL kernel, must sit between them, as the model's other layers do. This op lets its dependents launch only once +its head CTAs have passed their own wait, i.e. once the launch before it has completed, so one launch of this op (or of +`ssm/k3_kda_decode_attn`) on other pools in between is enough. A kernel that lets its dependents launch before its +wait (`ssm/k3_kda_verify` does, in its prologue) is not. + +**Why the test drives call sequences.** The state is carried by the pools from round to round: `pending` selects the +next round's starting state and conv window, and a round's per-draft states are only read by the next one. The test +therefore runs two requests one after the other, each for 9 rounds with pending running through 0..7, on every layer +of a real manager, against the unfused path bit for bit and layer 0 also against a float64 verify over the request's +committed history; then layers x rounds on one shared set against a set per layer, and a round of every layer captured +once and replayed with rewritten inputs and records. + +**What a wrong order does.** Two rounds of one request in swapped order (the test's negative control): nothing raises, +and the second round's output and the slot's state are those of a different history. Measured on sm_100: the second round's output is off by 1.00 and +the state by 1.01 (the largest absolute difference over the in-order result's largest magnitude). + +## Metadata consumed + +`slots` and `pending`. Both are read before the grid-dependency wait, so under CUDA graphs they must be written before +the graph runs (by the step's preparation and the previous step's acceptance), never by a kernel of the same step +ahead of this launch. + +## Preconditions + +- sm_100: the kernel uses tcgen05, TMA and clusters. +- The op checks `x`, `w`, `w_fb` (shape, dtype, contiguity), `num_spec` = 7, `ssm` / `state_tok` (shape, fp32, + `ssm.stride()[1:] == (16384, 128, 1)`, `state_tok` contiguous), `cs_q`'s window width, the fp32 dtypes, `slots` + (one int32 element), `pending` (int32) and the set's sizes (`epoch` of 128), and raises `ValueError` before a launch. + A conv cache that is not dense raises `ValueError` too. The other shapes in the table (the fp32 weights, `cs_k` / + `cs_v`, `pending` covering the slot, the slot inside the pool) are the caller's obligation. +- Every tensor the kernel addresses is 16-byte aligned except `slots` / `pending` (see *Notes*); the DSL checks the + assumed alignment at the call. +- The first call of each configuration (`lower_bound`, `scale`, `eps` and whether PDL is on) compiles the kernel and + must run outside CUDA-graph capture; under capture it raises `RuntimeError`. PDL follows `TRTLLM_ENABLE_PDL` (default + on), read at every call. +- The slot is not in use by another launch on the same pools (one request per launch; a batch is a sequence of + launches). + +## Notes + +- `slots` and `pending` may be slices of longer index tensors starting at any element (the mixer passes + `state_indices[num_prefills:]`): the kernel reads them at their element's alignment. The op's own test certifies + the bits at offsets 0..3. +- The pools are addressed at 64-bit slot offsets, so pools past 2 GiB are fine (the op's own test, + `test_k3_kda_pools_past_2g.py`). +- Repeated runs, and graph replays against eager runs, are bit-identical. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_attn.py new file mode 100644 index 000000000000..fce495f9baf5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_attn.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3's KDA projection and speculative verify of one request's 8 tokens in one launch, over a caller-owned +:class:`K3KdaBuffers`; and ``k3_kda_qkvg``, the projection stream alone.""" + +import torch + +from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_attn import ( # noqa: F401 — registers the ops + op as _k3_kda_attn_op, +) +from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_verify import ( # noqa: F401 — k3_kda_attn's helpers + op as _k3_kda_verify_op, +) + +from .k3_kda_buffers import K3KdaBuffers + + +def k3_kda_attn( + x: torch.Tensor, + w: torch.Tensor, + w_fb: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + a_log: torch.Tensor, + dt_bias: torch.Tensor, + onorm_w: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + ssm: torch.Tensor, + state_tok: torch.Tensor, + slots: torch.Tensor, + pending: torch.Tensor, + buffers: K3KdaBuffers, + num_spec: int, + lower_bound: float, + scale: float, + eps: float, +) -> torch.Tensor: + """Return the gated-norm core output bf16 [8, 6, 128] of one request's golden token and 7 drafts ``x`` bf16 + [8, 7168]: the projection ``x w^T`` and the verify of ``ssm/k3_kda_verify`` on it, in one launch. Updates the + slot's conv caches, state and per-draft states in place; advances ``buffers`` by one launch.""" + return torch.ops.trtllm.k3_kda_attn( + x, + w, + w_fb, + w_q, + w_k, + w_v, + a_log, + dt_bias, + onorm_w, + cs_q, + cs_k, + cs_v, + ssm, + state_tok, + slots, + pending, + buffers.p1, + buffers.part, + buffers.epoch, + num_spec, + lower_bound, + scale, + eps, + ) + + +def k3_kda_qkvg(x: torch.Tensor, w: torch.Tensor, buffers: K3KdaBuffers) -> None: + """The projection ``x w^T`` of T <= 8 tokens (``x`` bf16 [T, 7168], ``w`` bf16 [3208, 7168]) published into + ``buffers`` (made with ``ctas=CTAS``): buffer ``e = buffers.epoch[0]`` (before the call) holds the rows; advances + ``buffers`` by one launch. Returns None.""" + torch.ops.trtllm.k3_kda_qkvg(x, w, buffers.p1, buffers.part, buffers.epoch) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_buffers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_buffers.py new file mode 100644 index 000000000000..d2869f3c5afa --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_buffers.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Caller-owned state of Kimi K3's fused KDA projection: one set of its three Lamport buffers and per-CTA indices. + +A state type, not an entry: it launches nothing per call. Its constructor is eager; the target builds one per +device in ``post_load_weights`` (before any CUDA-graph capture) and passes it to every ``ssm/k3_kda_attn`` and +``ssm/k3_kda_decode_attn`` call on that device, of every KDA layer. ``ssm/k3_kda_attn``'s sibling +``k3_kda_qkvg`` (the projection stream alone) takes a set of its own, made with ``ctas=CTAS``. The contract is the +``## State`` section of ``k3_kda_attn.md``. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_attn import op as _op + +CTAS = _op.CTAS # k3_kda_qkvg's grid: 26 clusters of 4 CTAs +FUSED_CTAS = _op.FUSED_CTAS # k3_kda_attn's and k3_kda_decode_attn's grid + + +@dataclass(eq=False) +class K3KdaBuffers: + """The projection's Lamport set: three buffers of every published word, each word the sentinel (all ones) until a + launch writes it, and each CTA's buffer index. A launch writes buffer ``e = epoch[cta]``, re-arms buffer + ``(e + 1) % 3`` (the next launch's) to the sentinel and leaves ``epoch[cta] = (e + 1) % 3``, so launches on one + set run one at a time, in stream order, and every launch moves every index once (see the contract's + ``## State``).""" + + p1: torch.Tensor + """int16 [3 * 8 * 1664]: per buffer, 8 token rows of the q, k and f_a columns (bf16 bits).""" + part: torch.Tensor + """int32 [3 * 3 * 2 * 8 * 768]: per buffer, the fp32 bits of the two K-half partials of v, og and b.""" + epoch: torch.Tensor + """int32 [ctas]: each CTA's buffer index (its launch count mod 3).""" + ctas: int + + @classmethod + def create(cls, device, ctas: int = FUSED_CTAS) -> "K3KdaBuffers": + """Allocate and arm a set for ``ctas`` CTAs (``FUSED_CTAS`` for k3_kda_attn / k3_kda_decode_attn, ``CTAS`` + for k3_kda_qkvg) on ``device``: every buffer word the sentinel, every index 0. Eager: it allocates, so it + refuses to run under CUDA-graph capture.""" + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("K3KdaBuffers.create allocates: call it before CUDA-graph capture") + if ctas not in (CTAS, FUSED_CTAS): + raise ValueError( + f"K3KdaBuffers: ctas must be {CTAS} (k3_kda_qkvg) or {FUSED_CTAS}, got {ctas}" + ) + p1, part, epoch = _op.make_buffers(torch.device(device), ctas) + return cls(p1=p1, part=part, epoch=epoch, ctas=ctas) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_decode_attn.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_decode_attn.md new file mode 100644 index 000000000000..a65d9032c18d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_decode_attn.md @@ -0,0 +1,128 @@ +--- +receipts: + sm_100: {status: passed, tests: 14} +--- + +# k3_kda_decode_attn + +**Wraps** `torch.ops.trtllm.k3_kda_decode_attn` (one call), over a caller-owned `K3KdaBuffers` +(`catalog/ssm/k3_kda_buffers.py`), the same object `ssm/k3_kda_attn` takes. + +Kimi K3's KDA layer for plain (non-speculative) decode, at the TP16 rank slice: the fused input projection of one +token of each of R <= 8 requests and the one-token KDA decode of `ssm/kda_decode` on it, in one launch. + +## Semantics + +`x` bf16 `[R, 7168]`, one token of each of R requests. One launch computes: + +``` +# 1. projection: k3_kda_attn's stream (26 clusters of 4 CTAs, three Lamport phases), rows >= R read as zero +y = x @ w^T # q, k, f_a, v, og, b rounded to bf16 as in ssm/k3_kda_attn +# 2. per request r, on slot s = slots[r] (six head clusters of 4 CTAs, one V quarter per CTA): +g = bf16(f_a @ w_fb^T) +q, k, v = SiLU(conv4(conv[s] window, new raw)) # q, k L2-normalized, q *= scale +conv[s] = the window shifted by one: its last two raw inputs, then the new one +beta = sigmoid(b); decay = exp(lower_bound * sigmoid(exp(a_log) * (g + dt_bias))) +S = ssm[s] * decay (per key); S += beta (v - S k) k^T; ssm[s] = S; o = S q +out[r] = o * rsqrt(mean(o^2) + eps) * onorm_w * sigmoid(og) +``` + +and returns `out` bf16 `[R, 6, 128]`, the gated-norm core output. The decode is `ssm/kda_decode`'s arithmetic, row and +key layout. Certified against the model's unfused path (the projection stream alone, f_b as a bf16 `F.linear`, then +`ssm/kda_decode` on a copy of the pools) and a float64 decode, both at fp32 tolerance for the output (2e-2 relative, +bf16 outputs) and the state rows (1e-3); the conv windows bit for bit (raw bf16 inputs). Repeated runs are +bit-identical. + +## Signature + +```python +def k3_kda_decode_attn( + x: torch.Tensor, + w: torch.Tensor, + w_fb: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + a_log: torch.Tensor, + dt_bias: torch.Tensor, + onorm_w: torch.Tensor, + conv: torch.Tensor, + ssm: torch.Tensor, + slots: torch.Tensor, + buffers: K3KdaBuffers, + lower_bound: float, + scale: float, + eps: float, +) -> torch.Tensor +``` + +| Argument | Shape | Dtype | Layout | Device | +|---|---|---|---|---| +| `x` | `[R, 7168]`, 1 <= R <= 8 | bf16 | contiguous | CUDA | +| `w` | `[3208, 7168]` (as `ssm/k3_kda_attn`) | bf16 | contiguous | CUDA | +| `w_fb` | `[768, 128]` | bf16 | contiguous | CUDA | +| `w_q`, `w_k`, `w_v` | `[768, 4]` conv taps, oldest input first | fp32 | dense | CUDA | +| `a_log` | `[6]` | fp32 | dense | CUDA | +| `dt_bias` | `[768]` | fp32 | dense | CUDA | +| `onorm_w` | `[128]` | fp32 | dense | CUDA | +| `conv` | `[slots, 2304, 3]`: q, k, v channels, the last three raw inputs oldest first | bf16 | each slot dense; slot stride a multiple of 8 elements; 16-byte aligned base | CUDA | +| `ssm` | `[slots, 6, 128, 128]` (V rows, K contiguous) | fp32 | each slot dense; slot stride a multiple of 4; 16-byte aligned base | CUDA | +| `slots` | `[R]` | int32 | contiguous | CUDA | +| `buffers` | `K3KdaBuffers` made with `ctas=FUSED_CTAS` (128) | — | — | CUDA | +| `lower_bound`, `scale`, `eps` | scalars (Kimi K3: -5.0, 128^-0.5, 1e-5) | Python float | — | — | +| returns | `[R, 6, 128]` | bf16 | contiguous, fresh | CUDA | + +`mutates_args`: `conv`, `ssm`, and the set's `p1`, `part`, `epoch`. + +## State + +**Object.** `K3KdaBuffers`, shared with `ssm/k3_kda_attn`: one per device for both ops and every KDA layer. Its +contents, creation, re-arming, the 128-CTA requirement and the call-order invariant (launches on one set run one at a +time, in stream order; each launch moves every CTA's index once) are in `k3_kda_attn.md`'s `## State`. Certified here: +plain-decode launches and `ssm/k3_kda_attn` verify launches interleaved on one set give the bits of the same launches +on sets of their own; layers x steps on one set, a set per layer and two sets alternating agree bit for bit; the index +preset to 2^31 - 2 gives the bits of a fresh set and stays in 0..2; a 104-CTA set raises `ValueError` before anything +is written; `create` under capture raises `RuntimeError`. + +**The pools.** The cache manager's plain-decode states, not part of the object: with the V2 hybrid manager +(`conv_state_layout="q_k_v"`, bf16 conv states, fp32 SSM states), layer `l`'s `get_conv_states(l)` (`conv`) and +`get_ssm_states(l)` (`ssm`), their slots strided by the manager's per-slot coalescing; `slots` from +`get_state_indices`. A launch on layer `l` writes only layer `l`'s views, at its slots: certified on a real manager, +with the other slots of the layer and every slot of the other layers bit-unchanged. A batch names each slot once. + +**The PDL rule.** As `ssm/k3_kda_attn`: the head CTAs read the slots' pools (the conv windows, 32 state rows per CTA) +before their grid-dependency wait, so a launch must not directly follow another launch on the same pools in the +stream. A kernel that waits, or a non-PDL kernel, sits between them in the model; one launch of this op (or of +`ssm/k3_kda_attn`) on other pools is also enough, since its dependents launch only after its own wait. + +**Why the test drives call sequences.** Each step's state and conv window are the next step's input, so the test +runs layers x steps of a real manager, and captures one step of every layer once and replays it with rewritten +inputs: the replays give the bits of the same steps run eagerly on a copy of the pools. + +**What a wrong order does.** Two steps of one request in swapped order (the negative control): nothing raises, and the +second step's output and the final state are those of a different history. Measured on sm_100: the second step's output is off by 0.81 and the state by +0.89 (the largest absolute difference over the in-order result's largest magnitude). + +## Metadata consumed + +`slots`, read before the grid-dependency wait: under CUDA graphs it must be written before the graph runs. + +## Preconditions + +- sm_100: the kernel uses tcgen05, TMA and clusters. +- The op checks `x` (rank 2, 1..8 rows, 7168 columns, bf16, contiguous), `w`, `w_fb`, `conv` (bf16, `[*, 2304, 3]`, + strides `(*, 3, 1)`, slot stride a multiple of 8, 16-byte base), `ssm` (fp32, `[*, 6, 128, 128]`, dense slots, slot + stride a multiple of 4, 16-byte base), `slots` (R int32 elements, contiguous), the set's sizes (`epoch` of 128) and + the fp32 dtypes, and raises `ValueError` before a launch. The fp32 weights' shapes and the slots' range are the + caller's obligation. The V2 hybrid manager's per-layer views meet the layout (certified on a real manager). +- The first call of each configuration (`lower_bound`, `scale`, `eps` and whether PDL is on) compiles the kernel and + must run outside CUDA-graph capture; under capture it raises `RuntimeError`. PDL follows `TRTLLM_ENABLE_PDL` + (default on), read at every call. + +## Notes + +- The pools are addressed at 64-bit slot offsets, so pools past 2 GiB are fine (the op's own test, + `test_k3_kda_pools_past_2g.py`). `slots` may be a slice of a longer index tensor starting at any element; the op's + own test certifies offsets 0..3. +- Within the fp32 tolerances above, not bitwise, against `ssm/kda_decode`: the two read the same rows but sum the + decode in their own orders. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_decode_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_decode_attn.py new file mode 100644 index 000000000000..1d29d58f23d8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_decode_attn.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3's KDA projection and plain one-token decode of up to 8 requests in one launch, over a caller-owned +:class:`K3KdaBuffers`.""" + +import torch + +from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_attn import ( # noqa: F401 — registers the op + op as _k3_kda_attn_op, +) +from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_verify import ( # noqa: F401 — the op's helpers + op as _k3_kda_verify_op, +) + +from .k3_kda_buffers import K3KdaBuffers + + +def k3_kda_decode_attn( + x: torch.Tensor, + w: torch.Tensor, + w_fb: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + a_log: torch.Tensor, + dt_bias: torch.Tensor, + onorm_w: torch.Tensor, + conv: torch.Tensor, + ssm: torch.Tensor, + slots: torch.Tensor, + buffers: K3KdaBuffers, + lower_bound: float, + scale: float, + eps: float, +) -> torch.Tensor: + """Return the gated-norm core output bf16 [R, 6, 128] of one token of each of R <= 8 requests ``x`` bf16 + [R, 7168]: the projection ``x w^T`` and the plain decode of ``ssm/kda_decode`` on it, in one launch. Updates the + slots' conv and state pools in place; advances ``buffers`` by one launch.""" + return torch.ops.trtllm.k3_kda_decode_attn( + x, + w, + w_fb, + w_q, + w_k, + w_v, + a_log, + dt_bias, + onorm_w, + conv, + ssm, + slots, + buffers.p1, + buffers.part, + buffers.epoch, + lower_bound, + scale, + eps, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_verify.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_verify.md new file mode 100644 index 000000000000..a65197c0634f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_verify.md @@ -0,0 +1,145 @@ +--- +receipts: + sm_100: {status: passed, tests: 7} +--- + +# k3_kda_verify + +**Wraps** `torch.ops.trtllm.k3_kda_verify` (one call). + +Kimi K3's KDA speculative verify of N requests of 1 + num_spec tokens (a golden token and its drafts), from the fused +projection rows to the gated-norm core output, committing the state after every verify token so that the next round +starts from the drafts the sampler accepted instead of replaying them. + +## Semantics + +`proj` bf16 `[N (1 + num_spec), cols]` holds each request's 1 + num_spec rows of the fused projection +`[q | k | v | og | f_a | b | pad]` (H K, H K, H K, H K, K, H columns, then padding). Grid (H, N, 8): a cluster of 8 +CTAs per (head, request), each owning 16 V rows. For request n on slot `s = slots[n]`, with `P = pending[s]`: + +``` +S = ssm[s] if P == 0 else state_tok[s, P - 1] # the state after the last accepted token +window = conv caches cs_*[s] columns P..P+2 # raw inputs at positions -3..-1 +per token t = 0..num_spec: + g = bf16(f_a[t] @ w_fb^T) (or g_ext[t]: the unfused f_b output) + q, k, v = SiLU(conv4(window, raw[t])); q, k L2-normalized (q *= scale) + beta = sigmoid(b[t]); decay = exp(lower_bound * sigmoid(exp(a_log) * (g + dt_bias))) + S *= decay (per key); S += beta (v - S k) k^T; o = S q + out[t] = o * rsqrt(mean(o^2) + eps) * onorm_w * sigmoid(og[t]) +ssm[s] = the state after t = 0 (the golden token); state_tok[s, t - 1] = the state after draft t +cs_*[s] = the raw inputs at positions -2..num_spec around the golden token +``` + +Returns `out` bf16 `[N (1 + num_spec), H, 128]`. The recurrence is `kda_mtp_decode`'s V-split arithmetic over the +verify tokens, unrolled: with `g_ext` the op is bit-exact against `kda_mtp_decode` fed the same gate and replaying the +same accepted drafts (the op's own test). Certified here against a float64 verify over each request's committed +history: every round's outputs within 2e-2 relative (bf16) and the committed states within 1e-3, the per-draft states +and conv caches through the next round, which starts from them. Repeated runs, and graph replays against eager runs, +are bit-identical. + +## Signature + +```python +def k3_kda_verify( + proj: torch.Tensor, + w_fb: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + a_log: torch.Tensor, + dt_bias: torch.Tensor, + onorm_w: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + ssm: torch.Tensor, + state_tok: torch.Tensor, + slots: torch.Tensor, + pending: torch.Tensor, + num_spec: int, + lower_bound: float, + scale: float, + eps: float, + g_ext: Optional[torch.Tensor] = None, +) -> torch.Tensor +``` + +| Argument | Shape | Dtype | Layout | Device | +|---|---|---|---|---| +| `proj` | `[N (1 + num_spec), cols]`, cols a multiple of 8 and >= 4 H K + K + H | bf16 | contiguous | CUDA | +| `w_fb` | `[H K, K]` (f_b, out x in) | bf16 | contiguous | CUDA | +| `w_q`, `w_k`, `w_v` | `[H K, 4]` conv taps, oldest input first | fp32 | dense | CUDA | +| `a_log` | `[H]` | fp32 | dense | CUDA | +| `dt_bias` | `[H K]` | fp32 | dense | CUDA | +| `onorm_w` | `[K]` | fp32 | dense | CUDA | +| `cs_q`, `cs_k`, `cs_v` | `[pool, H K, 3 + num_spec]` | fp32 | channel stride 1 (dim-contiguous) | CUDA | +| `ssm` | `[pool, H, K, K]` (V rows, K contiguous) | fp32 | each slot dense, slots at any stride | CUDA | +| `state_tok` | `[pool, num_spec, H, K, K]` | fp32 | contiguous | CUDA | +| `slots` | `[N]` | int32 | any element offset | CUDA | +| `pending` | `[pool]`: drafts the sampler accepted last round, per slot | int32 | any element offset | CUDA | +| `num_spec` | drafts per request | Python int | — | — | +| `lower_bound`, `scale`, `eps` | scalars (Kimi K3: -5.0, 128^-0.5, 1e-5) | Python float | — | — | +| `g_ext` | `[N (1 + num_spec), H K]` or None (fold f_b into the kernel) | bf16 | contiguous | CUDA | +| returns | `[N (1 + num_spec), H, K]` | bf16 | contiguous, fresh | CUDA | + +K = V = 128, conv width 4. Kimi K3's TP16 rank slice is H = 6, num_spec = 7. `mutates_args`: `cs_q`, `cs_k`, `cs_v`, +`ssm`, `state_tok`. + +## State + +**Object.** None of its own (stateful kinds P and R): the op updates the caller's pools in place and reads the +record the sampler wrote after the previous round. No buffer persists in the op. + +**The pools and the record.** The cache manager's: with the V2 hybrid manager built with the KDA replay caches and +`kda_token_states` (`kda_replay_num_spec` = num_spec), layer `l`'s `mamba_layer_cache(l).kda_conv_q` / `kda_conv_k` +/ `kda_conv_v` (`cs_*`), `get_ssm_states(l)` (`ssm`, its slots strided by the manager's per-slot coalescing) and +`mamba_layer_cache(l).kda_state_tok`. `pending` is the manager's `prev_num_accepted_tokens`: one record shared by +every layer, written between steps by the sampler's acceptance. A call on layer `l` writes only layer `l`'s views, at +its slots: certified on a real manager, the other slots bit-unchanged. + +**Call-order invariant.** Per layer, one call per verify round, in round order, with `pending` updated between rounds +(after the sampler accepts) and not during one. A call reads what the previous round's call on that layer left at the +slot: the state after the last accepted token (`ssm` if none was accepted, else `state_tok[s, P - 1]`) and the conv +window starting at column P. The per-draft states of a round are read only by the next round. + +**The PDL rule.** The kernel reads the slot, P, the starting state and the conv window before its grid-dependency +wait, and lets its dependents launch in its prologue, before that wait. So a launch must not follow another launch on +the same pools without a kernel that waits (or a non-PDL kernel) in between, and the launch in between must not be +one that lets its dependents launch before its own wait either: one `k3_kda_verify` launch on other pools is not +enough. In the model, a layer's other kernels sit between its launches. + +**Why the test drives call sequences.** The record and the per-draft states make each round depend on the previous +one's acceptance, which no single call can check. The test therefore runs 6 rounds of every layer with a pending count +per request that changes every round (every count 0..num_spec over the rounds), against the float64 history, and +captures a round of every layer once, replaying it with rewritten rows and records: the replays give the bits of the +same rounds run eagerly, and the schedule run twice gives the same bits. + +**What a wrong order does.** Two rounds of one request in swapped order (the negative control): nothing raises, and +the second round's outputs and the slot's state are those of a different history. Measured on sm_100: the second round's outputs are off by 0.98 and the +state by 0.98 (the largest absolute difference over the in-order result's largest magnitude). + +## Metadata consumed + +`slots` and `pending`, both read before the grid-dependency wait: under CUDA graphs they must be written before the +graph runs (the step's preparation and the previous step's acceptance). + +## Preconditions + +- sm_100: the kernel uses tcgen05, TMA and clusters. +- The op checks `proj` (bf16, rank 2, N (1 + num_spec) rows, columns), `w_fb`, V = K = 128, `ssm` / `state_tok` + (shapes, fp32, `ssm.stride()[1:] == (K K, K, 1)`, `state_tok` contiguous), `cs_q`'s window width, the int32 index + tensors and the fp32 dtypes, and raises `ValueError` before a launch. A conv cache that is not dense raises + `ValueError` too. The fp32 weights' shapes, `cs_k` / `cs_v`, `pending` covering the slots and the slots' range are + the caller's obligation. +- The first call of each configuration (H, num_spec, `lower_bound`, `scale`, `eps`, with or without `g_ext`, and + whether PDL is on) compiles the kernel and must run outside CUDA-graph capture; under capture it raises + `RuntimeError`. PDL follows `TRTLLM_ENABLE_PDL` (default on), read at every call. +- A call names each slot once. + +## Notes + +- `slots` and `pending` may be slices of longer index tensors starting at any element (the mixer passes + `state_indices[num_prefills:]`); the op's own test certifies the bits at offsets 0..3. +- The pools are addressed at 64-bit slot offsets, so pools past 2 GiB are fine (the op's own test, + `test_k3_kda_pools_past_2g.py`). +- `ssm/k3_kda_attn` computes this verify fused with the layer's projection, for one request of 8 tokens. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_verify.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_verify.py new file mode 100644 index 000000000000..822f0d29e23c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/k3_kda_verify.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3's KDA speculative verify of N requests of 1 + num_spec tokens, from the fused projection rows to the +gated-norm core output, committing the state after every verify token.""" + +from typing import Optional + +import torch + +from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_verify import ( # noqa: F401 — registers the op + op as _k3_kda_verify_op, +) + + +def k3_kda_verify( + proj: torch.Tensor, + w_fb: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + a_log: torch.Tensor, + dt_bias: torch.Tensor, + onorm_w: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + ssm: torch.Tensor, + state_tok: torch.Tensor, + slots: torch.Tensor, + pending: torch.Tensor, + num_spec: int, + lower_bound: float, + scale: float, + eps: float, + g_ext: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Return the gated-norm core output bf16 [N (1 + num_spec), H, 128] of the verify tokens whose fused + projection rows are ``proj``; updates the slots' conv caches, state and per-draft states in place.""" + return torch.ops.trtllm.k3_kda_verify( + proj, + w_fb, + w_q, + w_k, + w_v, + a_log, + dt_bias, + onorm_w, + cs_q, + cs_k, + cs_v, + ssm, + state_tok, + slots, + pending, + num_spec, + lower_bound, + scale, + eps, + g_ext, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/kda_decode.md b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/kda_decode.md new file mode 100644 index 000000000000..8974bb788a54 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/kda_decode.md @@ -0,0 +1,139 @@ +--- +receipts: + sm_100: {status: passed, tests: 11} +--- + +# kda_decode + +**Wraps** `torch.ops.trtllm.kda_decode` (one call). + +One-token KDA (Kimi Delta Attention) decode of B requests: the causal depthwise conv over each request's conv window, +the gated delta rule on its recurrent state, and the gated RMS norm, with the state (and optionally the conv windows) +updated in place at each request's slot. Certified here at Kimi K3's TP16 rank slice (6 heads, K = V = 128, conv +width 4) on a real cache manager's pools. + +## Semantics + +For request b (slot `s = ssm_state_indices[b]`), head h, with `x_q`, `x_k`, `x_v` the new raw q, k, v inputs: + +``` +q, k, v = SiLU(conv4(window[s], new raw) + bias) # taps w_*_t[j], oldest input first; bias_* added before SiLU +q = q / sqrt(sum(q^2) + 1e-6) * scale; k = k / sqrt(sum(k^2) + 1e-6) +beta = sigmoid(beta) +decay = exp(lower_bound * sigmoid(exp(a_log) * (g + dt_bias))) +S = state[s] * decay (per key); S += beta (v - S k) k^T; state[s] = S; o = S q +output = o * rsqrt(mean(o^2) + onorm_eps) * onorm_weight * sigmoid(onorm_g) +``` + +`apply_onorm`, `use_lower_bound` and `apply_beta_sigmoid` must all be True: the op supports only that combination +(*Preconditions*), which is the one above and Kimi K3's. + +With `update_conv_cache` the conv windows at the slot shift by one (the last two raw inputs, then the new one). The +float64 reference of the test is this arithmetic; the op matches it within fp32 tolerance (outputs 2e-2 relative, +bf16; state rows 1e-3), and the conv windows bit for bit (raw bf16 inputs). + +**Kernels.** On sm_100 and sm_103 the dispatcher picks by the workload B x H: up to 32 the four-CTA cluster kernel, up +to 144 the legacy compact-heads kernel, above that per-architecture choices among the optimized single-CTA and bulk +kernels and the legacy many-heads kernel. At Kimi K3's 6 heads: B <= 5 runs the cluster kernel and 6 <= B <= 24 the +legacy compact-heads kernel (whose block reduction #19830 orders with `__syncwarp`). The test runs B = 1..8, so both. +Other architectures run the legacy kernels. + +## Signature + +```python +def kda_decode( + 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, + bias_k: torch.Tensor, + bias_v: torch.Tensor, + conv_state_q: torch.Tensor, + conv_state_k: torch.Tensor, + conv_state_v: torch.Tensor, + a_log: torch.Tensor, + g: torch.Tensor, + dt_bias: torch.Tensor, + beta: torch.Tensor, + onorm_g: torch.Tensor, + onorm_weight: torch.Tensor, + ssm_state_indices: Optional[torch.Tensor], + state: torch.Tensor, + apply_onorm: bool, + update_conv_cache: bool, + use_lower_bound: bool, + apply_beta_sigmoid: bool, + lower_bound: float, + scale: float, + onorm_eps: float, + output: torch.Tensor, +) -> None +``` + +| Argument | Shape | Dtype | Layout | Device | +|---|---|---|---|---| +| `x_q`, `x_k`, `x_v` | `[1, B, H, 128]` | bf16 | head and channel axes packed; any row stride | CUDA | +| `w_q_t`, `w_k_t`, `w_v_t` | `[4, H * 128]` conv taps, oldest input first | bf16 | contiguous | CUDA | +| `bias_q`, `bias_k`, `bias_v` | `[H * 128]` (zeros: no bias) | bf16 | contiguous | CUDA | +| `conv_state_q/k/v` | with `update_conv_cache`: `[slots, H * 128, 3]` section views of one packed `[slots, 3 H 128, 3]` pool (`q | k | v`); else `[B, H * 128, 3]` | bf16 | packed: equal slot strides >= `3 H 128 * 3`, strides `(*, 3, 1)`; else contiguous | CUDA | +| `a_log` | `[H]` | fp32 | contiguous | CUDA | +| `g` | `[1, B, H, 128]` (the gate before `dt_bias`) | bf16 | head and channel axes packed | CUDA | +| `dt_bias` | `[H * 128]` | fp32 | contiguous | CUDA | +| `beta` | `[1, B, H]` | bf16 | head axis stride 1 | CUDA | +| `onorm_g` | `[1, B, H, 128]` (the output gate) | bf16 | head and channel axes packed | CUDA | +| `onorm_weight` | `[128]` | fp32 | contiguous | CUDA | +| `ssm_state_indices` | `[B]` slots, or None (state rows 0..B-1) | int32 | contiguous | CUDA | +| `state` | `[slots >= B, H, 128, 128]` (V rows, K contiguous) | fp32 | each slot dense; slot stride a multiple of 4; 16-byte aligned base | CUDA | +| `apply_onorm`, `update_conv_cache`, `use_lower_bound`, `apply_beta_sigmoid` | scalars (Kimi K3: all True) | Python bool | — | — | +| `lower_bound`, `scale`, `onorm_eps` | scalars (Kimi K3: -5.0, 128^-0.5, 1e-5) | Python float | — | — | +| `output` | `[B, 1, H, 128]` | bf16 | contiguous | CUDA | + +H is one of 1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, 96 and the same for q / k and v. Returns None; writes `output`. +Schema mutations (`Tensor(a!)`): `conv_state_q`, `conv_state_k`, `conv_state_v`, `state`, `output`. + +## State + +**Object.** None of its own (stateful kind P): the op updates the caller's pools in place, the recurrent `state` at +each request's slot and, with `update_conv_cache`, the conv windows there. Nothing persists in the op between calls. + +**The pools.** The cache manager's: with the V2 hybrid manager (`conv_state_layout="q_k_v"`, bf16 conv states, fp32 +SSM states), layer `l`'s `get_conv_states(l)` split into its q, k, v sections and `get_ssm_states(l)`, their slots +strided by the manager's per-slot coalescing; `ssm_state_indices` from `get_state_indices`. A call on layer `l` writes +only layer `l`'s views, at its slots: certified on a real manager, with the other slots of the layer and every slot of +the other layers bit-unchanged. + +**Call-order invariant.** One call per layer per decode step, in step order; a batch names each slot once. A call +reads the state and conv window that the previous step's call on that layer left at the slot. + +**Why the test drives call sequences.** Each step's state is the next step's input, so the test runs layers x steps +on a real manager and captures a step of every layer once, replaying it with rewritten inputs: the replays give the +bits of the same steps run eagerly on a copy of the pools. + +**What a wrong order does.** Two steps of one request in swapped order (the negative control): nothing raises, and the +second step's output and the final state are those of a different history. Measured at Kimi K3's shape on sm_100: the second step's output is off by 1.38 and the state +by 1.27 (the largest absolute difference over the in-order result's largest magnitude). + +## Metadata consumed + +`ssm_state_indices` (the requests' slots). None of the attention metadata. + +## Preconditions + +- Every check in the table is the op's (`TORCH_CHECK`): a violation raises `RuntimeError` before the launch, with the + pools unchanged. Among them: the state base must be 16-byte aligned and its slot stride a multiple of 4 floats (the + kernels move state with 16-byte accesses at `slot * stride(0)`); `ssm_state_indices` must be int32. +- `apply_onorm`, `use_lower_bound` and `apply_beta_sigmoid` must all be True, on every architecture and batch size: + the op checks them first and otherwise raises `RuntimeError` ("KDA decode only supports apply_onorm=true, + use_lower_bound=true, and apply_beta_sigmoid=true") before the launch. Certified with the lower bound off at B = 2, + H = 6, the pools unchanged. +- The conv inputs, gates and outputs are bf16 and the state fp32; K = V = 128 and conv width 4 only. + +## Notes + +- Kimi K3's KDA layer calls this op through `run_kda_decode_fusion_cuda` + (`_torch/modules/kimi_kda/_kda_decode.py`), which fills the optional bias, gate and norm arguments with cached + dummies; this entry exposes the schema raw. +- `ssm/k3_kda_decode_attn` computes the same decode fused with the layer's projection, for up to 8 requests. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/kda_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/kda_decode.py new file mode 100644 index 000000000000..09dd96c8d95f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/ssm/kda_decode.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One-token KDA decode of B requests: the causal conv, the gated delta rule on each request's state slot and the +gated RMS norm, the conv and state pools updated in place.""" + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* + + +def kda_decode( + 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, + bias_k: torch.Tensor, + bias_v: torch.Tensor, + conv_state_q: torch.Tensor, + conv_state_k: torch.Tensor, + conv_state_v: torch.Tensor, + a_log: torch.Tensor, + g: torch.Tensor, + dt_bias: torch.Tensor, + beta: torch.Tensor, + onorm_g: torch.Tensor, + onorm_weight: torch.Tensor, + ssm_state_indices: Optional[torch.Tensor], + state: torch.Tensor, + apply_onorm: bool, + update_conv_cache: bool, + use_lower_bound: bool, + apply_beta_sigmoid: bool, + lower_bound: float, + scale: float, + onorm_eps: float, + output: torch.Tensor, +) -> None: + """Writes the decode output into ``output`` [B, 1, HV, 128] bf16; updates ``state`` at ``ssm_state_indices`` + and, with ``update_conv_cache``, the conv pools there. Returns None.""" + torch.ops.trtllm.kda_decode( + x_q, + x_k, + x_v, + w_q_t, + w_k_t, + w_v_t, + bias_q, + bias_k, + bias_v, + conv_state_q, + conv_state_k, + conv_state_v, + a_log, + g, + dt_bias, + beta, + onorm_g, + onorm_weight, + ssm_state_indices, + state, + apply_onorm, + update_conv_cache, + use_lower_bound, + apply_beta_sigmoid, + lower_bound, + scale, + onorm_eps, + output=output, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/__init__.py new file mode 100644 index 000000000000..8d4f159338fc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KimiK3ForConditionalGeneration targets.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/__init__.py new file mode 100644 index 000000000000..474e865afce4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 (MXFP4) / sm_100 / tp16 attention, routed experts moe_tp 4 x moe_ep 4.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_gemv.py new file mode 100644 index 000000000000..b384010b65dd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_gemv.py @@ -0,0 +1,411 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode path's GEMVs, LM head and embedding on the catalog's single-GPU Kimi K3 entries. + +* **Per-site GEMVs** (`K3DecodeGemvs.project`): a projection of a decode step runs on the kernel measured fastest + at its call site's weight shape (`SITES`): at most `MAX_ROWS` rows on `gemm/k3_decode_gemv`, + `gemm/k3_ctm_gemv_wide` or `gemm/k3_ctm_gemv_long`, and, where the site lists it, up to `WIDE_ROWS` rows on + `gemm/k3_ctm_gemv_wide`. The sites are the decode kernels' fused projections (MLA's [W_a; W_g] with the gate rows + through a sigmoid, KDA's [q | k | v | g | f_a | b]), the attention output projection, and the built-in MLA path's + q_a / kv_a, q_b and gate projections. +* **LM head** (`K3LogitsProcessor`): at most `MAX_ROWS` rows of this rank's vocabulary shard on + `gemm/k3_head_gemv` over the target's `K3HeadGemvWorkspace`, then the shards gathered (`comm/allgather`) as the + stock head gathers them. It is the shell's logits processor, so the speculative worker's target logits and the + drafter's logits on the same head take it too. +* **Embedding** (`K3DecodeGemvs.embed_norm`): a decode step's embedding rows written into the attention-residual + bank's slot 0 (layer 0's first snapshot) and layer 0's input RMSNorm applied, in one `norm/k3_embed_norm` launch. +* **Dense MLP** (`K3DecodeGemvs.dense_mlp`): layer 0's MLP at most `MAX_ROWS` rows, split over the whole TP group: + gate_up on `gemm/k3_ctm_gemv_long`, `activation/k3_situ_mul`, down on `gemm/k3_ctm_gemv_long`; the caller then + runs the down projection's all-reduce. + +Each returns None where its kernel does not take the call, and the caller then runs the generic path's module. + +A kernel compiles on its first call for a shape, which must not happen under CUDA-graph capture. +`K3DecodeGemvs.create`, run once the weights are final, runs every site's kernel and the head once, eagerly. The +embedding kernel compiles per token count, on the eager warm-up step before each capture. Under capture, a call +whose kernel has not run eagerly is refused. + +The head's workspace serves every `k3_head_gemv` call of its weight shape, so those calls must be ordered on one +stream: the logits are computed on the model's stream (see the `gemm/k3_head_gemv` contract's State section). +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Dict, Iterable, Optional, Set + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.k3_situ_mul import k3_situ_mul +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.allgather import allgather +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_long import ( + k3_ctm_gemv_long, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_wide import ( + k3_ctm_gemv_wide, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_decode_gemv import k3_decode_gemv +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_head_gemv import ( + K3HeadGemvWorkspace, + k3_head_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.k3_embed_norm import k3_embed_norm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.concat import concat +from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.split import split + +# The kernels' support predicates: metadata reads only. +from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import op as _ctm_op +from tensorrt_llm._torch.cute_dsl_kernels.k3_decode_gemv import op as _decode_op +from tensorrt_llm._torch.cute_dsl_kernels.k3_embed import op as _embed_op +from tensorrt_llm._torch.cute_dsl_kernels.k3_head_gemv import op as _head_op +from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE +from tensorrt_llm._utils import mpi_disabled + +# The row limit of the decode GEMV and head kernels (one token tile), and of k3_ctm_gemv_wide (a decode step of 8 +# requests of 8 tokens). +MAX_ROWS = 8 +WIDE_ROWS = 64 + + +@dataclass(frozen=True) +class Site: + """A call site's weight shape (this target's per-rank shapes) and its kernels: ``small`` at 1..MAX_ROWS rows + ("decode", "wide" or "long"), and k3_ctm_gemv_wide at MAX_ROWS+1..WIDE_ROWS rows where ``wide``. Output columns + from ``sig_col0`` on are stored through a sigmoid. ``split`` / ``ring`` / ``push``: k3_ctm_gemv_long's CTAs per + 128-row weight tile, weight-ring stages, and whether the partial sums are pushed to each row's owner.""" + + n: int + k: int + small: str + wide: bool = False + sig_col0: int = -1 + split: int = 0 + ring: int = 0 + push: bool = False + + +SITES: Dict[str, Site] = { + # MLA's [W_a; W_g] on a decode step: [q_a 1536 | kv_a 512 | k_pe 64] then the output gate (6 heads x 128), + # the gate rows through a sigmoid. + "mla_ag": Site(2880, 7168, "long", wide=True, sig_col0=2112, split=6, ring=6, push=True), + # KDA's [q | k | v | g | f_a | b] on a decode step (6 heads x 128 each, then 128 and 6, padded to 3208 rows). + "kda_proj": Site(3208, 7168, "long", wide=True, split=5, ring=6, push=True), + # The attention output projection (row parallel; 6 heads x 128 in). + "o_proj": Site(7168, 768, "decode", wide=True), + # The built-in MLA path's projections: kv_a_proj_with_mqa, q_b_proj and the output gate. + "kv_a": Site(2112, 7168, "decode"), + "q_b": Site(1152, 1536, "wide"), + "g_proj": Site(768, 7168, "wide"), + # Layer 0's dense MLP split over the 16-way TP group: gate_up [gate 2112 | up 2112] and down. + "dense_gate_up": Site(4224, 7168, "long", split=4, ring=5), + "dense_down": Site(7168, 2112, "long", split=2, ring=6), +} + + +def _wide_tile(rows: int) -> int: + from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import k3_ctm_gemv_kernel + + return k3_ctm_gemv_kernel.wide_tile(rows) + + +def _run( + spec: Site, kernel: str, x2d: torch.Tensor, weight: torch.Tensor +) -> Optional[torch.Tensor]: + """``spec``'s ``kernel`` on dense rows ``x2d``, or None where it does not take them.""" + if kernel == "decode": + if spec.sig_col0 >= 0 or not _decode_op.supports(x2d, weight): + return None + return k3_decode_gemv(x2d, weight) + if kernel == "wide": + if not _ctm_op.supports_wide(x2d, weight, spec.sig_col0, False): + return None + return k3_ctm_gemv_wide(x2d, weight, sig_col0=spec.sig_col0) + # One wave of the GPU's SMs: beyond it the long GEMV loses to the others. + sms = torch.cuda.get_device_properties(x2d.device).multi_processor_count + if math.ceil(spec.n / 128) * spec.split > sms or not _ctm_op.supports_long( + x2d, weight, spec.split, spec.ring + ): + return None + return k3_ctm_gemv_long( + x2d, + weight, + sig_col0=spec.sig_col0, + split=spec.split, + ring=spec.ring, + trigger_early=True, + push=spec.push, + ) + + +def _capturing() -> bool: + return not torch.compiler.is_compiling() and torch.cuda.is_current_stream_capturing() + + +def _dense_rows(x2d: torch.Tensor) -> torch.Tensor: + """``x2d`` itself, or a dense copy: the kernels' TMA descriptors need dense, 16-byte-aligned rows.""" + if x2d.stride() != (x2d.shape[1], 1) or x2d.data_ptr() % 16: + return x2d.clone(memory_format=torch.contiguous_format) + return x2d + + +def _head_takes_module(lm_head: nn.Module) -> bool: + """Whether ``lm_head(rows)`` is a plain vocabulary-parallel GEMM of a bf16 weight whose shards it gathers along + the vocabulary in rank order, with nothing else applied: the stock head this path reproduces.""" + weight = getattr(lm_head, "weight", None) + mapping = getattr(lm_head, "mapping", None) + return ( + isinstance(weight, torch.Tensor) + and weight.dim() == 2 + and weight.dtype == torch.bfloat16 + and weight.is_contiguous() + and getattr(getattr(lm_head, "tp_mode", None), "name", None) == "COLUMN" + and getattr(lm_head, "gather_output", False) + and getattr(lm_head, "gather_output_sizes", None) is None + and getattr(lm_head, "padding_size", None) == 0 + and getattr(lm_head, "bias", None) is None + and not getattr(lm_head, "has_any_quant", True) + and mapping is not None + and not mapping.enable_attention_dp + ) + + +def _plain_rmsnorm(norm: nn.Module) -> bool: + """Whether ``norm`` is the stock RMSNorm that runs flashinfer's kernel, which ``k3_embed_norm`` reproduces bit for + bit. An unknown module fails closed.""" + weight = getattr(norm, "weight", None) + return ( + IS_FLASHINFER_AVAILABLE + and isinstance(weight, torch.Tensor) + and weight.dtype == torch.bfloat16 + and not getattr(norm, "use_gemma", True) + and not getattr(norm, "is_nvfp4", True) + and not getattr(norm, "use_cuda_tile", True) + and not getattr(norm, "return_hp_output", True) + and getattr(norm, "nvfp4_scale", None) is None + and hasattr(norm, "variance_epsilon") + ) + + +class K3DecodeGemvs: + """The decode GEMVs' state for one target: the LM head's `K3HeadGemvWorkspace`, and the calls whose kernels ran + eagerly (compiled), which are the only ones a CUDA-graph capture may take. Built by `create` once the weights are + final; owned by the target.""" + + def __init__(self, head_workspace: Optional[K3HeadGemvWorkspace] = None) -> None: + self.head_workspace = head_workspace + self._ran: Set[tuple] = set() + + @classmethod + def create( + cls, + lm_head: Optional[nn.Module] = None, + sites: Iterable[str] = tuple(SITES), + device: Optional[torch.device] = None, + ) -> "K3DecodeGemvs": + """The state for ``lm_head`` and ``sites`` on ``device`` (default: the head's, else the current one). Eager: + it allocates the head's workspace and runs every kernel of every site once (each wide row class once) on + zero rows of a zero weight of the site's shape, so they compile here and not under a capture. A site or + head whose kernel does not take its shape keeps the generic path.""" + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "K3DecodeGemvs.create allocates and compiles: run it before CUDA-graph capture" + ) + head_weight = getattr(lm_head, "weight", None) if lm_head is not None else None + if device is None: + device = ( + head_weight.device + if isinstance(head_weight, torch.Tensor) + else torch.device("cuda", torch.cuda.current_device()) + ) + state = cls() + for site in sites: + spec = SITES[site] + weight = torch.zeros(spec.n, spec.k, dtype=torch.bfloat16, device=device) + for rows in (1, 16, 32, 64) if spec.wide else (1,): + state._project(site, weight.new_zeros(rows, spec.k), weight, warm=True) + del weight + if "dense_gate_up" in sites: + gu = torch.zeros(1, SITES["dense_gate_up"].n, dtype=torch.bfloat16, device=device) + if _ctm_op.supports_situ_mul(gu): + for linear_beta in (None, 1.0): + k3_situ_mul(gu, 1.0, linear_beta) + state._ran.add(("situ_mul", linear_beta is not None)) + if lm_head is not None and _head_takes_module(lm_head): + x = head_weight.new_zeros(1, head_weight.shape[1]) + if _head_op.supports(x, head_weight): + workspace = K3HeadGemvWorkspace.create( + head_weight.shape[0], head_weight.shape[1], head_weight.device + ) + k3_head_gemv(x, head_weight, workspace) + state.head_workspace = workspace + state._ran.add(("lm_head",)) + torch.cuda.synchronize(device) + return state + + def project(self, site: str, x: torch.Tensor, weight: torch.Tensor) -> Optional[torch.Tensor]: + """``x @ weight.T`` (bf16 ``[..., N]``, the site's sigmoid columns through the sigmoid) for ``site``'s weight + on its decode kernel, or None where none takes the call: more rows than the site's kernels take, another + shape or dtype, or, under capture, a kernel that has not run eagerly. The caller then runs its GEMM.""" + return self._project(site, x, weight, warm=False) + + def _project( + self, site: str, x: torch.Tensor, weight: torch.Tensor, warm: bool + ) -> Optional[torch.Tensor]: + spec = SITES[site] + if ( + weight.dtype != torch.bfloat16 + or tuple(weight.shape) != (spec.n, spec.k) + or not weight.is_contiguous() + or x.dtype != torch.bfloat16 + or x.dim() < 1 + or x.shape[-1] != spec.k + ): + return None + rows = x.numel() // spec.k + if 0 < rows <= MAX_ROWS: + kernel = spec.small + elif spec.wide and MAX_ROWS < rows <= WIDE_ROWS: + kernel = "wide" + else: + return None + key = (site, kernel, _wide_tile(rows) if kernel == "wide" else 0) + capturing = _capturing() + if capturing and not warm and key not in self._ran: + return None + y = _run(spec, kernel, _dense_rows(x.reshape(rows, spec.k)), weight) + if y is None: + return None + if not capturing: + self._ran.add(key) + return y.view(*x.shape[:-1], spec.n) + + def lm_head_logits(self, rows: torch.Tensor, lm_head: nn.Module) -> Optional[torch.Tensor]: + """``lm_head(rows)``, the gathered bf16 logits ``[M, vocab]``, with this rank's shard on + ``gemm/k3_head_gemv``; None where it does not take the call (more than `MAX_ROWS` rows, another head, ...).""" + workspace = self.head_workspace + if workspace is None or not _head_takes_module(lm_head): + return None + weight = lm_head.weight + if ( + tuple(weight.shape) != (workspace.n_out, workspace.k_in) + or weight.device != workspace.partials.device + or rows.dim() != 2 + or rows.dtype != torch.bfloat16 + or not 0 < rows.shape[0] <= MAX_ROWS + or rows.shape[1] != workspace.k_in + ): + return None + if _capturing() and ("lm_head",) not in self._ran: + return None + group = lm_head.mapping.tp_group + if len(group) > 1 and mpi_disabled(): + return None + x = _dense_rows(rows) + if not _head_op.supports(x, weight): + return None + local = k3_head_gemv(x, weight, workspace) + if len(group) == 1: + return local + gathered = allgather(local, None, group) + return concat(list(split(gathered, rows.shape[0], dim=0)), dim=-1) + + def dense_mlp( + self, + x: torch.Tensor, + gate_up_weight: torch.Tensor, + down_weight: torch.Tensor, + beta: float, + linear_beta: Optional[float], + ) -> Optional[torch.Tensor]: + """Layer 0's dense MLP of at most `MAX_ROWS` rows, before its down projection's all-reduce: gate_up on + ``gemm/k3_ctm_gemv_long``, ``SituAndMul(beta, linear_beta)`` on ``activation/k3_situ_mul``, down on + ``gemm/k3_ctm_gemv_long``. None, with nothing launched, where a kernel does not take the call; the caller then + runs the module.""" + gate_up, down = SITES["dense_gate_up"], SITES["dense_down"] + rows = x.shape[0] if x.dim() == 2 else 0 + if ( + not 0 < rows <= MAX_ROWS + or linear_beta == 0.0 + or tuple(gate_up_weight.shape) != (gate_up.n, gate_up.k) + or tuple(down_weight.shape) != (down.n, down.k) + or gate_up.n != 2 * down.k + ): + return None + if ( + _capturing() + and not { + ("dense_gate_up", "long", 0), + ("dense_down", "long", 0), + ("situ_mul", linear_beta is not None), + } + <= self._ran + ): + return None + gu = self.project("dense_gate_up", x, gate_up_weight) + if gu is None or not _ctm_op.supports_situ_mul(gu): + return None + return self.project("dense_down", k3_situ_mul(gu, beta, linear_beta), down_weight) + + def embed_norm( + self, + input_ids: torch.Tensor, + table: torch.Tensor, + norm: nn.Module, + bank: torch.Tensor, + ) -> Optional[torch.Tensor]: + """Layer 0's normed input for the step's tokens, with their embedding rows written into ``bank[0]``, in one + ``norm/k3_embed_norm`` launch: bit-identical to the embedding followed by ``norm``. None where it does not + apply (another norm, a token count or table the kernel does not take, or, under capture, a token count whose + kernel has not run eagerly).""" + if not _plain_rmsnorm(norm): + return None + ids = input_ids.reshape(-1) + if not _embed_op.supports_norm(ids, table, norm.weight, bank[0]): + return None + key = ("embed_norm", ids.numel(), ids.dtype) + capturing = _capturing() + if capturing and key not in self._ran: + return None + normed = k3_embed_norm(ids, table, norm.weight, norm.variance_epsilon, bank[0]) + if not capturing: + self._ran.add(key) + return normed + + +class _K3Head: + """``lm_head(rows)``, with the rows on ``k3_head_gemv`` where it takes them.""" + + __slots__ = ("_gemvs", "_lm_head") + + def __init__(self, gemvs: K3DecodeGemvs, lm_head: nn.Module) -> None: + self._gemvs = gemvs + self._lm_head = lm_head + + def __call__(self, rows: torch.Tensor) -> torch.Tensor: + logits = self._gemvs.lm_head_logits(rows, self._lm_head) + return self._lm_head(rows) if logits is None else logits + + +class K3LogitsProcessor(nn.Module): + """The shell's logits processor, with its LM head call on ``gemm/k3_head_gemv`` where ``gemvs`` takes the rows. + + It wraps the stock processor instead of repeating it: the row selection and the fp32 conversion stay the stock + processor's. ``gemvs`` is set once the weights are final; until then every call is the stock one. + """ + + def __init__(self, stock: nn.Module) -> None: + super().__init__() + self.stock = stock + self.gemvs: Optional[K3DecodeGemvs] = None + + def forward( + self, + hidden_states: torch.Tensor, + lm_head: nn.Module, + attn_metadata, + return_context_logits: bool = False, + ) -> torch.Tensor: + head = lm_head if self.gemvs is None else _K3Head(self.gemvs, lm_head) + return self.stock.forward(hidden_states, head, attn_metadata, return_context_logits) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py new file mode 100644 index 000000000000..7f50539b2ea1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py @@ -0,0 +1,2332 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ModelingV2 target: Kimi K3 (MXFP4) / sm_100 / tp16 attention, routed experts moe_tp 4 x moe_ep 4. + +Kimi K3's language model: 93 layers, hidden 7168. Every fourth layer (3, 7, ..., 91) is MLA attention with 96 +query heads; the others are Kimi Delta Attention (KDA), a gated linear-attention recurrence with per-request state. +Layer 0's MLP is dense. Every other layer's MLP is a latent MoE: 896 routed experts, top-16, expert width 3584. +Attention residuals mix each sublayer's input from a bank of earlier outputs. The routed experts are MXFP4 (the +checkpoint's compressed-tensors default, read as W4A8_MXFP4_MXFP8); everything else is bf16, and so is the KV pool. + +`tp16_moetp4ep4` is `tensor_parallel_size: 16` with `moe_tensor_parallel_size: 4` and +`moe_expert_parallel_size: 4`, no attention data parallelism, all 16 GPUs in one NVLink domain. Attention is +head-split (6 MLA query heads per rank), and every rank holds a quarter of the width of a quarter of the experts. + +**Each step is classified once, on the host, from its shape** (`decode_step`, a `DecodeStep` or None), and the +classification decides which kernels each module runs: + +* **small**: at most 8 tokens (`DECODE_MAX_TOKENS`, one token tile), context requests included. The token-count + kernels take it: decode GEMVs, the MoE front and routed experts, the sandwiches, the embedding and residual + epilogues. +* **decode**: a pure decode step of R <= 8 generation requests with the same T <= 8 tokens each (one token without + speculation, 1 + 7 drafts with DSpark) and no context request, so at most 64 tokens. The request-aware kernels + take it: MLA attention and its KV store, the KDA verify, the drafter's attention. +* **wide**: a decode step of more than one token tile (DSpark verify of several requests). It keeps the decode + layout's MoE head and tail, on M-general ops. + +Every other step (prefill, mixed steps, decode steps above those bounds) runs the **generic path**: this target's +text model (`KimiLinearModel` below: decoder layers, attention residuals, the MLA / KDA / MoE runtimes), computed +exactly as the built-in Kimi K3 text model computes it, on stock modules and ops that have no catalog entries yet. +`UNCERTIFIED_GENERIC_CALLS` names them. + +The text model hands each step's classification to its attention modules, which run a **decode step** on the K3 +decode kernels' catalog entries: + +* KDA (`K3DecodeKDA`): one token per request, the fused input projection and the plain decode in one + `ssm/k3_kda_decode_attn` launch. Verify tokens (DFlash / DSpark at an even verify width up to 8): the cache manager + then keeps the KDA state after every verify token, and every verify of the layer, on any step, runs the kernels + that keep it: `ssm/k3_kda_attn` for one request of 8 tokens (the projection fused in), else `ssm/k3_kda_verify`. +* MLA (`K3DecodeMLA`): `attention/k3_mla_qkv` (the query path and the step's latent KV rows into the paged cache), + then `attention/k3_mla_attn_vb_out` (the attention, v_b and the output gate in one launch). + +The projections around them run on the decode GEMV sites of `decode_gemv.py` (the [W_a; W_g] and KDA verify-row +projections, `o_proj` on every classified step), as do the LM head, the embedding and layer 0's dense MLP. The state +those kernels share (the KDA projection's Lamport buffers, the MLA attention workspace, the decode GEMVs' state) +lives in typed objects this target creates in `post_load_weights`, before any graph capture. The MoE front and routed +experts, the sandwiches and the residual epilogues come with their own entries; until then they run the generic path +on every step. + +**What this target asserts rather than adapts**: SM 10.0; the topology above; the MXFP4 checkpoint's quantization +(W4A16_MXFP4 with no per-layer declarations, so the routed experts run the W4A8_MXFP4_MXFP8 default and the excluded +modules stay bf16); bf16 weights and a bf16 KV pool; +tokens_per_block 64 (the MLA generation kernels K3's 96 heads reach exist only at 64); the V2 hybrid KV / state +manager, which holds the KDA states, with block reuse off and fp32 recurrent states; an all-reduce strategy of AUTO or +MNNVL. The construction-time ones fail in `__init__`, the per-engine ones on the first forward, each naming the +setting. A layer the decode kernels do not take fails the weight load. + +**Text only.** The checkpoint is the vision-language wrapper. This target builds and loads no vision tower (its +weights are a predicted non-load, `weights.py`), and a step carrying multimodal input raises. + +**Speculative decoding** goes through the stock one-engine shell: DSpark or DFlash with an external drafter +checkpoint, and SA. The worker and its kernels stay upstream code; this target does not own a worker. +""" + +from __future__ import annotations + +import copy +import math +import os +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, NamedTuple, Optional, Tuple + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_attn_vb_out import ( + k3_mla_attn_vb_out, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_attn_workspace import ( + K3MlaAttnWorkspace, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_qkv import k3_mla_qkv +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_attn import k3_kda_attn +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_buffers import K3KdaBuffers +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_decode_attn import ( + k3_kda_decode_attn, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_verify import k3_kda_verify +from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.attention.backends.fmha.cute_dsl_mla import k3_mla_decode_view + +# Registers trtllm::kda_mtp_decode (REQUIRED_TRTLLM_OPS), which the built-in KDA module loads on its first verify. +from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_kda_mtp_ops # noqa: F401 +from tensorrt_llm._torch.distributed import AllReduce +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM +from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM +from tensorrt_llm._torch.models.modeling_utils import DecoderModel, register_auto_model +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.kimi_k3_mla import KimiK3MLAAttention +from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.modules.situ import SituAndMul +from tensorrt_llm._torch.moe.fused_moe import ( + ConfigurableMoE, + SiTuActivation, + TRTLLMGenFusedMoE, + create_moe, +) +from tensorrt_llm._torch.moe.fused_moe.interface import MoESchedulerKind +from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 +from tensorrt_llm._torch.utils import AuxStreamType +from tensorrt_llm.functional import AllReduceStrategy +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + +from . import decode_gemv as _decode_gemv +from . import weights as _weights + +if TYPE_CHECKING: + from transformers import PretrainedConfig + + +# The GPU architecture this target IS. Routing will not send another one here, but a direct instantiation could, +# and the certification is per arch. +_SM = (10, 0) + +#: Every trtllm op this target reaches for, in its forward and in the weight load. Declared here, asserted in +#: tests/unittest/_torch/modeling_v2: the K3-specific ops of the generic path (attention residuals, KDA, the router +#: and fused-A GEMMs), then the decode kernels'. +REQUIRED_TRTLLM_OPS = ( + "attn_res_fwd", + "attn_res_rmsnorm_fwd", + "attn_res_add_rmsnorm_fwd", + "attn_res_add_rmsnorm_persistent_fwd", + "kda_prefill", + "kda_decode", + "kda_mtp_decode", + "dsv3_router_gemm_op", + "dsv3_fused_a_gemm_op", + "k3_kda_decode_attn", + "k3_kda_attn", + "k3_kda_verify", + "k3_mla_qkv", + "k3_mla_attn_vb_out", + # The decode path's GEMVs, the dense MLP's activation, the LM head and the embedding (decode_gemv.py). + "k3_decode_gemv", + "k3_ctm_gemv_wide", + "k3_ctm_gemv_long", + "k3_situ_mul", + "k3_head_gemv", + "k3_embed_norm", + "allgather", +) + +#: The engine surface the first forward checks before this target relies on it: per object, the attributes read. +#: A renamed field upstream fails here, loudly, instead of reading as a default. +REQUIRED_ENGINE_FIELDS = { + "attn_metadata": ( + "num_contexts", + "num_generations", + "num_seqs", + "num_tokens", + "seq_lens", + "tokens_per_block", + "kv_cache_manager", + "mamba_metadata", + ), + "kv_cache_manager": ("enable_block_reuse", "mamba_layer_cache"), +} + +#: Stock code the generic path runs outside the catalog, declared so it is not consumed silently: every +#: tensorrt_llm import of this module that computes (test_modeling_v2_claims.py checks the list both ways). An entry +#: leaves the list when a catalog entry replaces it. +UNCERTIFIED_GENERIC_CALLS = ( + # The checkpoint load and the engine hooks this target inherits, and the causal LM around the text model. + "tensorrt_llm._torch.models.modeling_kimi_linear.KimiLinearForCausalLM", + "tensorrt_llm._torch.models.modeling_speculative.SpecDecOneEngineForCausalLM", + "tensorrt_llm._torch.models.modeling_utils.DecoderModel", + # The text model's stock modules. + "tensorrt_llm._torch.modules.kimi_kda.KimiKDALinearAttention", + "tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops", + "tensorrt_llm._torch.modules.kimi_k3_mla.KimiK3MLAAttention", + "tensorrt_llm._torch.moe.fused_moe.create_moe", + "tensorrt_llm._torch.moe.fused_moe.ConfigurableMoE", + "tensorrt_llm._torch.moe.fused_moe.TRTLLMGenFusedMoE", + "tensorrt_llm._torch.moe.fused_moe.routing.DeepSeekV3MoeRoutingMethod", + "tensorrt_llm._torch.modules.gated_mlp.GatedMLP", + "tensorrt_llm._torch.modules.situ.SituAndMul", + "tensorrt_llm._torch.modules.rms_norm.RMSNorm", + "tensorrt_llm._torch.distributed.AllReduce", + "tensorrt_llm._torch.modules.multi_stream_utils.maybe_execute_in_parallel", +) + +# The K3 decode kernels' bounds: the token-count kernels take one tile of DECODE_MAX_TOKENS rows; the request-aware +# kernels take MAX_REQUESTS generation requests of at most MAX_TOKENS_PER_REQUEST tokens (1 + 7 drafts with DSpark). +DECODE_MAX_TOKENS = 8 +MAX_REQUESTS = 8 +MAX_TOKENS_PER_REQUEST = 8 + +# The MLA generation kernels for K3's 96 query heads exist only at a 64-token page (the built-in model's own +# get_model_defaults sets it for the same reason). +_TOKENS_PER_BLOCK = 64 + +_LANG_PREFIX = "language_model." + + +# ---------------------------------------------------------------------------------------------------------------------- +# The text model: Kimi K3's decoder (93 layers: KDA / MLA attention, attention residuals, the dense layer-0 MLP and +# the latent MoE), its generic path. Each step's classification goes to the attention modules (`K3DecodeKDA`, +# `K3DecodeMLA` below). +# ---------------------------------------------------------------------------------------------------------------------- + +# 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" +) + + +# 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" + + +class KimiK3MoEGate(nn.Module): + """Kimi K3 gate weights and routing method for ``ConfigurableMoE``.""" + + 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 + + assert self.moe_router_activation_func in ("sigmoid", "softmax"), ( + "K3 MoE gate supports sigmoid or softmax scoring only" + ) + + # 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, + ) + + @property + def routing_method(self) -> DeepSeekV3MoeRoutingMethod: + """Return the shared DeepSeek-V3 router used by ``ConfigurableMoE``.""" + if self.moe_router_activation_func != "sigmoid": + raise ValueError("Kimi K3 ConfigurableMoE routing requires sigmoid scores.") + if not self.moe_renormalize: + raise ValueError( + "Kimi K3 ConfigurableMoE routing requires top-k weight renormalization." + ) + return DeepSeekV3MoeRoutingMethod( + top_k=self.top_k, + n_group=self.num_expert_group, + topk_group=self.topk_group, + routed_scaling_factor=self.routed_scaling_factor, + callable_e_score_correction_bias=lambda: self.e_score_correction_bias, + is_fused=True, + ) + + +class KimiK3RMSNorm(nn.Module): + """RMSNorm matching the Kimi checkpoint implementation's rounding.""" + + def __init__( + self, + hidden_size: int, + eps: float = 1e-6, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, + ) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype, device=device)) + self.eps = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states_float = hidden_states.to(torch.float32) + variance = hidden_states_float.pow(2).mean(-1, keepdim=True) + hidden_states_float = hidden_states_float * torch.rsqrt(variance + self.eps) + return self.weight * hidden_states_float.to(input_dtype) + + +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) + + +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"] + + +KIMI_K3_AUX_ATTN_RES_STREAM_ENV = "KIMI_K3_AUX_ATTN_RES_STREAM" + + +_AUX_ATTN_RES_STREAM_ENABLED = os.environ.get(KIMI_K3_AUX_ATTN_RES_STREAM_ENV, "1") == "1" + + +KIMI_K3_FUSED_ATTN_RES_ENV = "KIMI_K3_FUSED_ATTN_RES" + + +_FUSED_ATTN_RES_ENABLED = os.environ.get(KIMI_K3_FUSED_ATTN_RES_ENV, "1") == "1" + + +KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV = "KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS" + + +KIMI_K3_ATTN_RES_TOPOLOGY_ENV = "KIMI_K3_ATTN_RES_TOPOLOGY" + + +_ATTN_RES_TOPOLOGIES = ("per_token", "persistent", "split") + + +# Turning the port on should not require knowing which of the three topologies +# is the right one: "1" selects the measured policy (``split``), so enabling the +# feature and choosing the policy are one action through one variable. +_ATTN_RES_TOPOLOGY_ON = "1" + + +def _read_attn_res_topology() -> str: + """Default ``per_token``: the persistent kernel is opt-in. + + ``1`` is the only accepted on-value and resolves to ``split``. The named + topologies stay available for measurement: ``persistent`` uses the + persistent kernel at every shape it implements, ``per_token`` at none. + """ + raw = os.environ.get(KIMI_K3_ATTN_RES_TOPOLOGY_ENV, "per_token") + if raw == _ATTN_RES_TOPOLOGY_ON: + return "split" + if raw not in _ATTN_RES_TOPOLOGIES: + # Loudly, for the same reason as the token ceiling below: a mistyped A/B + # arm that silently fell back to the default would measure one side + # twice and report no difference. + raise ValueError( + f"{KIMI_K3_ATTN_RES_TOPOLOGY_ENV} must be one of " + f"{_ATTN_RES_TOPOLOGIES} or {_ATTN_RES_TOPOLOGY_ON!r} " + f"(which means 'split'), got {raw!r}" + ) + return raw + + +def _read_fused_attn_res_max_tokens() -> int: + """Resolved after the topology, because its default follows it. + + With the persistent port off -- the default -- the ceiling is 1, the + pre-existing gate: the fused epilogue is taken at the single-token decode + shape and nowhere else. Enabling the port raises it to 32, the top of the + measured range, so that "off" keeps meaning "unchanged". + """ + default = "1" if _ATTN_RES_TOPOLOGY == "per_token" else "32" + raw = os.environ.get(KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV, default) + try: + value = int(raw) + except ValueError: + # Failing loudly matters more than usual here: a mistyped A/B arm that + # silently fell back to the default would measure the candidate twice + # and report no difference. + raise ValueError( + f"{KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV} must be a positive integer, got {raw!r}" + ) from None + if value < 1: + raise ValueError(f"{KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV} must be >= 1, got {value}") + return value + + +_ATTN_RES_TOPOLOGY = _read_attn_res_topology() + + +_FUSED_ATTN_RES_MAX_TOKENS = _read_fused_attn_res_max_tokens() + + +def _persistent_attn_res_applicable(M: int, H: int, N: int) -> bool: + """Shape gate for the persistent kernel: H == 7168 and 2 <= N <= 9. + + No token ceiling: the persistent grid is sized by the SM count, not by the + token count, so prefill is the case it exists for. + """ + del M # deliberately unused; see above + return H == 7168 and 2 <= N <= 9 + + +def _use_persistent_attn_res(M: int, H: int, N: int) -> bool: + """Pick between the two fused kernels for this call site. + + ``persistent`` takes the persistent kernel at every shape it implements; + ``split`` takes it only above ``KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS`` tokens, + which stands in for the prefill/decode boundary. Shapes the persistent + kernel does not implement fall through to the caller's existing gate and + land on the unfused path. + """ + if not _persistent_attn_res_applicable(M, H, N): + return False + if _ATTN_RES_TOPOLOGY == "persistent": + return True + return _ATTN_RES_TOPOLOGY == "split" and M > _FUSED_ATTN_RES_MAX_TOKENS + + +def _apply_attn_res_fused( + prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: KimiK3RMSNorm +) -> 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 M > 16384 or not (4096 <= H <= 8192 and H % 1024 == 0): + return None + try: + attn_res_op = torch.ops.trtllm.attn_res_fwd + except (AttributeError, RuntimeError): + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + output, _rsigma, _probs, _logits = attn_res_op( + layer_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + ) + return output.reshape(M, H) + + +def _rms_norm_eps(norm: nn.Module) -> float: + if hasattr(norm, "eps"): + return float(norm.eps) + return float(norm.variance_epsilon) + + +def _note_attn_res_fusion(site: str, fused: bool, M: int, H: int, N: int) -> None: + """Report whether the fused path was actually reached, once per shape. + + ``_FUSED_ATTN_RES_ENABLED`` only says the feature is switched on, not that + the shape gate let the call through, and a rejected call looks exactly like + a disabled one in the logs. 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}", + ) + + +def _apply_attn_res_rmsnorm_fused( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, +) -> 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. + if _use_persistent_attn_res(M, H, N): + try: + persistent_op = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd + except (AttributeError, RuntimeError): + return None + _, output = persistent_op( + 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), + _rms_norm_eps(output_norm), + ) + _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 + try: + attn_res_rmsnorm_op = torch.ops.trtllm.attn_res_rmsnorm_fwd + except (AttributeError, RuntimeError): + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + output = attn_res_rmsnorm_op( + 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), + _rms_norm_eps(output_norm), + ) + _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: KimiK3RMSNorm, + output_norm: nn.Module, +) -> 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. + if _use_persistent_attn_res(M, H, N): + try: + persistent_op = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd + except (AttributeError, RuntimeError): + return None + updated_prefix_sum, output = persistent_op( + 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), + _rms_norm_eps(output_norm), + ) + _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 + try: + attn_res_add_rmsnorm_op = torch.ops.trtllm.attn_res_add_rmsnorm_fwd + except (AttributeError, RuntimeError): + 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 = attn_res_add_rmsnorm_op( + 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), + _rms_norm_eps(output_norm), + ) + _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: KimiK3RMSNorm +) -> 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]`` + + Unless ``KIMI_K3_FUSED_ATTN_RES=0``, 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. + """ + if _FUSED_ATTN_RES_ENABLED: + fused = _apply_attn_res_fused(prefix_sum, block_residual, proj, norm) + if fused is not None: + return fused + 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: KimiK3RMSNorm, + output_norm: nn.Module, +) -> torch.Tensor: + """Apply attention-residual selection and the next RMSNorm.""" + if _FUSED_ATTN_RES_ENABLED: + fused = _apply_attn_res_rmsnorm_fused(prefix_sum, block_residual, proj, norm, output_norm) + if fused is not None: + return fused + return output_norm(_apply_attn_res(prefix_sum, block_residual, proj, norm)) + + +def _apply_attn_res_add_and_rmsnorm( + prefix_sum: torch.Tensor, + addend: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Add an attention output to the running residual, then select and norm.""" + if _FUSED_ATTN_RES_ENABLED: + 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 + ) + + +_K3_ROUTED_EXPERT_KEY_SUFFIXES = ("block_sparse_moe.experts", "mlp.experts") + + +# The routed experts' module path prefixes. ``exclude_modules`` matches with +# wildcards and walks ancestor prefixes, so an empty prefix would widen what +# matches instead of just missing. +_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 + + +def _load_packed_mxfp4_expert(backend, base, expert_idx, local_slot_id, get_tensor) -> None: + backend.quant_method.load_packed_mxfp4_expert( + backend, + global_expert_id=expert_idx, + local_slot_id=local_slot_id, + w1_weight=get_tensor(f"{base}.{expert_idx}.w1.weight_packed"), + w1_weight_scale=get_tensor(f"{base}.{expert_idx}.w1.weight_scale"), + w2_weight=get_tensor(f"{base}.{expert_idx}.w2.weight_packed"), + w2_weight_scale=get_tensor(f"{base}.{expert_idx}.w2.weight_scale"), + w3_weight=get_tensor(f"{base}.{expert_idx}.w3.weight_packed"), + w3_weight_scale=get_tensor(f"{base}.{expert_idx}.w3.weight_scale"), + ) + + +class _K3ExpertCkptSpec(NamedTuple): + """How one routed-expert quantization is spelled and loaded.""" + + # Per-``w{1,2,3}`` checkpoint tensor suffixes this layout stores. + kinds: Tuple[str, ...] + loader: Callable[..., None] + # Set of filled slots the loader maintains, checked after the load. + loaded_slots_attr: str + # Whether the layer is finalized after its experts load (the MXFP4 loaders write through). + needs_layer_finalize: bool + + +_K3_EXPERT_CKPT_SPECS = { + QuantAlgo.W4A8_MXFP4_MXFP8: _K3ExpertCkptSpec( + kinds=("weight_packed", "weight_scale"), + loader=_load_packed_mxfp4_expert, + loaded_slots_attr="_packed_mxfp4_loaded_slots", + needs_layer_finalize=False, + ), +} + + +def _k3_expert_ckpt_spec(quant_algo: Optional[QuantAlgo]) -> _K3ExpertCkptSpec: + spec = _K3_EXPERT_CKPT_SPECS.get(quant_algo) + if spec is None: + raise NotImplementedError( + f"Kimi K3 routed experts are quantized as {quant_algo}, for which " + "no per-expert checkpoint layout is known. Supported: " + f"{sorted(a.name for a in _K3_EXPERT_CKPT_SPECS)}." + ) + return spec + + +class KimiK3MoERuntime(nn.Module): + """Kimi K3 latent MoE block backed by ConfigurableMoE.""" + + 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.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; 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). + # KIMI_K3_ROUTER_BF16=0 forces the upcast-to-fp32 GEMM. + _router_bf16_env = os.environ.get("KIMI_K3_ROUTER_BF16") + _router_bf16 = _router_bf16_env == "1" if _router_bf16_env is not None else True + self.gate = KimiK3MoEGate(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.quant_algo) + routed_moe_kwargs = dict( + routing_method=self.gate.routing_method, + num_experts=self.num_experts, + hidden_size=self.moe_hidden_size, + intermediate_size=cfg.moe_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, + model_config=routed_moe_model_config, + override_quant_config=routed_quant_config, + layer_idx=layer_idx, + aux_stream_dict=aux_stream_dict, + # Let CommunicationFactory select the best available strategy. + communication_method=None, + activation=SiTuActivation( + gate_softcap=situ_beta, + linear_softcap=situ_linear_beta, + ), + # A request that silently degraded to CUTLASS would be benchmarked + # as if it were the backend that was asked for, and the decline is + # easy to trigger: MegaMoE has its own token / top-k limits and is + # EP-only, and CuteDSL declines on activation shape, SM version and + # the CuTe DSL dependency. As measured once: a CUTEDSL request + # was turned down on every one of the 92 MoE layers, on all 16 + # ranks, and still produced correct text and a zero exit -- the + # only trace was a warning line per layer. Fail in the resolver + # instead, which reports the rejection trail. + # + # CUTLASS is absent on purpose: it is the fallback target, so + # "degraded to CUTLASS" is not a thing that can happen to it. + allow_backend_degradation=routed_moe_model_config.moe_backend + not in ("MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL", "CUTEDSL"), + ) + self._check_trtllm_situ_quant( + routed_moe_model_config.moe_backend, routed_quant_config.quant_algo + ) + + self.routed_experts = create_moe(**routed_moe_kwargs) + if not isinstance(self.routed_experts, ConfigurableMoE): + raise RuntimeError( + "Kimi K3 requires ConfigurableMoE; ENABLE_CONFIGURABLE_MOE must not be disabled." + ) + if self.routed_experts.layer_load_balancer is not None: + raise NotImplementedError( + "Kimi K3 packed-checkpoint streaming does not yet support " + "dynamic EPLB or replicated expert slots." + ) + 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 + shared_model_config = copy.copy(model_config) + shared_model_config.quant_config = QuantConfig() + # Direct MoE-TP leaves both branches as partials for one concatenated + # all-reduce. + use_shared_tp = model_config.mapping.tp_size > 1 + self._reduce_routed_output = ( + use_shared_tp + and self.routed_experts.backend.scheduler_kind != MoESchedulerKind.FUSED_COMM + ) + if self._reduce_routed_output and self.routed_experts.all_reduce is None: + raise RuntimeError( + "Kimi K3 direct MoE tensor parallelism requires the " + "ConfigurableMoE all-reduce even when reduce_results=False." + ) + self.shared_experts = GatedMLP( + hidden_size=cfg.hidden_size, + intermediate_size=shared_intermediate, + bias=False, + activation=SituAndMul( + beta=situ_beta, + linear_beta=situ_linear_beta, + use_fused_activation=True, + ), + dtype=dtype, + config=shared_model_config, + reduce_output=use_shared_tp, + layer_idx=layer_idx, + is_shared_expert=True, + ) + # 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.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 + ) + # Stock fused RMSNorm (flashinfer kernel; the no-flashinfer + # fallback is the same fp32-variance eager math as KimiK3RMSNorm). + self.routed_expert_norm = RMSNorm( + hidden_size=self.moe_hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + + @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 + ) + + @staticmethod + def _select_moe_tp_ep(mapping: Mapping) -> Tuple[int, int]: + """The routed-expert ``(moe_tp, moe_ep)`` split: the user config's explicit + ``moe_tensor_parallel_size`` / ``moe_expert_parallel_size`` (4 x 4, asserted at + construction).""" + return mapping.moe_tp_size, mapping.moe_ep_size + + @staticmethod + def _resolve_routed_quant_config(model_config: ModelConfig, layer_idx: int) -> QuantConfig: + """Routed-expert quantization for ``layer_idx``: the MXFP4 checkpoint declares nothing per layer (asserted at + construction), so the experts keep the ``W4A8_MXFP4_MXFP8`` default. + + An exclusion outranks the default: ``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) + + 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 _check_trtllm_situ_quant(moe_backend: str, quant_algo: Optional[QuantAlgo]) -> None: + """Reject a routed-expert format trtllm-gen has no fused SiTu cubin for. + + trtllm-gen has fused SiTu FC1 cubins for two input formats and no + standalone SiTu activation kernel, so anything else has to die here + rather than in a cubin lookup deep inside the runner. Checked against + the resolved backend, not the K3 architecture branch, because the + generic FP8_BLOCK_SCALES fallback in ``resolve_moe_backend`` can also + land on TRTLLM. + + The admitted set is read off the backend rather than restated here, + because restating it is what broke. This guard was written in #17865 + when MXFP4 was the only fused SiTu drop; #17940 then added the NVFP4 + (group-16 ``Bmm_E2m1_E2m1E2m1_..._siTuGlu_*``) cubins and updated + ``TRTLLMGenFusedMoE``'s set without touching this copy. For the week + in between, an NVFP4 K3 checkpoint could not start at all -- and not + only when TRTLLM was asked for by name, because + ``ModelConfig.resolve_moe_backend`` sends every K3 architecture to + TRTLLM, so the default AUTO configuration hit this raise too. The unit + tests did not catch it: they call ``create_moe`` directly and never + reach this guard, so the kernel path stayed green while the model path + was closed. + + A staticmethod, not an inline block, so that the invariant is + reachable from a test without constructing the whole runtime. + """ + situ_supported = TRTLLMGenFusedMoE.situ_supported_quant_algos() + if moe_backend != "TRTLLM" or quant_algo in situ_supported: + return + supported = ", ".join(sorted(algo.name for algo in situ_supported)) + raise ValueError( + f"Kimi K3 routed experts are quantized as {quant_algo}, which the " + "TRTLLM (trtllm-gen) MoE backend cannot serve: fused SiTu cubins " + f"exist only for {supported}. Set moe_config.backend to CUTLASS " + "or MEGAMOE_CUTEDSL." + ) + + @staticmethod + def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: + """Build a private routed-expert mapping without mutating the shared + config, with the split of ``_select_moe_tp_ep``.""" + # Every backend here declares ``ActivationType.SiTu`` in its + # ``activation_support``; the list is not a preference order. CUTEDSL + # joined once its act-fusion kernel grew the SiTU epilogue. + supported_backends = { + "CUTLASS", + "TRTLLM", + "CUTEDSL", + "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL", + } + if model_config.moe_backend not in supported_backends: + raise ValueError( + "Kimi K3 SiTU routed experts only support the CUTLASS, TRTLLM, " + "CUTEDSL, MEGAMOE_DEEPGEMM, and MEGAMOE_CUTEDSL 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 + + moe_tp, moe_ep = KimiK3MoERuntime._select_moe_tp_ep(mapping) + 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 in { + "MEGAMOE_DEEPGEMM", + "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) -> torch.Tensor: + """``hidden_states``: ``[num_tokens, hidden_size]`` bf16.""" + identity = hidden_states + router_logits = self.gate.compute_logits(hidden_states) + moe_all_reduce = self.routed_experts.all_reduce if self._reduce_routed_output else None + + def _routed_output(): + # 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). + routed_in = self._routed_projection(hidden_states, self.routed_expert_down_proj) + y = self.routed_experts( + routed_in, + router_logits, + all_rank_num_tokens=all_rank_num_tokens, + ) + 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) + + # 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 GatedMLP includes its output all-reduce on + # the auxiliary stream. The join below must precede the routed + # all-reduce: concurrent collectives on different streams can corrupt + # SYMM_MEM all-reduce state. + routed_out, shared_out = maybe_execute_in_parallel( + _routed_output, + lambda: self.shared_experts(identity), + self.moe_main_event, + self.moe_shared_event, + self.shared_expert_stream, + disable_on_compile=True, + ) + if self._reduce_routed_output: + routed_latent = moe_all_reduce(routed_out) + routed_latent = self.routed_expert_norm(routed_latent) + routed_out = self._routed_projection(routed_latent, self.routed_expert_up_proj) + return routed_out + shared_out + + +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 + + +class KimiMLARuntime(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], + ) -> None: + super().__init__() + + max_positions = int( + os.environ.get( + _KIMI_K3_MLA_MAX_POSITIONS_ENV, + cfg.max_position_embeddings, + ) + ) + self.layer_idx = layer_idx + # KimiK3MLAAttention 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. + mapping = model_config.mapping + reduce_output = mapping.tp_size > 1 + self._o_allreduce = ( + AllReduce( + mapping=mapping, + strategy=model_config.allreduce_strategy, + dtype=torch.bfloat16, + ) + if reduce_output + else None + ) + 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", + ) + } + attention_config.quant_config = QuantConfig( + kv_cache_quant_algo=model_config.quant_config.kv_cache_quant_algo + if model_config.quant_config is not None + else None + ) + attention_config._frozen = model_config._frozen + self.mixer = K3DecodeMLA( + 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, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + step: Optional[DecodeStep] = None, + ) -> torch.Tensor: + # MLA.forward takes position_ids first; K3 is NoPE, so pass None. + out = self.mixer(None, hidden_states, attn_metadata, step=step) + if 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 + + +class KimiLinearDecoderLayer(nn.Module): + 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_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 = K3DecodeKDA( + cfg, + layer_idx, + mapping=model_config.mapping, + allreduce_strategy=model_config.allreduce_strategy, + aux_stream=aux_stream_dict[AuxStreamType.Attention], + model_config=attention_config, + ) + else: + self.self_attn = KimiMLARuntime( + cfg, + layer_idx, + model_config=model_config, + aux_stream_dict=aux_stream_dict, + ) + + 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 + ) + if self.is_moe: + self.block_sparse_moe = KimiK3MoERuntime(model_config, cfg, layer_idx, aux_stream_dict) + else: + situ_beta = getattr(cfg, "activation_situ_beta", None) or 1.0 + situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) + self.mlp_tp_size = math.gcd(cfg.intermediate_size, model_config.mapping.tp_size) + # Over MNNVL (one NVLink domain across the nodes, where a cross-node all-reduce costs what a node's + # does) the MLP stays split over the whole TP group, the per-rank shapes its decode GEMVs take + # (decode_gemv.SITES); otherwise it stays within one node. + spans_nodes = self._mnnvl_allreduce() is not None + if self.mlp_tp_size > model_config.mapping.gpus_per_node and not spans_nodes: + self.mlp_tp_size = math.gcd(self.mlp_tp_size, model_config.mapping.gpus_per_node) + mlp_model_config = copy.copy(model_config) + mlp_model_config.quant_config = QuantConfig() + # K3's dense layer is BF16, so a unit block size gives the same + # subgroup selection as DeepSeek-V3. + self.mlp = GatedMLP( + hidden_size=cfg.hidden_size, + intermediate_size=cfg.intermediate_size, + bias=False, + activation=SituAndMul( + beta=situ_beta, + linear_beta=situ_linear_beta, + use_fused_activation=True, + ), + dtype=dtype, + config=mlp_model_config, + overridden_tp_size=self.mlp_tp_size, + reduce_output=self.mlp_tp_size > 1, + layer_idx=layer_idx, + ) + self._situ = (situ_beta, situ_linear_beta) + # The decode GEMVs' state (decode_gemv.py), set by the target's cache_derived_state. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + # Stock fused RMSNorm for the plain (whole-tensor) norms; numerics + # are drop-in for KimiK3RMSNorm (fp32 variance, weight applied + # after downcast, use_gemma=False). + self.input_layernorm = RMSNorm( + hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + self.post_attention_layernorm = RMSNorm( + hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + + # Attention residual scheme (always on for K3). The res norms stay + # KimiK3RMSNorm: they are consumed field-wise (.weight/.eps) by + # _apply_attn_res and the fused attn_res op, never called as + # modules. + 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 = KimiK3RMSNorm( + cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + self.mlp_res_norm = KimiK3RMSNorm(cfg.hidden_size, eps=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 forward( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + num_snapshots: int, + attn_metadata: AttentionMetadata, + capture: Optional[Tuple[Any, int]] = None, + step: Optional[DecodeStep] = None, + prenormed: bool = False, + ) -> 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. + + ``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. + + ``step`` is the step's classification (``decode_step``), handed to the + attention module. + + ``prenormed`` (layer 0 on a decode step): ``hidden_states`` already is + this layer's input norm, and the layer's input, the step's embedding, + already is in ``block_residual[0]`` (``K3DecodeGemvs.embed_norm``). + """ + prefix_sum = hidden_states + valid_block_residual = block_residual[:num_snapshots] + + if prenormed: + assert num_snapshots == 0 and self.layer_idx % self.attn_res_block_size == 0 + elif capture is not None: + # The mixture tap needs the PRE-norm value, which the fused + # attn-res + RMSNorm kernel does not expose. Keep the two steps + # split on captured layers only and fuse everywhere else. + if num_snapshots > 0: + hidden_states = _apply_attn_res( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ) + # A property of the DRAFTER checkpoint, not a knob: a mismatch only lowers + # acceptance, silently. hidden_states is the pre-norm attn_res mixture; + # prefix_only wants the running prefix, already in hand as prefix_sum. + tapped = hidden_states if _AUX_ATTN_RES_STREAM_ENABLED else prefix_sum + capture[0].maybe_capture_hidden_states(capture[1], tapped, None) + hidden_states = self.input_layernorm(hidden_states) + 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 = self.input_layernorm(hidden_states) + + if self.layer_idx % self.attn_res_block_size == 0: + if not prenormed: + block_residual[num_snapshots].copy_(prefix_sum) + num_snapshots += 1 + valid_block_residual = block_residual[:num_snapshots] + prefix_sum = None + if self.is_kda: + hidden_states = self.linear_attn(hidden_states, attn_metadata, step=step) + else: + hidden_states = self.self_attn(hidden_states, attn_metadata, step=step) + + 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: + 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, + ) + if self.is_moe: + hidden_states = self.block_sparse_moe( + hidden_states, getattr(attn_metadata, "all_rank_num_tokens", None) + ) + else: + hidden_states = self._dense_mlp(hidden_states, step) + + prefix_sum = prefix_sum + hidden_states + return prefix_sum, num_snapshots + + def _dense_mlp(self, hidden_states: torch.Tensor, step: Optional[DecodeStep]) -> torch.Tensor: + """The dense MLP: on a step of at most DECODE_MAX_TOKENS tokens, its GEMVs and activation on the decode + kernels (``K3DecodeGemvs.dense_mlp``), then the down projection's all-reduce; else the module.""" + gemvs = self.decode_gemvs + if gemvs is not None and step is not None and step.small: + out = gemvs.dense_mlp( + hidden_states, + self.mlp.gate_up_proj.weight, + self.mlp.down_proj.weight, + *self._situ, + ) + if out is not None: + return self.mlp.down_proj.all_reduce(out) if self.mlp_tp_size > 1 else out + return self.mlp(hidden_states) + + def _mnnvl_allreduce(self): + """The MNNVL all-reduce of this layer's attention output, or None.""" + attention = self.linear_attn if self.is_kda else self.self_attn + return getattr(getattr(attention, "_o_allreduce", None), "mnnvl_allreduce", None) + + 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 + + +class KimiLinearModel(DecoderModel): + 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() for _ in range(4)] + 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], + } + + self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size, dtype=dtype) + self.layers = nn.ModuleList( + [ + KimiLinearDecoderLayer(model_config, cfg, layer_idx, self.aux_stream_dict) + for layer_idx in range(cfg.num_hidden_layers) + ] + ) + self.norm = RMSNorm(hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype) + + # KimiK3RMSNorm (not RMSNorm): consumed field-wise (.weight/.eps) + # by _apply_attn_res and the fused attn_res op. + self.output_attn_res_norm = KimiK3RMSNorm( + cfg.hidden_size, eps=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 + # The decode path's GEMVs and embedding (decode_gemv.py), built by the target's cache_derived_state once the + # weights are final. None: every step embeds on the generic path. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + # 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", + ) + + @property + def kda_token_states(self) -> bool: + """Whether the hybrid cache manager keeps the KDA state after every verify token, the protocol of + ``ssm/k3_kda_verify`` and ``ssm/k3_kda_attn``. The engine reads it once the weights are loaded, to build the + manager: DFlash / DSpark drafts of an even verify width up to 8, every KDA layer taking the K3 kernels. + Otherwise the KDA verify replays the accepted drafts (the built-in verify).""" + spec_config = getattr(self.model_config, "spec_config", None) + return bool( + spec_config is not None + and (spec_config.spec_dec_mode.is_dflash() or spec_config.spec_dec_mode.is_dspark()) + and spec_config.tokens_per_gen_step in (2, 4, 6, 8) + and all(layer.linear_attn.takes_k3_kernels for layer in self.layers if layer.is_kda) + ) + + 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") + + num_tokens = (input_ids if inputs_embeds is None else inputs_embeds).shape[0] + step = decode_step(attn_metadata, num_tokens) + # A decode step embeds and norms for layer 0 in one launch, the embedding written as layer 0's first snapshot. + prenormed = None + if ( + inputs_embeds is None + and step is not None + and self.decode_gemvs is not None + and len(self.layers) > 0 + and self.num_attn_res_snapshots > 0 + ): + table = self.embed_tokens.weight + block_residual = table.new_empty( + self.num_attn_res_snapshots, num_tokens, table.shape[1] + ) + prenormed = self.decode_gemvs.embed_norm( + input_ids, table, self.layers[0].input_layernorm, block_residual + ) + if prenormed is not None: + hidden_states = prenormed + else: + hidden_states = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds + 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, + step=step, + prenormed=i == 0 and prenormed is not None, + ) + + # 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) + + return _apply_attn_res_and_rmsnorm( + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + ) + + +# ---------------------------------------------------------------------------------------------------------------------- +# The attention modules: the built-in KDA and MLA modules, with the steps the decode kernels take on those kernels. +# ---------------------------------------------------------------------------------------------------------------------- + + +class K3DecodeKDA(KimiKDALinearAttention): + """Kimi K3's KDA attention: the built-in module, with the decode kernels on the steps they take. + + * A decode step of one token per request runs the fused input projection and the plain decode in one + ``ssm/k3_kda_decode_attn`` launch. + * With the cache manager's per-token states (``KimiLinearModel.kda_token_states``), every verify of the layer, on + any step, runs the kernels that keep them: ``ssm/k3_kda_attn`` for one request of 8 tokens (the projection fused + in), else ``ssm/k3_kda_verify`` on the projection's rows (the ``kda_proj`` decode GEMV site where it takes + them). The built-in verify replays drafts from caches these kernels do not fill, so the two never run on one + manager. + * On every step ``decode_step`` classifies, ``o_proj`` runs on the ``o_proj`` decode GEMV site where it takes the + rows, then the module's all-reduce. + + Every other step runs the built-in module. The kernels read one ``[q | k | v | g | f_a | b]`` weight, built at the + checkpoint load from the module's own, and the device's ``K3KdaBuffers`` and decode GEMVs' state, which the + target sets in ``post_load_weights``. + """ + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + # The fused [q | k | v | g | f_a | b | pad] projection weight; the six projections' weights and the built-in + # fused [q | k | v | g] and [f_a | b | pad] ones are views of it. + self.k3_proj_weight: Optional[torch.Tensor] = None + # The fused projection's Lamport buffers: one set per device, shared by every KDA layer. + self.k3_buffers: Optional[K3KdaBuffers] = None + # The decode GEMVs' state (decode_gemv.py), shared by the target's layers. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + @property + def takes_k3_kernels(self) -> bool: + """Whether the decode kernels can run this layer: its fused projection weight is built.""" + return self.k3_proj_weight is not None + + def finalize_decode_weights(self) -> None: + """The built-in fused weights, then one ``[q | k | v | g | f_a | b | pad]`` weight of both.""" + super().finalize_decode_weights() + assert ( + self.use_full_rank_gate + and self.gate_lower_bound is not None + and self._qkvg_proj_weight is not None + and self._bfa_proj_weight is not None + and self._qkvg_proj_weight.dtype == self._bfa_proj_weight.dtype == torch.bfloat16 + ), ( + f"Kimi K3 KDA layer {self.layer_idx}: the decode kernels read the bf16 fused projections the built-in " + "module builds on CUDA at head dim 128, with a full-rank output gate and a gate lower bound" + ) + rows = self._qkvg_proj_weight.shape[0] + with torch.no_grad(): + fused = self._merge_projection_weights( + (self.q_proj, self.k_proj, self.v_proj, self.g_proj, self.f_a_proj, self.b_proj), + pad_rows_to=8, + ) + self.k3_proj_weight = fused + self._qkvg_proj_weight, self._bfa_proj_weight = fused[:rows], fused[rows:] + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + step: Optional[DecodeStep] = None, + ) -> torch.Tensor: + """The built-in forward on a step ``decode_step`` does not classify, and under a breakable CUDA graph. On the + others: the plain decode on ``ssm/k3_kda_decode_attn`` on a decode step of one token per request, else the + built-in dispatch; then ``o_proj`` on its decode GEMV site.""" + if step is None or is_in_breakable_cuda_graph(): + return super().forward(hidden_states, attn_metadata) + if ( + step.decode + and step.tokens_per_request == 1 + and self.takes_k3_kernels + and self.k3_buffers is not None + ): + core = self._k3_decode(hidden_states[: step.num_tokens], attn_metadata) + else: + core = self._forward_impl(hidden_states, attn_metadata) + return self._k3_project_output(core) + + def _k3_project_output(self, core: torch.Tensor) -> torch.Tensor: + """``o_proj`` on the ``o_proj`` decode GEMV site where it takes the rows (else the module), then the TP + all-reduce.""" + out = None + if self.decode_gemvs is not None: + out = self.decode_gemvs.project( + "o_proj", core.reshape(-1, self.proj_size), self.o_proj.weight + ) + if out is None: + return self._project_output(core) + return out if self._o_allreduce is None else self._o_allreduce(out) + + def _k3_decode(self, x: torch.Tensor, attn_metadata: AttentionMetadata) -> torch.Tensor: + """``ssm/k3_kda_decode_attn``: the core output ``[R, H, 128]`` of one token of each of the step's R requests; + each slot's conv window and state advance in place.""" + mamba_metadata = attn_metadata.mamba_metadata + slots = getattr(mamba_metadata, "generation_state_indices", None) + if slots is None: + slots = mamba_metadata.state_indices[: x.shape[0]] + layer_cache = attn_metadata.kv_cache_manager.mamba_layer_cache(self.layer_idx) + w_q, w_k, w_v = self._get_mtp_conv_weights() + core = k3_kda_decode_attn( + x.contiguous(), + self.k3_proj_weight, + self.f_b_proj.weight, + w_q, + w_k, + w_v, + self._A_log_f32, + self._dt_bias_f32, + self._onorm_w_f32, + layer_cache.conv, + layer_cache.temporal, + slots, + self.k3_buffers, + float(self.gate_lower_bound), + self.head_k_dim**-0.5, + float(self.o_norm.eps), + ) + # Speculative decoding's replay caches keep their committed conv window in step with the pool's. + self._sync_kda_replay_conv_window(layer_cache, slots, layer_cache.conv) + return core + + def forward_verify( + self, + x2d, + num_steps, + layer_cache, + conv_pool, + ssm_pool, + slot_indices, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """The built-in verify, unless the cache manager keeps the per-token states; then the decode kernels'.""" + if not (layer_cache.has_kda_replay_caches and layer_cache.kda_state_tok is not None): + return super().forward_verify( + x2d, num_steps, layer_cache, conv_pool, ssm_pool, slot_indices, output=output + ) + assert self.takes_k3_kernels and self.k3_buffers is not None, ( + f"Kimi K3 KDA layer {self.layer_idx}: the cache manager keeps the per-token verify states, which only " + "the decode kernels write, and this layer has no fused projection weight or buffers" + ) + core = self._k3_verify(x2d, num_steps, layer_cache, ssm_pool, slot_indices) + return self._store_core(core, output) + + def _k3_verify(self, x, num_steps, layer_cache, ssm_pool, slot_indices) -> torch.Tensor: + """The core output ``[N T, H, 128]`` of N requests of T verify tokens: ``ssm/k3_kda_attn`` for one request of + 8 tokens, else ``ssm/k3_kda_verify`` on the projection's rows (the ``kda_proj`` decode GEMV site where it + takes them). Each slot's state after the golden token, its drafts' states and its conv window are written in + place.""" + num_requests = x.shape[0] // num_steps + slots = slot_indices[:num_requests] + w_q, w_k, w_v = self._get_mtp_conv_weights() + constants = (float(self.gate_lower_bound), self.head_k_dim**-0.5, float(self.o_norm.eps)) + if num_requests == 1 and num_steps == 8: + out = k3_kda_attn( + x.contiguous(), + self.k3_proj_weight, + self.f_b_proj.weight, + w_q, + w_k, + w_v, + self._A_log_f32, + self._dt_bias_f32, + self._onorm_w_f32, + layer_cache.kda_conv_q, + layer_cache.kda_conv_k, + layer_cache.kda_conv_v, + ssm_pool, + layer_cache.kda_state_tok, + slots, + layer_cache.prev_num_accepted_tokens, + self.k3_buffers, + num_steps - 1, + *constants, + ) + else: + rows = None + if self.decode_gemvs is not None: + rows = self.decode_gemvs.project("kda_proj", x, self.k3_proj_weight) + if rows is None: + rows = torch.nn.functional.linear(x, self.k3_proj_weight) + out = k3_kda_verify( + rows, + self.f_b_proj.weight, + w_q, + w_k, + w_v, + self._A_log_f32, + self._dt_bias_f32, + self._onorm_w_f32, + layer_cache.kda_conv_q, + layer_cache.kda_conv_k, + layer_cache.kda_conv_v, + ssm_pool, + layer_cache.kda_state_tok, + slots, + layer_cache.prev_num_accepted_tokens, + num_steps - 1, + *constants, + ) + return out.view(-1, self.num_heads, self.head_dim) + + +class K3DecodeMLA(KimiK3MLAAttention): + """Kimi K3's MLA attention: the built-in module, with a decode step's attention on the decode kernels. + + A decode step runs ``x [W_a; W_g]^T`` with the gate columns through a sigmoid (the ``mla_ag`` decode GEMV site + where it takes the rows, else one GEMM), then ``attention/k3_mla_qkv`` (the q_a / kv_a RMSNorms, q_b and the k_b + absorption into the fused query, the step's latent rows stored into the paged cache) and + ``attention/k3_mla_attn_vb_out`` (the attention over the paged cache, v_b and the output gate in one launch), then + ``o_proj`` (the ``o_proj`` site where it takes the rows, else the module). Every other step, and a decode step whose + cache the kernels do not read (``k3_mla_decode_view`` says why), runs the built-in module. + + ``[W_a; W_g]`` is built at load from the module's weights, which become views of it. The attention workspace is + the device's ``K3MlaAttnWorkspace``; it and the decode GEMVs' state are set by the target in + ``post_load_weights``. + """ + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + # [W_a; W_g]: the fused q_a / kv_a projection's rows, then the output gate's. + self.k3_ag_weight: Optional[torch.Tensor] = None + # The decode attention's workspace: one per device, shared by every MLA layer. + self.k3_workspace: Optional[K3MlaAttnWorkspace] = None + # The decode GEMVs' state (decode_gemv.py), shared by the target's layers. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + def post_load_weights(self) -> None: + """The built-in post-load, then ``[W_a; W_g]`` (once: CUDA graphs captured since read it).""" + super().post_load_weights() + if self.k3_ag_weight is not None: + return + gaps = self._k3_layout_gaps() + assert not gaps, ( + f"Kimi K3 MLA layer {self.layer_idx}: the decode kernels do not take {'; '.join(gaps)}" + ) + qkv_a, gate = self.kv_a_proj_with_mqa, self.g_proj + rows = qkv_a.weight.shape[0] + with torch.no_grad(): + fused = torch.cat([qkv_a.weight, gate.weight]) + qkv_a.weight = nn.Parameter(fused[:rows], requires_grad=False) + gate.weight = nn.Parameter(fused[rows:], requires_grad=False) + self.k3_ag_weight = fused + + def _k3_layout_gaps(self) -> list: + """What of this layer the decode kernels do not take (empty when they take all of it).""" + if not (self.use_output_gate and self.fuse_qkv_a_proj and not self.is_lite): + return ["a layer without the output gate or the fused q_a / kv_a projection"] + linears = (self.kv_a_proj_with_mqa, self.g_proj, self.q_b_proj, self.o_proj) + checks = ( + (not self.mapping.has_cp_helix(), "helix context parallelism"), + (not self.apply_rotary_emb and not self.llama_4_scaling, "RoPE or llama-4 scaling"), + (self.sparse_attn_hooks is None, "sparse attention"), + ( + self.kv_cache_dtype != "fp8_ds_mla" + and not getattr(self.mqa, "has_fp8_kv_cache", False) + and not getattr(self.mqa, "has_fp4_kv_cache", False), + "a quantized KV cache", + ), + ( + all(m.weight.dtype == torch.bfloat16 and m.bias is None for m in linears) + and self.k_b_proj_trans.dtype == self.v_b_proj.dtype == torch.bfloat16, + "projections other than bf16 and unbiased", + ), + ( + not getattr(self.q_a_layernorm, "is_nvfp4", False) + and not getattr(self.kv_a_layernorm, "use_gemma", False), + "an NVFP4 q_a norm or a Gemma kv_a norm", + ), + ( + self.num_heads_tp % 6 == 0 + and self.kv_lora_rank == 512 + and self.qk_rope_head_dim == 64 + and self.q_lora_rank == 1536, + f"{self.num_heads_tp} heads, latent {self.kv_lora_rank}, rope {self.qk_rope_head_dim}, " + f"q_lora {self.q_lora_rank}", + ), + ) + return [why for ok, why in checks if not ok] + + def forward( + self, + position_ids: Optional[torch.Tensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + all_reduce_params=None, + latent_cache_gen: Optional[torch.Tensor] = None, + step: Optional[DecodeStep] = None, + ) -> torch.Tensor: + """The built-in forward, except on a decode step whose cache the decode kernels read.""" + view = None + if ( + step is not None + and step.decode + and latent_cache_gen is None + and self.k3_ag_weight is not None + and self.k3_workspace is not None + and not is_in_breakable_cuda_graph() + ): + view = self._k3_decode_view(attn_metadata, step.num_tokens) + if view is None: + return super().forward( + position_ids, hidden_states, attn_metadata, all_reduce_params, latent_cache_gen + ) + x = hidden_states[: step.num_tokens].contiguous() + rows = self.kv_a_proj_with_mqa.weight.shape[0] + gemvs = self.decode_gemvs + ag = None if gemvs is None else gemvs.project("mla_ag", x, self.k3_ag_weight) + if ag is None: + ag = torch.nn.functional.linear(x, self.k3_ag_weight) + ag[:, rows:].sigmoid_() + fused_q = k3_mla_qkv( + ag, + self.q_a_layernorm.weight, + float(self.q_a_layernorm.variance_epsilon), + self.q_b_proj.weight, + self.k_b_proj_trans, + self.kv_a_layernorm.weight, + float(self.kv_a_layernorm.variance_epsilon), + view["pool"], + view["row_stride"], + view["page_table"], + view["page_offset"], + view["seq_len"], + ) + attn_output = self.create_output(x, 0) + k3_mla_attn_vb_out( + fused_q, + view["pool"], + view["row_stride"], + view["page_table"], + view["page_offset"], + view["seq_len"], + view["softmax_scale"], + self.v_b_proj, + attn_output, + self.k3_workspace, + gate=ag, + gate_col0=rows, + ) + out = None if gemvs is None else gemvs.project("o_proj", attn_output, self.o_proj.weight) + if out is None: + out = self._project_output( + [attn_output], position_ids, attn_metadata, all_reduce_params + ) + return out + + def _k3_decode_view(self, attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[dict]: + """The paged latent cache as the decode kernels read it this step, or None when they do not read it (the + reason is logged once).""" + view = k3_mla_decode_view(self.mqa, attn_metadata, num_tokens) + if isinstance(view, str): + logger.info_once( + f"Kimi K3 MLA: the built-in path for a decode step the decode kernels do not read ({view})", + key=f"k3_mla_decode_view_{view}", + ) + return None + return view + + +# ---------------------------------------------------------------------------------------------------------------------- +# The target: step classification, the construction checks and the registration shell. +# ---------------------------------------------------------------------------------------------------------------------- + + +def _text_model_config(model_config: ModelConfig) -> ModelConfig: + """The language model's ModelConfig: the checkpoint's text_config, with quant exclusions renamed to match. + + The checkpoint names its language-model modules `language_model.`; the text model's are ``, with + `layers.*` under `model.`. + """ + config = model_config.pretrained_config + text = copy.copy(model_config) + text._frozen = False + text.pretrained_config = config.text_config + excluded = text.quant_config.exclude_modules + if excluded: + text.quant_config = copy.copy(text.quant_config) + renamed = [] + for name in excluded: + if name.startswith(_LANG_PREFIX): + name = name[len(_LANG_PREFIX) :] + if name.startswith("layers."): + name = "model." + name + renamed.append(name) + text.quant_config.exclude_modules = renamed + text.skip_create_weights_in_init = True + text._frozen = True + return text + + +@dataclass(frozen=True) +class DecodeStep: + """A step the Kimi K3 decode kernels take: ``num_tokens`` rows and, on a pure decode step, ``num_requests`` + generation requests of ``tokens_per_request`` tokens each (None on a step with context requests).""" + + num_tokens: int + num_requests: Optional[int] = None + tokens_per_request: Optional[int] = None + + @property + def small(self) -> bool: + """Whether the step fits one token tile of the token-count kernels.""" + return self.num_tokens <= DECODE_MAX_TOKENS + + @property + def decode(self) -> bool: + """Whether the step is a pure decode step the request-aware kernels take.""" + return self.num_requests is not None + + @property + def wide(self) -> bool: + """Whether the step is a pure decode step of more than one token tile: its token-count work keeps the decode + layout's MoE head and tail, on M-general ops.""" + return self.decode and not self.small + + +def decode_step(attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[DecodeStep]: + """The step's shape if any Kimi K3 decode kernel takes it, else None (the generic path runs). + + ``num_tokens`` is the step's token count (the rows of the model input). Read on the host from per-step integers + and the host copy of the sequence lengths only. A CUDA graph is captured per decode batch shape, and every input + here is fixed by that shape, so a captured step and its replays are classified alike. + """ + if num_tokens <= 0: + return None + requests = _decode_requests(attn_metadata, num_tokens) + if requests is not None: + return DecodeStep(num_tokens, requests, num_tokens // requests) + if num_tokens <= DECODE_MAX_TOKENS: + return DecodeStep(num_tokens) + return None + + +def _decode_requests(attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[int]: + """R when the step is R <= 8 generation requests of the same T <= 8 tokens and no context request, else None.""" + if attn_metadata.num_contexts != 0: + return None + requests = attn_metadata.num_generations + if not 0 < requests <= MAX_REQUESTS or num_tokens % requests != 0: + return None + tokens = num_tokens // requests + if not 0 < tokens <= MAX_TOKENS_PER_REQUEST: + return None + seq_lens = getattr(attn_metadata, "seq_lens", None) + if seq_lens is not None and seq_lens.device.type == "cpu": + lens = seq_lens[:requests] + if lens.numel() != requests or bool((lens != tokens).any()): + return None + return requests + + +def _check_construction(model_config: ModelConfig) -> None: + """The settings this target is built for that are fixed before the first step.""" + capability = torch.cuda.get_device_capability() + assert capability == _SM, ( + f"this target is certified on sm_{_SM[0]}{_SM[1]}, running on sm_{capability[0]}{capability[1]}" + ) + mapping = model_config.mapping + topology = ( + mapping.world_size, + mapping.tp_size, + mapping.pp_size, + mapping.moe_tp_size, + mapping.moe_ep_size, + mapping.enable_attention_dp, + ) + assert topology == (16, 16, 1, 4, 4, False), ( + "the tp16_moetp4ep4 target needs world_size 16, tensor_parallel_size 16, pipeline_parallel_size 1, " + "moe_tensor_parallel_size 4, moe_expert_parallel_size 4 and enable_attention_dp false; the engine built " + f"(world, tp, pp, moe_tp, moe_ep, attention_dp) = {topology}" + ) + assert model_config.torch_dtype == torch.bfloat16, ( + f"this target computes in bf16; the engine resolved dtype {model_config.torch_dtype}" + ) + kv_algo = model_config.quant_config.kv_cache_quant_algo + assert kv_algo is None, ( + f"this target's MLA kernels read a bf16 KV pool; kv_cache_config.dtype resolved to {kv_algo}" + ) + quant_algo = model_config.quant_config.quant_algo + assert quant_algo == QuantAlgo.W4A16_MXFP4 and not model_config.quant_config_dict, ( + "this target loads the MXFP4 checkpoint, whose compressed-tensors config the model config reads as " + "W4A16_MXFP4 with no per-layer declarations (its routed experts run the W4A8_MXFP4_MXFP8 default); the " + f"engine read {quant_algo} with {len(model_config.quant_config_dict or {})} per-layer declarations" + ) + strategy = model_config.allreduce_strategy + assert strategy in (AllReduceStrategy.AUTO, AllReduceStrategy.MNNVL), ( + f"this target runs its all-reduces over MNNVL; allreduce_strategy is {strategy.name}" + ) + + +@register_auto_model("ModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4") +class ModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4(KimiLinearForCausalLM): + """The registration shell: this target's text model (`KimiLinearModel` above) behind its checks. + + It inherits the built-in Kimi K3 causal LM for the checkpoint load and the engine hooks (`load_weights` through + weights.py, the KDA metadata class, the model defaults), which walk the model by its module names; the text model + keeps the built-in one's. + """ + + @classmethod + def get_preferred_kv_cache_manager_version(cls, pretrained_config: Any = None) -> Literal["V2"]: + """The V2 hybrid manager holds the KDA states; the step contract requires it.""" + return "V2" + + def __init__(self, model_config: ModelConfig): + config = model_config.pretrained_config + assert getattr(config, "text_config", None) is not None, ( + "this target loads the KimiK3ForConditionalGeneration checkpoint, whose language model is its " + "text_config" + ) + _check_construction(model_config) + text = _text_model_config(model_config) + spec_config = getattr(text, "spec_config", None) + 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" + # The inherited loader reads these: this target has neither helix context parallelism nor the fp8 + # weight-read conversion of the shared / latent MLPs. + self._fp8_weight_read_moe_mlp = False + self.mapping_with_cp = None + self._repurposed_tp_mapping = None + cfg = text.pretrained_config + SpecDecOneEngineForCausalLM.__init__( + self, + KimiLinearModel(text), + text, + hidden_size=cfg.hidden_size, + vocab_size=cfg.vocab_size, + ) + self._step_checked = False + # The LM head on gemm/k3_head_gemv at decode size, once cache_derived_state has built its state. It is the + # processor the speculative worker's target logits and a parallel drafter's logits call too. + stock_logits_processor = self.logits_processor + self.logits_processor = _decode_gemv.K3LogitsProcessor(stock_logits_processor) + draft_model = getattr(self, "draft_model", None) + if getattr(draft_model, "logits_processor", None) is stock_logits_processor: + draft_model.logits_processor = self.logits_processor + # The executor reads generation settings (eos_token_id, ...) off the model config the engine holds, which + # must therefore be the text config, as the built-in wrapper leaves it. + model_config._frozen = False + model_config.pretrained_config = self.config + model_config._frozen = True + + def load_weights(self, weights, *args, **kwargs): + _weights.load(self, weights) + + def cache_derived_state(self) -> None: + """Build the decode GEMVs' state once the weights are final: the LM head's workspace, and one eager call of + every decode GEMV kernel at its site's shape (decode_gemv.SITES), so none compiles under capture. Built once: + a later call keeps it, since CUDA graphs captured in between hold its workspace.""" + super().cache_derived_state() + if self.model.decode_gemvs is not None: + return + gemvs = _decode_gemv.K3DecodeGemvs.create(self.lm_head) + self.model.decode_gemvs = gemvs + self.logits_processor.gemvs = gemvs + for layer in self.model.layers: + if not layer.is_moe: + layer.decode_gemvs = gemvs + + def post_load_weights(self) -> None: + """The state the decode kernels share, built once per device before any CUDA-graph capture and handed to + every layer of its kind: the KDA projection's Lamport buffers and the MLA decode attention's workspace; and + the decode GEMVs' state (built by ``cache_derived_state``) handed to every attention module.""" + super().post_load_weights() + kda = [layer.linear_attn for layer in self.model.layers if layer.is_kda] + mla = [layer.self_attn.mixer for layer in self.model.layers if not layer.is_kda] + if kda[0].k3_buffers is not None: + return # built by an earlier call; CUDA graphs captured since hold it + device = self.model.embed_tokens.weight.device + buffers = K3KdaBuffers.create(device) + workspace = K3MlaAttnWorkspace.create(device, mla[0].num_heads_tp // 6) + for module in kda: + module.k3_buffers = buffers + for module in mla: + module.k3_workspace = workspace + for module in kda + mla: + module.decode_gemvs = self.model.decode_gemvs + logger.info( + "Kimi K3 decode kernels: KDA on k3_kda_decode_attn, k3_kda_attn and k3_kda_verify " + f"({sum(m.takes_k3_kernels for m in kda)} / {len(kda)} layers take them), MLA on k3_mla_qkv and " + f"k3_mla_attn_vb_out ({len(mla)} layers)" + ) + + def _check_step_contract(self, attn_metadata: AttentionMetadata) -> None: + """First-forward checks of the engine surface and the per-engine settings.""" + objects = { + "attn_metadata": attn_metadata, + "kv_cache_manager": attn_metadata.kv_cache_manager, + } + missing = [ + f"{owner}.{name}" + for owner, names in REQUIRED_ENGINE_FIELDS.items() + for name in names + if not hasattr(objects[owner], name) + ] + assert not missing, f"engine fields this target reads are missing: {missing}" + assert attn_metadata.tokens_per_block == _TOKENS_PER_BLOCK, ( + f"this target needs kv_cache_config.tokens_per_block {_TOKENS_PER_BLOCK}; the engine built " + f"{attn_metadata.tokens_per_block}" + ) + manager = attn_metadata.kv_cache_manager + assert isinstance(manager, MambaHybridCacheManagerV2), ( + "this target needs the V2 hybrid KV / state cache manager " + "(kv_cache_config.use_kv_cache_manager_v2); the engine built " + f"{type(manager).__name__}" + ) + assert not manager.enable_block_reuse, ( + "this target runs with kv_cache_config.enable_block_reuse false; the engine enabled it" + ) + kda_layer = next(layer.layer_idx for layer in self.model.layers if layer.is_kda) + state_dtype = manager.mamba_layer_cache(kda_layer).temporal.dtype + assert state_dtype == torch.float32, ( + "this target's KDA kernels keep fp32 recurrent states (kv_cache_config.mamba_ssm_cache_dtype float32 or " + f"auto); the engine built a {state_dtype} state pool" + ) + self._step_checked = True + + 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: bool = False, + spec_metadata=None, + resource_manager=None, + **kwargs, + ) -> torch.Tensor: + if kwargs.pop("multimodal_params", None): + raise ValueError( + "this Kimi K3 target is text only: it loads no vision tower, and a request carried image input" + ) + if not self._step_checked: + self._check_step_contract(attn_metadata) + return super().forward( + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + return_context_logits, + spec_metadata, + resource_manager, + **kwargs, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/weights.py new file mode 100644 index 000000000000..b4177f698d8b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/weights.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weight loading: Kimi K3 (MXFP4) / sm_100 / tp16_moetp4ep4. + +The checkpoint is the vision-language wrapper's. `language_model.*` holds the language model; `vision_tower.*` and +`mm_projector.*` hold the vision tower and its projector. This target is text only, so the keys split three ways: + +* `language_model.*` goes to the built-in text model's loader with the prefix stripped. That loader streams the + routed experts one at a time, keeps this rank's slice (a quarter of the width of a quarter of the experts under + moe_tp 4 x moe_ep 4), and checks its own key coverage. +* The vision tower and the projector are a predicted non-load: listed here, never read. +* Any other key fails the load, naming it, rather than being dropped. +""" + +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM +from tensorrt_llm._torch.models.modeling_utils import filter_weights + +_LANG_PREFIX = "language_model." + +# The vision tower's and the projector's key families: in the checkpoint, not loaded by this text-only target. +PREDICTED_NON_LOAD = ("vision_tower.", "mm_projector.") + + +def load(model, weights) -> None: + """Load the language model's weights into `model`, the target shell, and check every other key is predicted.""" + unknown = sorted( + k + for k in weights.keys() + if not k.startswith(_LANG_PREFIX) and not k.startswith(PREDICTED_NON_LOAD) + ) + assert not unknown, ( + f"{len(unknown)} checkpoint key(s) are neither language-model weights nor a predicted non-load, " + f"first {unknown[:5]}" + ) + lm_weights = ConsumableWeightsDict(filter_weights(_LANG_PREFIX[:-1], weights)) + assert len(lm_weights), f"the checkpoint has no {_LANG_PREFIX}* keys" + checkpoint_dir = getattr(weights, "checkpoint_dir", None) + if checkpoint_dir is not None: + lm_weights.checkpoint_dir = checkpoint_dir + lm_weights.checkpoint_prefix = _LANG_PREFIX + KimiLinearForCausalLM.load_weights(model, 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 new file mode 100644 index 000000000000..c406423e441d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/routing.py @@ -0,0 +1,127 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Where a KimiK3ForConditionalGeneration config lands. Read this file and you know. + +One forward-reading decision tree per architecture family: the criteria are +evaluated in the order a reader would ask them, and every branch that does not +end in a target returns None (in ``auto`` the engine then uses the built-in +Kimi K3 implementation; in ``require`` it raises, quoting the trace below). + +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. +""" + +from __future__ import annotations + +from typing import Any, Optional + +from tensorrt_llm.quantization.mode import QuantAlgo + +from ..._router_index import NULL_TRACE, ModelingV2Context, Trace + +# The one GPU architecture these targets are written for. sm is part of a +# target's identity, not a knob: a different SM is a different target. The K3 +# decode kernels (the tcgen05 GEMVs, the fused MoE) are built for GB200 only. +_SM = (10, 0) + +# Config-shape fingerprint -> checkpoint identity; see the note in the sibling +# gpt_oss routing module for what shape-sniffing does and does not pin. +# +# (num_hidden_layers, hidden_size, num_experts, routed_expert_hidden_size), all +# read from ``text_config``. The NVFP4 requant of the same checkpoint has this +# shape too; the quantization criterion below is what keeps it out. +_CHECKPOINTS = { + (93, 7168, 896, 3584): "kimi_k3_mxfp4", +} + +_TARGETS = { + ("kimi_k3_mxfp4", "tp16_moetp4ep4"): "ModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4", +} + +# Synthetic architecture name -> the module whose import registers it. +TARGET_MODULES = { + "ModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4": ( + "models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4.modeling" + ), +} + + +def _field(config: Any, name: str) -> Any: + """Read ``name`` off a sub-config that may be an object or a plain dict.""" + if isinstance(config, dict): + return config.get(name) + return getattr(config, name, None) + + +def _mxfp4(quant_config: Any) -> bool: + """Whether the checkpoint is the MXFP4 one these targets load. + + The MXFP4 checkpoint declares its quantization inside + ``text_config.quantization_config`` (compressed-tensors, + ``mxfp4-pack-quantized``); ``KimiK3Config`` surfaces it and the model + config reads it as W4A16_MXFP4. The NVFP4 requant ships + ``hf_quant_config.json`` and reads as MIXED_PRECISION, and a checkpoint + without quantization reads as none; the experts of either would land in a + loader that reads packed MXFP4 tensors. + """ + return quant_config is not None and quant_config.quant_algo == QuantAlgo.W4A16_MXFP4 + + +def _parallel(m) -> Optional[str]: + """Name the parallel topology, or None if no target implements it. + + Attention and the dense layers are sharded 16 ways; the routed experts are + split 4 ways by tensor and 4 ways by expert. The expert split decides which + expert weights each rank loads and into which shapes, so it selects a + target rather than a runtime branch. + """ + if ( + m.world_size == 16 + and m.tp_size == 16 + and m.pp_size == 1 + and m.moe_tp_size == 4 + and m.moe_ep_size == 4 + and not m.enable_attention_dp + ): + return "tp16_moetp4ep4" + return None + + +def route(ctx: ModelingV2Context, trace: Trace = NULL_TRACE) -> Optional[str]: + c, m = ctx.pretrained_config, ctx.mapping + + if not trace.check("sm", ctx.sm, ctx.sm == _SM): + return None + + t = getattr(c, "text_config", None) + if not trace.check("text_config", type(t).__name__, t is not None): + return None + + shape = tuple( + _field(t, name) + for name in ( + "num_hidden_layers", + "hidden_size", + "num_experts", + "routed_expert_hidden_size", + ) + ) + ckpt = trace.resolve("shape", shape, _CHECKPOINTS.get(shape)) + if ckpt is None: + return None + + quant_algo = getattr(ctx.quant_config, "quant_algo", None) + if not trace.check("quant", quant_algo, _mxfp4(ctx.quant_config)): + return None + + parallel = trace.resolve( + "parallel", + f"ws={m.world_size} tp={m.tp_size} pp={m.pp_size} moe_tp={m.moe_tp_size} " + f"moe_ep={m.moe_ep_size} attention_dp={m.enable_attention_dp}", + _parallel(m), + ) + if parallel is None: + return None + + return _TARGETS.get((ckpt, parallel)) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py index ac97de1d550c..4f977580943d 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py @@ -40,6 +40,92 @@ _LOG2_E = math.log2(math.e) +def k3_mla_decode_view(attn, meta, num_tokens: int): + """The per-layer inputs of trtllm::k3_mla_qkv and trtllm::k3_mla_attn(_vb)_out for a bf16 decode step of R + generation requests of T = num_tokens / R tokens each (no context requests, R <= 8, T <= 8), or the reason it + does not apply (a string; never raises). The dict: ``pool`` (flat bf16 pool), ``row_stride``, ``page_table`` + (int32 [R, W], request i's pages in row i: a strided view of ``kv_cache_block_offsets``, row stride + ``page_table.stride(0)``), ``page_offset``, ``seq_len`` (int32 [R], the kv lengths including the step's tokens), + ``softmax_scale``, ``num_requests`` (R) and ``tokens_per_request`` (T). The page table and lengths are the + host-filled metadata buffers (``kv_cache_block_offsets``, ``kv_lens_cuda_runtime``), which the kernels may read + before their grid wait. T is taken as uniform, as the engine pads every generation request to the same number of + tokens.""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_mla.op import MAX_REQUEST_TOKENS, MAX_REQUESTS + + num_requests = meta.num_generations + if ( + meta.num_contexts != 0 + or not 0 < num_requests <= MAX_REQUESTS + or num_tokens % num_requests + or not 0 < num_tokens // num_requests <= MAX_REQUEST_TOKENS + ): + return ( + f"{meta.num_contexts} context / {num_requests} generation requests, {num_tokens} tokens" + ) + if ( + meta.beam_width != 1 + or getattr(meta, "is_spec_dec_tree", False) + or getattr(meta, "is_spec_dec_dynamic_tree", False) + ): + return f"beam width {meta.beam_width} or a tree speculative mask" + if attn.num_heads % 6 or attn.kv_lora_rank != 512 or attn.qk_rope_head_dim != 64: + return f"heads {attn.num_heads}, latent {attn.kv_lora_rank}, rope {attn.qk_rope_head_dim}" + if meta.tokens_per_block != 64 or meta.helix_position_offsets is not None: + return f"page {meta.tokens_per_block}, helix {meta.helix_position_offsets is not None}" + if meta.kv_cache_manager is None or meta.kv_cache_block_offsets is None: + return "no KV cache manager / block offsets" + kv_pool = meta.kv_cache_manager.get_buffers(attn.layer_idx) + if kv_pool.dtype != torch.bfloat16: + return f"kv cache {kv_pool.dtype}" + packed_block = 1 + for size in kv_pool.shape[1:]: + packed_block *= size + block_stride = kv_pool.stride(0) + layers_in_pool = block_stride // packed_block if packed_block else 1 + layer_in_pool = 0 + if layers_in_pool > 1 and block_stride == layers_in_pool * packed_block: + layer_in_pool = kv_pool.storage_offset() // packed_block + kv_pool = kv_pool.as_strided( + (kv_pool.shape[0] * layers_in_pool, *kv_pool.shape[1:]), + (packed_block, *kv_pool.stride()[1:]), + 0, + ) + if ( + kv_pool.dim() != 5 + or kv_pool.shape[1] != 1 + or kv_pool.shape[3] != 1 + or kv_pool.shape[2] != 64 + ): + return f"pool layout {tuple(kv_pool.shape)}" + if not (kv_pool.is_contiguous() and kv_pool.stride(2) >= 576 and kv_pool.stride(2) % 8 == 0): + return f"pool strides {kv_pool.stride()}" + pool_idx = int(meta.host_kv_cache_pool_mapping[attn.get_local_layer_idx(meta), 0]) + gen = slice(meta.num_contexts, meta.num_contexts + num_requests) + page_table = meta.kv_cache_block_offsets[pool_idx, gen, 0, :] + seq_len = meta.kv_lens_cuda_runtime[gen] + if page_table.dtype != torch.int32 or seq_len.dtype != torch.int32: + return f"page table {page_table.dtype} / length {seq_len.dtype} not int32" + if ( + page_table.shape[0] != num_requests + or seq_len.shape[0] != num_requests + or page_table.stride(1) != 1 + ): + return f"page table {tuple(page_table.shape)} strides {page_table.stride()} / lengths {tuple(seq_len.shape)}" + softmax_scale = float( + 1.0 / (math.sqrt(attn.qk_nope_head_dim + attn.qk_rope_head_dim) * attn.q_scaling) + ) + return dict( + pool=kv_pool.view(-1), + row_stride=kv_pool.stride(2), + page_table=page_table, + page_offset=int(layer_in_pool), + seq_len=seq_len, + softmax_scale=softmax_scale, + num_requests=num_requests, + tokens_per_request=num_tokens // num_requests, + ) + + class CuteDslMlaFmha(PhasedFmha): """Blackwell CuTe DSL FMHA library for decode-only MLA.""" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_ctm_gemv/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_ctm_gemv/__init__.py new file mode 100644 index 000000000000..71b0dc23eedd --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_ctm_gemv/__init__.py @@ -0,0 +1,19 @@ +# 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 CTM decode GEMV (``trtllm::k3_ctm_gemv``, ``_tail``, ``_gated``). + +Importing :mod:`.op` registers the torch ops; nothing is imported eagerly here so that +the CuTe DSL dependency stays optional for every other model. +""" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_ctm_gemv/k3_ctm_gemv_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_ctm_gemv/k3_ctm_gemv_kernel.py new file mode 100644 index 000000000000..a405e008f725 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_ctm_gemv/k3_ctm_gemv_kernel.py @@ -0,0 +1,1681 @@ +# 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 decode GEMV -- CTM (prims/cute) kernel, M <= 8 tokens, bf16 in/out +# ============================================================================= +# +# y[t, n] = sum_k B[t, k] * W[n, k] (A @ B^T, fp32 accumulation in TMEM) +# A = W (N, K <= 768) K-major bf16 (the streamed weight, 128 rows per cluster) +# B (M <= 8, K) K-major bf16 (8 token columns; rows past M are zero) +# +# Variants (compile-time): +# PLAIN B = x (KDA / MLA o_proj) +# TAIL y = rsqrt(mean(latent^2) + eps) * (latent_slice @ W_lat^T) + act @ W_act^T +# (two TMA maps, two TMEM accumulators; the row-parallel MoE tail) +# GATED B = bf16(a * bf16(sigmoid(g))) (MLA o_proj with its output gate) +# GATED_S B = bf16(a * s) with s = bf16(sigmoid(g)) precomputed (the long-K GEMV's sigmoid rows) +# SWIGLU B = bf16(g * sigmoid(g) * a) in fp32, g and a the two halves of a gate_up output (silu_and_mul folded +# into the down projection) +# a and g land by TMA in two identically swizzled tiles; the epilogue warps rewrite B in +# place (same offsets in both tiles), fence the generic writes to the async proxy and +# release the MMA through a 128-arrival barrier (the resident-B prolog of +# fuse_rmsnorm_qkv_rope). +# +# Geometry: one 128-row weight tile per cluster of SPLIT CTAs (1, 2 or 4). Rank r takes the interleaved +# k-tiles r, r + SPLIT, ... (when SPLIT does not divide the k-tiles, the first k-tiles % SPLIT ranks take one +# more) and owns output rows [R r, R r + R) of the tile, R = 128 / SPLIT (the TMEM +# lanes of epilogue warps r * 4 / SPLIT ...); every other rank's warps holding those rows send their +# fp32 partials to slot [source rank] of the owner's mailbox; the owner adds the SPLIT partials in rank +# order, then rounds once. Two ways to send (compile-time PUSH, the same sums): DSMEM stores and a release +# arrive at cluster scope on the owner's barrier, acquired by a try_wait at cluster scope; or 16-byte +# st.async stores that complete the owner's barrier by bytes, which the owner spins on with test_wait (acquire at +# cluster scope). +# +# Every k-tile of the CTA has its own shared-memory stage, and the whole weight slice is TMA'd +# (EVICT_FIRST) before griddepcontrol.wait: launched early under PDL the kernel streams its weight +# while the predecessor runs. Only the activation loads follow the wait. +# +# Warps: 0 weight TMA (+ early dependent trigger), 1 activation TMA after the grid dependency, +# 2 TMEM allocation + tcgen05 MMA (M 128, N 8, K 16), 3 idle, 4-7 epilogue (gate prologue, TMEM -> +# registers, split-K reduce, bf16 staging, 16-byte coalesced stores). +# ============================================================================= +"""CTM decode GEMV for Kimi K3 (``y = B @ W^T``, M <= 8): plain, MoE-tail and output-gated variants.""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +CTA_M = 128 # weight rows (output features) per cluster = the MMA's M +MMA_N = 8 # token columns +CTA_K = 128 # one k-tile: two 64-element halves of the 128-byte swizzle +MMA_K = 16 +TMA_K_BOX = 64 +TMA_COPY_ITERS = CTA_K // TMA_K_BOX +K_BLOCKS_PER_HALF = TMA_K_BOX // MMA_K +MAX_K_TILES = 6 # 6 x (32 KB weight + 2 KB B [+ 2 KB gate]) of shared memory per CTA +THREADS = 256 +EPI_THREADS = 128 +TMEM_COLS = 32 +ELEM_BYTES = 2 +VEC = 8 # bf16 per 16-byte vector +EVICT_FIRST = 0x12F0000000000000 # createpolicy.fractional.L2::evict_first, fraction 1.0 (sm_100) + +# Shared-memory descriptor strides for the 128-byte swizzle, in 16-byte units. +LEADING = 16 +STRIDE = 8 * TMA_K_BOX * ELEM_BYTES +A_HALF_ELEMS = CTA_M * TMA_K_BOX +B_HALF_ELEMS = MMA_N * TMA_K_BOX +STEP = (MMA_K * ELEM_BYTES) >> 4 +A_BOX = A_HALF_ELEMS >> 3 +B_BOX = B_HALF_ELEMS >> 3 +STAGE_A = (CTA_M * CTA_K * ELEM_BYTES) >> 4 +STAGE_B = (MMA_N * CTA_K * ELEM_BYTES) >> 4 + +PLAIN = 0 +TAIL = 1 +GATED = 2 # B = bf16(a * bf16(sigmoid(g))), g given +GATED_S = 3 # B = bf16(a * s), s = bf16(sigmoid(g)) given (e.g. by k3_ctm_gemv_long's sigmoid rows) +SWIGLU = 4 # B = bf16(g * sigmoid(g) * a) in fp32: silu_and_mul of a gate_up output (g the first half, a the second) + +io_dtype = cutlass.BFloat16 + + +def num_k_tiles(k_in: int) -> int: + return (k_in + CTA_K - 1) // CTA_K + + +def supports(n_out: int, k_in: int, split: int = 1) -> bool: + """Shapes the kernel runs: whole 128-row tiles, whole k-tiles, at least one and at most MAX_K_TILES k-tiles per + CTA (the first k-tiles % split ranks take one more).""" + k_tiles = num_k_tiles(k_in) + return ( + n_out % CTA_M == 0 + and k_in % CTA_K == 0 + and split in (1, 2, 4) + and split <= k_tiles + and (k_tiles + split - 1) // split <= MAX_K_TILES + ) + + +@dsl_user_op +def _div_rn(a, b, *, loc=None, ip=None): + """IEEE fp32 division (div.rn.f32), as torch's ``one / (one + exp(-x))``.""" + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + "div.rn.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, + ) + ) # fmt: skip + + +@dsl_user_op +def _mapa_u32(smem_ptr, peer, *, loc=None, ip=None): + """The shared::cluster address of this CTA's shared-memory location in cluster CTA ``peer``.""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [smem_ptr.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(peer).ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_async_v4(dst, a, b, c, d, mbar, *, loc=None, ip=None): + """st.async of four fp32 (16 bytes) to a shared::cluster address, completing ``mbar`` (a shared::cluster address) + by 16 bytes.""" + _llvm.inline_asm( + None, + [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Float32(a).ir_value(loc=loc, ip=ip), + cutlass.Float32(b).ir_value(loc=loc, ip=ip), cutlass.Float32(c).ir_value(loc=loc, ip=ip), + cutlass.Float32(d).ir_value(loc=loc, ip=ip), cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.f32 [$0], {$1, $2, $3, $4}, [$5];", "r,f,f,f,f,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _test_wait_cluster(mbar, parity, *, loc=None, ip=None): + """mbarrier.test_wait.parity.acquire.cluster on a barrier of this CTA (``mbar``, its shared-memory pointer): whether + phase ``parity`` has completed, acquiring at cluster scope. The barriers it is used on are completed by other CTAs + (st.async complete_tx, remote arrives), which release at cluster scope.""" + done = cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tmbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" + "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + return done != cutlass.Int32(0) + + +def _bf16_rn(v): + """fp32 -> bf16 precision (round to nearest even) in the integer domain, kept as fp32. + + Bit operations, not a float cast, so no fptrunc/fpext pair exists that the compiler could + fold into the following multiply and skip this rounding. + """ + u = v.bitcast(cutlass.Int32) + u = u + (((u >> 16) & 1) + 0x7FFF) + u = (u >> 16) << 16 + return u.bitcast(cutlass.Float32) + + +def _sigmoid_bf16(g): + """bf16(sigmoid(g)) for a bf16 value g held in fp32, as torch computes it for bf16 tensors.""" + one = cutlass.Float32(1.0) + return _bf16_rn(_div_rn(one, one + cute.math.exp(-g, fastmath=False))) + + +@cute.kernel +def k3_ctm_gemv_kernel( + tma_desc_w: cutlass.GridConstant[cuda.TensorMap], # W [N, K] bf16, 5-D, one call per k-tile + tma_desc_x: cutlass.GridConstant[ + cuda.TensorMap + ], # plain: x; tail: latent; gated: a. Box 64 x 8 + tma_desc_x2: cutlass.GridConstant[ + cuda.TensorMap + ], # tail: act; gated: the tensor holding g. Box 64 x 8 + y: cutlass.Array, # [M * N] bf16, token-major + rms_src: cutlass.Array, # tail: int32 words of the bf16 [M, rms_cols] latent rows + num_tokens: cutlass.Int32, + x_col0: cutlass.Int32, # tail: column of the latent where k-tile 0 starts + x2_col0: cutlass.Int32, # gated: column of g in its tensor + eps: cutlass.Float32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + variant: cutlass.Constexpr[int], + x_tiles: cutlass.Constexpr[ + int + ], # tail: latent k-tiles (the rest accumulate from x2 into accumulator 1) + rms_cols: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + push: cutlass.Constexpr[ + bool + ], # split-K partials by st.async (test_wait owner) instead of stores + arrive +): + """One weight tile of 128 rows per cluster; the CTA's k-tiles all resident; B by TMA (gated: rewritten).""" + k_tiles = num_k_tiles(k_in) + extra = k_tiles % split + my_tiles = (k_tiles + split - 1) // split # stages (the last one unused on ranks >= extra) + rows_owned = CTA_M // split # output rows reduced and stored by this CTA + tx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp_id = cute.arch.warp_idx() + tma_ptr_w = tma_desc_w.get_ptr() + tma_ptr_x = tma_desc_x.get_ptr() + tma_ptr_x2 = tma_desc_x2.get_ptr() + rank = cutlass.Int32(0) + if cutlass.const_expr(split > 1): + rank = cute.arch.block_idx_in_cluster() + m_offset = (bx // cutlass.Int32(split)) * cutlass.Int32(CTA_M) + my_count = cutlass.Int32(my_tiles) + if cutlass.const_expr(extra > 0): + my_count = cutlass.Int32(k_tiles // split) + cutlass.Int32( + cutlass.select_(rank < cutlass.Int32(extra), cutlass.Int32(1), cutlass.Int32(0)) + ) + + # Allocation order is the same in every CTA, so mapa() addresses the peer's mailbox and barrier. + smem_a = cutlass.Array( + io_dtype, my_tiles * CTA_M * CTA_K, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_b = cutlass.Array( + io_dtype, my_tiles * MMA_N * CTA_K, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_g = smem_b + if cutlass.const_expr(variant >= GATED): + smem_g = cutlass.Array( + io_dtype, my_tiles * MMA_N * CTA_K, space=cutlass.AddressSpace.smem, alignment=1024 + ) + stage = cutlass.Array( + io_dtype, MMA_N * rows_owned, space=cutlass.AddressSpace.smem, alignment=16 + ) + weight_full = cutlass.Array( + cutlass.Int64, my_tiles, space=cutlass.AddressSpace.smem, alignment=8 + ) + act_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + b_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + acc_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + mail_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + if cutlass.const_expr(variant == TAIL): + row_scale = cutlass.Array( + cutlass.Float32, MMA_N, space=cutlass.AddressSpace.smem, alignment=16 + ) + if cutlass.const_expr(split > 1): + # [source rank][row within the owned block][token] fp32 partials (the own rank's slot stays unused). + mailbox = cutlass.Array( + cutlass.Float32, + split * rows_owned * MMA_N, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + + if warp_id == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_x) + if cutlass.const_expr(variant != PLAIN): + prims.prefetch_tensormap(tma_ptr_x2) + if prims.elect_sync(): + for i in cutlass.range_constexpr(my_tiles): + prims.mbarrier_init(weight_full.subview(i), 1) + prims.mbarrier_init(act_full, 1) + prims.mbarrier_init(acc_done, 1) + if cutlass.const_expr(variant >= GATED): + prims.mbarrier_init(b_ready, EPI_THREADS) + if cutlass.const_expr(split > 1): + if cutlass.const_expr(push): + # The other ranks' partials of this CTA's rows arrive by st.async (16-byte stores completing + # this barrier's transaction count); expected here, before cluster formation. + prims.mbarrier_init(mail_full, 1) + prims.mbarrier_arrive_expect_tx(mail_full, (split - 1) * rows_owned * MMA_N * 4) + else: + # Every lane of the other ranks' epilogue warps holding this CTA's rows arrives once. + prims.mbarrier_init(mail_full, (split - 1) * rows_owned) + if warp_id == 2: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + if cutlass.const_expr(split > 1): + # Cluster formation: the peer's shared memory and barriers are addressable from here on. + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + tmem_ptr = cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32) + + if warp_id == 0: + # ===================================================================== + # Weight TMA: every k-tile of this CTA, ahead of the grid dependency. + # ===================================================================== + if prims.elect_sync(): + for i in cutlass.range_constexpr(my_tiles): + k = rank + cutlass.Int32(i * split) + if cutlass.Int32(i) < my_count: + prims.mbarrier_arrive_expect_tx( + weight_full.subview(i), CTA_M * CTA_K * ELEM_BYTES + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(i * CTA_M * CTA_K), + tma_ptr_w, + ( + cutlass.Int32(0), + m_offset, + k * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), + cutlass.Int32(0), + ), + weight_full.subview(i), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(trigger_early): + # Dependents may launch now; they wait for this whole grid before reading y. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + elif warp_id == 1: + # ===================================================================== + # Activation TMA: B (and g), written by predecessors -> after the wait. + # ===================================================================== + prims.griddepcontrol(prims.GridDepAction.WAIT) + if prims.elect_sync(): + tile_bytes = MMA_N * CTA_K * ELEM_BYTES + if cutlass.const_expr(variant >= GATED): + tile_bytes = 2 * tile_bytes + prims.mbarrier_arrive_expect_tx(act_full, my_count * cutlass.Int32(tile_bytes)) + for i in cutlass.range_constexpr(my_tiles): + k = rank + cutlass.Int32(i * split) + k_c = cutlass.Int32(cutlass.select_(cutlass.Int32(i) < my_count, k, rank)) + for half in cutlass.range_constexpr(TMA_COPY_ITERS): + b_off = i * MMA_N * CTA_K + half * B_HALF_ELEMS + col = k_c * cutlass.Int32(CTA_K) + cutlass.Int32(half * TMA_K_BOX) + if cutlass.const_expr(variant == TAIL): + if cutlass.const_expr(i < x_tiles): + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(b_off), + tma_ptr_x, + (x_col0 + col, cutlass.Int32(0)), + act_full, + ) + else: + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(b_off), + tma_ptr_x2, + ( + cutlass.Int32((i - x_tiles) * CTA_K + half * TMA_K_BOX), + cutlass.Int32(0), + ), + act_full, + ) + else: + if cutlass.Int32(i) < my_count: + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(b_off), + tma_ptr_x, + (x_col0 + col, cutlass.Int32(0)), + act_full, + ) + if cutlass.const_expr(variant >= GATED): + prims.cp_async_bulk_tensor_shared_cta_global( + smem_g.subview(b_off), + tma_ptr_x2, + (x2_col0 + col, cutlass.Int32(0)), + act_full, + ) + elif warp_id == 2: + # ===================================================================== + # MMA: this CTA's k-tiles in ascending order into one TMEM accumulator + # (tail: the act k-tiles into a second one). + # ===================================================================== + idesc = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, a_dtype=io_dtype, b_dtype=io_dtype, n_dim=MMA_N, m_dim=CTA_M + ) + desc_a_base = prims.Tcgen05SmemDesc.build( + start_address=smem_a, leading_byte_offset=LEADING, stride_byte_offset=STRIDE, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + desc_b_base = prims.Tcgen05SmemDesc.build( + start_address=smem_b, leading_byte_offset=LEADING, stride_byte_offset=STRIDE, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + split_acc = variant == TAIL + tmem_acc1 = tmem_ptr + if cutlass.const_expr(split_acc): + tmem_acc1 = cutlass.inttoptr( + tmem_ptr_i32.load() + cutlass.Int32(MMA_N), 6, cutlass.Int32 + ) + # B is ready when its TMA lands (gated: when the epilogue warps have rewritten it). + if cutlass.const_expr(variant >= GATED): + while not cute.arch.mbarrier_try_wait(b_ready.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + for i in cutlass.range_constexpr(my_tiles): + second = split_acc and i >= x_tiles + first_tile = i == 0 or (split_acc and i == x_tiles) + if cutlass.Int32(i) < my_count: + while not cute.arch.mbarrier_try_wait(weight_full.subview(i).data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for kb in cutlass.range_constexpr(TMA_COPY_ITERS * K_BLOCKS_PER_HALF): + box = kb // K_BLOCKS_PER_HALF + within = kb % K_BLOCKS_PER_HALF + desc_a = desc_a_base + (i * STAGE_A + box * A_BOX + within * STEP) + desc_b = desc_b_base + (i * STAGE_B + box * B_BOX + within * STEP) + if prims.elect_sync(): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, prims.CTAGroup.CTA_1, tmem_acc1 if second else tmem_ptr, + desc_a, desc_b, idesc, not (first_tile and kb == 0), + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + elif warp_id >= 4: + # ===================================================================== + # Epilogue warps: [gate prologue] -> TMEM -> registers -> [split-K + # reduce] -> bf16 staging -> 16-byte stores. + # warp 4 + w reads TMEM lanes 32 w .. 32 w + 31 = tile rows 32 w + lane. + # ===================================================================== + tid = tx - cutlass.Int32(EPI_THREADS) + lane = tx % 32 + w = warp_id - 4 + if cutlass.const_expr(variant >= GATED): + # B = bf16(a * bf16(sigmoid(g))): a and g share one swizzled layout, so chunk c of B is + # chunk c of both inputs. 16-byte chunks, consecutive threads -> no bank conflicts. + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + for j in cutlass.range_constexpr(my_tiles * MMA_N * CTA_K // (VEC * EPI_THREADS)): + c = (tid + cutlass.Int32(j * EPI_THREADS)) * cutlass.Int32(VEC) + av = smem_b.load(idx=c, vector_size=VEC, alignment=16) + gv = smem_g.load(idx=c, vector_size=VEC, alignment=16) + outs = [] + for e in cutlass.range_constexpr(VEC): + if cutlass.const_expr(variant == SWIGLU): + # silu_and_mul's order in fp32: g * sigmoid(g), then times the up value, one rounding. + g = cutlass.Float32(gv[e]) + one = cutlass.Float32(1.0) + sig = _div_rn(one, one + cute.math.exp(-g, fastmath=False)) + outs.append(((g * sig) * cutlass.Float32(av[e])).to(io_dtype)) + else: + if cutlass.const_expr(variant == GATED): + gate = _sigmoid_bf16(cutlass.Float32(gv[e])) + else: + gate = cutlass.Float32(gv[e]) + outs.append((cutlass.Float32(av[e]) * gate).to(io_dtype)) + smem_b.store( + cutlass.Vector.from_elements(tuple(outs), io_dtype), + idx=c, + vector_size=VEC, + alignment=16, + ) + # Generic-proxy writes of B -> the tensor core's async-proxy reads, then release the MMA. + prims.fence_proxy(prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) + prims.mbarrier_arrive(b_ready) + if cutlass.const_expr(variant == TAIL): + # Per-token RMS of the latent rows while the MMA runs: epilogue warp w owns tokens 2w, 2w+1. + prims.griddepcontrol(prims.GridDepAction.WAIT) + for j in cutlass.range_constexpr(MMA_N // 4): + t = w * cutlass.Int32(MMA_N // 4) + cutlass.Int32(j) + sum_sq = cutlass.Float32(0.0) + if t < num_tokens: + for v in cutlass.range_constexpr(rms_cols // (8 * 32)): + words = rms_src.load( + idx=t * cutlass.Int32(rms_cols // 2) + + cutlass.Int32(v * 32 * 4) + + lane * cutlass.Int32(4), + vector_size=4, + alignment=16, + ) + for q in cutlass.range_constexpr(4): + lo = (words[q] << cutlass.Int32(16)).bitcast(cutlass.Float32) + hi = (words[q] & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + sum_sq = sum_sq + lo * lo + hi * hi + for offset in (16, 8, 4, 2, 1): + sum_sq = sum_sq + cute.arch.shuffle_sync_bfly(sum_sq, offset=offset) + if lane == 0: + row_scale.store( + cute.math.rsqrt(sum_sq * cutlass.Float32(1.0 / rms_cols) + eps), idx=t + ) + prims.barrier_cta_sync(1, thread_count=EPI_THREADS) + + while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc0 = prims.tcgen05_ld( + "32x32b", cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), num=MMA_N + ) + acc1 = acc0 + if cutlass.const_expr(variant == TAIL): + acc1 = prims.tcgen05_ld( + "32x32b", + cutlass.inttoptr(tmem_ptr_i32.load() + cutlass.Int32(MMA_N), 6, cutlass.Float32), + num=MMA_N, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + vals = [] + for t in cutlass.range_constexpr(MMA_N): + value = cutlass.Float32(acc0[t]) + if cutlass.const_expr(variant == TAIL): + value = value * row_scale.load(idx=t) + cutlass.Float32(acc1[t]) + vals.append(value) + + row = w * cutlass.Int32(32) + lane # tile row of this thread + local_row = row + if cutlass.const_expr(split > 1): + # Rows [R r, R r + R) belong to rank r = w // (4 / SPLIT); the other ranks push, the owner adds. + owner = w // cutlass.Int32(rows_owned // 32) + local_row = row - owner * cutlass.Int32(rows_owned) + if owner != rank: + slot = (rank * cutlass.Int32(rows_owned) + local_row) * cutlass.Int32(MMA_N) + if cutlass.const_expr(push): + peer_slot = _mapa_u32(mailbox.subview(slot).data_ptr(), owner) + mbar_peer = _mapa_u32(mail_full.data_ptr(), owner) + _st_async_v4(peer_slot, vals[0], vals[1], vals[2], vals[3], mbar_peer) + _st_async_v4( + peer_slot + cutlass.Int32(16), vals[4], vals[5], vals[6], vals[7], mbar_peer + ) + else: + for t in cutlass.range_constexpr(MMA_N): + prims.mapa(mailbox.subview(slot + cutlass.Int32(t)), owner).store(vals[t]) + # Release at cluster scope: the partials above are visible to the owner's acquire. + prims.mbarrier_arrive( + prims.mapa(mail_full, owner), scope=prims.MemScope.CLUSTER + ) + else: + if cutlass.const_expr(push): + # test_wait spin: a warp suspended in try_wait on a barrier completed by peers wakes late. + # Acquire at cluster scope: the bytes come from other CTAs' st.async (complete_tx releases at + # cluster scope). + while not _test_wait_cluster(mail_full.data_ptr(), 0): + pass + else: + while not prims.mbarrier_try_wait_parity( + mail_full, 0, scope=prims.MBarrierScope.CLUSTER + ): + pass + # The SPLIT partials in rank order (this rank's own from registers). + summed = [cutlass.Float32(0.0)] * MMA_N + for q in cutlass.range_constexpr(split): + slot = (cutlass.Int32(q * rows_owned) + local_row) * cutlass.Int32(MMA_N) + for t in cutlass.range_constexpr(MMA_N): + part = mailbox.load(idx=slot + cutlass.Int32(t)) + summed[t] = summed[t] + cutlass.Float32( + cutlass.select_(rank == cutlass.Int32(q), vals[t], part) + ) + for t in cutlass.range_constexpr(MMA_N): + stage.store( + summed[t].to(io_dtype), idx=cutlass.Int32(t * rows_owned) + local_row + ) + else: + for t in cutlass.range_constexpr(MMA_N): + stage.store(vals[t].to(io_dtype), idx=cutlass.Int32(t * rows_owned) + local_row) + # Every TMEM reader has waited for its load; the staging tile is complete. + prims.barrier_cta_sync(1, thread_count=EPI_THREADS) + if warp_id == 4: + prims.tcgen05_dealloc(tmem_ptr, TMEM_COLS) + # [token][rows_owned] bf16 -> y[token, rows]: 16 bytes per thread, rows contiguous per token. + chunks_per_token = rows_owned // VEC + ct = tid // cutlass.Int32(chunks_per_token) + cc = tid % cutlass.Int32(chunks_per_token) + if ct < num_tokens: + if ct < cutlass.Int32(MMA_N): + v = stage.load( + idx=ct * cutlass.Int32(rows_owned) + cc * cutlass.Int32(VEC), + vector_size=VEC, + alignment=16, + ) + y.store( + v, + idx=ct * cutlass.Int32(n_out) + + m_offset + + rank * cutlass.Int32(rows_owned) + + cc * cutlass.Int32(VEC), + vector_size=VEC, + alignment=16, + ) + + +# ============================================================================= +# Long-K variant (K > 6 k-tiles, e.g. the MLA [W_a; W_g] projection 2880 x 7168): one 128-row tile per cluster +# of SPLIT >= 4 CTAs, rank r streaming the interleaved k-tiles r, r + SPLIT, ... (the first `extra` ranks one +# more when SPLIT does not divide the k-tiles) through a RING-stage ring filled before griddepcontrol.wait, +# the rest of its k-tiles prefetched into L2 at the same time, so a latency-bound predecessor (an all-reduce) +# hides the whole weight read. B (all of the rank's k-tiles of x) is resident. Rows [32 w, 32 w + 32) of the +# tile belong to rank w < 4 (epilogue warp w's TMEM lanes; with SPLIT 2, rank w // 2 owns two such blocks); every +# other rank sends its fp32 partials of them to slot [rank] of the owner's mailbox (PUSH as in the short kernel: +# DSMEM stores + a release arrive at cluster scope, or st.async completing the barrier by bytes); the owner adds +# the SPLIT partials in rank order and rounds once. Output rows >= sig_row0 store bf16(sigmoid(bf16(acc))) instead +# of bf16(acc): the MLA output gate, ready for GATED_S. +# ============================================================================= +LONG_ROWS_PER_WARP = 32 + + +def long_supports(n_out: int, k_in: int, split: int, ring: int) -> bool: + """K in whole 64-column halves: a last k-tile of 64 columns has its second half past K, which the TMA fills with + zeros in both W and x (e.g. the dense MLP's down projection, K = 2112 at TP16).""" + k_tiles = num_k_tiles(k_in) + return ( + k_in % TMA_K_BOX == 0 + and (split == 2 or 4 <= split <= 8) + and 1 <= ring <= k_tiles // split + and ring * CTA_M * CTA_K * ELEM_BYTES + (k_tiles // split + 1) * MMA_N * CTA_K * ELEM_BYTES + <= 216 * 1024 + and n_out > 0 + ) + + +@cute.kernel +def k3_ctm_gemv_long_kernel( + tma_desc_w: cutlass.GridConstant[cuda.TensorMap], # W [N, K] bf16, 5-D, one call per k-tile + tma_desc_x: cutlass.GridConstant[cuda.TensorMap], # x [M, K] bf16, box 64 x 8 + y: cutlass.Array, # [M * N] bf16, token-major + num_tokens: cutlass.Int32, + sig_row0: cutlass.Int32, # rows >= sig_row0 store bf16(sigmoid(bf16(acc))); n_out: none + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + push: cutlass.Constexpr[ + bool + ], # split-K partials by st.async (test_wait owner) instead of stores + arrive +): + """One 128-row tile per cluster of `split` CTAs; weight ring + L2 prefetch before the grid wait.""" + k_tiles = num_k_tiles(k_in) + extra = k_tiles % split + max_tiles = k_tiles // split + (1 if extra > 0 else 0) + tx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp_id = cute.arch.warp_idx() + rank = cute.arch.block_idx_in_cluster() + tma_ptr_w = tma_desc_w.get_ptr() + tma_ptr_x = tma_desc_x.get_ptr() + m_offset = (bx // cutlass.Int32(split)) * cutlass.Int32(CTA_M) + my_tiles = cutlass.Int32(k_tiles // split) + if cutlass.const_expr(extra > 0): + if rank < cutlass.Int32(extra): + my_tiles = my_tiles + cutlass.Int32(1) + + smem_a = cutlass.Array( + io_dtype, ring * CTA_M * CTA_K, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_b = cutlass.Array( + io_dtype, max_tiles * MMA_N * CTA_K, space=cutlass.AddressSpace.smem, alignment=1024 + ) + tma_full = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + mma_done = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + act_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + acc_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + mail_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + # [source rank][owned block][row in the block][token] fp32 partials (the own rank's slots stay unused). A rank owns + # one 32-row block with SPLIT >= 4, two with SPLIT 2. + blocks = 4 // split if split < 4 else 1 + mailbox = cutlass.Array( + cutlass.Float32, + split * blocks * LONG_ROWS_PER_WARP * MMA_N, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + + if warp_id == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_x) + if prims.elect_sync(): + for s in cutlass.range_constexpr(ring): + prims.mbarrier_init(tma_full.subview(s), 1) + prims.mbarrier_init(mma_done.subview(s), 1) + prims.mbarrier_init(act_full, 1) + prims.mbarrier_init(acc_done, 1) + if cutlass.const_expr(push): + # Owner ranks (< 4): the other ranks' partials of the owned rows arrive by st.async (completing + # the transaction count expected here, before cluster formation); other ranks never wait on it. + prims.mbarrier_init(mail_full, 1) + prims.mbarrier_arrive_expect_tx( + mail_full, (split - 1) * blocks * LONG_ROWS_PER_WARP * MMA_N * 4 + ) + else: + # Owner ranks: every lane of the other ranks' epilogue warps of the owned blocks arrives once. + prims.mbarrier_init(mail_full, (split - 1) * blocks * LONG_ROWS_PER_WARP) + if warp_id == 2: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + # Cluster formation: the peers' shared memory and barriers are addressable from here on. + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + tmem_ptr = cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32) + + if warp_id == 0: + # ===================================================================== + # Weight TMA: fill the ring and prefetch the rest into L2 before the + # grid dependency; then refill each stage when its MMAs are done. + # ===================================================================== + if prims.elect_sync(): + for i in cutlass.range_constexpr(ring): + k = rank + cutlass.Int32(i * split) + prims.mbarrier_arrive_expect_tx(tma_full.subview(i), CTA_M * CTA_K * ELEM_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(i * CTA_M * CTA_K), + tma_ptr_w, + ( + cutlass.Int32(0), + m_offset, + k * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), + cutlass.Int32(0), + ), + tma_full.subview(i), + l2_cache_hint=EVICT_FIRST, + ) + for i in range(ring, my_tiles): + k = rank + i * cutlass.Int32(split) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, + [ + cutlass.Int32(0), + m_offset, + k * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), + cutlass.Int32(0), + ], + [], # tile mode: no im2col offsets + ) + if cutlass.const_expr(trigger_early): + # Dependents may launch now; they wait for this whole grid before reading y. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if prims.elect_sync(): + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + for i in range(ring, my_tiles): + while not cute.arch.mbarrier_try_wait(mma_done.subview(stage).data_ptr(), phase): + pass + k = rank + i * cutlass.Int32(split) + prims.mbarrier_arrive_expect_tx(tma_full.subview(stage), CTA_M * CTA_K * ELEM_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(stage * cutlass.Int32(CTA_M * CTA_K)), + tma_ptr_w, + ( + cutlass.Int32(0), + m_offset, + k * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), + cutlass.Int32(0), + ), + tma_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + stage = stage + cutlass.Int32(1) + if stage == cutlass.Int32(ring): + stage = cutlass.Int32(0) + phase = phase ^ cutlass.Int32(1) + elif warp_id == 1: + # ===================================================================== + # Activation TMA: the rank's k-tiles of x, resident, after the wait. + # ===================================================================== + prims.griddepcontrol(prims.GridDepAction.WAIT) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + act_full, my_tiles * cutlass.Int32(MMA_N * CTA_K * ELEM_BYTES) + ) + for i in range(my_tiles): + k = rank + i * cutlass.Int32(split) + for half in cutlass.range_constexpr(TMA_COPY_ITERS): + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview( + i * cutlass.Int32(MMA_N * CTA_K) + cutlass.Int32(half * B_HALF_ELEMS) + ), + tma_ptr_x, + ( + k * cutlass.Int32(CTA_K) + cutlass.Int32(half * TMA_K_BOX), + cutlass.Int32(0), + ), + act_full, + ) + elif warp_id == 2: + # ===================================================================== + # MMA: the rank's k-tiles in order, one TMEM accumulator (partial sum). + # ===================================================================== + idesc = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, a_dtype=io_dtype, b_dtype=io_dtype, n_dim=MMA_N, m_dim=CTA_M + ) + desc_a_base = prims.Tcgen05SmemDesc.build( + start_address=smem_a, leading_byte_offset=LEADING, stride_byte_offset=STRIDE, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + desc_b_base = prims.Tcgen05SmemDesc.build( + start_address=smem_b, leading_byte_offset=LEADING, stride_byte_offset=STRIDE, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + for i in range(my_tiles): + while not cute.arch.mbarrier_try_wait(tma_full.subview(stage).data_ptr(), phase): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for kb in cutlass.range_constexpr(TMA_COPY_ITERS * K_BLOCKS_PER_HALF): + box = kb // K_BLOCKS_PER_HALF + within = kb % K_BLOCKS_PER_HALF + desc_a = desc_a_base + ( + stage * cutlass.Int32(STAGE_A) + cutlass.Int32(box * A_BOX + within * STEP) + ) + desc_b = desc_b_base + ( + i * cutlass.Int32(STAGE_B) + cutlass.Int32(box * B_BOX + within * STEP) + ) + accumulate = cutlass.Boolean(True) + if cutlass.const_expr(kb == 0): + accumulate = i > cutlass.Int32(0) + if prims.elect_sync(): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, + prims.CTAGroup.CTA_1, + tmem_ptr, + desc_a, + desc_b, + idesc, + accumulate, + ) + if prims.elect_sync(): + prims.tcgen05_commit(mma_done.subview(stage)) + stage = stage + cutlass.Int32(1) + if stage == cutlass.Int32(ring): + stage = cutlass.Int32(0) + phase = phase ^ cutlass.Int32(1) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + elif warp_id >= 4: + # ===================================================================== + # Epilogue: TMEM -> registers, push to / reduce at the row owner, store. + # ===================================================================== + lane = tx % 32 + w = warp_id - 4 # TMEM lanes 32 w .. 32 w + 31: tile rows owned by rank w + while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), num=MMA_N + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + slot_row = lane * cutlass.Int32(MMA_N) + owner = w // cutlass.Int32(blocks) + block_slot = (w % cutlass.Int32(blocks)) * cutlass.Int32( + LONG_ROWS_PER_WARP * MMA_N + ) + slot_row + if owner != rank: + base = rank * cutlass.Int32(blocks * LONG_ROWS_PER_WARP * MMA_N) + block_slot + if cutlass.const_expr(push): + peer_slot = _mapa_u32(mailbox.subview(base).data_ptr(), owner) + mbar_peer = _mapa_u32(mail_full.data_ptr(), owner) + _st_async_v4(peer_slot, acc[0], acc[1], acc[2], acc[3], mbar_peer) + _st_async_v4( + peer_slot + cutlass.Int32(16), acc[4], acc[5], acc[6], acc[7], mbar_peer + ) + else: + for t in cutlass.range_constexpr(MMA_N): + prims.mapa(mailbox.subview(base + cutlass.Int32(t)), owner).store( + cutlass.Float32(acc[t]) + ) + # Release at cluster scope: this lane's partials are visible to the owner's acquire. + prims.mbarrier_arrive(prims.mapa(mail_full, owner), scope=prims.MemScope.CLUSTER) + else: + if cutlass.const_expr(push): + # test_wait spin: a warp suspended in try_wait on a barrier completed by peers wakes late. + # Acquire at cluster scope: the bytes come from other CTAs' st.async (complete_tx releases at + # cluster scope). + while not _test_wait_cluster(mail_full.data_ptr(), 0): + pass + else: + while not prims.mbarrier_try_wait_parity( + mail_full, 0, scope=prims.MBarrierScope.CLUSTER + ): + pass + n = m_offset + w * cutlass.Int32(LONG_ROWS_PER_WARP) + lane + for t in cutlass.range_constexpr(MMA_N): + total = cutlass.Float32(0.0) + for q in cutlass.range_constexpr(split): + part = mailbox.load( + idx=cutlass.Int32(q * blocks * LONG_ROWS_PER_WARP * MMA_N) + + block_slot + + cutlass.Int32(t) + ) + total = total + cutlass.Float32( + cutlass.select_(rank == cutlass.Int32(q), cutlass.Float32(acc[t]), part) + ) + out = total.to(io_dtype) + if n >= sig_row0: + out = _sigmoid_bf16(_bf16_rn(total)).to(io_dtype) + if cutlass.Int32(t) < num_tokens: + if n < cutlass.Int32(n_out): + y.store(out, idx=cutlass.Int32(t * n_out) + n) + prims.barrier_cta_sync(1, thread_count=EPI_THREADS) + if warp_id == 4: + prims.tcgen05_dealloc(tmem_ptr, TMEM_COLS) + + +def _weight_tensor_map(w, n_out, k_in): + """W as five TMA dimensions (64-element column chunk, row, 64-element chunk index, 1, 1) so one call + per k-tile lands both 128-byte-swizzled halves; strides in 16-byte units.""" + return cuda.create_tensor_map_tiled( + global_address=w.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[TMA_K_BOX, n_out, k_in // TMA_K_BOX, 1, 1], + global_strides=[ + (k_in * ELEM_BYTES) // 16, + (TMA_K_BOX * ELEM_BYTES) // 16, + (n_out * k_in * ELEM_BYTES) // 16, + (n_out * k_in * ELEM_BYTES) // 16, + ], + box_dims=[TMA_K_BOX, CTA_M, TMA_COPY_ITERS, 1, 1], + swizzle=cuda.TensorMapSwizzle.s128b, + ) + + +def _activation_tensor_map(x, cols, num_tokens): + """x [M, cols] (rows dense) as (cols, M) with an 8-row box: rows past num_tokens arrive as zeros.""" + return cuda.create_tensor_map_tiled( + global_address=x.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[cols, num_tokens], + global_strides=[(cols * ELEM_BYTES) // 16], + box_dims=[TMA_K_BOX, MMA_N], + swizzle=cuda.TensorMapSwizzle.s128b, + ) + + +def _launch(kernel, n_out, split, use_pdl, stream): + kernel.launch( + grid=((n_out // CTA_M) * split, 1, 1), + block=(THREADS, 1, 1), + cluster=(split, 1, 1) if split > 1 else None, + stream=stream, + use_pdl=use_pdl, + ) + + +@cute.jit +def k3_ctm_gemv( + w: cute.Tensor, # [N, K] bf16, K contiguous + x: cute.Tensor, # [M, K] bf16, K contiguous, M <= 8 + y: cute.Tensor, # [M * N] bf16 + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + push: cutlass.Constexpr[bool], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + """``y = x @ w^T``.""" + tma_desc_w = _weight_tensor_map(w, n_out, k_in) + tma_desc_x = _activation_tensor_map(x, k_in, num_tokens) + kernel = k3_ctm_gemv_kernel( + tma_desc_w, tma_desc_x, tma_desc_x, y, y, num_tokens, cutlass.Int32(0), cutlass.Int32(0), + cutlass.Float32(0.0), n_out, k_in, PLAIN, num_k_tiles(k_in), 0, split, trigger_early, push, + ) # fmt: skip + _launch(kernel, n_out, split, use_pdl, stream) + + +@cute.jit +def k3_ctm_gemv_tail( + w: cute.Tensor, # [N, K_lat + K_act] bf16: latent-up columns of this rank's slice (zero-padded) | shared down + latent: cute.Tensor, # [M, rms_cols] bf16, the whole reduced latent row + latent_words: cute.Tensor, # the same memory as int32 words, for the RMS + act: cute.Tensor, # [M, K_act] bf16, the shared-expert activation + y: cute.Tensor, # [M * N] bf16 + num_tokens: cutlass.Int32, + lat_col0: cutlass.Int32, # first latent column of this rank's slice + eps: cutlass.Float32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + lat_tiles: cutlass.Constexpr[int], + rms_cols: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + """``y = rmsnorm(latent)[:, slice] @ W_lat^T + act @ W_act^T`` (the RMS applied to the latent accumulator).""" + tma_desc_w = _weight_tensor_map(w, n_out, k_in) + tma_desc_lat = _activation_tensor_map(latent, rms_cols, num_tokens) + tma_desc_act = _activation_tensor_map(act, k_in - lat_tiles * CTA_K, num_tokens) + kernel = k3_ctm_gemv_kernel( + tma_desc_w, tma_desc_lat, tma_desc_act, y, latent_words, num_tokens, lat_col0, cutlass.Int32(0), eps, + n_out, k_in, TAIL, lat_tiles, rms_cols, 1, trigger_early, False, + ) # fmt: skip + _launch(kernel, n_out, 1, use_pdl, stream) + + +@cute.jit +def k3_ctm_gemv_gated( + w: cute.Tensor, # [N, K] bf16 (the o_proj weight), K contiguous + a: cute.Tensor, # [M, K] bf16, the attention output (o_proj's input before the gate) + gsrc: cute.Tensor, # [M, gsrc_cols] bf16, dense rows; g = gsrc[:, g_col0 : g_col0 + K] + y: cute.Tensor, # [M * N] bf16 + num_tokens: cutlass.Int32, + g_col0: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + gsrc_cols: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + gate_sigmoid: cutlass.Constexpr[bool], # True: gsrc holds g; False: it holds bf16(sigmoid(g)) + trigger_early: cutlass.Constexpr[bool], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + """``y = (a * sigmoid(g)) @ w^T`` with torch's bf16 roundings of the sigmoid and of the product.""" + tma_desc_w = _weight_tensor_map(w, n_out, k_in) + tma_desc_a = _activation_tensor_map(a, k_in, num_tokens) + tma_desc_g = _activation_tensor_map(gsrc, gsrc_cols, num_tokens) + kernel = k3_ctm_gemv_kernel( + tma_desc_w, tma_desc_a, tma_desc_g, y, y, num_tokens, cutlass.Int32(0), g_col0, cutlass.Float32(0.0), + n_out, k_in, GATED if gate_sigmoid else GATED_S, num_k_tiles(k_in), 0, split, trigger_early, False, + ) # fmt: skip + _launch(kernel, n_out, split, use_pdl, stream) + + +@cute.jit +def k3_ctm_gemv_swiglu( + w: cute.Tensor, # [N, K] bf16 (the down projection), K contiguous + gu: cute.Tensor, # [M, 2 K] bf16, dense rows: the gate_up output (gate columns first) + y: cute.Tensor, # [M * N] bf16 + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + push: cutlass.Constexpr[bool], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + """``y = (silu(gu[:, :K]) * gu[:, K:]) @ w^T`` with silu_and_mul's fp32 arithmetic and bf16 rounding.""" + tma_desc_w = _weight_tensor_map(w, n_out, k_in) + tma_desc_gu = _activation_tensor_map(gu, 2 * k_in, num_tokens) + kernel = k3_ctm_gemv_kernel( + tma_desc_w, tma_desc_gu, tma_desc_gu, y, y, num_tokens, cutlass.Int32(k_in), cutlass.Int32(0), + cutlass.Float32(0.0), n_out, k_in, SWIGLU, num_k_tiles(k_in), 0, split, trigger_early, push, + ) # fmt: skip + _launch(kernel, n_out, split, use_pdl, stream) + + +@cute.jit +def k3_ctm_gemv_long( + w: cute.Tensor, # [N, K] bf16, K contiguous + x: cute.Tensor, # [M, K] bf16, K contiguous, M <= 8 + y: cute.Tensor, # [M * N] bf16 + num_tokens: cutlass.Int32, + sig_row0: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + push: cutlass.Constexpr[bool], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + """``y = x @ w^T``; output rows >= sig_row0 hold bf16(sigmoid(bf16(.))) instead.""" + tma_desc_w = _weight_tensor_map(w, n_out, k_in) + tma_desc_x = _activation_tensor_map(x, k_in, num_tokens) + k3_ctm_gemv_long_kernel( + tma_desc_w, + tma_desc_x, + y, + num_tokens, + sig_row0, + n_out, + k_in, + split, + ring, + trigger_early, + push, + ).launch( + grid=(((n_out + CTA_M - 1) // CTA_M) * split, 1, 1), + block=(THREADS, 1, 1), + cluster=(split, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) + + +# ============================================================================= +# Wide variant (up to 64 tokens: DSpark verify of several requests): the long kernel's geometry, weight ring and L2 +# prefetch, with all the call's token columns in one MMA of N = N_TILE (16, 32 or 64). x no longer fits resident +# beside the ring, so each ring stage holds a weight k-tile and the matching x k-tile: the weight half is loaded +# before griddepcontrol.wait (the first RING stages; the rest prefetched into L2), the x half after it, each refilled +# once the stage's MMAs are done. +# Split-K reduction by token: rank r owns tokens [r T, r T + T) of the tile (T = N_TILE / SPLIT rounded up to 4) for +# all 128 rows. Every epilogue thread (one row) sends each other owner its partials of that owner's tokens by st.async +# (16-byte chunks of 4 tokens, chunk j at j ^ f(lane) so a quarter-warp's 8 rows fall in 8 bank groups), writes its +# own into its own slot, then sums its row's T tokens over the SPLIT slots in rank order: every token's sums are the +# long kernel's (same k-tile order per rank, same rank order). Each warp stages its 32 rows as [token][row] and stores +# them with one TMA store (tokens >= M clipped; the block holding row N - 1 when N is not a multiple of 32 has its own +# box): bf16 in the x ring (free once the MMAs are done; rows >= sig_row0 as bf16(sigmoid(bf16(acc)))), or fp32 +# (OUT_FP32: the MoE head's router logits) in place of its own slot. +# ============================================================================= +WIDE_TILES = (16, 32, 64) +WIDE_SMEM_BYTES = 220 * 1024 # the ring, the x stages and the mailbox; the barriers fit in the rest + + +def wide_tile(num_tokens: int) -> int: + """The token tile (MMA N) that holds ``num_tokens`` columns.""" + for n_tile in WIDE_TILES: + if num_tokens <= n_tile: + return n_tile + raise ValueError(f"k3_ctm_gemv_wide: {num_tokens} tokens exceed {WIDE_TILES[-1]}") + + +def wide_owned_tokens(n_tile: int, split: int) -> int: + """Tokens of the tile a rank reduces and stores (the last owner may have fewer, later ranks none).""" + return (n_tile + split - 1) // split + 3 & ~3 + + +def wide_smem_bytes(split: int, ring: int, n_tile: int, x_ring: int) -> int: + return ( + ring * CTA_M + x_ring * n_tile + ) * CTA_K * ELEM_BYTES + split * CTA_M * wide_owned_tokens(n_tile, split) * 4 + + +def wide_supports(n_out: int, k_in: int, split: int, ring: int, n_tile: int, x_ring: int) -> bool: + """K in whole 64-column halves (a last half k-tile reads zeros past K, as in the long kernel); N a multiple of 8 + (the output rows of a token are a whole number of 16-byte units for the TMA store).""" + return ( + k_in % TMA_K_BOX == 0 + and n_out % 8 == 0 + and n_tile in WIDE_TILES + and (split == 2 or 4 <= split <= 8) + and 1 <= ring <= num_k_tiles(k_in) // split + and 1 <= x_ring <= ring + and wide_smem_bytes(split, ring, n_tile, x_ring) <= WIDE_SMEM_BYTES + and n_out > 0 + ) + + +def _chunk_swizzle(lane, chunks: int): + """XOR for the 16-byte chunk index of a row of ``chunks`` chunks: the 8 rows of a quarter-warp in 8 bank groups.""" + if chunks == 2: + return (lane >> cutlass.Int32(2)) & cutlass.Int32(1) + if chunks == 4: + return (lane >> cutlass.Int32(1)) & cutlass.Int32(3) + if chunks >= 8: + return lane & cutlass.Int32(7) + return cutlass.Int32(0) # 1 or 3 chunks: rows already in distinct bank groups + + +@cute.kernel +def k3_ctm_gemv_wide_kernel( + tma_desc_w: cutlass.GridConstant[cuda.TensorMap], # W [N, K] bf16, 5-D, one call per k-tile + tma_desc_x: cutlass.GridConstant[cuda.TensorMap], # x [M, K] bf16, box 64 x n_tile + tma_desc_y: cutlass.GridConstant[ + cuda.TensorMap + ], # y [M, N] bf16 (fp32 with out_fp32), box 32 x T + tma_desc_y_tail: cutlass.GridConstant[cuda.TensorMap], # the same, box (N % 32 or 32) x T + sig_row0: cutlass.Int32, # bf16 output: rows >= sig_row0 store bf16(sigmoid(bf16(acc))); n_out: none + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + x_ring: cutlass.Constexpr[int], + n_tile: cutlass.Constexpr[int], + out_fp32: cutlass.Constexpr[bool], + trigger_early: cutlass.Constexpr[bool], +): + """One 128-row tile per cluster of `split` CTAs; a weight ring filled before the grid wait, an x ring after it.""" + k_tiles = num_k_tiles(k_in) + extra = k_tiles % split + tmem_cols = max(TMEM_COLS, n_tile) + b_stage = n_tile * CTA_K # elements of one x k-tile + b_half = n_tile * TMA_K_BOX + tail_rows = n_out % LONG_ROWS_PER_WARP or LONG_ROWS_PER_WARP + owned = wide_owned_tokens(n_tile, split) # T + slot_elems = CTA_M * owned # one source rank's partials of the owner's tokens: [row][T] + tx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp_id = cute.arch.warp_idx() + rank = cute.arch.block_idx_in_cluster() + tma_ptr_w = tma_desc_w.get_ptr() + tma_ptr_x = tma_desc_x.get_ptr() + tma_ptr_y = tma_desc_y.get_ptr() + tma_ptr_y_tail = tma_desc_y_tail.get_ptr() + m_offset = (bx // cutlass.Int32(split)) * cutlass.Int32(CTA_M) + my_tiles = cutlass.Int32(k_tiles // split) + if cutlass.const_expr(extra > 0): + if rank < cutlass.Int32(extra): + my_tiles = my_tiles + cutlass.Int32(1) + # The tokens this rank owns: [rank T, rank T + my_owned). + my_owned = cutlass.Int32(n_tile) - rank * cutlass.Int32(owned) + if my_owned > cutlass.Int32(owned): + my_owned = cutlass.Int32(owned) + + smem_a = cutlass.Array( + io_dtype, ring * CTA_M * CTA_K, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_b = cutlass.Array( + io_dtype, x_ring * b_stage, space=cutlass.AddressSpace.smem, alignment=1024 + ) + a_full = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + b_full = cutlass.Array(cutlass.Int64, x_ring, space=cutlass.AddressSpace.smem, alignment=8) + mma_done = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + x_done = cutlass.Array(cutlass.Int64, x_ring, space=cutlass.AddressSpace.smem, alignment=8) + acc_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + mail_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + # [source rank][row][T tokens, chunks swizzled] fp32 partials of this rank's tokens (its own slot written locally). + mailbox = cutlass.Array( + cutlass.Float32, split * slot_elems, space=cutlass.AddressSpace.smem, alignment=128 + ) + + if warp_id == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_x) + prims.prefetch_tensormap(tma_ptr_y) + prims.prefetch_tensormap(tma_ptr_y_tail) + if prims.elect_sync(): + for s in cutlass.range_constexpr(ring): + prims.mbarrier_init(a_full.subview(s), 1) + prims.mbarrier_init(mma_done.subview(s), 1) + for s in cutlass.range_constexpr(x_ring): + prims.mbarrier_init(b_full.subview(s), 1) + prims.mbarrier_init(x_done.subview(s), 1) + prims.mbarrier_init(acc_done, 1) + # The other ranks' partials of this rank's tokens arrive by st.async, completing the transaction count + # expected here, before cluster formation (a rank owning no tokens never waits on it). + prims.mbarrier_init(mail_full, 1) + if my_owned > cutlass.Int32(0): + prims.mbarrier_arrive_expect_tx( + mail_full, cutlass.Int32((split - 1) * CTA_M * 4) * my_owned + ) + if warp_id == 2: + prims.tcgen05_alloc(tmem_ptr_i32, tmem_cols) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + # Cluster formation: the peers' shared memory and barriers are addressable from here on. + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + tmem_ptr = cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32) + + if warp_id == 0: + # ===================================================================== + # Weight TMA: fill the ring and prefetch the rest into L2 before the + # grid dependency; then refill each stage when its MMAs are done. + # ===================================================================== + if prims.elect_sync(): + for i in cutlass.range_constexpr(ring): + k = rank + cutlass.Int32(i * split) + prims.mbarrier_arrive_expect_tx(a_full.subview(i), CTA_M * CTA_K * ELEM_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(i * CTA_M * CTA_K), + tma_ptr_w, + ( + cutlass.Int32(0), + m_offset, + k * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), + cutlass.Int32(0), + ), + a_full.subview(i), + l2_cache_hint=EVICT_FIRST, + ) + for i in range(ring, my_tiles): + k = rank + i * cutlass.Int32(split) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, + [ + cutlass.Int32(0), + m_offset, + k * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), + cutlass.Int32(0), + ], + [], # tile mode: no im2col offsets + ) + if cutlass.const_expr(trigger_early): + # Dependents may launch now; they wait for this whole grid before reading y. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if prims.elect_sync(): + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + for i in range(ring, my_tiles): + while not cute.arch.mbarrier_try_wait(mma_done.subview(stage).data_ptr(), phase): + pass + k = rank + i * cutlass.Int32(split) + prims.mbarrier_arrive_expect_tx(a_full.subview(stage), CTA_M * CTA_K * ELEM_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(stage * cutlass.Int32(CTA_M * CTA_K)), + tma_ptr_w, + ( + cutlass.Int32(0), + m_offset, + k * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), + cutlass.Int32(0), + ), + a_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + stage = stage + cutlass.Int32(1) + if stage == cutlass.Int32(ring): + stage = cutlass.Int32(0) + phase = phase ^ cutlass.Int32(1) + elif warp_id == 1: + # ===================================================================== + # Activation TMA after the wait: the x k-tile of each weight k-tile + # into the x ring, refilled once the stage's MMAs are done. + # ===================================================================== + prims.griddepcontrol(prims.GridDepAction.WAIT) + if prims.elect_sync(): + for i in cutlass.range_constexpr(x_ring): + k = rank + cutlass.Int32(i * split) + prims.mbarrier_arrive_expect_tx(b_full.subview(i), b_stage * ELEM_BYTES) + for half in cutlass.range_constexpr(TMA_COPY_ITERS): + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(i * b_stage + half * b_half), + tma_ptr_x, + ( + k * cutlass.Int32(CTA_K) + cutlass.Int32(half * TMA_K_BOX), + cutlass.Int32(0), + ), + b_full.subview(i), + ) + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + for i in range(x_ring, my_tiles): + while not cute.arch.mbarrier_try_wait(x_done.subview(stage).data_ptr(), phase): + pass + k = rank + i * cutlass.Int32(split) + prims.mbarrier_arrive_expect_tx(b_full.subview(stage), b_stage * ELEM_BYTES) + for half in cutlass.range_constexpr(TMA_COPY_ITERS): + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview( + stage * cutlass.Int32(b_stage) + cutlass.Int32(half * b_half) + ), + tma_ptr_x, + ( + k * cutlass.Int32(CTA_K) + cutlass.Int32(half * TMA_K_BOX), + cutlass.Int32(0), + ), + b_full.subview(stage), + ) + stage = stage + cutlass.Int32(1) + if stage == cutlass.Int32(x_ring): + stage = cutlass.Int32(0) + phase = phase ^ cutlass.Int32(1) + elif warp_id == 2: + # ===================================================================== + # MMA: the rank's k-tiles in order, one TMEM accumulator (partial sum) + # of N = n_tile token columns. + # ===================================================================== + idesc = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, a_dtype=io_dtype, b_dtype=io_dtype, n_dim=n_tile, m_dim=CTA_M + ) + desc_a_base = prims.Tcgen05SmemDesc.build( + start_address=smem_a, leading_byte_offset=LEADING, stride_byte_offset=STRIDE, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + desc_b_base = prims.Tcgen05SmemDesc.build( + start_address=smem_b, leading_byte_offset=LEADING, stride_byte_offset=STRIDE, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + x_stage = cutlass.Int32(0) + x_phase = cutlass.Int32(0) + for i in range(my_tiles): + while not cute.arch.mbarrier_try_wait(a_full.subview(stage).data_ptr(), phase): + pass + while not cute.arch.mbarrier_try_wait(b_full.subview(x_stage).data_ptr(), x_phase): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for kb in cutlass.range_constexpr(TMA_COPY_ITERS * K_BLOCKS_PER_HALF): + box = kb // K_BLOCKS_PER_HALF + within = kb % K_BLOCKS_PER_HALF + desc_a = desc_a_base + ( + stage * cutlass.Int32(STAGE_A) + cutlass.Int32(box * A_BOX + within * STEP) + ) + desc_b = desc_b_base + ( + x_stage * cutlass.Int32(b_stage >> 3) + + cutlass.Int32(box * (b_half >> 3) + within * STEP) + ) + accumulate = cutlass.Boolean(True) + if cutlass.const_expr(kb == 0): + accumulate = i > cutlass.Int32(0) + if prims.elect_sync(): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, + prims.CTAGroup.CTA_1, + tmem_ptr, + desc_a, + desc_b, + idesc, + accumulate, + ) + if prims.elect_sync(): + prims.tcgen05_commit(mma_done.subview(stage)) + prims.tcgen05_commit(x_done.subview(x_stage)) + stage = stage + cutlass.Int32(1) + if stage == cutlass.Int32(ring): + stage = cutlass.Int32(0) + phase = phase ^ cutlass.Int32(1) + x_stage = x_stage + cutlass.Int32(1) + if x_stage == cutlass.Int32(x_ring): + x_stage = cutlass.Int32(0) + x_phase = x_phase ^ cutlass.Int32(1) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + elif warp_id >= 4: + # ===================================================================== + # Epilogue: TMEM -> registers; partials of each owner's tokens to that + # owner (own into the own slot); sum this rank's tokens in rank order; + # stage the warp's rows, one TMA store. + # ===================================================================== + lane = tx % 32 + w = warp_id - 4 # TMEM lanes 32 w .. 32 w + 31: tile rows 32 w + lane + row = w * cutlass.Int32(LONG_ROWS_PER_WARP) + lane + while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), num=n_tile + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + swz = _chunk_swizzle(lane, owned // 4) + row_slot = row * cutlass.Int32(owned) + peer_slots = [] + peer_mbars = [] + for r in cutlass.range_constexpr(split): + peer_slots.append( + _mapa_u32( + mailbox.subview(rank * cutlass.Int32(slot_elems) + row_slot).data_ptr(), + cutlass.Int32(r), + ) + ) + peer_mbars.append(_mapa_u32(mail_full.data_ptr(), cutlass.Int32(r))) + # Chunk j for every owner in turn, so the owners' incoming traffic interleaves; the own chunks to the own slot. + for j in cutlass.range_constexpr(owned // 4): + for r in cutlass.range_constexpr(split): + if cutlass.const_expr(r * owned + 4 * j < n_tile): + t = r * owned + 4 * j + if rank == cutlass.Int32(r): + own = ( + cutlass.Int32(r * slot_elems) + + row_slot + + ((cutlass.Int32(j) ^ swz) << cutlass.Int32(2)) + ) + mailbox.store( + cutlass.Vector.from_elements( + tuple(cutlass.Float32(acc[t + e]) for e in range(4)), cutlass.Float32 + ), + idx=own, + vector_size=4, + alignment=16, + ) # fmt: skip + else: + _st_async_v4( + peer_slots[r] + ((cutlass.Int32(j) ^ swz) << cutlass.Int32(4)), + acc[t], acc[t + 1], acc[t + 2], acc[t + 3], peer_mbars[r], + ) # fmt: skip + if my_owned > cutlass.Int32(0): + # test_wait spin: a warp suspended in try_wait on a barrier completed by peers wakes late. + # Acquire at cluster scope: the bytes come from other CTAs' st.async (complete_tx releases at + # cluster scope). + while not _test_wait_cluster(mail_full.data_ptr(), 0): + pass + # This row's owned tokens over the SPLIT slots in rank order, as the long kernel. + totals = [cutlass.Float32(0.0)] * owned + for q in cutlass.range_constexpr(split): + for j in cutlass.range_constexpr(owned // 4): + part = mailbox.load( + idx=cutlass.Int32(q * slot_elems) + row_slot + ((cutlass.Int32(j) ^ swz) << cutlass.Int32(2)), + vector_size=4, + alignment=16, + ) # fmt: skip + for e in cutlass.range_constexpr(4): + totals[4 * j + e] = totals[4 * j + e] + cutlass.Float32(part[e]) + # [token][rows] staging of the warp's rows below N, one TMA store (tokens >= M clipped). + n0 = m_offset + w * cutlass.Int32(LONG_ROWS_PER_WARP) + tail_block = n0 + cutlass.Int32(LONG_ROWS_PER_WARP) > cutlass.Int32(n_out) + pitch = cutlass.Int32(LONG_ROWS_PER_WARP) + if tail_block: + pitch = cutlass.Int32(tail_rows) + t0 = rank * cutlass.Int32(owned) + if cutlass.const_expr(out_fp32): + # In place of this rank's own slot, rows of this warp only (read above by this warp alone). + staged = mailbox.subview( + rank * cutlass.Int32(slot_elems) + w * cutlass.Int32(LONG_ROWS_PER_WARP * owned) + ) + cute.arch.sync_warp() + else: + staged = smem_b.subview(w * cutlass.Int32(LONG_ROWS_PER_WARP * owned)) + if n0 < cutlass.Int32(n_out): + if lane < pitch: + if cutlass.const_expr(out_fp32): + for i in cutlass.range_constexpr(owned): + staged.store(totals[i], idx=cutlass.Int32(i) * pitch + lane) + else: + if n0 + lane >= sig_row0: + for i in cutlass.range_constexpr(owned): + staged.store( + _sigmoid_bf16(_bf16_rn(totals[i])).to(io_dtype), idx=cutlass.Int32(i) * pitch + lane + ) # fmt: skip + else: + for i in cutlass.range_constexpr(owned): + staged.store( + totals[i].to(io_dtype), idx=cutlass.Int32(i) * pitch + lane + ) + # Generic-proxy writes of the staged rows -> the TMA store's async-proxy reads. + prims.fence_proxy(prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + if lane == 0: + if tail_block: + prims.cp_async_bulk_tensor_global_shared_cta( + tma_ptr_y_tail, staged, [n0, t0] + ) + else: + prims.cp_async_bulk_tensor_global_shared_cta(tma_ptr_y, staged, [n0, t0]) + prims.cp_async_bulk_commit_group() + prims.cp_async_bulk_wait_group(0, read=True) + prims.barrier_cta_sync(1, thread_count=EPI_THREADS) + if warp_id == 4: + prims.tcgen05_dealloc(tmem_ptr, tmem_cols) + + +def _activation_tensor_map_rows(x, cols, num_tokens, rows): + """x [M, cols] (rows dense) as (cols, M) with a ``rows``-row box: rows past num_tokens arrive as zeros.""" + return cuda.create_tensor_map_tiled( + global_address=x.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[cols, num_tokens], + global_strides=[(cols * ELEM_BYTES) // 16], + box_dims=[TMA_K_BOX, rows], + swizzle=cuda.TensorMapSwizzle.s128b, + ) + + +def _output_tensor_map(y, n_out, num_tokens, tokens, out_fp32, rows=LONG_ROWS_PER_WARP): + """y [M, N] (rows dense) as (N, M) with a rows x tokens box, no swizzle: a warp's staged [token][rows].""" + elem_bytes = 4 if out_fp32 else ELEM_BYTES + return cuda.create_tensor_map_tiled( + global_address=y.iterator.toint(), + dtype=cutlass.Float32 if out_fp32 else cutlass.BFloat16, + global_dims=[n_out, num_tokens], + global_strides=[(n_out * elem_bytes) // 16], + box_dims=[rows, tokens], + swizzle=cuda.TensorMapSwizzle.none, + ) + + +@cute.jit +def k3_ctm_gemv_wide( + w: cute.Tensor, # [N, K] bf16, K contiguous + x: cute.Tensor, # [M, K] bf16, K contiguous, M <= n_tile + y: cute.Tensor, # [M * N] bf16 (fp32 with out_fp32) + num_tokens: cutlass.Int32, + sig_row0: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + x_ring: cutlass.Constexpr[int], + n_tile: cutlass.Constexpr[int], + out_fp32: cutlass.Constexpr[bool], + trigger_early: cutlass.Constexpr[bool], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + """``y = x @ w^T`` for up to ``n_tile`` tokens; bf16 output rows >= sig_row0 hold bf16(sigmoid(bf16(.))).""" + owned = wide_owned_tokens(n_tile, split) + tma_desc_w = _weight_tensor_map(w, n_out, k_in) + tma_desc_x = _activation_tensor_map_rows(x, k_in, num_tokens, n_tile) + tma_desc_y = _output_tensor_map(y, n_out, num_tokens, owned, out_fp32) + tma_desc_y_tail = _output_tensor_map( + y, n_out, num_tokens, owned, out_fp32, n_out % LONG_ROWS_PER_WARP or LONG_ROWS_PER_WARP + ) + k3_ctm_gemv_wide_kernel( + tma_desc_w, tma_desc_x, tma_desc_y, tma_desc_y_tail, sig_row0, n_out, k_in, split, ring, x_ring, n_tile, + out_fp32, trigger_early, + ).launch( + grid=(((n_out + CTA_M - 1) // CTA_M) * split, 1, 1), + block=(THREADS, 1, 1), + cluster=(split, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) # fmt: skip + + +# ============================================================================= +# SiTU-and-mul with PDL: a = bf16(situ(g) * situ_lin(u)) for a gate_up output [M, 2K] (g its first K columns), in +# SituAndMul's fp32 order (beta tanh(g / beta) sigmoid(g), times linear_beta tanh(u / linear_beta) or u), one 16-byte +# vector of 8 columns per thread. The dependents launch at once, so a following projection streams its weights while +# this kernel and its predecessor run; the activation reads wait for the grid. +# ============================================================================= +SITU_THREADS = 128 +SITU_COLS_PER_CTA = SITU_THREADS * VEC + + +def _situ_mul(g, u, beta, linear_beta, has_linear: bool): + """bf16(situ(g) * situ_lin(u)) in fp32, SituAndMul's order.""" + one = cutlass.Float32(1.0) + sig = _div_rn(one, one + cute.math.exp(-g, fastmath=False)) + a = beta * cute.math.tanh(_div_rn(g, beta), fastmath=False) * sig + if cutlass.const_expr(has_linear): + u = linear_beta * cute.math.tanh(_div_rn(u, linear_beta), fastmath=False) + return (a * u).to(io_dtype) + + +@cute.kernel +def k3_situ_mul_kernel( + gu: cutlass.Array, # bf16 [M * 2K] + out: cutlass.Array, # bf16 [M * K] + num_tokens: cutlass.Int32, + beta: cutlass.Float32, + linear_beta: cutlass.Float32, + k_in: cutlass.Constexpr[int], + has_linear: cutlass.Constexpr[bool], +): + tx, _, _ = cute.arch.thread_idx() + bx, by, _ = cute.arch.block_idx() + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + prims.griddepcontrol(prims.GridDepAction.WAIT) + col = bx * cutlass.Int32(SITU_COLS_PER_CTA) + tx * cutlass.Int32(VEC) + if col < cutlass.Int32(k_in): + if by < num_tokens: + row = by * cutlass.Int32(2 * k_in) + gv = gu.load(idx=row + col, vector_size=VEC, alignment=16) + uv = gu.load(idx=row + cutlass.Int32(k_in) + col, vector_size=VEC, alignment=16) + outs = [ + _situ_mul( + cutlass.Float32(gv[e]), cutlass.Float32(uv[e]), beta, linear_beta, has_linear + ) + for e in range(VEC) + ] + out.store( + cutlass.Vector.from_elements(tuple(outs), io_dtype), + idx=by * cutlass.Int32(k_in) + col, vector_size=VEC, alignment=16, + ) # fmt: skip + + +@cute.jit +def k3_situ_mul( + gu: cute.Tensor, # bf16 [M * 2K] + out: cute.Tensor, # bf16 [M * K] + num_tokens: cutlass.Int32, + beta: cutlass.Float32, + linear_beta: cutlass.Float32, + max_tokens: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + has_linear: cutlass.Constexpr[bool], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + """``out = SituAndMul(beta, linear_beta)(gu)`` for rows < num_tokens (grid over max_tokens rows).""" + k3_situ_mul_kernel(gu, out, num_tokens, beta, linear_beta, k_in, has_linear).launch( + grid=((k_in + SITU_COLS_PER_CTA - 1) // SITU_COLS_PER_CTA, max_tokens, 1), + block=(SITU_THREADS, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_ctm_gemv/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_ctm_gemv/op.py new file mode 100644 index 000000000000..97ff2483e831 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_ctm_gemv/op.py @@ -0,0 +1,544 @@ +# 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. +"""Torch ops of the CTM decode GEMV (M <= 8, bf16, fp32 accumulation, K <= 768). + +``trtllm::k3_ctm_gemv``: ``x @ weight^T`` (same results as ``trtllm::k3_decode_gemv`` with ``split=1``). +``trtllm::k3_ctm_gemv_tail``: the row-parallel MoE tail (same results as ``trtllm::k3_decode_gemv_tail``). +``trtllm::k3_ctm_gemv_gated``: ``(a * sigmoid(g)) @ weight^T`` with torch's bf16 roundings, where ``g`` is a +column window of a dense bf16 matrix (the MLA output gate inside the fused q_a/kv_a/gate projection output), or +``a * s`` when that window already holds ``s = bf16(sigmoid(g))``. +``trtllm::k3_ctm_gemv_long``: ``x @ weight^T`` for long K (split-K over a 4-8 CTA cluster, weight ring and L2 +prefetch before the grid-dependency wait); output columns >= ``sig_col0`` hold ``bf16(sigmoid(bf16(.)))``. +``trtllm::k3_ctm_gemv_swiglu``: ``silu_and_mul(gu) @ weight^T`` with the activation in the B prologue. +``trtllm::k3_ctm_gemv_wide``: ``x @ weight^T`` for up to 64 tokens in one MMA of 16, 32 or 64 token columns, on the +long kernel's split-K clusters (bf16 output with optional sigmoid columns, or fp32 output). + +``push`` (the split-K ops) chooses how a cluster's ranks send their fp32 partials to the row owner: DSMEM +stores + a release arrive (False) or 16-byte st.async completing the owner's barrier by bytes (True). The +sums and their order are the same; each call site takes the one measured faster at its shape. + +Each kernel is compiled on the first call for its shape and flags, which must happen outside CUDA-graph +capture. The whole weight slice is loaded before the grid-dependency wait. +""" + +from __future__ import annotations + +import os +import threading +from typing import Dict, Optional + +import torch + +MAX_TOKENS = 8 +WIDE_MAX_TOKENS = 64 + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + + +def _arg(t: torch.Tensor): + from cutlass.cute.runtime import from_dlpack + + # detach(): DLPack refuses tensors that require grad (weights are parameters). + return from_dlpack(t.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +def _use_pdl() -> bool: + return os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +def _stream(t: torch.Tensor): + import cuda.bindings.driver as cuda_driver + + return cuda_driver.CUstream(torch.cuda.current_stream(t.device).cuda_stream) + + +def _compile(key, entry, *args): + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"{key[0]}: run once per shape outside CUDA-graph capture first (it compiles its kernel)." + ) + import cutlass.cute as cute + + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile(entry, *args) + return fn + + +def supports(x: torch.Tensor, weight: torch.Tensor, split: int = 1) -> bool: + """Whether ``k3_ctm_gemv`` runs ``x @ weight^T``.""" + from . import k3_ctm_gemv_kernel as kernel + + return ( + x.is_cuda + and x.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and x.dim() == 2 + and weight.dim() == 2 + and 0 < x.shape[0] <= MAX_TOKENS + and x.shape[1] == weight.shape[1] + and x.is_contiguous() + and weight.is_contiguous() + and kernel.supports(weight.shape[0], weight.shape[1], split) + ) + + +@torch.library.custom_op("trtllm::k3_ctm_gemv", mutates_args=()) +def k3_ctm_gemv( + x: torch.Tensor, + weight: torch.Tensor, + trigger_early: bool = True, + split: int = 1, + push: bool = False, +) -> torch.Tensor: + """``x @ weight.T`` for bf16 ``x`` [M <= 8, K <= 768] and ``weight`` [N, K]; returns bf16 [M, N]. + + ``split=2`` splits the k-tiles of each 128-row tile over a 2-CTA cluster (twice the CTAs streaming the + weight; the two fp32 partials are added before the bf16 rounding).""" + if not supports(x, weight, split): + raise ValueError( + f"k3_ctm_gemv: unsupported call x {tuple(x.shape)} {x.dtype}, weight {tuple(weight.shape)} " + f"{weight.dtype}, split {split}" + ) + from . import k3_ctm_gemv_kernel as kernel + + num_tokens, k_in = x.shape + n_out = weight.shape[0] + y = torch.empty(num_tokens, n_out, dtype=torch.bfloat16, device=x.device) + args = (_arg(weight), _arg(x), _arg(y.view(-1))) + stream = _stream(x) + use_pdl = _use_pdl() + fn = _compile( + ("k3_ctm_gemv", n_out, k_in, split, trigger_early, push, use_pdl), + kernel.k3_ctm_gemv, *args, num_tokens, n_out, k_in, split, trigger_early, push, use_pdl, stream, + ) # fmt: skip + fn(*args, num_tokens, stream) + return y + + +@k3_ctm_gemv.register_fake +def _(x, weight, trigger_early=True, split=1, push=False): + return x.new_empty((x.shape[0], weight.shape[0]), dtype=torch.bfloat16) + + +def supports_tail( + latent: torch.Tensor, act: torch.Tensor, weight: torch.Tensor, width: int +) -> bool: + """Whether ``k3_ctm_gemv_tail`` runs the call (the conditions of ``k3_decode_gemv``'s tail).""" + from . import k3_ctm_gemv_kernel as kernel + + k_act = act.shape[1] if act.dim() == 2 else -1 + k_lat = weight.shape[1] - k_act + return ( + latent.is_cuda + and latent.dtype == act.dtype == weight.dtype == torch.bfloat16 + and latent.dim() == 2 + and act.dim() == 2 + and weight.dim() == 2 + and 0 < latent.shape[0] <= MAX_TOKENS + and act.shape[0] == latent.shape[0] + and latent.is_contiguous() + and act.is_contiguous() + and weight.is_contiguous() + and k_lat % kernel.CTA_K == 0 + and k_act % kernel.CTA_K == 0 + and 0 < width <= k_lat + and latent.shape[1] % (8 * 32) == 0 + and kernel.supports(weight.shape[0], weight.shape[1]) + ) + + +@torch.library.custom_op("trtllm::k3_ctm_gemv_tail", mutates_args=()) +def k3_ctm_gemv_tail( + latent: torch.Tensor, + act: torch.Tensor, + weight: torch.Tensor, + lo: int, + width: int, + eps: float, + trigger_early: bool = True, +) -> torch.Tensor: + """``[rmsnorm(latent)[:, lo:lo+width] | act] @ weight.T`` with the RMS applied to the fp32 latent + accumulator (``weight`` = [latent-up columns of the slice, zero-padded to 128 | shared down]).""" + if not supports_tail(latent, act, weight, width): + raise ValueError( + f"k3_ctm_gemv_tail: unsupported call latent {tuple(latent.shape)}, act {tuple(act.shape)}, " + f"weight {tuple(weight.shape)}, width {width}" + ) + from . import k3_ctm_gemv_kernel as kernel + + num_tokens, rms_cols = latent.shape + n_out, k_in = weight.shape + lat_tiles = (k_in - act.shape[1]) // kernel.CTA_K + y = torch.empty(num_tokens, n_out, dtype=torch.bfloat16, device=latent.device) + args = ( + _arg(weight), + _arg(latent), + _arg(latent.view(-1).view(torch.int32)), + _arg(act), + _arg(y.view(-1)), + ) + stream = _stream(latent) + use_pdl = _use_pdl() + fn = _compile( + ("k3_ctm_gemv_tail", n_out, k_in, lat_tiles, rms_cols, trigger_early, use_pdl), + kernel.k3_ctm_gemv_tail, *args, num_tokens, lo, float(eps), n_out, k_in, lat_tiles, rms_cols, + trigger_early, use_pdl, stream, + ) # fmt: skip + fn(*args, num_tokens, lo, float(eps), stream) + return y + + +@k3_ctm_gemv_tail.register_fake +def _(latent, act, weight, lo, width, eps, trigger_early=True): + return latent.new_empty((latent.shape[0], weight.shape[0]), dtype=torch.bfloat16) + + +def supports_gated( + a: torch.Tensor, gsrc: torch.Tensor, g_col0: int, weight: torch.Tensor, split: int = 2 +) -> bool: + """Whether ``k3_ctm_gemv_gated`` runs ``(a * sigmoid(gsrc[:, g_col0 : g_col0 + K])) @ weight^T``.""" + from . import k3_ctm_gemv_kernel as kernel + + k_in = weight.shape[1] if weight.dim() == 2 else -1 + return ( + a.is_cuda + and a.dtype == gsrc.dtype == weight.dtype == torch.bfloat16 + and a.dim() == 2 + and gsrc.dim() == 2 + and weight.dim() == 2 + and 0 < a.shape[0] <= MAX_TOKENS + and gsrc.shape[0] == a.shape[0] + and a.shape[1] == k_in + and a.is_contiguous() + and gsrc.is_contiguous() + and weight.is_contiguous() + and g_col0 % 8 == 0 + and 0 <= g_col0 + and g_col0 + k_in <= gsrc.shape[1] + and gsrc.shape[1] % 8 == 0 + and kernel.supports(weight.shape[0], k_in, split) + ) + + +@torch.library.custom_op("trtllm::k3_ctm_gemv_gated", mutates_args=()) +def k3_ctm_gemv_gated( + a: torch.Tensor, + gsrc: torch.Tensor, + g_col0: int, + weight: torch.Tensor, + trigger_early: bool = True, + split: int = 2, + gate_sigmoid: bool = True, +) -> torch.Tensor: + """``(a * gsrc[:, g_col0:g_col0 + K].sigmoid()) @ weight.T``: bf16 ``a`` [M <= 8, K], ``gsrc`` [M, C] with dense + rows, ``weight`` [N, K]; the sigmoid and the product each round to bf16 as the unfused torch ops do. With + ``gate_sigmoid=False`` the window already holds the sigmoid: ``(a * gsrc[:, g_col0:g_col0 + K]) @ weight.T``.""" + if not supports_gated(a, gsrc, g_col0, weight, split): + raise ValueError( + f"k3_ctm_gemv_gated: unsupported call a {tuple(a.shape)}, gsrc {tuple(gsrc.shape)} at {g_col0}, " + f"weight {tuple(weight.shape)}, split {split}" + ) + from . import k3_ctm_gemv_kernel as kernel + + num_tokens, k_in = a.shape + n_out = weight.shape[0] + gsrc_cols = gsrc.shape[1] + y = torch.empty(num_tokens, n_out, dtype=torch.bfloat16, device=a.device) + args = (_arg(weight), _arg(a), _arg(gsrc), _arg(y.view(-1))) + stream = _stream(a) + use_pdl = _use_pdl() + fn = _compile( + ("k3_ctm_gemv_gated", n_out, k_in, gsrc_cols, split, gate_sigmoid, trigger_early, use_pdl), + kernel.k3_ctm_gemv_gated, *args, num_tokens, g_col0, n_out, k_in, gsrc_cols, split, gate_sigmoid, + trigger_early, use_pdl, stream, + ) # fmt: skip + fn(*args, num_tokens, g_col0, stream) + return y + + +@k3_ctm_gemv_gated.register_fake +def _(a, gsrc, g_col0, weight, trigger_early=True, split=2, gate_sigmoid=True): + return a.new_empty((a.shape[0], weight.shape[0]), dtype=torch.bfloat16) + + +def supports_swiglu(gu: torch.Tensor, weight: torch.Tensor, split: int = 2) -> bool: + """Whether ``k3_ctm_gemv_swiglu`` runs ``(silu(gu[:, :K]) * gu[:, K:]) @ weight^T``.""" + from . import k3_ctm_gemv_kernel as kernel + + k_in = weight.shape[1] if weight.dim() == 2 else -1 + return ( + gu.is_cuda + and gu.dtype == weight.dtype == torch.bfloat16 + and gu.dim() == 2 + and weight.dim() == 2 + and 0 < gu.shape[0] <= MAX_TOKENS + and gu.shape[1] == 2 * k_in + and gu.is_contiguous() + and weight.is_contiguous() + and kernel.supports(weight.shape[0], k_in, split) + ) + + +@torch.library.custom_op("trtllm::k3_ctm_gemv_swiglu", mutates_args=()) +def k3_ctm_gemv_swiglu( + gu: torch.Tensor, + weight: torch.Tensor, + trigger_early: bool = True, + split: int = 2, + push: bool = False, +) -> torch.Tensor: + """``silu_and_mul(gu) @ weight.T`` for a bf16 gate_up output ``gu`` [M <= 8, 2 K] (gate columns first) and + ``weight`` [N, K]: the activation is computed in the GEMV's B prologue with silu_and_mul's fp32 arithmetic and + one bf16 rounding. ``split`` k-tile ranks per 128-row tile (need not divide the k-tiles).""" + if not supports_swiglu(gu, weight, split): + raise ValueError( + f"k3_ctm_gemv_swiglu: unsupported call gu {tuple(gu.shape)} {gu.dtype}, weight {tuple(weight.shape)} " + f"{weight.dtype}, split {split}" + ) + from . import k3_ctm_gemv_kernel as kernel + + num_tokens = gu.shape[0] + n_out, k_in = weight.shape + y = torch.empty(num_tokens, n_out, dtype=torch.bfloat16, device=gu.device) + args = (_arg(weight), _arg(gu), _arg(y.view(-1))) + stream = _stream(gu) + use_pdl = _use_pdl() + fn = _compile( + ("k3_ctm_gemv_swiglu", n_out, k_in, split, trigger_early, push, use_pdl), + kernel.k3_ctm_gemv_swiglu, *args, num_tokens, n_out, k_in, split, trigger_early, push, use_pdl, stream, + ) # fmt: skip + fn(*args, num_tokens, stream) + return y + + +@k3_ctm_gemv_swiglu.register_fake +def _(gu, weight, trigger_early=True, split=2, push=False): + return gu.new_empty((gu.shape[0], weight.shape[0]), dtype=torch.bfloat16) + + +def supports_long(x: torch.Tensor, weight: torch.Tensor, split: int, ring: int) -> bool: + """Whether ``k3_ctm_gemv_long`` runs ``x @ weight^T`` with this split and ring.""" + from . import k3_ctm_gemv_kernel as kernel + + return ( + x.is_cuda + and x.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and x.dim() == 2 + and weight.dim() == 2 + and 0 < x.shape[0] <= MAX_TOKENS + and x.shape[1] == weight.shape[1] + and x.is_contiguous() + and weight.is_contiguous() + and kernel.long_supports(weight.shape[0], weight.shape[1], split, ring) + ) + + +@torch.library.custom_op("trtllm::k3_ctm_gemv_long", mutates_args=()) +def k3_ctm_gemv_long( + x: torch.Tensor, + weight: torch.Tensor, + sig_col0: int = -1, + split: int = 6, + ring: int = 5, + trigger_early: bool = True, + push: bool = False, +) -> torch.Tensor: + """``x @ weight.T`` for bf16 ``x`` [M <= 8, K] and ``weight`` [N, K] (long K); columns >= ``sig_col0`` (if >= 0) + hold ``bf16(sigmoid(bf16(x @ weight.T)))``. Each 128-row weight tile is split over a cluster of ``split`` CTAs, + each streaming its k-tiles through a ``ring``-stage ring filled before the grid-dependency wait.""" + return _launch_long(x, weight, sig_col0, split, ring, trigger_early, push=push) + + +def _launch_long(x, weight, sig_col0, split, ring, trigger_early, push=False) -> torch.Tensor: + if not supports_long(x, weight, split, ring): + raise ValueError( + f"k3_ctm_gemv_long: unsupported call x {tuple(x.shape)} {x.dtype}, weight {tuple(weight.shape)} " + f"{weight.dtype}, split {split}, ring {ring}" + ) + from . import k3_ctm_gemv_kernel as kernel + + num_tokens, k_in = x.shape + n_out = weight.shape[0] + y = torch.empty(num_tokens, n_out, dtype=torch.bfloat16, device=x.device) + args = (_arg(weight), _arg(x), _arg(y.view(-1))) + stream = _stream(x) + use_pdl = _use_pdl() + sig_row0 = sig_col0 if sig_col0 >= 0 else n_out + fn = _compile( + ("k3_ctm_gemv_long", n_out, k_in, split, ring, trigger_early, push, use_pdl), + kernel.k3_ctm_gemv_long, *args, num_tokens, sig_row0, n_out, k_in, split, ring, trigger_early, push, use_pdl, + stream, + ) # fmt: skip + fn(*args, num_tokens, sig_row0, stream) + return y + + +@k3_ctm_gemv_long.register_fake +def _(x, weight, sig_col0=-1, split=6, ring=5, trigger_early=True, push=False): + return x.new_empty((x.shape[0], weight.shape[0]), dtype=torch.bfloat16) + + +def wide_config(n_out: int, k_in: int, n_tile: int, num_sms: int) -> Optional[tuple]: + """``(split, ring, x_ring)`` of a ``k3_ctm_gemv_wide`` call, or None: the largest cluster whose CTAs fit one wave, + then the deepest weight ring that fits shared memory beside an x ring of up to 3 stages. At the call sites that run + ``k3_ctm_gemv_long`` at most 8 tokens (MLA [W_a; W_g] 6, dense gate_up 4, dense down 2, drafter 8) this is their + split, so a token's sums are the same.""" + from . import k3_ctm_gemv_kernel as kernel + + tiles = (n_out + kernel.CTA_M - 1) // kernel.CTA_M + for split in (8, 7, 6, 5, 4, 2): + if tiles * split > num_sms: + continue + for ring in range(kernel.num_k_tiles(k_in) // split, 0, -1): + if kernel.wide_supports(n_out, k_in, split, ring, n_tile, min(ring, 3)): + return split, ring, min(ring, 3) + return None + + +def _num_sms(device: torch.device) -> int: + return torch.cuda.get_device_properties(device).multi_processor_count + + +def supports_wide( + x: torch.Tensor, weight: torch.Tensor, sig_col0: int = -1, out_fp32: bool = False +) -> bool: + """Whether ``k3_ctm_gemv_wide`` runs ``x @ weight^T``.""" + from . import k3_ctm_gemv_kernel as kernel + + if not ( + x.is_cuda + and x.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and x.dim() == 2 + and weight.dim() == 2 + and 0 < x.shape[0] <= WIDE_MAX_TOKENS + and x.shape[1] == weight.shape[1] + and x.is_contiguous() + and weight.is_contiguous() + and x.data_ptr() % 16 == 0 + and weight.data_ptr() % 16 == 0 + and sig_col0 < weight.shape[0] + and (sig_col0 < 0 or not out_fp32) + ): + return False + n_tile = kernel.wide_tile(x.shape[0]) + return wide_config(weight.shape[0], weight.shape[1], n_tile, _num_sms(x.device)) is not None + + +@torch.library.custom_op("trtllm::k3_ctm_gemv_wide", mutates_args=()) +def k3_ctm_gemv_wide( + x: torch.Tensor, weight: torch.Tensor, sig_col0: int = -1, out_fp32: bool = False +) -> torch.Tensor: + """``x @ weight.T`` for bf16 ``x`` [M <= 64, K] and ``weight`` [N, K], all M tokens in one MMA of N = 16, 32 or 64 + columns (the long kernel's split-K clusters and weight ring, the partials reduced by token; split and ring from + ``wide_config``). Returns bf16 + [M, N], whose columns >= ``sig_col0`` (if >= 0) hold ``bf16(sigmoid(bf16(x @ weight.T)))``, or fp32 [M, N] with + ``out_fp32``.""" + if not supports_wide(x, weight, sig_col0, out_fp32): + raise ValueError( + f"k3_ctm_gemv_wide: unsupported call x {tuple(x.shape)} {x.dtype}, weight {tuple(weight.shape)} " + f"{weight.dtype}, sig_col0 {sig_col0}, out_fp32 {out_fp32}" + ) + from . import k3_ctm_gemv_kernel as kernel + + n_tile = kernel.wide_tile(x.shape[0]) + split, ring, x_ring = wide_config(weight.shape[0], weight.shape[1], n_tile, _num_sms(x.device)) + return _launch_wide(x, weight, sig_col0, out_fp32, split, ring, x_ring) + + +def _launch_wide( + x, weight, sig_col0, out_fp32, split, ring, x_ring, trigger_early=True +) -> torch.Tensor: + from . import k3_ctm_gemv_kernel as kernel + + num_tokens, k_in = x.shape + n_out = weight.shape[0] + n_tile = kernel.wide_tile(num_tokens) + if not kernel.wide_supports(n_out, k_in, split, ring, n_tile, x_ring): + raise ValueError( + f"k3_ctm_gemv_wide: split {split}, rings {ring} / {x_ring} do not fit weight {tuple(weight.shape)} at " + f"{n_tile} tokens" + ) + y = torch.empty( + num_tokens, n_out, dtype=torch.float32 if out_fp32 else torch.bfloat16, device=x.device + ) + args = (_arg(weight), _arg(x), _arg(y.view(-1))) + stream = _stream(x) + use_pdl = _use_pdl() + sig_row0 = sig_col0 if sig_col0 >= 0 else n_out + fn = _compile( + ("k3_ctm_gemv_wide", n_out, k_in, split, ring, x_ring, n_tile, out_fp32, trigger_early, use_pdl), + kernel.k3_ctm_gemv_wide, *args, num_tokens, sig_row0, n_out, k_in, split, ring, x_ring, n_tile, out_fp32, + trigger_early, use_pdl, stream, + ) # fmt: skip + fn(*args, num_tokens, sig_row0, stream) + return y + + +@k3_ctm_gemv_wide.register_fake +def _(x, weight, sig_col0=-1, out_fp32=False): + return x.new_empty( + (x.shape[0], weight.shape[0]), dtype=torch.float32 if out_fp32 else torch.bfloat16 + ) + + +def supports_situ_mul(gu: torch.Tensor) -> bool: + """Whether ``k3_situ_mul`` runs SituAndMul on ``gu`` [M <= 8, 2 K] (K a multiple of 8, rows dense).""" + return ( + gu.is_cuda + and gu.dtype == torch.bfloat16 + and gu.dim() == 2 + and 0 < gu.shape[0] <= MAX_TOKENS + and gu.shape[1] % 16 == 0 + and gu.is_contiguous() + and gu.data_ptr() % 16 == 0 + ) + + +@torch.library.custom_op("trtllm::k3_situ_mul", mutates_args=()) +def k3_situ_mul( + gu: torch.Tensor, beta: float = 1.0, linear_beta: Optional[float] = None +) -> torch.Tensor: + """``SituAndMul(beta, linear_beta)(gu)`` for a bf16 gate_up output ``gu`` [M <= 8, 2 K] (gate columns first), with + programmatic dependent launch: the next kernel launches at once and may stream its weights meanwhile.""" + if not supports_situ_mul(gu): + raise ValueError(f"k3_situ_mul: unsupported call gu {tuple(gu.shape)} {gu.dtype}") + from . import k3_ctm_gemv_kernel as kernel + + num_tokens, two_k = gu.shape + k_in = two_k // 2 + out = torch.empty(num_tokens, k_in, dtype=torch.bfloat16, device=gu.device) + args = (_arg(gu.view(-1)), _arg(out.view(-1))) + stream = _stream(gu) + use_pdl = _use_pdl() + has_linear = linear_beta is not None + fn = _compile( + ("k3_situ_mul", k_in, has_linear, use_pdl), + kernel.k3_situ_mul, *args, num_tokens, float(beta), float(linear_beta or 1.0), MAX_TOKENS, k_in, has_linear, + use_pdl, stream, + ) # fmt: skip + fn(*args, num_tokens, float(beta), float(linear_beta or 1.0), stream) + return out + + +@k3_situ_mul.register_fake +def _(gu, beta=1.0, linear_beta=None): + return gu.new_empty((gu.shape[0], gu.shape[1] // 2)) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_decode_gemv/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_decode_gemv/__init__.py new file mode 100644 index 000000000000..cb3670f60918 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_decode_gemv/__init__.py @@ -0,0 +1,19 @@ +# 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 decode GEMV in CuTe DSL (``trtllm::k3_decode_gemv``). + +Importing :mod:`.op` registers the torch op; nothing is imported eagerly here so that +the CuTe DSL dependency stays optional for every other model. +""" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_decode_gemv/k3_decode_gemv_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_decode_gemv/k3_decode_gemv_kernel.py new file mode 100644 index 000000000000..430dd42d7cd3 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_decode_gemv/k3_decode_gemv_kernel.py @@ -0,0 +1,599 @@ +# 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. +"""Decode GEMV for Kimi K3 projections: ``y[M, N] = x[M, K] @ W[N, K]^T``, M <= 8, bf16 in and out. + +tcgen05 with the weight as the MMA's A operand (128 output features per CTA) and the activation as B +(8 token columns), fp32 accumulation in TMEM. For the short-K projections (K <= 6 * 128: the MoE +tail, the attention o_proj) one CTA owns one 128-row weight tile over the whole K extent, and every +weight k-tile of that CTA gets its own shared-memory stage. The weight TMA for all stages is issued +at launch, before ``griddepcontrol.wait``: launched early under PDL, the kernel streams its whole +weight while its predecessor still runs, and after the wait only the activation (8 x K) remains to +load. The activation box is always 8 rows; rows past M are zero-filled by the TMA, so one compiled +kernel serves every M <= 8. + +Warps: 0 weight TMA, 1 activation TMA (after the grid dependency), 2 TMEM allocation and MMA, 3 idle, +4-7 epilogue (TMEM -> registers -> bf16 stores; warp 4 + i reads TMEM lanes 32i..32i+31). +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass.experimental import primitives as prims + +CTA_M = 128 # weight rows (output features) per CTA = the MMA's M +MMA_N = 8 # token columns +CTA_K = 128 # one k-tile: two 64-element halves of the 128-byte swizzle +MMA_K = 16 +TMA_K_BOX = 64 +TMA_COPY_ITERS = CTA_K // TMA_K_BOX +K_BLOCKS_PER_HALF = TMA_K_BOX // MMA_K +MAX_K_TILES = 6 # 6 x (32 KB weight + 2 KB activation) of shared memory +THREADS = 256 +TMEM_COLS = 32 +ELEM_BYTES = 2 +EVICT_FIRST = 0x12F0000000000000 # createpolicy.fractional.L2::evict_first, fraction 1.0 (sm_100) + +# Shared-memory descriptor strides for the 128-byte swizzle, in 16-byte units. +LEADING = 16 +STRIDE = 8 * TMA_K_BOX * ELEM_BYTES +A_HALF_ELEMS = CTA_M * TMA_K_BOX +B_HALF_ELEMS = MMA_N * TMA_K_BOX +STEP = (MMA_K * ELEM_BYTES) >> 4 +A_BOX = A_HALF_ELEMS >> 3 +B_BOX = B_HALF_ELEMS >> 3 +STAGE_A = (CTA_M * CTA_K * ELEM_BYTES) >> 4 +STAGE_B = (MMA_N * CTA_K * ELEM_BYTES) >> 4 + +io_dtype = cutlass.BFloat16 + + +def num_k_tiles(k_in: int) -> int: + return (k_in + CTA_K - 1) // CTA_K + + +def supports(n_out: int, k_in: int) -> bool: + """Shapes this kernel runs: whole 128-row tiles and at most MAX_K_TILES k-tiles.""" + return n_out % CTA_M == 0 and k_in % TMA_K_BOX == 0 and num_k_tiles(k_in) <= MAX_K_TILES + + +@cute.kernel +def k3_decode_gemv_kernel( + tma_desc_w: cutlass.GridConstant[cuda.TensorMap], # W [N, K] bf16, 5-D, one call per k-tile + tma_desc_x: cutlass.GridConstant[ + cuda.TensorMap + ], # x [M, *] bf16, box 8 x 64: k-tiles [0, x_tiles) + tma_desc_x2: cutlass.GridConstant[ + cuda.TensorMap + ], # x2 [M, *] bf16, box 8 x 64: k-tiles [x_tiles, K/128) + y: cutlass.Array, # [M * N] bf16, token-major + rms_src: cutlass.Array, # int32 words of the bf16 [M, rms_cols] rows whose RMS scales accumulator 0 + num_tokens: cutlass.Int32, + x_col0: cutlass.Int32, # column of x where k-tile 0 starts + eps: cutlass.Float32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + x_tiles: cutlass.Constexpr[int], + split_acc: cutlass.Constexpr[bool], # the x2 k-tiles accumulate into a second TMEM accumulator + rms_cols: cutlass.Constexpr[int], # > 0: y = rsqrt(mean(rms_src row^2) + eps) * acc0 + acc1 + trigger_early: cutlass.Constexpr[bool], +): + k_tiles = num_k_tiles(k_in) + tx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp_id = cute.arch.warp_idx() + tma_ptr_w = tma_desc_w.get_ptr() + tma_ptr_x = tma_desc_x.get_ptr() + tma_ptr_x2 = tma_desc_x2.get_ptr() + m_offset = bx * cutlass.Int32(CTA_M) + + smem_a = cutlass.Array( + io_dtype, k_tiles * CTA_M * CTA_K, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_b = cutlass.Array( + io_dtype, k_tiles * MMA_N * CTA_K, space=cutlass.AddressSpace.smem, alignment=1024 + ) + weight_full = cutlass.Array( + cutlass.Int64, k_tiles, space=cutlass.AddressSpace.smem, alignment=8 + ) + act_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + acc_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + row_scale = cutlass.Array(cutlass.Float32, MMA_N, space=cutlass.AddressSpace.smem, alignment=16) + + if warp_id == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_x) + if cutlass.const_expr(x_tiles < k_tiles): + prims.prefetch_tensormap(tma_ptr_x2) + if prims.elect_sync(): + for k in cutlass.range_constexpr(k_tiles): + prims.mbarrier_init(weight_full.subview(k), 1) + prims.mbarrier_init(act_full, 1) + prims.mbarrier_init(acc_done, 1) + if warp_id == 2: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + prims.barrier_cta_sync(0) + tmem_ptr = cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32) + + if warp_id == 0: + # The whole weight slice of this CTA, ahead of the grid dependency. + if prims.elect_sync(): + for k in cutlass.range_constexpr(k_tiles): + prims.mbarrier_arrive_expect_tx(weight_full.subview(k), CTA_M * CTA_K * ELEM_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(k * CTA_M * CTA_K), + tma_ptr_w, + ( + cutlass.Int32(0), + m_offset, + cutlass.Int32(k * TMA_COPY_ITERS), + cutlass.Int32(0), + cutlass.Int32(0), + ), + weight_full.subview(k), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(trigger_early): + # Dependents may launch now; they wait for this whole grid before reading y. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + elif warp_id == 1: + # The activations are written by the predecessor. + prims.griddepcontrol(prims.GridDepAction.WAIT) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(act_full, k_tiles * MMA_N * CTA_K * ELEM_BYTES) + for k in cutlass.range_constexpr(k_tiles): + for half in cutlass.range_constexpr(TMA_COPY_ITERS): + if cutlass.const_expr(k < x_tiles): + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(k * MMA_N * CTA_K + half * B_HALF_ELEMS), + tma_ptr_x, + ( + x_col0 + cutlass.Int32(k * CTA_K + half * TMA_K_BOX), + cutlass.Int32(0), + ), + act_full, + ) + else: + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(k * MMA_N * CTA_K + half * B_HALF_ELEMS), + tma_ptr_x2, + ( + cutlass.Int32((k - x_tiles) * CTA_K + half * TMA_K_BOX), + cutlass.Int32(0), + ), + act_full, + ) + elif warp_id == 2: + idesc = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, a_dtype=io_dtype, b_dtype=io_dtype, n_dim=MMA_N, m_dim=CTA_M + ) + desc_a_base = prims.Tcgen05SmemDesc.build( + start_address=smem_a, leading_byte_offset=LEADING, stride_byte_offset=STRIDE, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + desc_b_base = prims.Tcgen05SmemDesc.build( + start_address=smem_b, leading_byte_offset=LEADING, stride_byte_offset=STRIDE, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + tmem_acc1 = tmem_ptr + if cutlass.const_expr(split_acc): + tmem_acc1 = cutlass.inttoptr( + tmem_ptr_i32.load() + cutlass.Int32(MMA_N), 6, cutlass.Int32 + ) + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + for k in cutlass.range_constexpr(k_tiles): + second = split_acc and k >= x_tiles + first_tile = k == 0 or (split_acc and k == x_tiles) + while not cute.arch.mbarrier_try_wait(weight_full.subview(k).data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for kb in cutlass.range_constexpr(TMA_COPY_ITERS * K_BLOCKS_PER_HALF): + box = kb // K_BLOCKS_PER_HALF + within = kb % K_BLOCKS_PER_HALF + desc_a = desc_a_base + (k * STAGE_A + box * A_BOX + within * STEP) + desc_b = desc_b_base + (k * STAGE_B + box * B_BOX + within * STEP) + if prims.elect_sync(): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, prims.CTAGroup.CTA_1, tmem_acc1 if second else tmem_ptr, + desc_a, desc_b, idesc, not (first_tile and kb == 0), + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + elif warp_id >= 4: + lane = tx % 32 + if cutlass.const_expr(rms_cols > 0): + # Per-token RMS of the rms_src rows while the MMA runs: epilogue warp w owns tokens 2w, 2w+1. + prims.griddepcontrol(prims.GridDepAction.WAIT) + for j in cutlass.range_constexpr(MMA_N // 4): + t = (warp_id % 4) * cutlass.Int32(MMA_N // 4) + cutlass.Int32(j) + sum_sq = cutlass.Float32(0.0) + if t < num_tokens: + for v in cutlass.range_constexpr(rms_cols // (8 * 32)): + words = rms_src.load( + idx=t * cutlass.Int32(rms_cols // 2) + + cutlass.Int32(v * 32 * 4) + + lane * cutlass.Int32(4), + vector_size=4, + alignment=16, + ) + for i in cutlass.range_constexpr(4): + lo = (words[i] << cutlass.Int32(16)).bitcast(cutlass.Float32) + hi = (words[i] & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + sum_sq = sum_sq + lo * lo + hi * hi + for offset in (16, 8, 4, 2, 1): + sum_sq = sum_sq + cute.arch.shuffle_sync_bfly(sum_sq, offset=offset) + if lane == 0: + row_scale.store( + cute.math.rsqrt(sum_sq * cutlass.Float32(1.0 / rms_cols) + eps), idx=t + ) + prims.barrier_cta_sync(1, thread_count=128) + while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc0 = prims.tcgen05_ld( + "32x32b", cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), num=MMA_N + ) + acc1 = acc0 + if cutlass.const_expr(split_acc): + acc1 = prims.tcgen05_ld( + "32x32b", + cutlass.inttoptr(tmem_ptr_i32.load() + cutlass.Int32(MMA_N), 6, cutlass.Float32), + num=MMA_N, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + row = (warp_id % 4) * 32 + lane + n = m_offset + row + for t in cutlass.range_constexpr(MMA_N): + if cutlass.Int32(t) < num_tokens: + value = cutlass.Float32(acc0[t]) + if cutlass.const_expr(rms_cols > 0): + value = value * row_scale.load(idx=t) + if cutlass.const_expr(split_acc): + value = value + cutlass.Float32(acc1[t]) + y.store(cutlass.BFloat16(value), idx=cutlass.Int32(t * n_out) + n) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.barrier_cta_sync(1, thread_count=128) + if warp_id == 4: + prims.tcgen05_dealloc(tmem_ptr, TMEM_COLS) + + +# ============================================================================= +# Long-K split-K variant (e.g. the KDA qkvg projection, 3208 x 7168): the fuse_o_linear_ra structure. +# A cluster of SPLIT CTAs shares one 128-row weight tile; rank r streams the interleaved k-tiles r, +# r + SPLIT, ... through a RING_STAGES ring. Before the grid dependency each CTA fills its ring from HBM and +# prefetches the rest of its k-tiles into L2. Rows [32r, 32r + 32) of the tile are owned by rank r: every lane +# of the other ranks' epilogue warp r stores its fp32 partials into rank r's shared memory and arrives on its +# barrier (release, cluster scope); the owner waits (acquire, cluster scope), adds the SPLIT partials in rank +# order and stores bf16. (fuse_o_linear_ra pushes with st.async, which this DSL build does not expose.) +# ============================================================================= +SPLIT = 4 +RING_STAGES = 6 +ROWS_PER_RANK = CTA_M // SPLIT + + +def splitk_supports(n_out: int, k_in: int) -> bool: + k_tiles = num_k_tiles(k_in) + return k_in % CTA_K == 0 and k_tiles > MAX_K_TILES and k_tiles % SPLIT == 0 and n_out > 0 + + +@cute.kernel +def k3_decode_gemv_splitk_kernel( + tma_desc_w: cutlass.GridConstant[cuda.TensorMap], + tma_desc_x: cutlass.GridConstant[cuda.TensorMap], + y: cutlass.Array, # [M * N] bf16, token-major + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], +): + k_tiles = num_k_tiles(k_in) + my_tiles = k_tiles // SPLIT + tx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp_id = cute.arch.warp_idx() + rank = cute.arch.block_idx_in_cluster() + tma_ptr_w = tma_desc_w.get_ptr() + tma_ptr_x = tma_desc_x.get_ptr() + m_offset = (bx // cutlass.Int32(SPLIT)) * cutlass.Int32(CTA_M) + + smem_a = cutlass.Array( + io_dtype, RING_STAGES * CTA_M * CTA_K, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_b = cutlass.Array( + io_dtype, RING_STAGES * MMA_N * CTA_K, space=cutlass.AddressSpace.smem, alignment=1024 + ) + tma_done = cutlass.Array( + cutlass.Int64, RING_STAGES, space=cutlass.AddressSpace.smem, alignment=8 + ) + mma_done = cutlass.Array( + cutlass.Int64, RING_STAGES, space=cutlass.AddressSpace.smem, alignment=8 + ) + acc_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + reduce_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + # Mailbox of the rows this rank owns: [source rank][row][token] fp32 (its own slot stays unused). + mailbox = cutlass.Array( + cutlass.Float32, + SPLIT * ROWS_PER_RANK * MMA_N, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + + if warp_id == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_x) + if prims.elect_sync(): + for s in cutlass.range_constexpr(RING_STAGES): + prims.mbarrier_init(tma_done.subview(s), 2) # the weight and the activation TMA + prims.mbarrier_init(mma_done.subview(s), 1) + prims.mbarrier_init(acc_done, 1) + prims.mbarrier_init(reduce_ready, (SPLIT - 1) * 32) # every lane of the 3 pushing warps + if warp_id == 2: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + # Cluster formation: makes the peers' shared memory and barriers addressable. + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + tmem_ptr = cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32) + + if warp_id == 0: + if prims.elect_sync(): + # k-tiles past the ring go to L2 now, so the refills after the dependency read L2. + for i in cutlass.range_constexpr(RING_STAGES, my_tiles): + k_global = rank + cutlass.Int32(i * SPLIT) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, + [ + cutlass.Int32(0), + m_offset, + k_global * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), + cutlass.Int32(0), + ], + [], # tile mode: no im2col offsets + ) + for i in cutlass.range_constexpr(my_tiles): + stage = i % RING_STAGES + if cutlass.const_expr(i >= RING_STAGES): + while not cute.arch.mbarrier_try_wait( + mma_done.subview(stage).data_ptr(), (i // RING_STAGES - 1) % 2 + ): + pass + k_global = rank + cutlass.Int32(i * SPLIT) + prims.mbarrier_arrive_expect_tx(tma_done.subview(stage), CTA_M * CTA_K * ELEM_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(stage * CTA_M * CTA_K), + tma_ptr_w, + ( + cutlass.Int32(0), + m_offset, + k_global * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), + cutlass.Int32(0), + ), + tma_done.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(trigger_early and i == RING_STAGES - 1): + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if cutlass.const_expr(trigger_early and my_tiles < RING_STAGES): + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + elif warp_id == 1: + prims.griddepcontrol(prims.GridDepAction.WAIT) + if prims.elect_sync(): + for i in cutlass.range_constexpr(my_tiles): + stage = i % RING_STAGES + if cutlass.const_expr(i >= RING_STAGES): + while not cute.arch.mbarrier_try_wait( + mma_done.subview(stage).data_ptr(), (i // RING_STAGES - 1) % 2 + ): + pass + k_global = rank + cutlass.Int32(i * SPLIT) + prims.mbarrier_arrive_expect_tx(tma_done.subview(stage), MMA_N * CTA_K * ELEM_BYTES) + for half in cutlass.range_constexpr(TMA_COPY_ITERS): + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(stage * MMA_N * CTA_K + half * B_HALF_ELEMS), + tma_ptr_x, + ( + k_global * cutlass.Int32(CTA_K) + cutlass.Int32(half * TMA_K_BOX), + cutlass.Int32(0), + ), + tma_done.subview(stage), + ) + elif warp_id == 2: + idesc = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, a_dtype=io_dtype, b_dtype=io_dtype, n_dim=MMA_N, m_dim=CTA_M + ) + desc_a_base = prims.Tcgen05SmemDesc.build( + start_address=smem_a, leading_byte_offset=LEADING, stride_byte_offset=STRIDE, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + desc_b_base = prims.Tcgen05SmemDesc.build( + start_address=smem_b, leading_byte_offset=LEADING, stride_byte_offset=STRIDE, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + for i in cutlass.range_constexpr(my_tiles): + stage = i % RING_STAGES + while not cute.arch.mbarrier_try_wait( + tma_done.subview(stage).data_ptr(), (i // RING_STAGES) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for kb in cutlass.range_constexpr(TMA_COPY_ITERS * K_BLOCKS_PER_HALF): + box = kb // K_BLOCKS_PER_HALF + within = kb % K_BLOCKS_PER_HALF + desc_a = desc_a_base + (stage * STAGE_A + box * A_BOX + within * STEP) + desc_b = desc_b_base + (stage * STAGE_B + box * B_BOX + within * STEP) + if prims.elect_sync(): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, prims.CTAGroup.CTA_1, tmem_ptr, desc_a, desc_b, idesc, + i > 0 or kb > 0, + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_commit(mma_done.subview(stage)) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + elif warp_id >= 4: + lane = tx % 32 + owner = warp_id % 4 # TMEM lanes 32 * owner .. + 31 = the rows rank `owner` reduces + while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), num=MMA_N + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + slot_row = lane * MMA_N + if owner != rank: + base = rank * cutlass.Int32(ROWS_PER_RANK * MMA_N) + slot_row + for t in cutlass.range_constexpr(MMA_N): + prims.mapa(mailbox.subview(base + cutlass.Int32(t)), owner).store( + cutlass.Float32(acc[t]) + ) + # Release at cluster scope: this lane's partials are visible to the owner's acquire. + prims.mbarrier_arrive(prims.mapa(reduce_ready, owner), scope=prims.MemScope.CLUSTER) + else: + while not prims.mbarrier_try_wait_parity( + reduce_ready, 0, scope=prims.MBarrierScope.CLUSTER + ): + pass + n = m_offset + owner * cutlass.Int32(ROWS_PER_RANK) + lane + for t in cutlass.range_constexpr(MMA_N): + total = cutlass.Float32(0.0) + for q in cutlass.range_constexpr(SPLIT): + part = mailbox.load( + idx=cutlass.Int32(q * ROWS_PER_RANK * MMA_N) + slot_row + cutlass.Int32(t) + ) + total = total + cutlass.Float32( + cutlass.select_(rank == cutlass.Int32(q), cutlass.Float32(acc[t]), part) + ) + if cutlass.Int32(t) < num_tokens: + if n < cutlass.Int32(n_out): + y.store(cutlass.BFloat16(total), idx=cutlass.Int32(t * n_out) + n) + prims.barrier_cta_sync(1, thread_count=128) + if warp_id == 4: + prims.tcgen05_dealloc(tmem_ptr, TMEM_COLS) + + +def _weight_tensor_map(w, n_out, k_in): + """W as five TMA dimensions (64-element column chunk, row, 64-element chunk index, 1, 1) so one call + per k-tile lands both 128-byte-swizzled halves; strides in 16-byte units.""" + return cuda.create_tensor_map_tiled( + global_address=w.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[TMA_K_BOX, n_out, k_in // TMA_K_BOX, 1, 1], + global_strides=[ + (k_in * ELEM_BYTES) // 16, + (TMA_K_BOX * ELEM_BYTES) // 16, + (n_out * k_in * ELEM_BYTES) // 16, + (n_out * k_in * ELEM_BYTES) // 16, + ], + box_dims=[TMA_K_BOX, CTA_M, TMA_COPY_ITERS, 1, 1], + swizzle=cuda.TensorMapSwizzle.s128b, + ) + + +def _activation_tensor_map(x, cols, num_tokens): + """x [M, cols] as (cols, M) with an 8-row box: rows past num_tokens arrive as zeros.""" + return cuda.create_tensor_map_tiled( + global_address=x.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[cols, num_tokens], + global_strides=[(cols * ELEM_BYTES) // 16], + box_dims=[TMA_K_BOX, MMA_N], + swizzle=cuda.TensorMapSwizzle.s128b, + ) + + +@cute.jit +def k3_decode_gemv( + w: cute.Tensor, # [N, K] bf16, K contiguous + x: cute.Tensor, # [M, K] bf16, K contiguous, M <= 8 + y: cute.Tensor, # [M * N] bf16 + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + tma_desc_w = _weight_tensor_map(w, n_out, k_in) + tma_desc_x = _activation_tensor_map(x, k_in, num_tokens) + k3_decode_gemv_kernel( + tma_desc_w, tma_desc_x, tma_desc_x, y, y, num_tokens, cutlass.Int32(0), cutlass.Float32(0.0), + n_out, k_in, num_k_tiles(k_in), False, 0, trigger_early, + ).launch(grid=(n_out // CTA_M, 1, 1), block=(THREADS, 1, 1), stream=stream, use_pdl=use_pdl) # fmt: skip + + +@cute.jit +def k3_decode_gemv_tail( + w: cute.Tensor, # [N, K_lat + K_act] bf16: latent-up columns of this rank's slice (zero-padded) | shared down + latent: cute.Tensor, # [M, rms_cols] bf16, the whole reduced latent row + latent_words: cute.Tensor, # the same memory as int32 words, for the RMS + act: cute.Tensor, # [M, K_act] bf16, the shared-expert activation + y: cute.Tensor, # [M * N] bf16 + num_tokens: cutlass.Int32, + lat_col0: cutlass.Int32, # first latent column of this rank's slice + eps: cutlass.Float32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + lat_tiles: cutlass.Constexpr[int], + rms_cols: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + """``y = rmsnorm(latent)[:, slice] @ W_lat^T + act @ W_act^T`` (the RMS applied to the latent accumulator).""" + tma_desc_w = _weight_tensor_map(w, n_out, k_in) + tma_desc_lat = _activation_tensor_map(latent, rms_cols, num_tokens) + tma_desc_act = _activation_tensor_map(act, k_in - lat_tiles * CTA_K, num_tokens) + k3_decode_gemv_kernel( + tma_desc_w, tma_desc_lat, tma_desc_act, y, latent_words, num_tokens, lat_col0, eps, + n_out, k_in, lat_tiles, True, rms_cols, trigger_early, + ).launch(grid=(n_out // CTA_M, 1, 1), block=(THREADS, 1, 1), stream=stream, use_pdl=use_pdl) # fmt: skip + + +@cute.jit +def k3_decode_gemv_splitk( + w: cute.Tensor, # [N, K] bf16, K contiguous, K / 128 a multiple of SPLIT + x: cute.Tensor, # [M, K] bf16 + y: cute.Tensor, # [M * N] bf16 + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + tma_desc_w = _weight_tensor_map(w, n_out, k_in) + tma_desc_x = _activation_tensor_map(x, k_in, num_tokens) + k3_decode_gemv_splitk_kernel( + tma_desc_w, tma_desc_x, y, num_tokens, n_out, k_in, trigger_early + ).launch( + grid=(((n_out + CTA_M - 1) // CTA_M) * SPLIT, 1, 1), + block=(THREADS, 1, 1), + cluster=(SPLIT, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_decode_gemv/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_decode_gemv/op.py new file mode 100644 index 000000000000..1e816ca8d964 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_decode_gemv/op.py @@ -0,0 +1,209 @@ +# 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::k3_decode_gemv``: ``x[M, K] @ weight[N, K]^T`` for M <= 8 in CuTe DSL (bf16, fp32 accumulation). + +The kernel loads its whole weight slice before waiting for the producer of ``x``, so launched with +PDL behind a long predecessor it only pays for the activation load, the MMA and the store. It is +compiled on the first call for each (N, K, early trigger, PDL), which must happen outside CUDA-graph +capture. +""" + +from __future__ import annotations + +import os +import threading +from typing import Dict + +import torch + +MAX_TOKENS = 8 + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + + +def _arg(t: torch.Tensor): + from cutlass.cute.runtime import from_dlpack + + # detach(): DLPack refuses tensors that require grad (weights are parameters). + return from_dlpack(t.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +def _use_pdl() -> bool: + return os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +def supports(x: torch.Tensor, weight: torch.Tensor) -> bool: + """Whether the kernel runs this call (shapes, dtypes, layout).""" + from . import k3_decode_gemv_kernel as kernel + + return ( + x.is_cuda + and x.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and x.dim() == 2 + and weight.dim() == 2 + and 0 < x.shape[0] <= MAX_TOKENS + and x.shape[1] == weight.shape[1] + and x.is_contiguous() + and weight.is_contiguous() + and ( + kernel.supports(weight.shape[0], weight.shape[1]) + or kernel.splitk_supports(weight.shape[0], weight.shape[1]) + ) + ) + + +@torch.library.custom_op("trtllm::k3_decode_gemv", mutates_args=()) +def k3_decode_gemv( + x: torch.Tensor, weight: torch.Tensor, trigger_early: bool = True +) -> torch.Tensor: + """``x @ weight.T`` for bf16 ``x`` [M <= 8, K] and ``weight`` [N, K]; returns bf16 [M, N]. + + ``trigger_early`` lets the dependent grid launch once every CTA has issued its weight loads + (for dependents that wait for this whole grid before reading the output).""" + import cuda.bindings.driver as cuda_driver + + if not supports(x, weight): + raise ValueError( + f"k3_decode_gemv: unsupported call x {tuple(x.shape)} {x.dtype}, weight {tuple(weight.shape)} " + f"{weight.dtype}" + ) + num_tokens, k_in = x.shape + n_out = weight.shape[0] + y = torch.empty(num_tokens, n_out, dtype=torch.bfloat16, device=x.device) + args = (_arg(weight), _arg(x), _arg(y.view(-1))) + stream = cuda_driver.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + use_pdl = _use_pdl() + key = (n_out, k_in, trigger_early, use_pdl) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "trtllm::k3_decode_gemv must run once per shape outside CUDA-graph capture first " + "(it compiles its kernel on the first call)." + ) + import cutlass.cute as cute + + from . import k3_decode_gemv_kernel as kernel + + # Short K: one CTA per weight tile, whole slice resident; long K: split-K over a cluster. + entry = ( + kernel.k3_decode_gemv if kernel.supports(n_out, k_in) else kernel.k3_decode_gemv_splitk + ) + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + entry, *args, num_tokens, n_out, k_in, trigger_early, use_pdl, stream + ) + # The compiled function takes the runtime arguments only. + fn(*args, num_tokens, stream) + return y + + +@k3_decode_gemv.register_fake +def _(x, weight, trigger_early=True): + return x.new_empty((x.shape[0], weight.shape[0]), dtype=torch.bfloat16) + + +def supports_tail( + latent: torch.Tensor, act: torch.Tensor, weight: torch.Tensor, width: int +) -> bool: + """Whether the tail kernel runs this call: the latent part of the weight a whole number of k-tiles + covering the slice, the activation too, at most MAX_K_TILES in all, a latent row the RMS loop tiles.""" + from . import k3_decode_gemv_kernel as kernel + + k_act = act.shape[1] if act.dim() == 2 else -1 + k_lat = weight.shape[1] - k_act + return ( + latent.is_cuda + and latent.dtype == act.dtype == weight.dtype == torch.bfloat16 + and latent.dim() == 2 + and act.dim() == 2 + and weight.dim() == 2 + and 0 < latent.shape[0] <= MAX_TOKENS + and act.shape[0] == latent.shape[0] + and latent.is_contiguous() + and act.is_contiguous() + and weight.is_contiguous() + and k_lat % kernel.CTA_K == 0 + and k_act % kernel.CTA_K == 0 + and 0 < width <= k_lat + and latent.shape[1] % (8 * 32) == 0 + and kernel.supports(weight.shape[0], weight.shape[1]) + ) + + +@torch.library.custom_op("trtllm::k3_decode_gemv_tail", mutates_args=()) +def k3_decode_gemv_tail( + latent: torch.Tensor, + act: torch.Tensor, + weight: torch.Tensor, + lo: int, + width: int, + eps: float, + trigger_early: bool = True, +) -> torch.Tensor: + """The row-parallel MoE tail, ``[rmsnorm(latent)[:, lo:lo+width] | act] @ weight.T``: ``latent`` + bf16 [M, H] is the whole reduced latent row, whose RMS scales the latent part; ``weight`` is + ``[latent up columns of the slice, zero-padded to a multiple of 128 | shared down]``. The RMS is + applied to the fp32 latent accumulator, not to bf16 inputs.""" + import cuda.bindings.driver as cuda_driver + + if not supports_tail(latent, act, weight, width): + raise ValueError( + f"k3_decode_gemv_tail: unsupported call latent {tuple(latent.shape)}, act {tuple(act.shape)}, " + f"weight {tuple(weight.shape)}, width {width}" + ) + from . import k3_decode_gemv_kernel as kernel + + num_tokens, rms_cols = latent.shape + n_out, k_in = weight.shape + lat_tiles = (k_in - act.shape[1]) // kernel.CTA_K + y = torch.empty(num_tokens, n_out, dtype=torch.bfloat16, device=latent.device) + args = ( + _arg(weight), + _arg(latent), + _arg(latent.view(-1).view(torch.int32)), + _arg(act), + _arg(y.view(-1)), + ) + stream = cuda_driver.CUstream(torch.cuda.current_stream(latent.device).cuda_stream) + use_pdl = _use_pdl() + key = ("tail", n_out, k_in, lat_tiles, rms_cols, trigger_early, use_pdl) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "trtllm::k3_decode_gemv_tail must run once per shape outside CUDA-graph capture first " + "(it compiles its kernel on the first call)." + ) + import cutlass.cute as cute + + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kernel.k3_decode_gemv_tail, *args, num_tokens, lo, float(eps), n_out, k_in, lat_tiles, + rms_cols, trigger_early, use_pdl, stream, + ) # fmt: skip + fn(*args, num_tokens, lo, float(eps), stream) + return y + + +@k3_decode_gemv_tail.register_fake +def _(latent, act, weight, lo, width, eps, trigger_early=True): + return latent.new_empty((latent.shape[0], weight.shape[0]), dtype=torch.bfloat16) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_embed/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_embed/__init__.py new file mode 100644 index 000000000000..b34a787db3b4 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_embed/__init__.py @@ -0,0 +1,19 @@ +# 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. +"""Decode-size embedding row gather in CuTe DSL (``trtllm::k3_embed``). + +Importing :mod:`.op` registers the torch op; nothing is imported eagerly here so that +the CuTe DSL dependency stays optional for every other model. +""" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_embed/k3_embed_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_embed/k3_embed_kernel.py new file mode 100644 index 000000000000..ece6836f32cc --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_embed/k3_embed_kernel.py @@ -0,0 +1,224 @@ +# 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. +# ============================================================================= +# Decode-size embedding row gather -- CTM (prims/cute) kernel +# ============================================================================= +# +# out[t, :] = table[ids[t], :] t < N (a decode step's tokens), bf16 rows of H elements +# +# One 16-byte vector per thread: the grid covers N x H / 8 vectors, so every row is read by one wave of loads (one +# round trip after the grid dependency). The ids are the predecessor's output, read after griddepcontrol.wait; the +# dependents are launched at entry (they wait for this grid before reading out). An id outside [0, V) gives a zero +# row (the table is not read out of bounds). +# +# Norm mode (k3_embed_norm): the rows and the first layer's input RMSNorm in one launch +# +# raw[t, :] = table[ids[t], :] (the attention-residual snapshot bank's slot 0) +# out[t, :] = raw[t, :] * rsqrt(mean(raw[t, :]^2) + eps) * weight +# +# Bit-identical to k3_embed followed by flashinfer.norm.rmsnorm, whose GB200 kernel is the CuTe DSL RMSNormKernel +# (flashinfer/norm/kernels/rmsnorm.py). For a row of 6144 < H <= 16384 elements that kernel runs one 128-thread CTA +# per row, and thread t holds columns 8 t + v + 1024 k. This kernel keeps that geometry, the same tiled copies and +# fragments, and the same DSL expressions in the same order: +# - x * x, then the fragment's TensorSSA sum; +# - a butterfly over offsets 1 .. 16, the 4 warp sums through shared memory, a butterfly again; +# - sum / H, rsqrt(mean + eps) (fast-math), x * rstd * (w + 0). +# One CTA per token; no CTA reads another's writes. +# ============================================================================= +"""CTM decode embedding: the rows of a replicated embedding table for a step's token ids, one launch (optionally +with the first layer's input RMSNorm).""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +from cutlass.experimental import primitives as prims + +THREADS = 256 +VEC_WORDS = 4 # int32 words per 16-byte vector (8 bf16) +NORM_THREADS = 128 # flashinfer's threads per row for 6144 < H <= 16384 (one row per CTA) +NORM_VEC = 8 # bf16 per 16-byte copy +NORM_WARPS = NORM_THREADS // 32 + + +def norm_supports_hidden(hidden: int) -> bool: + """Rows this norm mode reproduces: flashinfer's one-CTA 128-thread geometry, and a tile that covers the row.""" + return 6144 < hidden <= 16384 and hidden % (NORM_VEC * NORM_THREADS) == 0 + + +@cute.kernel +def k3_embed_kernel( + ids: cutlass.Array, # int32 or int64 [N] + table: cutlass.Array, # int32 words of the bf16 table [V * H / 2] + out: cutlass.Array, # int32 words of the bf16 output [N * H / 2] + vocab: cutlass.Int32, + n_tokens: cutlass.Constexpr[int], + row_vecs: cutlass.Constexpr[int], # H / 8 +): + tx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + prims.griddepcontrol(prims.GridDepAction.WAIT) + v = bx * cutlass.Int32(THREADS) + tx + if v < cutlass.Int32(n_tokens * row_vecs): + t = v // cutlass.Int32(row_vecs) + c = v - t * cutlass.Int32(row_vecs) + tok = cutlass.Int64(ids.load(idx=t)) + valid = (tok >= cutlass.Int64(0)) & (tok < cutlass.Int64(vocab)) + row = cutlass.Int64(cutlass.select_(valid, tok, cutlass.Int64(0))) + e = table.load(idx=(row * cutlass.Int64(row_vecs) + cutlass.Int64(c)) * cutlass.Int64(VEC_WORDS), + vector_size=VEC_WORDS, alignment=16) # fmt: skip + zero = cutlass.Int32(0) + out.store( + (cutlass.Int32(cutlass.select_(valid, cutlass.Int32(e[0]), zero)), + cutlass.Int32(cutlass.select_(valid, cutlass.Int32(e[1]), zero)), + cutlass.Int32(cutlass.select_(valid, cutlass.Int32(e[2]), zero)), + cutlass.Int32(cutlass.select_(valid, cutlass.Int32(e[3]), zero))), + idx=v * cutlass.Int32(VEC_WORDS), + alignment=16, + ) # fmt: skip + + +@cute.jit +def k3_embed( + ids: cute.Tensor, + table: cute.Tensor, + out: cute.Tensor, + vocab: cutlass.Int32, + n_tokens: cutlass.Constexpr[int], + row_vecs: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + grid = (n_tokens * row_vecs + THREADS - 1) // THREADS + k3_embed_kernel(ids, table, out, vocab, n_tokens, row_vecs).launch( + grid=[grid, 1, 1], + block=[THREADS, 1, 1], + stream=stream, + use_pdl=use_pdl, + ) + + +@cute.kernel +def k3_embed_norm_kernel( + ids: cute.Tensor, # int32 or int64 [N] + table: cute.Tensor, # bf16 [V, H] + weight: cute.Tensor, # bf16 [H] + raw: cute.Tensor, # bf16 [N, H], out: the rows + out: cute.Tensor, # bf16 [N, H], out: the normed rows + vocab: cutlass.Int64, + eps: cutlass.Float32, + tv_layout: cute.Layout, + tiler_mn: cute.Shape, + hidden: cutlass.Constexpr[int], +): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + prims.griddepcontrol(prims.GridDepAction.WAIT) + + smem = cutlass.utils.SmemAllocator() + reduction_buffer = smem.allocate_tensor( + cutlass.Float32, cute.make_layout((1, NORM_WARPS)), byte_alignment=4 + ) + + tok = cutlass.Int64(ids[bidx]) + valid = (tok >= cutlass.Int64(0)) & (tok < vocab) + row = cutlass.Int64(cutlass.select_(valid, tok, cutlass.Int64(0))) + + gX = cute.local_tile(table, tiler_mn, (row, 0)) + gR = cute.local_tile(raw, tiler_mn, (bidx, 0)) + gY = cute.local_tile(out, tiler_mn, (bidx, 0)) + w_layout = cute.prepend(weight.layout, cute.make_layout((tiler_mn[0],), stride=(0,))) + gW = cute.local_tile(cute.make_tensor(weight.iterator, w_layout), tiler_mn, (0, 0)) + + copy_atom_load = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), table.element_type, num_bits_per_copy=128 + ) + copy_atom_store = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), out.element_type, num_bits_per_copy=128 + ) + thr_copy_X = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) + thr_copy_W = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) + thr_copy_O = cute.make_tiled_copy(copy_atom_store, tv_layout, tiler_mn).get_slice(tidx) + + tXgX = thr_copy_X.partition_S(gX) + tXrX = cute.make_fragment_like(tXgX) + tWgW = thr_copy_W.partition_S(gW) + tWrW = cute.make_fragment_like(tWgW) + tXrW = thr_copy_X.retile(tWrW) + tXgR = thr_copy_O.partition_D(gR) + tXgO = thr_copy_O.partition_D(gY) + tXrO = cute.make_fragment_like(tXgO) + + # The row (zero for an id outside [0, V)), straight to the snapshot slot; the weight. + tXrX.store(cute.zeros_like(tXrX, dtype=table.element_type)) + if valid: + cute.copy(copy_atom_load, tXgX, tXrX) + cute.copy(copy_atom_load, tWgW, tWrW) + cute.copy(copy_atom_store, tXrX, tXgR) + + # flashinfer's RMSNormKernel arithmetic, in its order (see the header). + x = tXrX.load().to(cutlass.Float32) + x_sq = x * x + sum_sq = x_sq.reduce(cute.ReductionOp.ADD, init_val=cutlass.Float32(0.0), reduction_profile=0) + for i in cutlass.range_constexpr(5): + sum_sq = sum_sq + cute.arch.shuffle_sync_bfly(sum_sq, offset=1 << i) + lane = cute.arch.lane_idx() + warp = cute.arch.warp_idx() + if lane == 0: + reduction_buffer[0, warp] = sum_sq + cute.arch.barrier() + total = cutlass.Float32(0.0) + if lane < NORM_WARPS: + total = reduction_buffer[0, lane] + for i in cutlass.range_constexpr(5): + total = total + cute.arch.shuffle_sync_bfly(total, offset=1 << i) + mean_sq = total / cutlass.Float32(hidden) + rstd = cute.math.rsqrt(mean_sq + eps, fastmath=True) + w = tXrW.load().to(cutlass.Float32) + y = x * rstd * (w + cutlass.Float32(0.0)) + tXrO.store(y.to(out.element_type)) + cute.copy(copy_atom_store, tXrO, tXgO) + + +@cute.jit +def k3_embed_norm( + ids: cute.Tensor, + table: cute.Tensor, + weight: cute.Tensor, + raw: cute.Tensor, + out: cute.Tensor, + vocab: cutlass.Int64, + eps: cutlass.Float32, + n_tokens: cutlass.Constexpr[int], + hidden: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + # flashinfer's thread-value layout for one row: thread t, value (v, k) -> column 8 t + v + 1024 k. + blocks = hidden // (NORM_VEC * NORM_THREADS) + tv_layout = cute.make_layout(((NORM_THREADS, 1), (NORM_VEC, blocks)), + stride=((NORM_VEC, 1), (1, NORM_VEC * NORM_THREADS))) # fmt: skip + k3_embed_norm_kernel( + ids, table, weight, raw, out, vocab, eps, tv_layout, (1, hidden), hidden + ).launch( + grid=[n_tokens, 1, 1], + block=[NORM_THREADS, 1, 1], + smem=NORM_WARPS * 4, + stream=stream, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_embed/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_embed/op.py new file mode 100644 index 000000000000..420d619fc2ee --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_embed/op.py @@ -0,0 +1,181 @@ +# 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. +"""Torch ops of the decode-size embedding gather: ``trtllm::k3_embed`` (``table[ids]`` for up to ``MAX_TOKENS`` int32 +or int64 ids and a bf16 table whose rows are whole 16-byte vectors, one vector per thread) and ``trtllm::k3_embed_norm`` +(the same rows written into a caller's buffer, plus their RMSNorm, bit-identical to ``flashinfer.norm.rmsnorm``). +Compiled on the first call for its shape, which must happen outside CUDA-graph capture.""" + +from __future__ import annotations + +import os +import threading +from typing import Dict + +import torch + +MAX_TOKENS = 64 + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + + +def _arg(t: torch.Tensor): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +def supports(ids: torch.Tensor, table: torch.Tensor) -> bool: + """Whether ``k3_embed`` gathers ``table[ids]``.""" + return ( + ids.is_cuda + and ids.dim() == 1 + and ids.dtype in (torch.int32, torch.int64) + and 0 < ids.numel() <= MAX_TOKENS + and table.dtype == torch.bfloat16 + and table.dim() == 2 + and table.is_contiguous() + and table.shape[1] % 8 == 0 + and table.data_ptr() % 16 == 0 + ) + + +@torch.library.custom_op("trtllm::k3_embed", mutates_args=()) +def k3_embed(ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: + """``table[ids]`` for int32 / int64 ``ids`` [N <= 64] and a contiguous bf16 ``table`` [V, H] (H % 8 == 0); ids + outside [0, V) give zero rows.""" + import cuda.bindings.driver as cuda_driver + + if not supports(ids, table): + raise ValueError( + f"k3_embed: unsupported call: ids {tuple(ids.shape)} {ids.dtype}, table {tuple(table.shape)} {table.dtype} " + f"(int32 / int64 [N <= {MAX_TOKENS}], contiguous bf16 [V, H % 8 == 0])" + ) + from . import k3_embed_kernel as kern + + n = ids.numel() + vocab, hidden = table.shape + out = torch.empty(n, hidden, dtype=torch.bfloat16, device=ids.device) + args = ( + _arg(ids.contiguous()), + _arg(table.view(-1).view(torch.int32)), + _arg(out.view(-1).view(torch.int32)), + ) + use_pdl = os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + stream = cuda_driver.CUstream(torch.cuda.current_stream(ids.device).cuda_stream) + key = (n, hidden // 8, ids.dtype, use_pdl) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "trtllm::k3_embed must run once per shape outside CUDA-graph capture first" + ) + import cutlass.cute as cute + + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kern.k3_embed, *args, int(vocab), n, hidden // 8, use_pdl, stream + ) + fn(*args, int(vocab), stream) + return out + + +@k3_embed.register_fake +def _(ids, table): + return table.new_empty((ids.numel(), table.shape[1])) + + +def _plain_arg(t: torch.Tensor, align: int = 16): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align) + + +def norm_supports_hidden(hidden: int) -> bool: + """Row widths ``k3_embed_norm`` reproduces flashinfer's RMSNorm for.""" + from . import k3_embed_kernel as kern + + return kern.norm_supports_hidden(hidden) + + +def supports_norm( + ids: torch.Tensor, table: torch.Tensor, weight: torch.Tensor, raw: torch.Tensor +) -> bool: + """Whether ``k3_embed_norm`` gathers ``table[ids]`` into ``raw`` and norms it with ``weight``.""" + hidden = table.shape[1] if table.dim() == 2 else 0 + return ( + supports(ids, table) + and norm_supports_hidden(hidden) + and weight.dtype == torch.bfloat16 + and tuple(weight.shape) == (hidden,) + and weight.is_contiguous() + and weight.data_ptr() % 16 == 0 + and raw.dtype == torch.bfloat16 + and tuple(raw.shape) == (ids.numel(), hidden) + and raw.is_contiguous() + and raw.data_ptr() % 16 == 0 + ) + + +@torch.library.custom_op("trtllm::k3_embed_norm", mutates_args=("raw",)) +def k3_embed_norm( + ids: torch.Tensor, table: torch.Tensor, weight: torch.Tensor, eps: float, raw: torch.Tensor +) -> torch.Tensor: + """``raw[:] = table[ids]`` (ids outside [0, V) give zero rows) and returns ``rmsnorm(raw, weight, eps)``, bf16 + [N, H]: bit-identical to ``k3_embed`` followed by ``flashinfer.norm.rmsnorm`` (see ``k3_embed_kernel``).""" + import cuda.bindings.driver as cuda_driver + + if not supports_norm(ids, table, weight, raw): + raise ValueError( + f"k3_embed_norm: unsupported call: ids {tuple(ids.shape)} {ids.dtype}, table {tuple(table.shape)} " + f"{table.dtype}, weight {tuple(weight.shape)} {weight.dtype}, raw {tuple(raw.shape)} {raw.dtype} " + f"(int32 / int64 [N <= {MAX_TOKENS}], contiguous bf16 [V, H] with 6144 < H <= 16384 and H % 1024 == 0, " + f"weight [H], raw [N, H])" + ) + from . import k3_embed_kernel as kern + + n = ids.numel() + vocab, hidden = table.shape + out = torch.empty(n, hidden, dtype=torch.bfloat16, device=ids.device) + ids = ids.contiguous() + args = (_plain_arg(ids, ids.element_size()),) + tuple( + _plain_arg(t) for t in (table, weight, raw, out) + ) + use_pdl = os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + stream = cuda_driver.CUstream(torch.cuda.current_stream(ids.device).cuda_stream) + key = ("norm", n, vocab, hidden, ids.dtype, use_pdl) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "trtllm::k3_embed_norm must run once per shape outside CUDA-graph capture first" + ) + import cutlass.cute as cute + + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kern.k3_embed_norm, *args, int(vocab), float(eps), n, hidden, use_pdl, stream + ) + fn(*args, int(vocab), float(eps), stream) + return out + + +@k3_embed_norm.register_fake +def _(ids, table, weight, eps, raw): + return table.new_empty((ids.numel(), table.shape[1])) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_head_gemv/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_head_gemv/__init__.py new file mode 100644 index 000000000000..db044cfd9984 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_head_gemv/__init__.py @@ -0,0 +1,19 @@ +# 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 persistent CTM head GEMV (``trtllm::k3_head_gemv``). + +Importing :mod:`.op` registers the torch op; nothing is imported eagerly here so that +the CuTe DSL dependency stays optional for every other model. +""" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_head_gemv/k3_head_gemv_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_head_gemv/k3_head_gemv_kernel.py new file mode 100644 index 000000000000..247b6c1f6ca8 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_head_gemv/k3_head_gemv_kernel.py @@ -0,0 +1,1137 @@ +# 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 vocab-shard head GEMV -- CTM (prims/cute) kernel, M <= 8 tokens, bf16 in/out, persistent +# ============================================================================= +# +# y[t, n] = sum_k x[t, k] * W[n, k] (fp32 accumulation in TMEM, one bf16 rounding) +# A = W (N, K) K-major bf16 (the streamed weight: an lm_head shard, N ~ 10^4) +# B = x (M <= 8, K) K-major bf16 (8 token columns; rows past M arrive as zeros) +# +# Work units: the weight is cut into 128-row tiles and every tile's K into CHUNKS chunks of CH k-tiles (128 each); +# unit u is (tile u // CHUNKS, chunk u % CHUNKS). One CTA per SM (GRID CTAs), each running units until the pool is +# empty: unit bx is static (its first stages are loaded before griddepcontrol.wait), the rest are claimed from a +# global counter (every CTA makes exactly one failing claim; the last ticket rolls the counter back to 0, so every +# launch and graph replay starts from 0). Units are claimed in tile-major order, so a tile's chunks finish together. +# +# Split-K combine (CHUNKS > 1): a unit's epilogue stores its fp32 partial [128 rows][8 tokens] to ws[u], then +# thread 0 of the epilogue warps counts the unit on cnt[tile] with an acq_rel atomic after a barrier of those warps +# (the barrier orders their stores before the release). The unit that brings the count to CHUNKS finalizes the +# tile: it adds the CHUNKS partials in chunk order (its own from registers), rounds once to bf16 and stores y's +# [M][128] block, and resets cnt[tile] to 0. A unit's partial does not depend on which CTA computes it and the +# combine order is fixed, so y is bit-identical from run to run whatever the claim order. +# +# Stage = one k-tile: A 32 KB (both 64-element halves of the 128-byte swizzle, one 5-D TMA call) + B 2 KB (two TMA +# calls from x). RING stages; the phases run straight through unit boundaries. TMEM: two 8-column fp32 +# accumulators, so the epilogue of one unit overlaps the MMAs of the next. +# +# L2: the weight loads are EVICT_FIRST, except for tiles < keep_tiles (normal priority), so that a later reader of +# the same shard (the drafter head after the target head) can find them in L2. +# +# Warps: 0 claimer + A/B TMA (one elected lane), 2 TMEM allocation + tcgen05 MMA (M 128, N 8, K 16), 4-7 epilogue +# (TMEM -> registers -> partial / fixup -> bf16 staging -> 16-byte stores), 1 and 3 idle. +# +# PDL: before the grid-dependency wait only barrier init, TMEM allocation, the tensormap prefetches and the static +# unit's A loads (a weight: nothing in the graph writes it). After it: B loads, claims, every global write. The +# dependents are released after the CTA's failing claim, so they launch only once every CTA of this grid has +# started (no dependent CTA can hold an SM that a static unit still needs). +# ============================================================================= +"""Persistent CTM head GEMV for Kimi K3 (``y = x @ W^T``, M <= 8, a large vocab shard): dynamic (tile, k-chunk) +units with a deterministic split-K combine, the weight streamed before the grid-dependency wait.""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +CTA_M = 128 # weight rows per tile = the MMA's M +MMA_N = 8 # token columns +CTA_K = 128 # one k-tile: two 64-element halves of the 128-byte swizzle +MMA_K = 16 +TMA_K_BOX = 64 +TMA_COPY_ITERS = CTA_K // TMA_K_BOX +K_BLOCKS_PER_HALF = TMA_K_BOX // MMA_K +THREADS = 256 +EPI_THREADS = 128 +ELEM_BYTES = 2 +VEC = 8 # bf16 per 16-byte vector +TMEM_COLS = 32 +UNIT_RING = ( + 16 # claimed-unit broadcast slots (the claimer runs at most RING k-tiles + 2 units ahead) +) +EVICT_FIRST = 0x12F0000000000000 # createpolicy.fractional.L2::evict_first, fraction 1.0 (sm_100) +SMEM_BYTES = 227 * 1024 +MAX_RING = 6 +FLAG_BACKOFF_NS = ( + 256 # between polls of a split tile's piece flags (the pieces finish long before the finalizer) +) + +# Shared-memory descriptor strides for the 128-byte swizzle, in 16-byte units. +LEADING = 16 +STRIDE = 8 * TMA_K_BOX * ELEM_BYTES +A_HALF_ELEMS = CTA_M * TMA_K_BOX +B_HALF_ELEMS = MMA_N * TMA_K_BOX +STEP = (MMA_K * ELEM_BYTES) >> 4 +A_BOX = A_HALF_ELEMS >> 3 +B_BOX = B_HALF_ELEMS >> 3 +STAGE_A = (CTA_M * CTA_K * ELEM_BYTES) >> 4 +STAGE_B = (MMA_N * CTA_K * ELEM_BYTES) >> 4 +A_BYTES = CTA_M * CTA_K * ELEM_BYTES +B_BYTES = MMA_N * CTA_K * ELEM_BYTES + +io_dtype = cutlass.BFloat16 + + +def num_k_tiles(k_in: int) -> int: + return k_in // CTA_K + + +def num_units(n_out: int, k_in: int, chunk_tiles: int) -> int: + return (n_out // CTA_M) * (num_k_tiles(k_in) // chunk_tiles) + + +def smem_bytes(ring: int) -> int: + """Shared memory of a CTA: the A and B rings, the bf16 output staging tile, barriers and slots.""" + return ring * (A_BYTES + B_BYTES) + MMA_N * CTA_M * ELEM_BYTES + 1024 + + +def supports(n_out: int, k_in: int, chunk_tiles: int, ring: int) -> bool: + """Shapes the kernel runs: whole 128-row tiles, whole k-tiles, chunks that tile K, a ring that fits.""" + k_tiles = num_k_tiles(k_in) + return ( + n_out > 0 + and n_out % CTA_M == 0 + and k_in > 0 + and k_in % CTA_K == 0 + and 0 < chunk_tiles <= k_tiles + and k_tiles % chunk_tiles == 0 + and 1 <= ring <= MAX_RING + and smem_bytes(ring) <= SMEM_BYTES + ) + + +@dsl_user_op +def _atomic_add_acq_rel(addr_i64, val, *, loc=None, ip=None): + """atom.acq_rel.gpu.global.add.u32, returning the old value.""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [addr_i64.ir_value(loc=loc, ip=ip), val.ir_value(loc=loc, ip=ip)], + "atom.acq_rel.gpu.global.add.u32 $0, [$1], $2;", "=r,l,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _atomic_add_relaxed(addr_i64, val, *, loc=None, ip=None): + """atom.relaxed.gpu.global.add.u32, returning the old value (the unit claim).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [addr_i64.ir_value(loc=loc, ip=ip), val.ir_value(loc=loc, ip=ip)], + "atom.relaxed.gpu.global.add.u32 $0, [$1], $2;", "=r,l,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _ld_acquire(addr_i64, *, loc=None, ip=None): + """ld.acquire.gpu.global.u32 (a flag written by another SM).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [addr_i64.ir_value(loc=loc, ip=ip)], + "ld.acquire.gpu.global.u32 $0, [$1];", "=r,l", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _flag_release(addr_i64, *, loc=None, ip=None): + """fence.acq_rel.gpu, then st.relaxed.gpu 1: the release of everything the thread (and, through a preceding + barrier, its CTA) stored before.""" + _llvm.inline_asm( + None, [addr_i64.ir_value(loc=loc, ip=ip)], + "fence.acq_rel.gpu;\n\tst.relaxed.gpu.global.u32 [$0], 1;", "l", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + """add.rn.f32: an fp32 add the compiler cannot reassociate or contract.""" + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + "add.rn.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, + ) + ) # fmt: skip + + +def _store_tile(stage_out, y, vals, row, tid, tile, n_out): + """The finalized [8 tokens] sums of tile row ``row`` -> bf16 -> y[token, tile rows] for all 8 token rows (y has 8 + rows; rows past M hold the zero-filled x rows' zeros): staged per token in shared memory, then 16-byte stores + (rows contiguous per token). The epilogue warps' barriers bracket the staging tile's reuse.""" + for t in range(MMA_N): + stage_out.store(vals[t].to(io_dtype), idx=cutlass.Int32(t * CTA_M) + row) + prims.barrier_cta_sync(1, thread_count=EPI_THREADS) + ct = tid // cutlass.Int32(CTA_M // VEC) + cc = tid % cutlass.Int32(CTA_M // VEC) + v = stage_out.load( + idx=ct * cutlass.Int32(CTA_M) + cc * cutlass.Int32(VEC), vector_size=VEC, alignment=16 + ) + y.store(v, idx=ct * cutlass.Int32(n_out) + tile * cutlass.Int32(CTA_M) + cc * cutlass.Int32(VEC), + vector_size=VEC, alignment=16) # fmt: skip + prims.barrier_cta_sync(1, thread_count=EPI_THREADS) + + +@cute.kernel +def k3_head_gemv_kernel( + tma_desc_w: cutlass.GridConstant[cuda.TensorMap], # W [N, K] bf16, 5-D, one call per k-tile + tma_desc_x: cutlass.GridConstant[cuda.TensorMap], # x [M, K] bf16, box 64 x 8 + y: cutlass.Array, # [8 * N] bf16, token-major (all 8 token rows are written) + ws: cutlass.Array, # fp32 [units * 128 * 8]: the units' partials (CHUNKS > 1) + cnt: cutlass.Array, # int32 [tiles]: units of the tile done (0 between launches) + claim: cutlass.Array, # int32 [1]: the unit pool's ticket counter (0 between launches) + num_tokens: cutlass.Int32, + keep_tiles: cutlass.Int32, # tiles [0, keep_tiles) load at normal L2 priority, the rest EVICT_FIRST + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + chunk_tiles: cutlass.Constexpr[int], # CH: k-tiles per unit + ring: cutlass.Constexpr[int], + grid: cutlass.Constexpr[int], +): + """Persistent: unit bx, then claimed units until the pool is empty; the last unit of a tile finalizes it.""" + k_tiles = num_k_tiles(k_in) + chunks = k_tiles // chunk_tiles + units = (n_out // CTA_M) * chunks + tickets = max( + units, grid + ) # claims made in one launch: units - grid succeed (if positive), grid fail + pre = min(ring, chunk_tiles) # stages of the static unit loaded before the grid-dependency wait + tx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp_id = cute.arch.warp_idx() + tma_ptr_w = tma_desc_w.get_ptr() + tma_ptr_x = tma_desc_x.get_ptr() + + smem_a = cutlass.Array( + io_dtype, ring * CTA_M * CTA_K, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_b = cutlass.Array( + io_dtype, ring * MMA_N * CTA_K, space=cutlass.AddressSpace.smem, alignment=1024 + ) + stage_out = cutlass.Array( + io_dtype, MMA_N * CTA_M, space=cutlass.AddressSpace.smem, alignment=16 + ) + full = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + empty = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + acc_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + unit_ready = cutlass.Array( + cutlass.Int64, UNIT_RING, space=cutlass.AddressSpace.smem, alignment=8 + ) + unit_slot = cutlass.Array( + cutlass.Int32, UNIT_RING, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_last = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + + if warp_id == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_x) + if prims.elect_sync(): + for s in cutlass.range_constexpr(ring): + prims.mbarrier_init(full.subview(s), 1) + prims.mbarrier_init(empty.subview(s), 1) + for b in cutlass.range_constexpr(2): + prims.mbarrier_init(acc_full.subview(b), 1) + prims.mbarrier_init(acc_empty.subview(b), 4) # one elected lane per epilogue warp + for s in cutlass.range_constexpr(UNIT_RING): + prims.mbarrier_init(unit_ready.subview(s), 1) + if warp_id == 2: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + prims.barrier_cta_sync(0) + + if warp_id == 0: + # ===================================================================== + # Claimer + TMA: the static unit's first stages before the wait, then + # every stage of every unit this CTA gets, A and B on one barrier. + # ===================================================================== + if prims.elect_sync(): + unit = cutlass.Int32(cutlass.select_(bx < cutlass.Int32(units), bx, cutlass.Int32(-1))) + unit_slot.store(unit, idx=0) + prims.mbarrier_arrive(unit_ready.subview(0)) + static_tile = unit // cutlass.Int32(chunks) + static_k0 = (unit % cutlass.Int32(chunks)) * cutlass.Int32(chunk_tiles) + if unit >= cutlass.Int32(0): + for jj in cutlass.range_constexpr(pre): + prims.mbarrier_arrive_expect_tx(full.subview(jj), A_BYTES + B_BYTES) + if static_tile < keep_tiles: + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(jj * CTA_M * CTA_K), + tma_ptr_w, + ( + cutlass.Int32(0), + static_tile * cutlass.Int32(CTA_M), + (static_k0 + cutlass.Int32(jj)) * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), + cutlass.Int32(0), + ), + full.subview(jj), + ) + else: + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(jj * CTA_M * CTA_K), + tma_ptr_w, + ( + cutlass.Int32(0), + static_tile * cutlass.Int32(CTA_M), + (static_k0 + cutlass.Int32(jj)) * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), + cutlass.Int32(0), + ), + full.subview(jj), + l2_cache_hint=EVICT_FIRST, + ) + prims.griddepcontrol(prims.GridDepAction.WAIT) + issued = cutlass.Int32(0) # k-tiles issued by this CTA (the ring position) + next_j = cutlass.Int32(0) # first k-tile of the current unit still to issue + if unit >= cutlass.Int32(0): + for jj in cutlass.range_constexpr(pre): + for half in cutlass.range_constexpr(TMA_COPY_ITERS): + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(jj * MMA_N * CTA_K + half * B_HALF_ELEMS), + tma_ptr_x, + ((static_k0 + cutlass.Int32(jj)) * cutlass.Int32(CTA_K) + cutlass.Int32(half * TMA_K_BOX), + cutlass.Int32(0)), + full.subview(jj), + ) # fmt: skip + issued = cutlass.Int32(pre) + next_j = cutlass.Int32(pre) + n_units = cutlass.Int32(0) # units published so far - 1 + pending = cutlass.Boolean(True) + while pending: + if unit >= cutlass.Int32(0): + tile = unit // cutlass.Int32(chunks) + k0 = (unit % cutlass.Int32(chunks)) * cutlass.Int32(chunk_tiles) + j = next_j + while j < cutlass.Int32(chunk_tiles): + stage = issued % cutlass.Int32(ring) + # The stage's previous MMAs committed (a fresh barrier passes parity 1 at once). + parity = ( + (issued // cutlass.Int32(ring)) & cutlass.Int32(1) + ) ^ cutlass.Int32(1) + while not cute.arch.mbarrier_try_wait( + empty.subview(stage).data_ptr(), parity + ): + pass + k = k0 + j + prims.mbarrier_arrive_expect_tx(full.subview(stage), A_BYTES + B_BYTES) + if tile < keep_tiles: + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(stage * cutlass.Int32(CTA_M * CTA_K)), tma_ptr_w, + (cutlass.Int32(0), tile * cutlass.Int32(CTA_M), k * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), cutlass.Int32(0)), + full.subview(stage), + ) # fmt: skip + else: + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(stage * cutlass.Int32(CTA_M * CTA_K)), tma_ptr_w, + (cutlass.Int32(0), tile * cutlass.Int32(CTA_M), k * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), cutlass.Int32(0)), + full.subview(stage), l2_cache_hint=EVICT_FIRST, + ) # fmt: skip + for half in cutlass.range_constexpr(TMA_COPY_ITERS): + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview( + stage * cutlass.Int32(MMA_N * CTA_K) + + cutlass.Int32(half * B_HALF_ELEMS) + ), + tma_ptr_x, + ( + k * cutlass.Int32(CTA_K) + cutlass.Int32(half * TMA_K_BOX), + cutlass.Int32(0), + ), + full.subview(stage), + ) + issued = issued + cutlass.Int32(1) + j = j + cutlass.Int32(1) + # Claim the next unit; the holder of the last ticket rolls the counter back to 0. + ticket = _atomic_add_relaxed(claim.data_ptr().toint(), cutlass.Int32(1)) + if ticket == cutlass.Int32(tickets - 1): + _atomic_add_relaxed(claim.data_ptr().toint(), cutlass.Int32(-tickets)) + claimed = ticket + cutlass.Int32(grid) + unit = cutlass.Int32( + cutlass.select_(claimed < cutlass.Int32(units), claimed, cutlass.Int32(-1)) + ) + n_units = n_units + cutlass.Int32(1) + slot = n_units % cutlass.Int32(UNIT_RING) + unit_slot.store(unit, idx=slot) + prims.mbarrier_arrive(unit_ready.subview(slot)) + next_j = cutlass.Int32(0) + pending = unit >= cutlass.Int32(0) + # Every CTA of this grid has started and made its last claim: the dependents may launch. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + elif warp_id == 2: + # ===================================================================== + # MMA: each unit's CH k-tiles into accumulator (unit count) % 2. + # ===================================================================== + idesc = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, a_dtype=io_dtype, b_dtype=io_dtype, n_dim=MMA_N, m_dim=CTA_M + ) + desc_a_base = prims.Tcgen05SmemDesc.build( + start_address=smem_a, leading_byte_offset=LEADING, stride_byte_offset=STRIDE, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + desc_b_base = prims.Tcgen05SmemDesc.build( + start_address=smem_b, leading_byte_offset=LEADING, stride_byte_offset=STRIDE, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + tmem_base = tmem_ptr_i32.load() + used = cutlass.Int32(0) # k-tiles consumed (the ring position) + count = cutlass.Int32(0) # units done by this CTA + while not cute.arch.mbarrier_try_wait(unit_ready.subview(0).data_ptr(), 0): + pass + unit = unit_slot.load(idx=0) + while unit >= cutlass.Int32(0): + buf = count & cutlass.Int32(1) + if count >= cutlass.Int32(2): + # The epilogue has read this accumulator's previous unit. + while not cute.arch.mbarrier_try_wait( + acc_empty.subview(buf).data_ptr(), + ((count >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + ): + pass + tmem_acc = cutlass.inttoptr(tmem_base + buf * cutlass.Int32(MMA_N), 6, cutlass.Int32) + for j in cutlass.range(chunk_tiles, unroll=1): + stage = used % cutlass.Int32(ring) + while not cute.arch.mbarrier_try_wait( + full.subview(stage).data_ptr(), (used // cutlass.Int32(ring)) & cutlass.Int32(1) + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for kb in cutlass.range_constexpr(TMA_COPY_ITERS * K_BLOCKS_PER_HALF): + box = kb // K_BLOCKS_PER_HALF + within = kb % K_BLOCKS_PER_HALF + desc_a = desc_a_base + ( + stage * cutlass.Int32(STAGE_A) + cutlass.Int32(box * A_BOX + within * STEP) + ) + desc_b = desc_b_base + ( + stage * cutlass.Int32(STAGE_B) + cutlass.Int32(box * B_BOX + within * STEP) + ) + accumulate = cutlass.Boolean(True) + if cutlass.const_expr(kb == 0): + accumulate = j > cutlass.Int32(0) + if prims.elect_sync(): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, + prims.CTAGroup.CTA_1, + tmem_acc, + desc_a, + desc_b, + idesc, + accumulate, + ) + if prims.elect_sync(): + prims.tcgen05_commit(empty.subview(stage)) + used = used + cutlass.Int32(1) + if prims.elect_sync(): + prims.tcgen05_commit(acc_full.subview(buf)) + count = count + cutlass.Int32(1) + slot = count % cutlass.Int32(UNIT_RING) + while not cute.arch.mbarrier_try_wait( + unit_ready.subview(slot).data_ptr(), + (count // cutlass.Int32(UNIT_RING)) & cutlass.Int32(1), + ): + pass + unit = unit_slot.load(idx=slot) + elif warp_id >= 4: + # ===================================================================== + # Epilogue: TMEM -> registers -> partial + count (or the whole sum when + # K is one chunk) -> the finalizer's ordered sum -> bf16 -> y. + # warp 4 + w reads TMEM lanes 32 w .. 32 w + 31 = tile rows 32 w + lane. + # ===================================================================== + tid = tx - cutlass.Int32(EPI_THREADS) + lane = tx % 32 + w = warp_id - 4 + row = w * cutlass.Int32(32) + lane + tmem_base = tmem_ptr_i32.load() + prims.griddepcontrol(prims.GridDepAction.WAIT) + count = cutlass.Int32(0) + while not cute.arch.mbarrier_try_wait(unit_ready.subview(0).data_ptr(), 0): + pass + unit = unit_slot.load(idx=0) + while unit >= cutlass.Int32(0): + buf = count & cutlass.Int32(1) + while not cute.arch.mbarrier_try_wait( + acc_full.subview(buf).data_ptr(), (count >> cutlass.Int32(1)) & cutlass.Int32(1) + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", + cutlass.inttoptr( + tmem_base + + ((w * cutlass.Int32(32)) << cutlass.Int32(16)) + + buf * cutlass.Int32(MMA_N), + 6, + cutlass.Float32, + ), + num=MMA_N, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if prims.elect_sync(): + prims.mbarrier_arrive(acc_empty.subview(buf)) + vals = [cutlass.Float32(acc[t]) for t in range(MMA_N)] + tile = unit // cutlass.Int32(chunks) + if cutlass.const_expr(chunks == 1): + _store_tile(stage_out, y, vals, row, tid, tile, n_out) + else: + chunk = unit % cutlass.Int32(chunks) + part = (unit * cutlass.Int32(CTA_M) + row) * cutlass.Int32(MMA_N) + ws.store(cutlass.Vector.from_elements((vals[0], vals[1], vals[2], vals[3]), cutlass.Float32), idx=part, + vector_size=4, alignment=16) # fmt: skip + ws.store(cutlass.Vector.from_elements((vals[4], vals[5], vals[6], vals[7]), cutlass.Float32), + idx=part + cutlass.Int32(4), vector_size=4, alignment=16) # fmt: skip + # Every epilogue thread's partial is stored; thread 0's release covers them. + prims.barrier_cta_sync(1, thread_count=EPI_THREADS) + if tid == cutlass.Int32(0): + done = _atomic_add_acq_rel( + cnt.subview(tile).data_ptr().toint(), cutlass.Int32(1) + ) + last = cutlass.Int32( + cutlass.select_( + done == cutlass.Int32(chunks - 1), cutlass.Int32(1), cutlass.Int32(0) + ) + ) + s_last.store(last, idx=0) + prims.barrier_cta_sync(1, thread_count=EPI_THREADS) + if s_last.load(idx=0) != cutlass.Int32(0): + # The tile's partials in chunk order (this unit's own from registers). Volatile loads: other + # SMs wrote them; thread 0's acquire and the barrier above order these reads after the writes. + total = [] + for c in cutlass.range_constexpr(chunks): + src = ( + (tile * cutlass.Int32(chunks) + cutlass.Int32(c)) * cutlass.Int32(CTA_M) + + row + ) * cutlass.Int32(MMA_N) + lo = ws.load(idx=src, vector_size=4, alignment=16, is_volatile=True) + hi = ws.load( + idx=src + cutlass.Int32(4), + vector_size=4, + alignment=16, + is_volatile=True, + ) + mine = chunk == cutlass.Int32(c) + for t in cutlass.range_constexpr(MMA_N): + got = cutlass.Float32(lo[t]) if t < 4 else cutlass.Float32(hi[t - 4]) + got = cutlass.Float32(cutlass.select_(mine, vals[t], got)) + if cutlass.const_expr(c == 0): + total.append(got) + else: + total[t] = _add_rn(total[t], got) + if tid == cutlass.Int32(0): + cnt.store(cutlass.Int32(0), idx=tile) + _store_tile(stage_out, y, total, row, tid, tile, n_out) + count = count + cutlass.Int32(1) + slot = count % cutlass.Int32(UNIT_RING) + while not cute.arch.mbarrier_try_wait( + unit_ready.subview(slot).data_ptr(), + (count // cutlass.Int32(UNIT_RING)) & cutlass.Int32(1), + ): + pass + unit = unit_slot.load(idx=slot) + prims.barrier_cta_sync(1, thread_count=EPI_THREADS) + if warp_id == 4: + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_base, 6, cutlass.Int32), TMEM_COLS) + + +# ============================================================================= +# Stream-K variant: the (tile, k-tile) space flattened tile-major and cut into GRID equal ranges, CTA c owning +# [c T / G, (c + 1) T / G). A CTA's range is a sequence of segments (one tile's contiguous k-tiles each); a tile +# split over several CTAs ("pieces", in k order = CTA order) is combined by the CTA of its piece 0 (the one holding +# k-tile 0), which is the end of that CTA's range: every other piece starts its CTA's range (or is all of it), so +# it is done early. A piece p > 0 stores its fp32 partial to ws[tile][p] and raises flag[tile][p] (fence.acq_rel + +# a relaxed store after a barrier of the epilogue warps); the finalizer acquires those flags and loads the partials +# while its own MMAs still run, then adds them after its accumulator in k order, rounds once, stores y and lowers +# the flags. The partition is fixed, so the sums are too. No claims: every warp walks the same static sequence, and +# the whole ring is loaded before the grid-dependency wait (optionally the next k-tiles are prefetched into L2). +# ============================================================================= +def streamk_cta_of(q: int, total: int, grid: int) -> int: + """The CTA whose range holds flat k-tile q (ranges [c T / G, (c + 1) T / G), some empty when G > T).""" + return ((q + 1) * grid - 1) // total + + +def streamk_max_pieces(n_out: int, k_in: int, grid: int) -> int: + """The most CTAs any tile is split over.""" + k_tiles = num_k_tiles(k_in) + tiles = n_out // CTA_M + total = tiles * k_tiles + return max( + streamk_cta_of((t + 1) * k_tiles - 1, total, grid) + - streamk_cta_of(t * k_tiles, total, grid) + + 1 + for t in range(tiles) + ) + + +def streamk_supports(n_out: int, k_in: int, ring: int) -> bool: + return ( + n_out > 0 + and n_out % CTA_M == 0 + and k_in > 0 + and k_in % CTA_K == 0 + and 1 <= ring <= MAX_RING + and (smem_bytes(ring) <= SMEM_BYTES) + ) + + +@cute.kernel +def k3_head_gemv_sk_kernel( + tma_desc_w: cutlass.GridConstant[cuda.TensorMap], # W [N, K] bf16, 5-D, one call per k-tile + tma_desc_x: cutlass.GridConstant[cuda.TensorMap], # x [M, K] bf16, box 64 x 8 + y: cutlass.Array, # [8 * N] bf16, token-major (all 8 token rows are written) + ws: cutlass.Array, # fp32 [tiles * max_pieces * 128 * 8]: the pieces' partials + flags: cutlass.Array, # int32 [tiles * max_pieces]: piece p > 0 of the tile stored its partial (0 between launches) + num_tokens: cutlass.Int32, + keep_tiles: cutlass.Int32, # tiles [0, keep_tiles) load at normal L2 priority, the rest EVICT_FIRST + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + grid: cutlass.Constexpr[int], + max_pieces: cutlass.Constexpr[int], + prefetch: cutlass.Constexpr[ + int + ], # k-tiles after the ring prefetched into L2 before the grid-dependency wait +): + """Stream-K: this CTA's equal share of the flat (tile, k-tile) space; split tiles combined by their last piece.""" + k_tiles = num_k_tiles(k_in) + total = (n_out // CTA_M) * k_tiles + tx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp_id = cute.arch.warp_idx() + tma_ptr_w = tma_desc_w.get_ptr() + tma_ptr_x = tma_desc_x.get_ptr() + lo = (bx * cutlass.Int32(total)) // cutlass.Int32(grid) + hi = ((bx + cutlass.Int32(1)) * cutlass.Int32(total)) // cutlass.Int32(grid) + count = hi - lo + + smem_a = cutlass.Array( + io_dtype, ring * CTA_M * CTA_K, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_b = cutlass.Array( + io_dtype, ring * MMA_N * CTA_K, space=cutlass.AddressSpace.smem, alignment=1024 + ) + stage_out = cutlass.Array( + io_dtype, MMA_N * CTA_M, space=cutlass.AddressSpace.smem, alignment=16 + ) + full = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + empty = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + acc_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + + if warp_id == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_x) + if prims.elect_sync(): + for s in cutlass.range_constexpr(ring): + prims.mbarrier_init(full.subview(s), 1) + prims.mbarrier_init(empty.subview(s), 1) + for b in cutlass.range_constexpr(2): + prims.mbarrier_init(acc_full.subview(b), 1) + prims.mbarrier_init(acc_empty.subview(b), 4) # one elected lane per epilogue warp + if warp_id == 2: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + prims.barrier_cta_sync(0) + + if warp_id == 0: + # ===================================================================== + # TMA: the range's first RING k-tiles' weight before the wait, their + # activation after it, then every later k-tile when its stage frees. + # ===================================================================== + if prims.elect_sync(): + for jj in cutlass.range_constexpr(ring): + if cutlass.Int32(jj) < count: + q0 = lo + cutlass.Int32(jj) + t0 = q0 // cutlass.Int32(k_tiles) + prims.mbarrier_arrive_expect_tx(full.subview(jj), A_BYTES + B_BYTES) + if t0 < keep_tiles: + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(jj * CTA_M * CTA_K), + tma_ptr_w, + ( + cutlass.Int32(0), + t0 * cutlass.Int32(CTA_M), + (q0 % cutlass.Int32(k_tiles)) * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), + cutlass.Int32(0), + ), + full.subview(jj), + ) + else: + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(jj * CTA_M * CTA_K), + tma_ptr_w, + ( + cutlass.Int32(0), + t0 * cutlass.Int32(CTA_M), + (q0 % cutlass.Int32(k_tiles)) * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), + cutlass.Int32(0), + ), + full.subview(jj), + l2_cache_hint=EVICT_FIRST, + ) + # The next k-tiles of the range into L2 while the predecessor runs (a weight; the loads below hit L2). + for jp in cutlass.range_constexpr(ring, ring + prefetch): + if cutlass.Int32(jp) < count: + qp = lo + cutlass.Int32(jp) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, + [ + cutlass.Int32(0), + (qp // cutlass.Int32(k_tiles)) * cutlass.Int32(CTA_M), + (qp % cutlass.Int32(k_tiles)) * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), + cutlass.Int32(0), + ], + [], # tile mode: no im2col offsets + ) + prims.griddepcontrol(prims.GridDepAction.WAIT) + for jj in cutlass.range_constexpr(ring): + if cutlass.Int32(jj) < count: + kb0 = (lo + cutlass.Int32(jj)) % cutlass.Int32(k_tiles) + for half in cutlass.range_constexpr(TMA_COPY_ITERS): + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(jj * MMA_N * CTA_K + half * B_HALF_ELEMS), + tma_ptr_x, + (kb0 * cutlass.Int32(CTA_K) + cutlass.Int32(half * TMA_K_BOX), cutlass.Int32(0)), + full.subview(jj), + ) # fmt: skip + n = cutlass.Int32(ring) + while n < count: + stage = n % cutlass.Int32(ring) + # The stage's previous MMAs committed. + while not cute.arch.mbarrier_try_wait( + empty.subview(stage).data_ptr(), + ((n // cutlass.Int32(ring)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + ): + pass + q = lo + n + tile = q // cutlass.Int32(k_tiles) + k = q % cutlass.Int32(k_tiles) + prims.mbarrier_arrive_expect_tx(full.subview(stage), A_BYTES + B_BYTES) + if tile < keep_tiles: + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(stage * cutlass.Int32(CTA_M * CTA_K)), tma_ptr_w, + (cutlass.Int32(0), tile * cutlass.Int32(CTA_M), k * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), cutlass.Int32(0)), + full.subview(stage), + ) # fmt: skip + else: + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(stage * cutlass.Int32(CTA_M * CTA_K)), tma_ptr_w, + (cutlass.Int32(0), tile * cutlass.Int32(CTA_M), k * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), cutlass.Int32(0)), + full.subview(stage), l2_cache_hint=EVICT_FIRST, + ) # fmt: skip + for half in cutlass.range_constexpr(TMA_COPY_ITERS): + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(stage * cutlass.Int32(MMA_N * CTA_K) + cutlass.Int32(half * B_HALF_ELEMS)), + tma_ptr_x, + (k * cutlass.Int32(CTA_K) + cutlass.Int32(half * TMA_K_BOX), cutlass.Int32(0)), + full.subview(stage), + ) # fmt: skip + n = n + cutlass.Int32(1) + # All of this CTA's loads are issued: the dependents may launch once every CTA gets here. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + elif warp_id == 2: + # ===================================================================== + # MMA: each segment's k-tiles into accumulator (segment count) % 2. + # ===================================================================== + idesc = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, a_dtype=io_dtype, b_dtype=io_dtype, n_dim=MMA_N, m_dim=CTA_M + ) + desc_a_base = prims.Tcgen05SmemDesc.build( + start_address=smem_a, leading_byte_offset=LEADING, stride_byte_offset=STRIDE, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + desc_b_base = prims.Tcgen05SmemDesc.build( + start_address=smem_b, leading_byte_offset=LEADING, stride_byte_offset=STRIDE, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + tmem_base = tmem_ptr_i32.load() + used = cutlass.Int32(0) + seg = cutlass.Int32(0) + q = lo + while q < hi: + seg_tile = q // cutlass.Int32(k_tiles) + seg_end = cutlass.Int32( + cutlass.select_( + (seg_tile + cutlass.Int32(1)) * cutlass.Int32(k_tiles) < hi, + (seg_tile + cutlass.Int32(1)) * cutlass.Int32(k_tiles), + hi, + ) + ) + buf = seg & cutlass.Int32(1) + if seg >= cutlass.Int32(2): + # The epilogue has read this accumulator's previous segment. + while not cute.arch.mbarrier_try_wait( + acc_empty.subview(buf).data_ptr(), + ((seg >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + ): + pass + tmem_acc = cutlass.inttoptr(tmem_base + buf * cutlass.Int32(MMA_N), 6, cutlass.Int32) + j = q + while j < seg_end: + stage = used % cutlass.Int32(ring) + while not cute.arch.mbarrier_try_wait( + full.subview(stage).data_ptr(), (used // cutlass.Int32(ring)) & cutlass.Int32(1) + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for kb in cutlass.range_constexpr(TMA_COPY_ITERS * K_BLOCKS_PER_HALF): + box = kb // K_BLOCKS_PER_HALF + within = kb % K_BLOCKS_PER_HALF + desc_a = desc_a_base + ( + stage * cutlass.Int32(STAGE_A) + cutlass.Int32(box * A_BOX + within * STEP) + ) + desc_b = desc_b_base + ( + stage * cutlass.Int32(STAGE_B) + cutlass.Int32(box * B_BOX + within * STEP) + ) + accumulate = cutlass.Boolean(True) + if cutlass.const_expr(kb == 0): + accumulate = j > q + if prims.elect_sync(): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, + prims.CTAGroup.CTA_1, + tmem_acc, + desc_a, + desc_b, + idesc, + accumulate, + ) + if prims.elect_sync(): + prims.tcgen05_commit(empty.subview(stage)) + used = used + cutlass.Int32(1) + j = j + cutlass.Int32(1) + if prims.elect_sync(): + prims.tcgen05_commit(acc_full.subview(buf)) + seg = seg + cutlass.Int32(1) + q = seg_end + elif warp_id >= 4: + # ===================================================================== + # Epilogue: per segment, TMEM -> registers -> y (a whole tile), or piece + # p > 0's partial + flag, or piece 0's ordered sum of all pieces -> y. + # ===================================================================== + tid = tx - cutlass.Int32(EPI_THREADS) + lane = tx % 32 + w = warp_id - 4 + row = w * cutlass.Int32(32) + lane + tmem_base = tmem_ptr_i32.load() + prims.griddepcontrol(prims.GridDepAction.WAIT) + seg = cutlass.Int32(0) + q = lo + while q < hi: + seg_tile = q // cutlass.Int32(k_tiles) + seg_end = cutlass.Int32( + cutlass.select_( + (seg_tile + cutlass.Int32(1)) * cutlass.Int32(k_tiles) < hi, + (seg_tile + cutlass.Int32(1)) * cutlass.Int32(k_tiles), + hi, + ) + ) + buf = seg & cutlass.Int32(1) + acc_addr = ( + tmem_base + + ((w * cutlass.Int32(32)) << cutlass.Int32(16)) + + buf * cutlass.Int32(MMA_N) + ) + acc_parity = (seg >> cutlass.Int32(1)) & cutlass.Int32(1) + # The tile's pieces: the CTAs whose ranges hold its first and last k-tiles, and everything between. + first_cta = ( + (seg_tile * cutlass.Int32(k_tiles) + cutlass.Int32(1)) * cutlass.Int32(grid) + - cutlass.Int32(1) + ) // cutlass.Int32(total) + last_cta = ( + (seg_tile + cutlass.Int32(1)) * cutlass.Int32(k_tiles) * cutlass.Int32(grid) + - cutlass.Int32(1) + ) // cutlass.Int32(total) + pieces = last_cta - first_cta + cutlass.Int32(1) + piece = bx - first_cta + slot0 = seg_tile * cutlass.Int32(max_pieces) + if pieces == cutlass.Int32(1): + while not cute.arch.mbarrier_try_wait(acc_full.subview(buf).data_ptr(), acc_parity): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", cutlass.inttoptr(acc_addr, 6, cutlass.Float32), num=MMA_N + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if prims.elect_sync(): + prims.mbarrier_arrive(acc_empty.subview(buf)) + _store_tile( + stage_out, + y, + [cutlass.Float32(acc[t]) for t in range(MMA_N)], + row, + tid, + seg_tile, + n_out, + ) + elif piece == cutlass.Int32(0): + # The finalizer. The other pieces' partials first, while this segment's MMAs run: thread 0 + # acquires the flags (backing off between polls, so the epilogue warps do not load the memory + # system while the weight streams), the barrier orders every thread's loads after its acquires, + # then each thread loads its row (a slot past the piece count reloads piece 1's and is not added). + if tid == cutlass.Int32(0): + for p in cutlass.range_constexpr(1, max_pieces): + if cutlass.Int32(p) < pieces: + while _ld_acquire( + flags.subview(slot0 + cutlass.Int32(p)).data_ptr().toint() + ) == cutlass.Int32(0): + prims.nanosleep(FLAG_BACKOFF_NS) + prims.barrier_cta_sync(1, thread_count=EPI_THREADS) + pre = [] + for p in cutlass.range_constexpr(1, max_pieces): + live = cutlass.Int32(p) < pieces + src = ( + ( + slot0 + + cutlass.Int32( + cutlass.select_(live, cutlass.Int32(p), cutlass.Int32(1)) + ) + ) + * cutlass.Int32(CTA_M) + + row + ) * cutlass.Int32(MMA_N) + pre.append( + ( + ws.load(idx=src, vector_size=4, alignment=16, is_volatile=True), + ws.load( + idx=src + cutlass.Int32(4), + vector_size=4, + alignment=16, + is_volatile=True, + ), + ) + ) + while not cute.arch.mbarrier_try_wait(acc_full.subview(buf).data_ptr(), acc_parity): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", cutlass.inttoptr(acc_addr, 6, cutlass.Float32), num=MMA_N + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if prims.elect_sync(): + prims.mbarrier_arrive(acc_empty.subview(buf)) + # Piece 0 (k from 0), then pieces 1, 2, ... in k order. + total_v = [cutlass.Float32(acc[t]) for t in range(MMA_N)] + for p in cutlass.range_constexpr(1, max_pieces): + live = cutlass.Int32(p) < pieces + lo_v, hi_v = pre[p - 1] + for t in cutlass.range_constexpr(MMA_N): + got = cutlass.Float32(lo_v[t]) if t < 4 else cutlass.Float32(hi_v[t - 4]) + total_v[t] = cutlass.Float32( + cutlass.select_(live, _add_rn(total_v[t], got), total_v[t]) + ) + _store_tile(stage_out, y, total_v, row, tid, seg_tile, n_out) + # Every epilogue thread's acquires are behind the store helper's barriers: lower the flags. + if tid == cutlass.Int32(0): + for p in cutlass.range_constexpr(1, max_pieces): + if cutlass.Int32(p) < pieces: + flags.store(cutlass.Int32(0), idx=slot0 + cutlass.Int32(p)) + else: + while not cute.arch.mbarrier_try_wait(acc_full.subview(buf).data_ptr(), acc_parity): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", cutlass.inttoptr(acc_addr, 6, cutlass.Float32), num=MMA_N + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if prims.elect_sync(): + prims.mbarrier_arrive(acc_empty.subview(buf)) + part = ((slot0 + piece) * cutlass.Int32(CTA_M) + row) * cutlass.Int32(MMA_N) + ws.store( + cutlass.Vector.from_elements( + ( + cutlass.Float32(acc[0]), + cutlass.Float32(acc[1]), + cutlass.Float32(acc[2]), + cutlass.Float32(acc[3]), + ), + cutlass.Float32, + ), + idx=part, + vector_size=4, + alignment=16, + ) + ws.store( + cutlass.Vector.from_elements( + ( + cutlass.Float32(acc[4]), + cutlass.Float32(acc[5]), + cutlass.Float32(acc[6]), + cutlass.Float32(acc[7]), + ), + cutlass.Float32, + ), + idx=part + cutlass.Int32(4), + vector_size=4, + alignment=16, + ) + # Every epilogue thread's partial is stored; thread 0's release covers them. + prims.barrier_cta_sync(1, thread_count=EPI_THREADS) + if tid == cutlass.Int32(0): + _flag_release(flags.subview(slot0 + piece).data_ptr().toint()) + seg = seg + cutlass.Int32(1) + q = seg_end + prims.barrier_cta_sync(1, thread_count=EPI_THREADS) + if warp_id == 4: + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_base, 6, cutlass.Int32), TMEM_COLS) + + +def _weight_tensor_map(w, n_out, k_in): + """W as five TMA dimensions (64-element column chunk, row, 64-element chunk index, 1, 1) so one call + per k-tile lands both 128-byte-swizzled halves; strides in 16-byte units.""" + return cuda.create_tensor_map_tiled( + global_address=w.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[TMA_K_BOX, n_out, k_in // TMA_K_BOX, 1, 1], + global_strides=[ + (k_in * ELEM_BYTES) // 16, + (TMA_K_BOX * ELEM_BYTES) // 16, + (n_out * k_in * ELEM_BYTES) // 16, + (n_out * k_in * ELEM_BYTES) // 16, + ], + box_dims=[TMA_K_BOX, CTA_M, TMA_COPY_ITERS, 1, 1], + swizzle=cuda.TensorMapSwizzle.s128b, + ) + + +def _activation_tensor_map(x, cols, num_tokens): + """x [M, cols] (rows dense) as (cols, M) with an 8-row box: rows past num_tokens arrive as zeros.""" + return cuda.create_tensor_map_tiled( + global_address=x.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[cols, num_tokens], + global_strides=[(cols * ELEM_BYTES) // 16], + box_dims=[TMA_K_BOX, MMA_N], + swizzle=cuda.TensorMapSwizzle.s128b, + ) + + +@cute.jit +def k3_head_gemv( + w: cute.Tensor, # [N, K] bf16, K contiguous + x: cute.Tensor, # [M, K] bf16, K contiguous, M <= 8 + y: cute.Tensor, # [8 * N] bf16 + ws: cute.Tensor, # fp32 [units * 128 * 8] + cnt: cute.Tensor, # int32 [tiles], zero + claim: cute.Tensor, # int32 [1], zero + num_tokens: cutlass.Int32, + keep_tiles: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + chunk_tiles: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + grid: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + """``y = x @ w^T`` over ``grid`` persistent CTAs.""" + tma_desc_w = _weight_tensor_map(w, n_out, k_in) + tma_desc_x = _activation_tensor_map(x, k_in, num_tokens) + k3_head_gemv_kernel( + tma_desc_w, + tma_desc_x, + y, + ws, + cnt, + claim, + num_tokens, + keep_tiles, + n_out, + k_in, + chunk_tiles, + ring, + grid, + ).launch( + grid=[grid, 1, 1], + block=[THREADS, 1, 1], + stream=stream, + use_pdl=use_pdl, + ) + + +@cute.jit +def k3_head_gemv_sk( + w: cute.Tensor, # [N, K] bf16, K contiguous + x: cute.Tensor, # [M, K] bf16, K contiguous, M <= 8 + y: cute.Tensor, # [8 * N] bf16 + ws: cute.Tensor, # fp32 [tiles * max_pieces * 128 * 8] + flags: cute.Tensor, # int32 [tiles * max_pieces], zero + num_tokens: cutlass.Int32, + keep_tiles: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + grid: cutlass.Constexpr[int], + max_pieces: cutlass.Constexpr[int], + prefetch: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + """``y = x @ w^T`` over ``grid`` stream-K CTAs.""" + tma_desc_w = _weight_tensor_map(w, n_out, k_in) + tma_desc_x = _activation_tensor_map(x, k_in, num_tokens) + k3_head_gemv_sk_kernel( + tma_desc_w, + tma_desc_x, + y, + ws, + flags, + num_tokens, + keep_tiles, + n_out, + k_in, + ring, + grid, + max_pieces, + prefetch, + ).launch( + grid=[grid, 1, 1], + block=[THREADS, 1, 1], + stream=stream, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_head_gemv/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_head_gemv/op.py new file mode 100644 index 000000000000..91e0bd132da1 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_head_gemv/op.py @@ -0,0 +1,268 @@ +# 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::k3_head_gemv``: ``x @ weight^T`` for a large weight (an lm_head vocab shard) and M <= 8 rows. + +One CTA per SM, and under stream-K at most one per (tile, k-tile), streams the weight, EVICT_FIRST except for the +first ``keep_tiles`` 128-row tiles, and adds each tile's k-split partials in a fixed order before one bf16 rounding: +the result is bit-identical from run to run (not bit-identical to cuBLAS, whose accumulation order differs). Two +schedules: ``streamk`` (every CTA an equal share of the flat (tile, k-tile) space; the default) and +``dynamic`` ((tile, k-chunk) units claimed from a counter, ``chunk_tiles`` k-tiles each). + +A call's cross-launch state is a caller-owned ``K3HeadGemvWorkspace`` (the partials, the per-tile flag or count words +and the unit counter, the counters zero between launches), created for one weight shape and schedule by its eager +``create()`` before any CUDA-graph capture. The calls that share one workspace must be ordered on one stream. Each +kernel compiles on the first call for its shape, which must also happen outside CUDA-graph capture. +""" + +from __future__ import annotations + +import os +import threading +from dataclasses import dataclass +from typing import Dict + +import torch + +MAX_TOKENS = 8 +UNITS_PER_CTA = 4 # the chunk choice keeps at least this many units per CTA for the dynamic balance + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + + +def _arg(t: torch.Tensor): + from cutlass.cute.runtime import from_dlpack + + # detach(): DLPack refuses tensors that require grad (weights are parameters). + return from_dlpack(t.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +def _use_pdl() -> bool: + return os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +def _stream(t: torch.Tensor): + import cuda.bindings.driver as cuda_driver + + return cuda_driver.CUstream(torch.cuda.current_stream(t.device).cuda_stream) + + +def _grid(device) -> int: + return torch.cuda.get_device_properties(device).multi_processor_count + + +def pick_chunk(n_out: int, k_in: int, grid: int) -> int: + """k-tiles per unit: the largest divisor of the k-tiles that leaves at least UNITS_PER_CTA units per CTA, or the + smallest divisor above 1 when none does (every divisor if the k-tiles are prime).""" + from . import k3_head_gemv_kernel as kernel + + k_tiles = kernel.num_k_tiles(k_in) + tiles = n_out // kernel.CTA_M + divisors = [d for d in range(1, k_tiles + 1) if k_tiles % d == 0] + fitting = [d for d in divisors if tiles * (k_tiles // d) >= UNITS_PER_CTA * grid] + if fitting: + return max(fitting) + return min([d for d in divisors if d > 1] or divisors) + + +def supports( + x: torch.Tensor, + weight: torch.Tensor, + chunk_tiles: int = 0, + ring: int = 6, + schedule: str = "streamk", +) -> bool: + """Whether ``k3_head_gemv`` runs ``x @ weight^T`` (``chunk_tiles`` 0: the op's choice).""" + from . import k3_head_gemv_kernel as kernel + + if not ( + x.is_cuda + and x.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and x.dim() == 2 + and weight.dim() == 2 + and 0 < x.shape[0] <= MAX_TOKENS + and x.shape[1] == weight.shape[1] + and x.is_contiguous() + and weight.is_contiguous() + and weight.shape[0] % kernel.CTA_M == 0 + and weight.shape[1] % kernel.CTA_K == 0 + ): + return False + n_out, k_in = weight.shape + if schedule == "streamk": + return kernel.streamk_supports(n_out, k_in, ring) + if schedule != "dynamic": + return False + chunk = chunk_tiles or pick_chunk(n_out, k_in, _grid(x.device)) + return kernel.supports(n_out, k_in, chunk, ring) + + +def _launch_config(n_out: int, k_in: int, device, chunk_tiles: int, schedule: str): + """(grid, units, flag words, pieces or k-tiles per unit) of a launch: the workspace sizes follow from them.""" + from . import k3_head_gemv_kernel as kernel + + tiles = n_out // kernel.CTA_M + grid = _grid(device) + if schedule == "streamk": + # At most one CTA per (tile, k-tile): every CTA's share is then non-empty, so each CTA between a split tile's + # first and last piece holds part of that tile and raises the flag its finalizer waits for. + grid = min(grid, tiles * kernel.num_k_tiles(k_in)) + pieces = kernel.streamk_max_pieces(n_out, k_in, grid) + return grid, tiles * pieces, tiles * pieces, pieces + chunk = chunk_tiles or pick_chunk(n_out, k_in, grid) + return grid, kernel.num_units(n_out, k_in, chunk), tiles, chunk + + +@dataclass(frozen=True) +class K3HeadGemvWorkspace: + """The cross-launch state of ``trtllm::k3_head_gemv`` for one weight shape and schedule. + + ``partials`` fp32 [units * 128 * 8] holds the fp32 partials of a split tile's pieces after the first (stream-K) + or of each unit (dynamic). ``flags`` int32 holds one word per stream-K piece, raised by the piece and lowered by + its tile's finalizer, or one unit count per tile (dynamic). ``claim`` int32 [1] is the dynamic schedule's unit + counter (unused by stream-K). Every launch leaves the words at zero for the next one. ``create()`` allocates them + zeroed, outside CUDA-graph capture. Calls of the same shape and schedule may share one workspace when they are + ordered on one stream. + """ + + n_out: int + k_in: int + schedule: str + chunk_tiles: int + partials: torch.Tensor + flags: torch.Tensor + claim: torch.Tensor + + @classmethod + def create( + cls, + n_out: int, + k_in: int, + device: torch.device, + schedule: str = "streamk", + chunk_tiles: int = 0, + ) -> "K3HeadGemvWorkspace": + """The workspace of [n_out, k_in] weights on ``device`` (``chunk_tiles``: the dynamic schedule's k-tiles per + unit, 0 for the op's choice).""" + if schedule not in ("streamk", "dynamic"): + raise ValueError(f"k3_head_gemv: unknown schedule {schedule!r}") + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("k3_head_gemv: create the workspace outside CUDA-graph capture") + device = torch.device(device) + _, units, words, per = _launch_config(n_out, k_in, device, chunk_tiles, schedule) + return cls( + n_out=n_out, + k_in=k_in, + schedule=schedule, + chunk_tiles=per if schedule == "dynamic" else 0, + partials=torch.empty(max(units, 1) * 128 * 8, dtype=torch.float32, device=device), + flags=torch.zeros(words, dtype=torch.int32, device=device), + claim=torch.zeros(1, dtype=torch.int32, device=device), + ) + + +def _compiled_fn(key, entry, *args): + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "k3_head_gemv: run once per shape outside CUDA-graph capture first (it compiles)." + ) + import cutlass.cute as cute + + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile(entry, *args) + return fn + + +@torch.library.custom_op("trtllm::k3_head_gemv", mutates_args=("partials", "flags", "claim")) +def k3_head_gemv( + x: torch.Tensor, + weight: torch.Tensor, + partials: torch.Tensor, + flags: torch.Tensor, + claim: torch.Tensor, + keep_tiles: int = 0, + chunk_tiles: int = 0, + ring: int = 6, + schedule: str = "streamk", + prefetch: int = 16, +) -> torch.Tensor: + """``x @ weight.T`` for bf16 ``x`` [M <= 8, K] and ``weight`` [N, K] (N and K multiples of 128); returns bf16 + [M, N]. ``partials`` / ``flags`` / ``claim`` are a ``K3HeadGemvWorkspace``'s, created for this shape and + schedule (and ``chunk_tiles``). Weight tiles (128 rows) below ``keep_tiles`` are read at normal L2 priority, the + rest EVICT_FIRST. ``schedule``: ``streamk`` or ``dynamic``; ``chunk_tiles`` (dynamic; 0: the op's choice) sets the + k-tiles per work unit; ``prefetch`` (stream-K) the k-tiles after the ring that each CTA prefetches into L2 before + the grid dependency.""" + if not supports(x, weight, chunk_tiles, ring, schedule): + raise ValueError( + f"k3_head_gemv: unsupported call x {tuple(x.shape)} {x.dtype}, weight {tuple(weight.shape)} " + f"{weight.dtype}, chunk_tiles {chunk_tiles}, ring {ring}, schedule {schedule}" + ) + from . import k3_head_gemv_kernel as kernel + + num_tokens, k_in = x.shape + n_out = weight.shape[0] + grid, units, words, per = _launch_config(n_out, k_in, x.device, chunk_tiles, schedule) + # A workspace of another shape or schedule is too small, or keeps its counters elsewhere: refused, not raced. + if not ( + partials.dtype == torch.float32 + and flags.dtype == torch.int32 + and claim.dtype == torch.int32 + and partials.device == x.device + and flags.device == x.device + and claim.device == x.device + and partials.numel() >= units * 128 * 8 + and flags.numel() == words + and claim.numel() == 1 + ): + raise ValueError( + f"k3_head_gemv: the workspace (partials {tuple(partials.shape)}, flags {tuple(flags.shape)}) is not " + f"one for weight {tuple(weight.shape)}, schedule {schedule}, chunk_tiles {chunk_tiles}" + ) + # All 8 token rows are written (the rows past M from TMA's zero-filled x rows); the result is the first M. + y = torch.empty(kernel.MMA_N, n_out, dtype=torch.bfloat16, device=x.device) + stream = _stream(x) + use_pdl = _use_pdl() + if schedule == "streamk": + args = (_arg(weight), _arg(x), _arg(y.view(-1)), _arg(partials), _arg(flags)) + fn = _compiled_fn(("k3_head_gemv_sk", n_out, k_in, ring, grid, per, prefetch, use_pdl), + kernel.k3_head_gemv_sk, *args, num_tokens, keep_tiles, n_out, k_in, ring, grid, per, + prefetch, use_pdl, stream) # fmt: skip + else: + args = (_arg(weight), _arg(x), _arg(y.view(-1)), _arg(partials), _arg(flags), _arg(claim)) + fn = _compiled_fn(("k3_head_gemv", n_out, k_in, per, ring, grid, use_pdl), kernel.k3_head_gemv, *args, + num_tokens, keep_tiles, n_out, k_in, per, ring, grid, use_pdl, stream) # fmt: skip + fn(*args, num_tokens, keep_tiles, stream) + return y[:num_tokens] + + +@k3_head_gemv.register_fake +def _( + x, + weight, + partials, + flags, + claim, + keep_tiles=0, + chunk_tiles=0, + ring=6, + schedule="streamk", + prefetch=16, +): + return x.new_empty((x.shape[0], weight.shape[0]), dtype=torch.bfloat16) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_attn/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_attn/__init__.py new file mode 100644 index 000000000000..cc433eba6f41 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_attn/__init__.py @@ -0,0 +1,19 @@ +# 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 attention projection as one CTM kernel (``trtllm::k3_kda_qkvg``). + +Importing :mod:`.op` registers the torch op; nothing is imported eagerly here so that +the CuTe DSL dependency stays optional for every other model. +""" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_attn/k3_kda_attn_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_attn/k3_kda_attn_kernel.py new file mode 100644 index 000000000000..3ff125d1c3b4 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_attn/k3_kda_attn_kernel.py @@ -0,0 +1,1561 @@ +# 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 fused projection y = x W^T (TP16 rank slice: W bf16 [3208, 7168], x bf16 [T <= 8, 7168]) streamed in +three priority phases, so that a consumer of the projection can start on the rows it needs first. + +Rows of W (= columns of y): q [0, 768) | k [768, 1536) | v [1536, 2304) | og [2304, 3072) | f_a [3072, 3200) | +b [3200, 3206) | pad [3206, 3208). + +Grid: 26 clusters of 4 CTAs, 256 threads. Cluster k, rank r: + phase 1 (q | k | f_a, 26 tiles of 64 rows): tile k, K chunks [28 r, 28 r + 28) (split-K 4 over the cluster); + phase 2 (v, 12 tiles, then b): cluster k < 24: v tile k // 2, K half k % 2, chunks 56 (k % 2) + [14 r, 14 r + 14) + (split-K 8 = two clusters per tile); clusters 24 and 25: the b tile (8 rows), K half k - 24, likewise; + phase 3 (og, 12 tiles): clusters k < 24 as phase 2; clusters 24 and 25 have no phase 3. +Warps: 0 weight TMA (issued from launch, before the grid dependency, EVICT_FIRST), 1 activation TMA (after it), +2 TMEM allocation and the M = 64, N = 8 MMAs (weight = A, activation = B; each phase accumulates into its own 8 +TMEM columns so a phase's epilogue runs beside the next phase's MMAs), 3 idle, 4-7 epilogue: warp 4 + w loads TMEM +rows [16 w, 16 w + 16) of the tile (16x256b), which rank w of the cluster owns. The other ranks st.async them into +the owner's mailbox; the owner adds the four partials in rank order and publishes them Lamport-style (the data words +are the flags: no fence, no counter, so nothing waits for the stores to be acknowledged under the weight stream): + phase 1: bf16 bits into ``p1`` [3 buffers][8 tokens][1664 rows (q | k | f_a)]; + phases 2 and 3: the cluster's fp32 partial bits into ``part`` [3 buffers][region (v, og, b)][half][8][768]; the + projection is bf16(half 0 + half 1), the consumer's job. +Every word of a buffer holds the sentinel (all ones: a NaN no finite GEMV produces; a computed all-ones word is +stored as the canonical NaN instead) until its producer writes it, so a consumer polls the words it needs until none +is the sentinel. A launch reads e = ``epoch[cta]`` (the CTA's buffer index, its launch count mod 3) after the grid +dependency, writes buffer e and stores the sentinel into the same words of buffer (e + 1) % 3, which the next launch +writes and the launch before last read, then writes (e + 1) % 3 back at its end. The index stays in 0..2: a raw +launch count would turn negative after 2^31 launches and its signed remainder would index before the buffers. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.base_dsl.array +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +try: + from cutlass.memory.smem import SmemAllocator +except ImportError: # older DSL layout + from cutlass.utils import SmemAllocator + +from ..k3_kda_verify.k3_kda_verify_kernel import ( + _bf16, + _butterfly, + _st_async_f32, + _store8, + _test_wait_cluster, +) + +K_IN = 7168 +HK = 768 # 6 local heads x 128 +PROJ_ROWS = 4 * HK + 128 + 8 # 3208 +V_ROW = 2 * HK +OG_ROW = 3 * HK +FA_ROW = 4 * HK +B_ROW = FA_ROW + 128 +TILE = 64 # weight rows per tile = the MMA's M +MMA_N = 8 +MMA_K = 16 +BOX_K = 64 # bf16 per 128-byte swizzled row +BOX_CH = 2 # 64-column chunks per weight box +BOX_ELEMS = TILE * BOX_K * BOX_CH # 16 KB +X_ELEMS = MMA_N * BOX_K * BOX_CH # 2 KB +STREAM_STAGES = 10 # weight ring depth of the projection alone (B1) +FUSED_STAGES = 12 # B2's stream CTAs: the verify role's buffers alias the same bytes (_SmemCarver) +P3_WINDOW = 3 # phase-3 weight boxes in flight per stream CTA +CLUSTER = 4 +FUSED_RAW_BYTES = ( + FUSED_STAGES * (BOX_ELEMS + X_ELEMS) * 2 + 3 * CLUSTER * 32 * 4 * 4 +) # the stream role's carve; the verify role's fits inside +STREAM_CLUSTERS = 26 +P1_BOXES = K_IN // (BOX_K * BOX_CH) // CLUSTER # 14 +P23_BOXES = P1_BOXES // 2 # 7 +THREADS = 256 +TMEM_COLS = 32 +PART_ROWS = HK +P1_ROWS = 2 * HK + 128 # q | k | f_a +P1_BUF = 8 * P1_ROWS +PART_BUF = 3 * 2 * 8 * PART_ROWS +BUFFERS = 3 +SENTINEL = -1 # all-ones bits +CANON_NAN16 = 0x7FC0 +CANON_NAN32 = 0x7FC00000 +EVICT_FIRST = 0x12F0000000000000 # createpolicy.fractional.L2::evict_first, fraction 1.0 + +# KDA verify role (stage B2): clusters 26-31 = local heads 0-5, cluster rank = V quarter (32 V rows, 4 per warp). +HD = 128 # key and value head dim +H_LOCAL = HK // HD +CONV_W = 4 +NT = 8 # verify tokens: 1 golden + 7 drafts (one request) +NUM_SPEC = NT - 1 +S_COLS = CONV_W - 1 + NUM_SPEC # conv-cache columns +ROWS_U = CONV_W - 1 + NT # raw conv inputs by position: 3 before token 0, then the tokens +V_CTA = HD // CLUSTER # 32 +REC_ROWS = V_CTA // 8 # V rows per warp in the recurrence +VEC = HD // 32 # keys per lane +# The recurrence's registers (one rmem array, static indices): the state [REC_ROWS][VEC], then the previous token's +# per-lane q-dot partials [REC_ROWS], then the current token's operands: v rows [REC_ROWS], q, decay * k, beta * k and +# decay [VEC] each. +R_Q = REC_ROWS * VEC +R_OP = R_Q + REC_ROWS +REC_REGS = R_OP + REC_ROWS + 4 * VEC +# The drafts' records, in each slot's per-token state region (fp32 words from the slot's start) in place of their +# full states: the row innovations vn [NUM_SPEC][H][V], then beta * k and the decay [NUM_SPEC][H][K]. The next launch +# rebuilds the state after the accepted drafts from the pool (the golden token's state) with the update's own +# arithmetic, S = fma(decay, S, vn * (beta * k)) draft by draft, so it is bit-identical to the one the drafts reached. +# k3_kda_verify writes and reads the same records. +CT_VN = 0 +CT_WB = NUM_SPEC * H_LOCAL * HD +CT_WD = CT_WB + NUM_SPEC * H_LOCAL * HD +CT_SLOT = NUM_SPEC * H_LOCAL * HD * HD # a slot's per-token state region +REC_CTA = ( + V_CTA + 2 * HD +) # one draft's records a verify CTA replays: vn of its 32 rows, beta * k and decay of 128 keys +PEND_WORDS = ( + 4 # pending counts a warp loads beside the slot index: pools up to 32 * PEND_WORDS slots +) +CP_CG = cutlass.base_dsl.array.LoadCacheModifier.CG # cp.async 16 B through L2 only +HEAD_CLUSTERS = H_LOCAL +FB_TMEM_COLS = 32 + +# Shared-memory descriptor units (16 bytes). +W_STAGE_U = (BOX_ELEMS * 2) >> 4 +W_CHUNK_U = (TILE * BOX_K * 2) >> 4 +X_STAGE_U = (X_ELEMS * 2) >> 4 +X_CHUNK_U = (MMA_N * BOX_K * 2) >> 4 +K_STEP_U = (MMA_K * 2) >> 4 + + +@dsl_user_op +def _mapa_u32(smem_ptr, peer, *, loc=None, ip=None): + """The shared::cluster address of this CTA's shared-memory location in cluster CTA ``peer``.""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [smem_ptr.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(peer).ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_async_v4_f32(dst, a, b, c, d, mbar, *, loc=None, ip=None): + """st.async of four fp32 to a shared::cluster address, completing ``mbar`` (shared::cluster) by 16 bytes.""" + _llvm.inline_asm( + None, + [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Float32(a).ir_value(loc=loc, ip=ip), + cutlass.Float32(b).ir_value(loc=loc, ip=ip), cutlass.Float32(c).ir_value(loc=loc, ip=ip), + cutlass.Float32(d).ir_value(loc=loc, ip=ip), cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.f32 [$0], {$1, $2, $3, $4}, [$5];", "r,f,f,f,f,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """The bf16 rounding (round to nearest even) of an fp32, as its 16 bits in the low half of an int32.""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], + "{ .reg .b16 h; cvt.rn.bf16.f32 h, $1; cvt.u32.u16 $0, h; }", "=r,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _f32_bits(x, *, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], "mov.b32 $0, $1;", "=r,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_b16(ptr, bits, *, loc=None, ip=None): + """st.global.b16 of the low 16 bits of an int32 at a global address (int64).""" + _llvm.inline_asm( + None, [cutlass.Int64(ptr).ir_value(loc=loc, ip=ip), cutlass.Int32(bits).ir_value(loc=loc, ip=ip)], + "{ .reg .b16 h; cvt.u16.u32 h, $1; st.global.b16 [$0], h; }", "l,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _st_f32(ptr, value, *, loc=None, ip=None): + """st.global.f32 of an fp32 at a global address (int64).""" + _llvm.inline_asm( + None, [cutlass.Int64(ptr).ir_value(loc=loc, ip=ip), cutlass.Float32(value).ir_value(loc=loc, ip=ip)], + "st.global.f32 [$0], $1;", "l,f", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _st_f32_if(ptr, value, pred, *, loc=None, ip=None): + """st.global.f32 of an fp32 at a global address (int64) where pred (int32) is non-zero.""" + _llvm.inline_asm( + None, [cutlass.Int64(ptr).ir_value(loc=loc, ip=ip), cutlass.Float32(value).ir_value(loc=loc, ip=ip), + cutlass.Int32(pred).ir_value(loc=loc, ip=ip)], + "{ .reg .pred p; setp.ne.s32 p, $2, 0; @p st.global.f32 [$0], $1; }", "l,f,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +def _lamport16(bits): + """bf16 bits with the sentinel pattern replaced by the canonical NaN.""" + return cutlass.select_(bits == cutlass.Int32(0xFFFF), cutlass.Int32(CANON_NAN16), bits) + + +def _lamport32(bits): + return cutlass.select_(bits == cutlass.Int32(SENTINEL), cutlass.Int32(CANON_NAN32), bits) + + +@cute.jit +def _publish( + p: cutlass.Constexpr[int], + mbox: cutlass.Array, + mbox_bar: cutlass.Array, + p1: cutlass.Array, + part: cutlass.Array, + a0, + a1, + a2, + a3, # this rank's partials: (r0, t0), (r0, t0 + 1), (r0 + 8, t0), (r0 + 8, t0 + 1) + rank, + lane, + cid, + half, + has_p3, + buf, + nbuf, +): + """Owner warp of rows [16 rank, 16 rank + 16) of the tile in phase p: wait for the three peers' partials, add + the four in rank order, and store them Lamport-style into buffer ``buf`` (phase 1: bf16 rows; phases 2-3: the + cluster's fp32 partial), the sentinel into the same words of buffer ``nbuf``. Clusters without rows in the phase + (24-25 in phase 3) store nothing.""" + while not _test_wait_cluster(mbox_bar.subview(p).data_ptr(), 0): + pass + s0 = cutlass.Float32(0.0) + s1 = cutlass.Float32(0.0) + s2 = cutlass.Float32(0.0) + s3 = cutlass.Float32(0.0) + for src in cutlass.range_constexpr(CLUSTER): + peer4 = mbox.load(idx=((p * CLUSTER + src) * 32 + lane) * 4, vector_size=4, alignment=16) + own = rank == cutlass.Int32(src) + s0 = s0 + cutlass.select_(own, a0, cutlass.Float32(peer4[0])) + s1 = s1 + cutlass.select_(own, a1, cutlass.Float32(peer4[1])) + s2 = s2 + cutlass.select_(own, a2, cutlass.Float32(peer4[2])) + s3 = s3 + cutlass.select_(own, a3, cutlass.Float32(peer4[3])) + r0 = rank * 16 + lane // 4 + t0 = (lane % 4) * 2 + if cutlass.const_expr(p == 0): + # p1 [buffer][token][q | k | f_a]: the tile's rows map to 64 cid + r (cid < 24: q, k; 24-25: f_a). + row_p = cid * cutlass.Int32(TILE) + r0 + cur = buf * cutlass.Int32(P1_BUF) + row_p + nxt = nbuf * cutlass.Int32(P1_BUF) + row_p + o00 = t0 * P1_ROWS + o10 = (t0 + 1) * P1_ROWS + _st_b16(p1.subview(cur + o00).data_ptr().toint(), _lamport16(_bf16_bits(s0))) + _st_b16(p1.subview(cur + o10).data_ptr().toint(), _lamport16(_bf16_bits(s1))) + _st_b16(p1.subview(cur + o00 + 8).data_ptr().toint(), _lamport16(_bf16_bits(s2))) + _st_b16(p1.subview(cur + o10 + 8).data_ptr().toint(), _lamport16(_bf16_bits(s3))) + _st_b16(p1.subview(nxt + o00).data_ptr().toint(), cutlass.Int32(0xFFFF)) + _st_b16(p1.subview(nxt + o10).data_ptr().toint(), cutlass.Int32(0xFFFF)) + _st_b16(p1.subview(nxt + o00 + 8).data_ptr().toint(), cutlass.Int32(0xFFFF)) + _st_b16(p1.subview(nxt + o10 + 8).data_ptr().toint(), cutlass.Int32(0xFFFF)) + else: + # Region 0 = v (phase 2, clusters < 24), 1 = og (phase 3), 2 = b (phase 2, clusters 24-25; rows 0-7). + if cutlass.const_expr(p == 1): + region = cutlass.select_(has_p3, cutlass.Int32(0), cutlass.Int32(2)) + lo_ok = has_p3 | (r0 < cutlass.Int32(8)) + else: + region = cutlass.Int32(1) + lo_ok = has_p3 + row_r = cutlass.select_(has_p3, (cid // cutlass.Int32(2)) * cutlass.Int32(TILE) + r0, r0) + slot = (region * cutlass.Int32(2) + half) * cutlass.Int32(8 * PART_ROWS) + row_r + cur_p = buf * cutlass.Int32(PART_BUF) + slot + nxt_p = nbuf * cutlass.Int32(PART_BUF) + slot + if lo_ok: + part.store(_lamport32(_f32_bits(s0)), idx=cur_p + t0 * PART_ROWS) + part.store(_lamport32(_f32_bits(s1)), idx=cur_p + (t0 + 1) * PART_ROWS) + part.store(cutlass.Int32(SENTINEL), idx=nxt_p + t0 * PART_ROWS) + part.store(cutlass.Int32(SENTINEL), idx=nxt_p + (t0 + 1) * PART_ROWS) + if has_p3: + part.store(_lamport32(_f32_bits(s2)), idx=cur_p + t0 * PART_ROWS + 8) + part.store(_lamport32(_f32_bits(s3)), idx=cur_p + (t0 + 1) * PART_ROWS + 8) + part.store(cutlass.Int32(SENTINEL), idx=nxt_p + t0 * PART_ROWS + 8) + part.store(cutlass.Int32(SENTINEL), idx=nxt_p + (t0 + 1) * PART_ROWS + 8) + + +@cute.jit +def _stream_role( + tma_w, + tma_x, + p1: cutlass.Array, + part: cutlass.Array, + epoch: cutlass.Array, + ring_w: cutlass.Array, + ring_x: cutlass.Array, + full: cutlass.Array, + empty: cutlass.Array, + acc_done: cutlass.Array, + mbox_bar: cutlass.Array, + mbox: cutlass.Array, + tmem_holder: cutlass.Array, + USE_PDL: cutlass.Constexpr[bool], + STAGES: cutlass.Constexpr[int], + lbx=None, +): + """One CTA of a stream cluster (cluster ids 0-25): the three projection phases of its tile rows. ``lbx`` is the + CTA's logical index in the fused grid (the block index when None).""" + tidx, _, _ = cute.arch.thread_idx() + bx = cute.arch.block_idx()[0] if lbx is None else lbx + lane = tidx % 32 + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + rank = cute.arch.block_idx_in_cluster() + # The f_a / b clusters (logical 24-25) take the grid's first cluster slots, so they are resident with the + # first CTAs even when a predecessor still holds SMs: every verify CTA needs their phase 1 (f_a). + cid = (bx // cutlass.Int32(CLUSTER) + cutlass.Int32(STREAM_CLUSTERS - 2)) % cutlass.Int32( + STREAM_CLUSTERS + ) + half = cid % cutlass.Int32(2) + has_p3 = cid < cutlass.Int32(2 * (HK // TILE)) + p1_row = cutlass.select_( + cid < cutlass.Int32(2 * HK // TILE), + cid * cutlass.Int32(TILE), + cutlass.Int32(FA_ROW) + (cid - cutlass.Int32(2 * HK // TILE)) * cutlass.Int32(TILE), + ) + p2_row = cutlass.select_( + has_p3, + cutlass.Int32(V_ROW) + (cid // cutlass.Int32(2)) * cutlass.Int32(TILE), + cutlass.Int32(B_ROW), + ) + p3_row = cutlass.Int32(OG_ROW) + (cid // cutlass.Int32(2)) * cutlass.Int32(TILE) + p1_chunk = rank * cutlass.Int32(P1_BOXES * BOX_CH) + p23_chunk = half * cutlass.Int32(CLUSTER * P23_BOXES * BOX_CH) + rank * cutlass.Int32( + P23_BOXES * BOX_CH + ) + + tma_ptr_w = tma_w.get_ptr() + tma_ptr_x = tma_x.get_ptr() + if warp == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_x) + if prims.elect_sync(): + for s in cutlass.range_constexpr(STAGES): + prims.mbarrier_init(full.subview(s), 2) # the weight and the activation TMA + prims.mbarrier_init(empty.subview(s), 1) + for p in cutlass.range_constexpr(3): + prims.mbarrier_init(acc_done.subview(p), 1) + elif warp == 2: + prims.tcgen05_alloc(tmem_holder, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + elif warp == 3: + if prims.elect_sync(): + for p in cutlass.range_constexpr(3): + prims.mbarrier_init(mbox_bar.subview(p), 1) + prims.mbarrier_arrive_expect_tx(mbox_bar.subview(p), (CLUSTER - 1) * 32 * 16) + prims.fence_mbarrier_init() + # Cluster formation: the peers' mailboxes and their barriers are initialized and addressable from here on. + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + + if warp == 0: + # Weight boxes: phase 1 (14), phase 2 (7), phase 3 (7, clusters 0-23). The first STAGES go out at launch. + if prims.elect_sync(): + for i in cutlass.range_constexpr(P1_BOXES + P23_BOXES): + stage_w = i % STAGES + if cutlass.const_expr(i == STAGES and STAGES < P1_BOXES): + # Phase 1's boxes beyond the ring go to L2 now, before the grid dependency, so that their loads + # after it (once the first stages are consumed) hit L2 instead of HBM. + for i_pf in cutlass.range_constexpr(STAGES, P1_BOXES): + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, + [cutlass.Int32(0), p1_row, p1_chunk + cutlass.Int32(i_pf * BOX_CH), cutlass.Int32(0), + cutlass.Int32(0)], + [], + ) # fmt: skip + if cutlass.const_expr(i >= STAGES): + while not cute.arch.mbarrier_test_wait( + empty.subview(stage_w).data_ptr(), (i // STAGES + 1) % 2 + ): + pass + if cutlass.const_expr(i < P1_BOXES): + row_w = p1_row + chunk_w = p1_chunk + cutlass.Int32(i * BOX_CH) + else: + row_w = p2_row + chunk_w = p23_chunk + cutlass.Int32((i - P1_BOXES) * BOX_CH) + prims.mbarrier_arrive_expect_tx(full.subview(stage_w), BOX_ELEMS * 2) + prims.cp_async_bulk_tensor_shared_cta_global( + ring_w.subview(stage_w * BOX_ELEMS), tma_ptr_w, + (cutlass.Int32(0), row_w, chunk_w, cutlass.Int32(0), cutlass.Int32(0)), full.subview(stage_w), + l2_cache_hint=EVICT_FIRST, + ) # fmt: skip + if has_p3: + for j in cutlass.range_constexpr(P23_BOXES): + i3 = P1_BOXES + P23_BOXES + j + stage_w3 = i3 % STAGES + # Phase 3 (the output gate, needed only after the recurrence) keeps at most P3_WINDOW boxes in + # flight: box i3 waits for box i3 - P3_WINDOW to be consumed, which also frees its own stage. A + # shallower HBM queue keeps the verify CTAs' phase-2 polls short while phase 3 streams. + jw = i3 - P3_WINDOW + while not cute.arch.mbarrier_test_wait( + empty.subview(jw % STAGES).data_ptr(), (jw // STAGES) % 2 + ): + pass + prims.mbarrier_arrive_expect_tx(full.subview(stage_w3), BOX_ELEMS * 2) + prims.cp_async_bulk_tensor_shared_cta_global( + ring_w.subview(stage_w3 * BOX_ELEMS), tma_ptr_w, + (cutlass.Int32(0), p3_row, p23_chunk + cutlass.Int32(j * BOX_CH), cutlass.Int32(0), + cutlass.Int32(0)), + full.subview(stage_w3), l2_cache_hint=EVICT_FIRST, + ) # fmt: skip + if cutlass.const_expr(USE_PDL): + # Dependents launch once the last weight box is in flight: they cannot take HBM bandwidth from the + # boxes this grid still needs. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + elif warp == 1: + if cutlass.const_expr(USE_PDL): + prims.griddepcontrol(prims.GridDepAction.WAIT) + if prims.elect_sync(): + for i in cutlass.range_constexpr(P1_BOXES + P23_BOXES): + stage_x = i % STAGES + if cutlass.const_expr(i >= STAGES): + while not cute.arch.mbarrier_test_wait( + empty.subview(stage_x).data_ptr(), (i // STAGES + 1) % 2 + ): + pass + if cutlass.const_expr(i < P1_BOXES): + chunk_x = p1_chunk + cutlass.Int32(i * BOX_CH) + else: + chunk_x = p23_chunk + cutlass.Int32((i - P1_BOXES) * BOX_CH) + prims.mbarrier_arrive_expect_tx(full.subview(stage_x), X_ELEMS * 2) + for c in cutlass.range_constexpr(BOX_CH): + prims.cp_async_bulk_tensor_shared_cta_global( + ring_x.subview(stage_x * X_ELEMS + c * MMA_N * BOX_K), tma_ptr_x, + ((chunk_x + cutlass.Int32(c)) * cutlass.Int32(BOX_K), cutlass.Int32(0)), full.subview(stage_x), + ) # fmt: skip + if has_p3: + for j in cutlass.range_constexpr(P23_BOXES): + i3x = P1_BOXES + P23_BOXES + j + stage_x3 = i3x % STAGES + while not cute.arch.mbarrier_test_wait( + empty.subview(stage_x3).data_ptr(), (i3x // STAGES + 1) % 2 + ): + pass + prims.mbarrier_arrive_expect_tx(full.subview(stage_x3), X_ELEMS * 2) + for c in cutlass.range_constexpr(BOX_CH): + prims.cp_async_bulk_tensor_shared_cta_global( + ring_x.subview(stage_x3 * X_ELEMS + c * MMA_N * BOX_K), tma_ptr_x, + ((p23_chunk + cutlass.Int32(j * BOX_CH + c)) * cutlass.Int32(BOX_K), cutlass.Int32(0)), + full.subview(stage_x3), + ) # fmt: skip + elif warp == 2: + idesc = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cutlass.BFloat16, + b_dtype=cutlass.BFloat16, + n_dim=MMA_N, + m_dim=TILE, + ) + desc_w = prims.Tcgen05SmemDesc.build( + start_address=ring_w, leading_byte_offset=16, stride_byte_offset=1024, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + desc_x = prims.Tcgen05SmemDesc.build( + start_address=ring_x, leading_byte_offset=16, stride_byte_offset=1024, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + tmem_raw = tmem_holder.load() + for i in cutlass.range_constexpr(P1_BOXES + P23_BOXES): + stage_m = i % STAGES + phase_m = 0 if i < P1_BOXES else 1 + first_m = i == 0 or i == P1_BOXES + while not cute.arch.mbarrier_test_wait( + full.subview(stage_m).data_ptr(), (i // STAGES) % 2 + ): + pass + for kb in cutlass.range_constexpr(BOX_CH * (BOX_K // MMA_K)): + c = kb // (BOX_K // MMA_K) + kk = kb % (BOX_K // MMA_K) + if prims.elect_sync(): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw + cutlass.Int32(MMA_N * phase_m), 6, cutlass.Int32), + desc_w + (stage_m * W_STAGE_U + c * W_CHUNK_U + kk * K_STEP_U), + desc_x + (stage_m * X_STAGE_U + c * X_CHUNK_U + kk * K_STEP_U), + idesc, not (first_m and kb == 0), + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_commit(empty.subview(stage_m)) + if cutlass.const_expr(i == P1_BOXES - 1 or i == P1_BOXES + P23_BOXES - 1): + prims.tcgen05_commit(acc_done.subview(phase_m)) + if has_p3: + for j in cutlass.range_constexpr(P23_BOXES): + i3m = P1_BOXES + P23_BOXES + j + stage_m3 = i3m % STAGES + while not cute.arch.mbarrier_test_wait( + full.subview(stage_m3).data_ptr(), (i3m // STAGES) % 2 + ): + pass + for kb3 in cutlass.range_constexpr(BOX_CH * (BOX_K // MMA_K)): + c3 = kb3 // (BOX_K // MMA_K) + kk3 = kb3 % (BOX_K // MMA_K) + if prims.elect_sync(): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw + cutlass.Int32(2 * MMA_N), 6, cutlass.Int32), + desc_w + (stage_m3 * W_STAGE_U + c3 * W_CHUNK_U + kk3 * K_STEP_U), + desc_x + (stage_m3 * X_STAGE_U + c3 * X_CHUNK_U + kk3 * K_STEP_U), + idesc, not (j == 0 and kb3 == 0), + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_commit(empty.subview(stage_m3)) + if cutlass.const_expr(j == P23_BOXES - 1): + prims.tcgen05_commit(acc_done.subview(2)) + else: + # No phase-3 rows: the phase completes empty (its epilogue reduces unused columns and stores nothing). + if prims.elect_sync(): + prims.tcgen05_commit(acc_done.subview(2)) + elif warp >= 4: + w = ( + warp - 4 + ) # TMEM lanes 32 w .. 32 w + 31 = tile rows 16 w .. 16 w + 15, owned by cluster rank w + if cutlass.const_expr(USE_PDL): + prims.griddepcontrol(prims.GridDepAction.WAIT) + e = epoch.load(idx=bx) + buf = e % cutlass.Int32(BUFFERS) + nbuf = (e + cutlass.Int32(1)) % cutlass.Int32(BUFFERS) + tmem_raw_e = tmem_holder.load() + lane_base = (tmem_raw_e >> 16) + w * 32 + col_base = tmem_raw_e & 0xFFFF + for p in cutlass.range_constexpr(3): + while not cute.arch.mbarrier_test_wait(acc_done.subview(p).data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "16x256b", + cutlass.inttoptr((lane_base << 16) | (col_base + p * MMA_N), 6, cutlass.Float32), + num=1, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if w != rank: + _st_async_v4_f32( + _mapa_u32(mbox.subview(((p * CLUSTER + rank) * 32 + lane) * 4).data_ptr(), w), + cutlass.Float32(acc[0]), cutlass.Float32(acc[1]), cutlass.Float32(acc[2]), cutlass.Float32(acc[3]), + _mapa_u32(mbox_bar.subview(p).data_ptr(), w), + ) # fmt: skip + else: + _publish(p, mbox, mbox_bar, p1, part, cutlass.Float32(acc[0]), cutlass.Float32(acc[1]), + cutlass.Float32(acc[2]), cutlass.Float32(acc[3]), rank, lane, cid, half, has_p3, buf, + nbuf) # fmt: skip + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.barrier_cta_sync(1, thread_count=128) + if warp == 4: + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_raw_e, 6, cutlass.Int32), TMEM_COLS) + if lane == 0: + epoch.store(nbuf, idx=bx) + + +def _rows_out(shq, lane): + """The warp's 4 row outputs from each lane's q-dot partials: two reduce-scatter rounds (bits 4, 3) then + butterflies (bits 2, 1, 0), the butterfly's pairs in its order (bit-identical); lanes 8 m hold row + 2 (m >> 1) + (m & 1).""" + up4 = (lane & cutlass.Int32(16)) != cutlass.Int32(0) + w0 = cutlass.select_(up4, shq[2], shq[0]) + cute.arch.shuffle_sync_bfly( + cutlass.select_(up4, shq[0], shq[2]), offset=16, mask=-1, mask_and_clamp=31 + ) + w1 = cutlass.select_(up4, shq[3], shq[1]) + cute.arch.shuffle_sync_bfly( + cutlass.select_(up4, shq[1], shq[3]), offset=16, mask=-1, mask_and_clamp=31 + ) + up3 = (lane & cutlass.Int32(8)) != cutlass.Int32(0) + xq = cutlass.select_(up3, w1, w0) + cute.arch.shuffle_sync_bfly( + cutlass.select_(up3, w0, w1), offset=8, mask=-1, mask_and_clamp=31 + ) + for offset in [4, 2, 1]: + xq = xq + cute.arch.shuffle_sync_bfly(xq, offset=offset, mask=-1, mask_and_clamp=31) + return xq + + +def _load_rec_operands(r_st: cutlass.Array, t, warp, lane, s_v, s_q, s_kd, s_bk, s_dec): + """Token t's recurrence operands into ``r_st[R_OP:]`` (v of the warp's rows, then the lane's keys of q, + decay * k, beta * k and decay).""" + v4 = s_v.load(idx=t * V_CTA + warp * REC_ROWS, vector_size=4, alignment=16) + for j in range(REC_ROWS): + r_st.store(cutlass.Float32(v4[j]), idx=R_OP + j) + for i in range(VEC): + c = t * HD + i * 32 + lane + r_st.store(s_q.load(idx=c), idx=R_OP + REC_ROWS + i) + r_st.store(s_kd.load(idx=c), idx=R_OP + REC_ROWS + VEC + i) + r_st.store(s_bk.load(idx=c), idx=R_OP + REC_ROWS + 2 * VEC + i) + r_st.store(s_dec.load(idx=c), idx=R_OP + REC_ROWS + 3 * VEC + i) + + +def _sent16(word): + """Whether either bf16 in an int32 word is the Lamport sentinel (all ones).""" + return ((word & cutlass.Int32(0xFFFF)) == cutlass.Int32(0xFFFF)) | ( + ((word >> 16) & cutlass.Int32(0xFFFF)) == cutlass.Int32(0xFFFF) + ) + + +def _sent32(word): + return word == cutlass.Int32(SENTINEL) + + +@cute.jit +def _poll_partials(part: cutlass.Array, idx, dst: cutlass.Array, dst_idx): + """Spin (relaxed, gpu scope, no sleep) until the four fp32 partial words at part[idx] are published, then store + them into dst[dst_idx : dst_idx + 4].""" + _poll_partials_from(part, idx, dst, dst_idx, cutlass.Int32(SENTINEL), cutlass.Int32(SENTINEL), + cutlass.Int32(SENTINEL), cutlass.Int32(SENTINEL)) # fmt: skip + + +@cute.jit +def _poll_partials_from(part: cutlass.Array, idx, dst: cutlass.Array, dst_idx, w0, w1, w2, w3): + """``_poll_partials`` starting from the words of an earlier read of part[idx]: re-reads only while one of them + is still the sentinel.""" + while _sent32(w0) | _sent32(w1) | _sent32(w2) | _sent32(w3): + vals = prims.load_ext( + part.subview(idx), dtype=cutlass.Int32, count=4, order="relaxed", scope="gpu" + ) + w0 = cutlass.Int32(vals[0]) + w1 = cutlass.Int32(vals[1]) + w2 = cutlass.Int32(vals[2]) + w3 = cutlass.Int32(vals[3]) + dst.store((w0.bitcast(cutlass.Float32), w1.bitcast(cutlass.Float32), w2.bitcast(cutlass.Float32), + w3.bitcast(cutlass.Float32)), idx=dst_idx, alignment=16) # fmt: skip + + +@cute.jit +def _head_role( + tma_wfb, + p1w: cutlass.Array, + part: cutlass.Array, + w_q: cutlass.Array, + w_k: cutlass.Array, + w_v: cutlass.Array, + a_log: cutlass.Array, + dt_bias: cutlass.Array, + onorm_w: cutlass.Array, + cs_q: cutlass.Array, + cs_k: cutlass.Array, + cs_v: cutlass.Array, + ssm: cutlass.Array, + state_tok: cutlass.Array, + slots: cutlass.Array, + pending: cutlass.Array, + out: cutlass.Array, + epoch: cutlass.Array, + smem_a: cutlass.Array, + smem_b: cutlass.Array, + bars: cutlass.Array, + tmem_holder: cutlass.Array, + s_uq: cutlass.Array, + s_uk: cutlass.Array, + s_uv: cutlass.Array, + s_wq: cutlass.Array, + s_wk: cutlass.Array, + s_wv: cutlass.Array, + s_dtb: cutlass.Array, + s_onw: cutlass.Array, + s_gr: cutlass.Array, + s_braw: cutlass.Array, + s_og: cutlass.Array, + s_q: cutlass.Array, + s_k: cutlass.Array, + s_dec: cutlass.Array, + s_kd: cutlass.Array, + s_bk: cutlass.Array, + s_beta: cutlass.Array, + s_v: cutlass.Array, + s_o: cutlass.Array, + s_ss: cutlass.Array, + s_rs: cutlass.Array, + s_vp: cutlass.Array, + s_ogp: cutlass.Array, + s_bp: cutlass.Array, + s_rec: cutlass.Array, + r_st: cutlass.Array, + ssm_stride, + pool_n, + lower_bound: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + eps: cutlass.Constexpr[float], + USE_PDL: cutlass.Constexpr[bool], + lbx, +): + """One CTA of a KDA verify cluster (cluster ids 26-31): stage A's verify for local head h = cid - 26 and V rows + [32 q, 32 q + 32) (q = cluster rank), fed by the stream clusters' Lamport buffers of this launch: q, k and f_a + (phase 1), v and b (phase 2), the output gate (phase 3). Arithmetic, rounding and state contract as + ``k3_kda_verify`` (V split 4 instead of 8; the output norm adds the same eight 16-row partial sums in the same + order).""" + tidx, _, _ = cute.arch.thread_idx() + bx = lbx + lane = tidx % 32 + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + vq = cute.arch.block_idx_in_cluster() + h = bx // cutlass.Int32(CLUSTER) - cutlass.Int32(STREAM_CLUSTERS) + v0 = vq * V_CTA + ch0 = h * HD + w_full = bars.subview(0) + acc_done = bars.subview(1) + ss_ready = bars.subview(2) + + tma_ptr_w = tma_wfb.get_ptr() + if warp == 0: + prims.prefetch_tensormap(tma_ptr_w) + if prims.elect_sync(): + prims.mbarrier_init(w_full, 1) + prims.mbarrier_init(acc_done, 1) + elif warp == 2: + prims.tcgen05_alloc(tmem_holder, FB_TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + elif warp == 3: + if prims.elect_sync(): + prims.mbarrier_init(ss_ready, 1) + prims.mbarrier_arrive_expect_tx(ss_ready, CLUSTER * 2 * NT * 4) + prims.fence_mbarrier_init() + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + if warp == 0: + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(w_full, HD * HD * 2) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a, tma_ptr_w, (cutlass.Int32(0), ch0, cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0)), + w_full, + ) # fmt: skip + + # ---- Before the grid dependency: nothing here is written by this launch's stream clusters or predecessors. + # Two round trips: first the slot index, every slot's accepted-draft count and what needs neither (conv weights, + # a_log, the norm and gate constants); then, as asynchronous copies into shared memory, everything the slot and + # its count select (the state, the conv caches' last positions and, below, the drafts' records). + slot = slots.load(idx=0) + pend_w = [] + for i_pw in cutlass.range_constexpr(PEND_WORDS): + pw_i = lane + cutlass.Int32(32 * i_pw) + pend_w.append( + cutlass.Int32(pending.load(idx=cutlass.select_(pw_i < pool_n, pw_i, cutlass.Int32(0)))) + ) + a_raw = a_log.load(idx=h) + QH = 3 * HD // 4 + VH = 3 * V_CTA // 4 + if tidx >= 2 * QH + VH + V_CTA: + c4_n = (tidx - (2 * QH + VH + V_CTA)) * 4 + prims.cp_async_shared_global(s_onw.subview(c4_n), onorm_w.subview(v0 + c4_n), 16, CP_CG) + if tidx < HD // 4: + prims.cp_async_shared_global( + s_dtb.subview(tidx * 4), dt_bias.subview(ch0 + tidx * 4), 16, CP_CG + ) + if tidx < HD: + wq4 = w_q.load(idx=(ch0 + tidx) * CONV_W, vector_size=CONV_W, alignment=16) + for w in cutlass.range_constexpr(CONV_W): + s_wq.store(cutlass.Float32(wq4[w]), idx=w * HD + tidx) + else: + ck_w = tidx - HD + wk4 = w_k.load(idx=(ch0 + ck_w) * CONV_W, vector_size=CONV_W, alignment=16) + for w in cutlass.range_constexpr(CONV_W): + s_wk.store(cutlass.Float32(wk4[w]), idx=w * HD + ck_w) + if tidx >= 2 * QH + VH: + if tidx < 2 * QH + VH + V_CTA: + cv_w = tidx - (2 * QH + VH) + wv4 = w_v.load(idx=(ch0 + v0 + cv_w) * CONV_W, vector_size=CONV_W, alignment=16) + for w in cutlass.range_constexpr(CONV_W): + s_wv.store(cutlass.Float32(wv4[w]), idx=w * V_CTA + cv_w) + pend_word = pend_w[PEND_WORDS - 1] + for i_pw in cutlass.range_constexpr(PEND_WORDS - 2, -1, -1): + pend_word = cutlass.Int32( + cutlass.select_(slot < cutlass.Int32(32 * (i_pw + 1)), pend_w[i_pw], pend_word) + ) + pend = cutlass.Int32(cute.arch.shuffle_sync(pend_word, offset=slot % cutlass.Int32(32))) + if pool_n > cutlass.Int32(32 * PEND_WORDS): + pend = cutlass.Int32(pending.load(idx=slot)) + pend = cutlass.Int32( + cutlass.select_(pend > cutlass.Int32(NUM_SPEC), cutlass.Int32(NUM_SPEC), pend) + ) + exp_a = cute.math.exp(a_raw, fastmath=True) + row_w = v0 + warp * REC_ROWS + # The slot's pool state and per-token region from their first element: the slot offset in 64 bits, once. + pool = ssm.subview(cutlass.Int64(slot) * ssm_stride) + tok = state_tok.subview(cutlass.Int64(slot) * CT_SLOT) + st_base = (h * HD + row_w) * HD + for r in cutlass.range_constexpr(REC_ROWS): + for i in cutlass.range_constexpr(VEC): + r_st.store(pool.load(idx=st_base + r * HD + i * 32 + lane), idx=r * VEC + i) + if tidx < QH: + m_q = tidx // (HD // 4) + c4_q = (tidx % (HD // 4)) * 4 + prims.cp_async_shared_global( + s_uq.subview(m_q * HD + c4_q), + cs_q.subview((slot * S_COLS + pend + m_q) * HK + ch0 + c4_q), + 16, + CP_CG, + ) + elif tidx < 2 * QH: + m_k = (tidx - QH) // (HD // 4) + c4_k = ((tidx - QH) % (HD // 4)) * 4 + prims.cp_async_shared_global( + s_uk.subview(m_k * HD + c4_k), + cs_k.subview((slot * S_COLS + pend + m_k) * HK + ch0 + c4_k), + 16, + CP_CG, + ) + elif tidx < 2 * QH + VH: + m_v = (tidx - 2 * QH) // (V_CTA // 4) + c4_v = ((tidx - 2 * QH) % (V_CTA // 4)) * 4 + prims.cp_async_shared_global( + s_uv.subview(m_v * V_CTA + c4_v), + cs_v.subview((slot * S_COLS + pend + m_v) * HK + ch0 + v0 + c4_v), + 16, + CP_CG, + ) + # The drafts the sampler accepted, replayed from their records onto the golden token's state. Every record this + # CTA may need (all NUM_SPEC drafts) comes into shared memory by one round trip first; a draft per iteration from + # global memory was one dependent round trip per accepted draft. + for it_rc in cutlass.range_constexpr((NUM_SPEC * REC_CTA // 4 + THREADS - 1) // THREADS): + q_rc = tidx + it_rc * THREADS + if q_rc < NUM_SPEC * REC_CTA // 4: + t_rc = q_rc // (REC_CTA // 4) + u_rc = (q_rc % (REC_CTA // 4)) * 4 + rec_rc = (t_rc * H_LOCAL + h) * HD + src_rc = cutlass.Int32( + cutlass.select_( + u_rc < V_CTA, + rec_rc + CT_VN + v0 + u_rc, + cutlass.select_( + u_rc < V_CTA + HD, + rec_rc + CT_WB + (u_rc - V_CTA), + rec_rc + CT_WD + (u_rc - V_CTA - HD), + ), + ) + ) + prims.cp_async_shared_global( + s_rec.subview(t_rc * REC_CTA + u_rc), tok.subview(src_rc), 16, CP_CG + ) + prims.cp_async_commit_group() + prims.cp_async_wait_group(0) + prims.barrier_cta_sync(0) + for t_acc in cutlass.range(pend, unroll=1): + rec_s = t_acc * REC_CTA + vn4_acc = s_rec.load(idx=rec_s + warp * REC_ROWS, vector_size=4, alignment=16) + vns_acc = [cutlass.Float32(vn4_acc[r]) for r in range(REC_ROWS)] + wbs_acc = [s_rec.load(idx=rec_s + V_CTA + i * 32 + lane) for i in range(VEC)] + wds_acc = [s_rec.load(idx=rec_s + V_CTA + HD + i * 32 + lane) for i in range(VEC)] + sts_acc = [r_st.load(idx=j) for j in range(REC_ROWS * VEC)] + for r in cutlass.range_constexpr(REC_ROWS): + for _pi in cutlass.range_constexpr(VEC // 2): + _p = _pi * 2 + vb0_acc, vb1_acc = cute.arch.mul_packed_f32x2( + (vns_acc[r], vns_acc[r]), (wbs_acc[_p], wbs_acc[_p + 1]) + ) + sts_acc[r * VEC + _p], sts_acc[r * VEC + _p + 1] = cute.arch.fma_packed_f32x2( + src_a=(wds_acc[_p], wds_acc[_p + 1]), src_b=(sts_acc[r * VEC + _p], sts_acc[r * VEC + _p + 1]), + src_c=(vb0_acc, vb1_acc), + ) # fmt: skip + for j in cutlass.range_constexpr(REC_ROWS * VEC): + r_st.store(sts_acc[j], idx=j) + # Every CTA of the head has read the head's q / k conv caches and per-key records once all four arrive here + # (release; the acquiring wait follows phase 1): from then on each CTA rewrites its quarter of both, in the idle + # and pre-recurrence slots instead of after the output norm. + prims.barrier_cluster_arrive() + + if cutlass.const_expr(USE_PDL): + prims.griddepcontrol(prims.GridDepAction.WAIT) + e = epoch.load(idx=bx) + buf = e % cutlass.Int32(BUFFERS) + if cutlass.const_expr(USE_PDL): + # The dependents may launch once the stream CTAs have also triggered (after their last weight box): they + # then take the SMs the stream leaves and run their pre-wait work beside the verify instead of after it. + if tidx == 0: + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + + # ---- Phase 1 (q, k, f_a of this launch): thread i polls 16-byte chunk i of the q | k rows, threads < 128 also + # chunk i of the f_a rows; q and k go to the conv inputs as fp32, f_a to the MMA's B operand (128-byte swizzle). + st_qk = tidx // HD + t_qk = (tidx % HD) // 16 + j_qk = tidx % 16 + a0 = cutlass.Int32(SENTINEL) + a1 = cutlass.Int32(SENTINEL) + a2 = cutlass.Int32(SENTINEL) + a3 = cutlass.Int32(SENTINEL) + qk_idx = ((buf * NT + t_qk) * P1_ROWS + st_qk * HK + ch0 + j_qk * 8) // 2 + while _sent16(a0) | _sent16(a1) | _sent16(a2) | _sent16(a3): + vqk = prims.load_ext( + p1w.subview(qk_idx), dtype=cutlass.Int32, count=4, order="relaxed", scope="gpu" + ) + a0 = cutlass.Int32(vqk[0]) + a1 = cutlass.Int32(vqk[1]) + a2 = cutlass.Int32(vqk[2]) + a3 = cutlass.Int32(vqk[3]) + if st_qk == 0: + _store8(s_uq, (a0, a1, a2, a3), (CONV_W - 1 + t_qk) * HD + j_qk * 8) + else: + _store8(s_uk, (a0, a1, a2, a3), (CONV_W - 1 + t_qk) * HD + j_qk * 8) + if tidx < HD: + t_fa = tidx // 16 + j_fa = tidx % 16 + f0 = cutlass.Int32(SENTINEL) + f1 = cutlass.Int32(SENTINEL) + f2 = cutlass.Int32(SENTINEL) + f3 = cutlass.Int32(SENTINEL) + fa_idx = ((buf * NT + t_fa) * P1_ROWS + 2 * HK + j_fa * 8) // 2 + while _sent16(f0) | _sent16(f1) | _sent16(f2) | _sent16(f3): + vfa = prims.load_ext( + p1w.subview(fa_idx), dtype=cutlass.Int32, count=4, order="relaxed", scope="gpu" + ) + f0 = cutlass.Int32(vfa[0]) + f1 = cutlass.Int32(vfa[1]) + f2 = cutlass.Int32(vfa[2]) + f3 = cutlass.Int32(vfa[3]) + # Row t_fa, 16-byte unit u of 64-column chunk c: byte 1024 c + 128 t_fa + 16 (u ^ t_fa). + sw_word = ((j_fa // 8) * 1024 + t_fa * 128 + ((j_fa % 8) ^ t_fa) * 16) // 4 + smem_b.store((f0, f1, f2, f3), idx=sw_word, alignment=16) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.barrier_cta_sync(0) + prims.barrier_cluster_wait() + + # ---- f_b on the tensor cores (warp 2 issues, warps 4-7 read TMEM), then the q/k pre-compute, one token per warp. + if warp == 2: + idesc = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cutlass.BFloat16, + b_dtype=cutlass.BFloat16, + n_dim=MMA_N, + m_dim=HD, + ) + desc_a_base = prims.Tcgen05SmemDesc.build( + start_address=smem_a, leading_byte_offset=16, stride_byte_offset=1024, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + desc_b_base = prims.Tcgen05SmemDesc.build( + start_address=smem_b, leading_byte_offset=16, stride_byte_offset=1024, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + tmem_acc = cutlass.inttoptr(tmem_holder.load(), 6, cutlass.Int32) + while not cute.arch.mbarrier_test_wait(w_full.data_ptr(), 0): + pass + prims.tcgen05_fence( + prims.Tcgen05Fence.AFTER_THREAD_SYNC + ) # f_a came from the other threads' smem stores + for kb in cutlass.range_constexpr(2 * (BOX_K // MMA_K)): + box = kb // (BOX_K // MMA_K) + within = kb % (BOX_K // MMA_K) + if prims.elect_sync(): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, prims.CTAGroup.CTA_1, tmem_acc, + desc_a_base + (box * ((HD * BOX_K * 2) >> 4) + within * K_STEP_U), + desc_b_base + (box * ((MMA_N * BOX_K * 2) >> 4) + within * K_STEP_U), + idesc, kb != 0, + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + if warp >= 4: + while not cute.arch.mbarrier_test_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", cutlass.inttoptr(tmem_holder.load(), 6, cutlass.Float32), num=MMA_N + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + c_fb = (warp - 4) * 32 + lane + for t_fb in cutlass.range_constexpr(NT): + s_gr.store(_bf16(cutlass.Float32(acc[t_fb])), idx=t_fb * HD + c_fb) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.barrier_cta_sync(1, thread_count=128) + if warp == 4: + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_holder.load(), 6, cutlass.Int32), FB_TMEM_COLS + ) + # q and k of token w, one token per warp (the arithmetic of kda_mtp_decode's pre-compute warps). + tk = warp + pq = [cutlass.Float32(0.0)] * VEC + for i in cutlass.range_constexpr(VEC): + c = i * 32 + lane + conv = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(CONV_W - 1): + conv += s_uq.load(idx=(tk + w) * HD + c) * s_wq.load(idx=w * HD + c) + conv += s_uq.load(idx=(tk + CONV_W - 1) * HD + c) * s_wq.load(idx=(CONV_W - 1) * HD + c) + ex = cute.math.exp(-conv, fastmath=True) + pq[i] = conv * cute.arch.rcp_approx(cutlass.Float32(1.0) + ex) + sum_q = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(VEC): + sum_q += pq[i] * pq[i] + sum_q = _butterfly(sum_q) + rnorm_q = cute.math.rsqrt(sum_q + 1e-06, fastmath=True) * scale + for i in cutlass.range_constexpr(VEC): + s_q.store(pq[i] * rnorm_q, idx=tk * HD + i * 32 + lane) + pk = [cutlass.Float32(0.0)] * VEC + for i in cutlass.range_constexpr(VEC): + c = i * 32 + lane + conv = s_uk.load(idx=tk * HD + c) * s_wk.load(idx=c) + for w in cutlass.range_constexpr(1, CONV_W - 1): + conv += s_uk.load(idx=(tk + w) * HD + c) * s_wk.load(idx=w * HD + c) + conv += s_uk.load(idx=(tk + CONV_W - 1) * HD + c) * s_wk.load(idx=(CONV_W - 1) * HD + c) + pk[i] = conv * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-conv, fastmath=True) + ) + sum_k = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(VEC): + sum_k += pk[i] * pk[i] + sum_k = _butterfly(sum_k) + rnorm_k = cute.math.rsqrt(sum_k + 1e-06, fastmath=True) + for i in cutlass.range_constexpr(VEC): + s_k.store(pk[i] * rnorm_k, idx=tk * HD + i * 32 + lane) + if warp < 4: + # Conv caches for the next round, q and k, this CTA's 32 channels: column s = position s - 2 = row s + 1 of + # the position-indexed inputs. + t_cq = tidx + for it_qk in cutlass.range_constexpr((2 * S_COLS * V_CTA + 127) // 128): + item_qk = t_cq + it_qk * 128 + if item_qk < 2 * S_COLS * V_CTA: + sc_qk = (item_qk % (S_COLS * V_CTA)) // V_CTA + ch_qk = vq * V_CTA + item_qk % V_CTA + dst_qk = (slot * S_COLS + sc_qk) * HK + ch0 + ch_qk + if item_qk < S_COLS * V_CTA: + cs_q.store(s_uq.load(idx=(sc_qk + 1) * HD + ch_qk), idx=dst_qk) + else: + cs_k.store(s_uk.load(idx=(sc_qk + 1) * HD + ch_qk), idx=dst_qk) + prims.barrier_cta_sync(0) + + # ---- Phase 2 (v, b): the two K-half partials of this CTA's 32 v rows and of this head's b; the projection is + # their bf16-rounded sum. + if tidx < 2 * NT * (V_CTA // 4): + half_v = tidx // (NT * (V_CTA // 4)) + t_v = (tidx // (V_CTA // 4)) % NT + j_v = tidx % (V_CTA // 4) + _poll_partials( + part, + ((buf * 3 + 0) * 2 + half_v) * (8 * PART_ROWS) + t_v * PART_ROWS + ch0 + v0 + j_v * 4, + s_vp, + (half_v * NT + t_v) * V_CTA + j_v * 4, + ) + elif tidx < 2 * NT * (V_CTA // 4) + 2 * NT: + ib = tidx - 2 * NT * (V_CTA // 4) + half_b = ib // NT + t_b = ib % NT + wb = cutlass.Int32(SENTINEL) + while _sent32(wb): + wb = prims.load_ext( + part.subview(((buf * 3 + 2) * 2 + half_b) * (8 * PART_ROWS) + t_b * PART_ROWS + h), + dtype=cutlass.Int32, + order="relaxed", + scope="gpu", + ) + s_bp.store(wb.bitcast(cutlass.Float32), idx=half_b * NT + t_b) + prims.barrier_cta_sync(0) + t_cv = tidx // V_CTA + c_cv = tidx % V_CTA + s_uv.store(_bf16(s_vp.load(idx=t_cv * V_CTA + c_cv) + s_vp.load(idx=(NT + t_cv) * V_CTA + c_cv)), + idx=(CONV_W - 1 + t_cv) * V_CTA + c_cv) # fmt: skip + if tidx < NT: + s_braw.store(_bf16(s_bp.load(idx=tidx) + s_bp.load(idx=NT + tidx)), idx=tidx) + prims.barrier_cta_sync(0) + for it_cs in cutlass.range_constexpr((S_COLS * V_CTA + THREADS - 1) // THREADS): + item_cs = tidx + it_cs * THREADS + if item_cs < S_COLS * V_CTA: + s_cs = item_cs // V_CTA + v_cs = item_cs % V_CTA + cs_v.store( + s_uv.load(idx=(s_cs + 1) * V_CTA + v_cs), + idx=(slot * S_COLS + s_cs) * HK + ch0 + v0 + v_cs, + ) + # v (this CTA's 32 channels, one (token, channel) per thread) and beta. + conv_v = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(CONV_W): + conv_v += s_uv.load(idx=(t_cv + w) * V_CTA + c_cv) * s_wv.load(idx=w * V_CTA + c_cv) + conv_v = conv_v * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-conv_v, fastmath=True) + ) + s_v.store(conv_v, idx=t_cv * V_CTA + c_cv) + if tidx < NT: + b_pre = s_braw.load(idx=tidx) + s_beta.store( + cute.arch.rcp_approx(cutlass.Float32(1.0) + cute.math.exp(-b_pre, fastmath=True)), + idx=tidx, + ) + prims.barrier_cta_sync(0) + + # ---- The recurrence's per-token operands: warp t, token t: decay = exp(gate), decay * k, beta * k. + tg = warp + r_beta_g = s_beta.load(idx=tg) + for i_pair in cutlass.range_constexpr(VEC // 2): + dks = [] + for i in (i_pair * 2, i_pair * 2 + 1): + c = i * 32 + lane + g_raw = s_gr.load(idx=tg * HD + c) + s_dtb.load(idx=c) + xg = exp_a * g_raw + sig = cute.arch.rcp_approx(cutlass.Float32(1.0) + cute.math.exp(-xg, fastmath=True)) + dks.append( + (cute.math.exp(lower_bound * sig, fastmath=True), s_k.load(idx=tg * HD + c), c) + ) + (d0, k0v, c0), (d1, k1v, c1) = dks + bk0, bk1 = cute.arch.mul_packed_f32x2((r_beta_g, r_beta_g), (k0v, k1v)) + kd0, kd1 = cute.arch.mul_packed_f32x2((d0, d1), (k0v, k1v)) + s_dec.store(d0, idx=tg * HD + c0) + s_dec.store(d1, idx=tg * HD + c1) + s_kd.store(kd0, idx=tg * HD + c0) + s_kd.store(kd1, idx=tg * HD + c1) + s_bk.store(bk0, idx=tg * HD + c0) + s_bk.store(bk1, idx=tg * HD + c1) + prims.barrier_cta_sync(0) + # The drafts' per-key records (beta * k and the decay of tokens 1-7), this CTA's 32 keys. + if tidx < NUM_SPEC * V_CTA: + t_rq = tidx // V_CTA + key_rq = vq * V_CTA + tidx % V_CTA + rec_q = h * HD + t_rq * H_LOCAL * HD + key_rq + tok.store(s_bk.load(idx=(t_rq + 1) * HD + key_rq), idx=rec_q + CT_WB) + tok.store(s_dec.load(idx=(t_rq + 1) * HD + key_rq), idx=rec_q + CT_WD) + + # ---- Recurrence over the verify tokens: warp w owns V rows row_w .. row_w + 3, lane l keys 32 i + l. A rolled + # loop: in the model the MoE's stream has evicted this kernel's code from L2, and the 8 tokens fully unrolled + # (~1,100 instructions, run once per launch) stall on instruction fetch; one token's body is fetched once. The + # loop is software-pipelined so that consecutive tokens overlap: iteration tr reduces token tr - 1's q-dot + # partials over the warp beside token tr's k-dot butterflies (the chain that carries the state) and loads token + # tr + 1's operands. Every sum is formed from the same values in the same order as before (bit-identical). + # This lane's first state element in the pool (the golden token's state) and in the first draft's per-token + # state; element (r, i) sits (r HD + 32 i) floats past either. + pool_ptr = pool.subview(st_base + lane).data_ptr().toint() + # Lanes 0-3 store the warp's row innovations, one row each (a lane's address beyond them is never stored to). + vn_ptr = tok.subview(CT_VN + h * HD + row_w + lane).data_ptr().toint() + _load_rec_operands(r_st, cutlass.Int32(0), warp, lane, s_v, s_q, s_kd, s_bk, s_dec) + for r in cutlass.range_constexpr(REC_ROWS): + r_st.store(cutlass.Float32(0.0), idx=R_Q + r) + for tr in cutlass.range(NT, unroll=1): + vrows = [r_st.load(idx=R_OP + j) for j in range(REC_ROWS)] + wq = [r_st.load(idx=R_OP + REC_ROWS + i) for i in range(VEC)] + wk = [r_st.load(idx=R_OP + REC_ROWS + VEC + i) for i in range(VEC)] + wb = [r_st.load(idx=R_OP + REC_ROWS + 2 * VEC + i) for i in range(VEC)] + wd = [r_st.load(idx=R_OP + REC_ROWS + 3 * VEC + i) for i in range(VEC)] + shq_prev = [r_st.load(idx=R_Q + r) for r in range(REC_ROWS)] + # Token tr + 1's operands (the last iteration reloads token NT - 1's and does not use them). + t_next = cutlass.Int32(cutlass.select_(tr + 1 < NT, tr + 1, cutlass.Int32(NT - 1))) + _load_rec_operands(r_st, t_next, warp, lane, s_v, s_q, s_kd, s_bk, s_dec) + stw = [r_st.load(idx=j) for j in range(REC_ROWS * VEC)] + xq_prev = _rows_out(shq_prev, lane) + shk = [] + for r in cutlass.range_constexpr(REC_ROWS): + p1a = cutlass.Float32(0.0) + p2a = cutlass.Float32(0.0) + for _pi in cutlass.range_constexpr(VEC // 2): + _p = _pi * 2 + p1a, p2a = cute.arch.fma_packed_f32x2( + src_a=(stw[r * VEC + _p], stw[r * VEC + _p + 1]), + src_b=(wk[_p], wk[_p + 1]), + src_c=(p1a, p2a), + ) + shk.append(p1a + p2a) + for offset in [16, 8, 4, 2, 1]: + for r in cutlass.range_constexpr(REC_ROWS): + shk[r] += cute.arch.shuffle_sync_bfly( + shk[r], offset=offset, mask=-1, mask_and_clamp=31 + ) + shq = [] + vns = [] + for r in cutlass.range_constexpr(REC_ROWS): + vn = vrows[r] - shk[r] + vns.append(vn) + q1 = cutlass.Float32(0.0) + q2 = cutlass.Float32(0.0) + for _pi in cutlass.range_constexpr(VEC // 2): + _p = _pi * 2 + vb0, vb1 = cute.arch.mul_packed_f32x2((vn, vn), (wb[_p], wb[_p + 1])) + stw[r * VEC + _p], stw[r * VEC + _p + 1] = cute.arch.fma_packed_f32x2( + src_a=(wd[_p], wd[_p + 1]), src_b=(stw[r * VEC + _p], stw[r * VEC + _p + 1]), src_c=(vb0, vb1), + ) # fmt: skip + q1, q2 = cute.arch.fma_packed_f32x2( + src_a=(stw[r * VEC + _p], stw[r * VEC + _p + 1]), + src_b=(wq[_p], wq[_p + 1]), + src_c=(q1, q2), + ) + shq.append(q1 + q2) + # The golden token's state to the pool (tr = 0); a draft's record: this warp's row innovations (lanes 0-3). + # Predicated stores, no branch: they issue between the update's FMAs. + first = cutlass.Int32(cutlass.select_(tr == 0, cutlass.Int32(1), cutlass.Int32(0))) + for r in cutlass.range_constexpr(REC_ROWS): + for i in cutlass.range_constexpr(VEC): + _st_f32_if(pool_ptr + cutlass.Int64((r * HD + i * 32) * 4), stw[r * VEC + i], first) + vn_lane = cutlass.Float32( + cutlass.select_( + lane == 0, + vns[0], + cutlass.select_(lane == 1, vns[1], cutlass.select_(lane == 2, vns[2], vns[3])), + ) + ) + draft_rec = cutlass.Int32( + cutlass.select_((tr > 0) & (lane < 4), cutlass.Int32(1), cutlass.Int32(0)) + ) + _st_f32_if( + vn_ptr + cutlass.Int64(tr - 1) * cutlass.Int64(H_LOCAL * HD * 4), vn_lane, draft_rec + ) + for j in cutlass.range_constexpr(REC_ROWS * VEC): + r_st.store(stw[j], idx=j) + for r in cutlass.range_constexpr(REC_ROWS): + r_st.store(shq[r], idx=R_Q + r) + # Token tr - 1's outputs, from the first lane of each group of 8 (which all hold the same sum); at tr = 0 a + # placeholder in token 0's slot, which iteration 1 overwrites from the same lanes. + o_tok = cutlass.Int32(cutlass.select_(tr > 0, tr - 1, cutlass.Int32(0))) + if (lane & cutlass.Int32(7)) == cutlass.Int32(0): + s_o.store( + _bf16(xq_prev), + idx=o_tok * V_CTA + warp * REC_ROWS + (lane >> 4) * 2 + ((lane >> 3) & 1), + ) + xq_last = _rows_out([r_st.load(idx=R_Q + r) for r in range(REC_ROWS)], lane) + if (lane & cutlass.Int32(7)) == cutlass.Int32(0): + s_o.store( + _bf16(xq_last), + idx=(NT - 1) * V_CTA + warp * REC_ROWS + (lane >> 4) * 2 + ((lane >> 3) & 1), + ) + + # ---- Phase 3 (the output gate), then the gated RMSNorm over V: every CTA sends its two 16-row sums of squares + # per token to all four CTAs (the eight partials of V split 8, added in the same order). Thread (token t, row v) + # reads its own gate's two K-half partials first, so their round trip overlaps the norm exchange. + t_out = tidx // V_CTA + v_y = tidx % V_CTA + og_at = (buf * 3 + 1) * 2 * (8 * PART_ROWS) + t_out * PART_ROWS + ch0 + v0 + v_y + og_h0 = cutlass.Int32( + prims.load_ext(part.subview(og_at), dtype=cutlass.Int32, order="relaxed", scope="gpu") + ) + og_h1 = cutlass.Int32(prims.load_ext(part.subview(og_at + 8 * PART_ROWS), dtype=cutlass.Int32, order="relaxed", + scope="gpu")) # fmt: skip + prims.barrier_cta_sync(0) + x_o = s_o.load(idx=tidx) + ss = x_o * x_o + for off_ss in [8, 4, 2, 1]: + ss = ss + cute.arch.shuffle_sync_bfly(ss, offset=off_ss, mask=-1, mask_and_clamp=31) + if lane % 16 == 0: + src_slot = vq * 2 + lane // 16 + for r in cutlass.range_constexpr(CLUSTER): + _st_async_f32( + _mapa_u32(s_ss.subview(src_slot * NT + t_out).data_ptr(), r), + ss, + _mapa_u32(ss_ready.data_ptr(), r), + ) + while not _test_wait_cluster(ss_ready.data_ptr(), 0): + pass + while _sent32(og_h0) | _sent32(og_h1): + og_h0 = cutlass.Int32( + prims.load_ext(part.subview(og_at), dtype=cutlass.Int32, order="relaxed", scope="gpu") + ) + og_h1 = cutlass.Int32(prims.load_ext(part.subview(og_at + 8 * PART_ROWS), dtype=cutlass.Int32, + order="relaxed", scope="gpu")) # fmt: skip + if tidx < NT: + total = cutlass.Float32(0.0) + for r in cutlass.range_constexpr(2 * CLUSTER): + total = total + s_ss.load(idx=r * NT + tidx) + s_rs.store(cute.math.rsqrt(total / HD + eps), idx=tidx) + prims.barrier_cta_sync(0) + z = _bf16(og_h0.bitcast(cutlass.Float32) + og_h1.bitcast(cutlass.Float32)) + gate = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.math.exp(-z, fastmath=True)) + y = s_o.load(idx=tidx) * s_rs.load(idx=t_out) * s_onw.load(idx=v_y) * gate + out.store(cutlass.BFloat16(y), idx=(t_out * H_LOCAL + h) * HD + v0 + v_y) + + if tidx == 0: + epoch.store((e + cutlass.Int32(1)) % cutlass.Int32(BUFFERS), idx=bx) + + +@cute.kernel +def k3_kda_attn_kernel( + tma_w: cutlass.GridConstant[ + cuda.TensorMap + ], # W [3208, 7168] bf16, 5-D, box 64 cols x 64 rows x 2 chunks + tma_x: cutlass.GridConstant[ + cuda.TensorMap + ], # x [T, 7168] bf16, box 64 cols x 8 rows (rows >= T read as 0) + p1: cutlass.Array, # int16 bits of bf16 [3][8][1664]: phase-1 rows (Lamport) + part: cutlass.Array, # int32 bits of fp32 [3][3][2][8][768]: v, og and b cluster partials (Lamport) + epoch: cutlass.Array, # int32 [CTAs]: each CTA's buffer index (launches completed mod 3) + USE_PDL: cutlass.Constexpr[bool], +): + ring_w = cutlass.Array( + cutlass.BFloat16, STREAM_STAGES * BOX_ELEMS, space=cutlass.AddressSpace.smem, alignment=1024 + ) + ring_x = cutlass.Array( + cutlass.BFloat16, STREAM_STAGES * X_ELEMS, space=cutlass.AddressSpace.smem, alignment=1024 + ) + full = cutlass.Array(cutlass.Int64, STREAM_STAGES, space=cutlass.AddressSpace.smem, alignment=8) + empty = cutlass.Array( + cutlass.Int64, STREAM_STAGES, space=cutlass.AddressSpace.smem, alignment=8 + ) + acc_done = cutlass.Array(cutlass.Int64, 3, space=cutlass.AddressSpace.smem, alignment=8) + mbox_bar = cutlass.Array(cutlass.Int64, 3, space=cutlass.AddressSpace.smem, alignment=8) + # [phase][source rank][lane][4]: the partials of the 16 rows this rank owns (its own slot stays unused). + mbox = cutlass.Array( + cutlass.Float32, 3 * CLUSTER * 32 * 4, space=cutlass.AddressSpace.smem, alignment=16 + ) + tmem_holder = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + + _stream_role( + tma_w, + tma_x, + p1, + part, + epoch, + ring_w, + ring_x, + full, + empty, + acc_done, + mbox_bar, + mbox, + tmem_holder, + USE_PDL, + STREAM_STAGES, + ) + + +@cute.jit +def k3_kda_qkvg( + w: cute.Tensor, # bf16 [3208, 7168] + x: cute.Tensor, # bf16 [T, 7168] + p1: cute.Tensor, # int16 [3 * 8 * 1664] + part: cute.Tensor, # int32 [3 * 3 * 2 * 8 * 768] + epoch: cute.Tensor, # int32 [104] + num_tokens: cutlass.Int32, + USE_PDL: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + """One launch of the three-phase projection (the stream role alone).""" + tma_w = cuda.create_tensor_map_tiled( + global_address=w.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[BOX_K, PROJ_ROWS, K_IN // BOX_K, 1, 1], + global_strides=[(K_IN * 2) // 16, (BOX_K * 2) // 16, (PROJ_ROWS * K_IN * 2) // 16, + (PROJ_ROWS * K_IN * 2) // 16], + box_dims=[BOX_K, TILE, BOX_CH, 1, 1], + swizzle=cuda.TensorMapSwizzle.s128b, + ) # fmt: skip + tma_x = cuda.create_tensor_map_tiled( + global_address=x.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[K_IN, num_tokens], + global_strides=[(K_IN * 2) // 16], + box_dims=[BOX_K, MMA_N], + swizzle=cuda.TensorMapSwizzle.s128b, + ) + k3_kda_attn_kernel(tma_w, tma_x, p1, part, epoch, USE_PDL).launch( + grid=(STREAM_CLUSTERS * CLUSTER, 1, 1), + block=(THREADS, 1, 1), + cluster=(CLUSTER, 1, 1), + stream=stream, + use_pdl=USE_PDL, + ) + + +class _SmemCarver: + """Typed views laid out one after another inside one raw shared-memory buffer: a CTA runs one role for its whole + life and every CTA of a cluster runs the same role, so the two roles share the bytes and a DSMEM address always + names the same view in the peer.""" + + def __init__(self, raw, raw_bytes: int): + self.raw = raw + self.cap = raw_bytes + self.off = 0 + + def view(self, dtype, n: int, align: int = 16) -> cutlass.Array: + self.off = -(-self.off // align) * align + base = self.raw if self.off == 0 else self.raw + self.off + ptr = cute.recast_ptr(base, dtype=dtype) + self.off += n * dtype.width // 8 + assert self.off <= self.cap, "a role's arrays exceed the shared raw buffer" + return cutlass.Array( + ptr, + shape=(n,), + dtype=dtype, + bounds_check=False, + addrspace=cutlass.AddressSpace.smem.value, + alignment=align, + ) + + +@cute.kernel +def k3_kda_attn_fused_kernel( + tma_w: cutlass.GridConstant[ + cuda.TensorMap + ], # W [3208, 7168] bf16, 5-D, box 64 cols x 64 rows x 2 chunks + tma_x: cutlass.GridConstant[cuda.TensorMap], # x [8, 7168] bf16, box 64 cols x 8 rows + tma_wfb: cutlass.GridConstant[ + cuda.TensorMap + ], # W_fb [768, 128] bf16, a head's 128 rows by one call + p1: cutlass.Array, # int16 bits of bf16 [3][8][1664] (Lamport, written by the stream role) + p1w: cutlass.Array, # the same memory as int32 words (polled by the verify role) + part: cutlass.Array, # int32 bits of fp32 [3][3][2][8][768] (Lamport) + epoch: cutlass.Array, # int32 [128]: each CTA's buffer index (launches completed mod 3) + w_q: cutlass.Array, # fp32 [768, 4] + w_k: cutlass.Array, + w_v: cutlass.Array, + a_log: cutlass.Array, # fp32 [6] + dt_bias: cutlass.Array, # fp32 [768] + onorm_w: cutlass.Array, # fp32 [128] + cs_q: cutlass.Array, # fp32 [pool][10][768]: raw conv inputs, column s = position s - 2 from the golden token + cs_k: cutlass.Array, + cs_v: cutlass.Array, + ssm: cutlass.Array, # fp32 [pool][6][128][128]: the state after the last golden token + state_tok: cutlass.Array, # fp32 [pool][7][6][128][128]: the state after each draft of the last round + slots: cutlass.Array, # int32 [1] + pending: cutlass.Array, # int32 [pool]: drafts the sampler accepted last round + out: cutlass.Array, # bf16 [8][6][128] + ssm_stride: cutlass.Int64, # fp32 elements between the pool's slots + pool_n: cutlass.Int32, # the pool's slots (entries of pending) + lower_bound: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + eps: cutlass.Constexpr[float], + USE_PDL: cutlass.Constexpr[bool], +): + # The stream CTAs (clusters 0-25) and the verify CTAs (26-31) never share a cluster, so their large arrays alias + # one raw buffer and the stream's weight ring takes the bytes the verify buffers need elsewhere. + raw = SmemAllocator().allocate(FUSED_RAW_BYTES, byte_alignment=1024) + stream_carve = _SmemCarver(raw, FUSED_RAW_BYTES) + ring_w = stream_carve.view(cutlass.BFloat16, FUSED_STAGES * BOX_ELEMS, 1024) + ring_x = stream_carve.view(cutlass.BFloat16, FUSED_STAGES * X_ELEMS, 1024) + mbox = stream_carve.view(cutlass.Float32, 3 * CLUSTER * 32 * 4, 16) + verify_carve = _SmemCarver(raw, FUSED_RAW_BYTES) + smem_a = verify_carve.view(cutlass.BFloat16, HD * HD, 1024) + smem_b = verify_carve.view(cutlass.Int32, MMA_N * HD // 2, 1024) + s_uq = verify_carve.view(cutlass.Float32, ROWS_U * HD) + s_uk = verify_carve.view(cutlass.Float32, ROWS_U * HD) + s_uv = verify_carve.view(cutlass.Float32, ROWS_U * V_CTA) + s_wq = verify_carve.view(cutlass.Float32, CONV_W * HD) + s_wk = verify_carve.view(cutlass.Float32, CONV_W * HD) + s_wv = verify_carve.view(cutlass.Float32, CONV_W * V_CTA) + s_dtb = verify_carve.view(cutlass.Float32, HD) + s_onw = verify_carve.view(cutlass.Float32, V_CTA) + s_gr = verify_carve.view(cutlass.Float32, NT * HD) + s_braw = verify_carve.view(cutlass.Float32, NT) + s_og = verify_carve.view(cutlass.Float32, NT * V_CTA) + s_q = verify_carve.view(cutlass.Float32, NT * HD) + s_k = verify_carve.view(cutlass.Float32, NT * HD) + s_dec = verify_carve.view(cutlass.Float32, NT * HD) + s_kd = verify_carve.view(cutlass.Float32, NT * HD) + s_bk = verify_carve.view(cutlass.Float32, NT * HD) + s_beta = verify_carve.view(cutlass.Float32, NT) + s_v = verify_carve.view(cutlass.Float32, NT * V_CTA) + s_o = verify_carve.view(cutlass.Float32, NT * V_CTA) + s_ss = verify_carve.view(cutlass.Float32, 2 * CLUSTER * NT) + s_rs = verify_carve.view(cutlass.Float32, NT) + s_vp = verify_carve.view(cutlass.Float32, 2 * NT * V_CTA) + s_ogp = verify_carve.view(cutlass.Float32, 2 * NT * V_CTA) + s_bp = verify_carve.view(cutlass.Float32, 2 * NT) + s_rec = verify_carve.view(cutlass.Float32, NUM_SPEC * REC_CTA) + # Barriers and the TMEM holder stay static (both roles initialize their own). + full = cutlass.Array(cutlass.Int64, FUSED_STAGES, space=cutlass.AddressSpace.smem, alignment=8) + empty = cutlass.Array(cutlass.Int64, FUSED_STAGES, space=cutlass.AddressSpace.smem, alignment=8) + acc_done = cutlass.Array(cutlass.Int64, 3, space=cutlass.AddressSpace.smem, alignment=8) + mbox_bar = cutlass.Array(cutlass.Int64, 3, space=cutlass.AddressSpace.smem, alignment=8) + tmem_holder = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + bars = cutlass.Array(cutlass.Int64, 3, space=cutlass.AddressSpace.smem, alignment=8) + r_st = cutlass.Array(cutlass.Float32, REC_REGS, space=cutlass.AddressSpace.rmem) + + # Launch order by cluster: the f_a / b stream clusters, the verify clusters, then the q/k/v/og stream clusters. + # The CTAs of the first clusters become resident while the predecessor still holds SMs, so every verify CTA + # runs its pre-wait prologue (state, records, the drafts' replay) before the grid dependency instead of after + # it; the stream clusters that come last load their boxes after it. Roles use the logical index: + # stream CTAs 0-103 (the f_a / b clusters first), verify CTAs 104-127. + pbx, _, _ = cute.arch.block_idx() + lbx_head = pbx + cutlass.Int32((STREAM_CLUSTERS - 2) * CLUSTER) + lbx_qkv = pbx - cutlass.Int32(HEAD_CLUSTERS * CLUSTER) + lbx = cutlass.Int32( + cutlass.select_( + pbx < cutlass.Int32(2 * CLUSTER), + pbx, + cutlass.select_(pbx < cutlass.Int32((2 + HEAD_CLUSTERS) * CLUSTER), lbx_head, lbx_qkv), + ) + ) + if lbx < cutlass.Int32(STREAM_CLUSTERS * CLUSTER): + _stream_role(tma_w, tma_x, p1, part, epoch, ring_w, ring_x, full, empty, acc_done, mbox_bar, mbox, tmem_holder, + USE_PDL, FUSED_STAGES, lbx) # fmt: skip + else: + _head_role(tma_wfb, p1w, part, w_q, w_k, w_v, a_log, dt_bias, onorm_w, cs_q, cs_k, cs_v, ssm, state_tok, + slots, pending, out, epoch, smem_a, smem_b, bars, tmem_holder, s_uq, s_uk, s_uv, s_wq, s_wk, s_wv, + s_dtb, s_onw, s_gr, s_braw, s_og, s_q, s_k, s_dec, s_kd, s_bk, s_beta, s_v, s_o, s_ss, s_rs, s_vp, + s_ogp, s_bp, s_rec, r_st, ssm_stride, pool_n, lower_bound, scale, eps, USE_PDL, lbx) # fmt: skip + + +@cute.jit +def k3_kda_attn( + w: cute.Tensor, # bf16 [3208, 7168] + x: cute.Tensor, # bf16 [8, 7168] + w_fb: cute.Tensor, # bf16 [768, 128] + p1: cute.Tensor, # int16 [3 * 8 * 1664] + p1w: cute.Tensor, # the same memory as int32 + part: cute.Tensor, # int32 [3 * 3 * 2 * 8 * 768] + epoch: cute.Tensor, # int32 [128] + w_q: cute.Tensor, + w_k: cute.Tensor, + w_v: cute.Tensor, + a_log: cute.Tensor, + dt_bias: cute.Tensor, + onorm_w: cute.Tensor, + cs_q: cute.Tensor, + cs_k: cute.Tensor, + cs_v: cute.Tensor, + ssm: cute.Tensor, + state_tok: cute.Tensor, + slots: cute.Tensor, + pending: cute.Tensor, + out: cute.Tensor, + ssm_stride: cutlass.Int64, + lower_bound: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + eps: cutlass.Constexpr[float], + USE_PDL: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + """One launch of the fused KDA projection + verify for one request of 8 verify tokens (stage B2).""" + tma_w = cuda.create_tensor_map_tiled( + global_address=w.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[BOX_K, PROJ_ROWS, K_IN // BOX_K, 1, 1], + global_strides=[(K_IN * 2) // 16, (BOX_K * 2) // 16, (PROJ_ROWS * K_IN * 2) // 16, + (PROJ_ROWS * K_IN * 2) // 16], + box_dims=[BOX_K, TILE, BOX_CH, 1, 1], + swizzle=cuda.TensorMapSwizzle.s128b, + ) # fmt: skip + tma_x = cuda.create_tensor_map_tiled( + global_address=x.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[K_IN, NT], + global_strides=[(K_IN * 2) // 16], + box_dims=[BOX_K, MMA_N], + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_wfb = cuda.create_tensor_map_tiled( + global_address=w_fb.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[BOX_K, HK, HD // BOX_K, 1, 1], + global_strides=[(HD * 2) // 16, (BOX_K * 2) // 16, (HK * HD * 2) // 16, (HK * HD * 2) // 16], + box_dims=[BOX_K, HD, HD // BOX_K, 1, 1], + swizzle=cuda.TensorMapSwizzle.s128b, + ) # fmt: skip + k3_kda_attn_fused_kernel( + tma_w, tma_x, tma_wfb, p1, p1w, part, epoch, w_q, w_k, w_v, a_log, dt_bias, onorm_w, cs_q, cs_k, cs_v, ssm, + state_tok, slots, pending, out, ssm_stride, cutlass.Int32(cute.size(pending)), lower_bound, scale, eps, USE_PDL, + ).launch( + grid=((STREAM_CLUSTERS + HEAD_CLUSTERS) * CLUSTER, 1, 1), block=(THREADS, 1, 1), cluster=(CLUSTER, 1, 1), + stream=stream, use_pdl=USE_PDL, + # One CTA per SM (the shared-memory carve); without it ptxas may pick an occupancy-driven register target. + min_blocks_per_mp=1, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_attn/k3_kda_decode_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_attn/k3_kda_decode_kernel.py new file mode 100644 index 000000000000..0955213cc57a --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_attn/k3_kda_decode_kernel.py @@ -0,0 +1,802 @@ +# 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 plain decode in one launch: the fused projection of ``k3_kda_attn`` (TP16 rank slice, its 26 stream +clusters and three Lamport phases, unchanged) and the gated-delta decode of R <= 8 requests of one token each, on the +pools of ``trtllm::kda_decode``. + +The projection's 8 token rows are the R requests' tokens (rows >= R read as zero). Six head clusters (local heads +0-5) of 4 CTAs, cluster rank = V quarter; CTA (h, q) owns V rows [32 q, 32 q + 32) of local head h for every request. +Before ``griddepcontrol.wait`` (nothing there is written by this launch or its predecessors): the slots, the head's +W_fb block (TMA), the conv weights and gate constants, each request's conv windows (the head's q / k channels, this +CTA's 32 v channels) and each request's 32 state rows (16 KB) into shared memory. +After it, per request (warp r = request r where a step is per request): + phase 1 q, k and f_a rows; f_b on the tensor cores (M = 128 gate channels, N = 8 rows); + q, k = L2norm(SiLU(conv4(window, new))) (q also scaled); the q / k windows of this CTA's quarter rewritten + once all four CTAs of the head have read the old ones (cluster barrier); + phase 2 v and b (bf16 of the two K-half partials): v = SiLU(conv4) of this CTA's 32 channels, its windows + rewritten; beta = sigmoid(b); decay = exp(lower_bound * sigmoid(exp(A_log) * (g + dt_bias))); + update ``kda_decode``'s arithmetic and row / key layout: lane l keys 4 l .. 4 l + 3, warp w rows w, w + 8, + w + 16, w + 24; S *= decay; r = (v - S k) beta; S += k r; o = S q; S back to the pool; + phase 3 the output gate; the gated RMSNorm over V (each CTA's two 16-row sums of squares per row to the four CTAs + of the head); y = o * rms * w_norm * sigmoid(gate), bf16. +Pools: conv bf16 [slots][3 * 768][3] (q | k | v channels, the last three raw inputs, oldest first; slot stride in +elements, a multiple of 8); state fp32 [slots][6][128][128] (V rows, K contiguous; slot stride a multiple of 4). +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +try: + from cutlass.memory.smem import SmemAllocator +except ImportError: # older DSL layout + from cutlass.utils import SmemAllocator + +from ..k3_kda_verify.k3_kda_verify_kernel import _bf16, _butterfly, _st_async_f32, _store8 +from .k3_kda_attn_kernel import ( + BOX_CH, + BOX_ELEMS, + BOX_K, + BUFFERS, + CLUSTER, + CONV_W, + CP_CG, + FB_TMEM_COLS, + FUSED_RAW_BYTES, + FUSED_STAGES, + HD, + HEAD_CLUSTERS, + HK, + K_IN, + K_STEP_U, + MMA_K, + MMA_N, + P1_ROWS, + PART_ROWS, + PROJ_ROWS, + SENTINEL, + STREAM_CLUSTERS, + THREADS, + TILE, + V_CTA, + X_ELEMS, + _bf16_bits, + _mapa_u32, + _poll_partials, + _sent16, + _sent32, + _SmemCarver, + _st_b16, + _stream_role, +) + +NR = MMA_N # request rows of a launch (the projection's token rows) +WIN = CONV_W - 1 # raw inputs a conv window keeps +WIN_QK = HD * WIN # bf16 of one request's q (or k) window for a head: 128 channels x 3 +WIN_V = V_CTA * WIN # bf16 of one request's v window for a CTA's 32 channels +QK_CHUNKS = WIN_QK // 8 # 16-byte copies per window +V_CHUNKS = WIN_V // 8 +REQ_CHUNKS = 2 * QK_CHUNKS + V_CHUNKS +ST_ROWS = V_CTA * HD # fp32 of one request's state rows owned by a CTA (16 KB) +ST_CHUNKS = ST_ROWS // 4 +LANE_KEYS = HD // 32 # keys per lane in the update (kda_decode's float4) +CW_ITEMS = 2 * V_CTA * WIN # q and k window entries a CTA rewrites per request + + +@dsl_user_op +def _ld_bf16(smem_ptr, *, loc=None, ip=None): + """fp32 of one bf16 in shared memory by ld.shared.b16: a channel's three window entries are only 2-byte aligned, + so their loads must not be merged into one wider load.""" + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [smem_ptr.toint(loc=loc, ip=ip).ir_value()], + "{ .reg .b16 h; ld.shared.b16 h, [$1]; cvt.f32.bf16 $0, h; }", "=f,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _test_wait_cluster(mbar, parity, *, loc=None, ip=None): + """mbarrier.test_wait.parity.acquire.cluster on a barrier of this CTA (``mbar``, its shared-memory pointer): whether + phase ``parity`` has completed, acquiring at cluster scope. The barrier is completed by other CTAs' st.async, whose + complete_tx releases at cluster scope.""" + done = cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tmbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" + "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + return done != cutlass.Int32(0) + + +@cute.jit +def _decode_head_role( + tma_wfb, + p1w: cutlass.Array, + part: cutlass.Array, + w_q: cutlass.Array, + w_k: cutlass.Array, + w_v: cutlass.Array, + a_log: cutlass.Array, + dt_bias: cutlass.Array, + onorm_w: cutlass.Array, + conv: cutlass.Array, + ssm: cutlass.Array, + slots: cutlass.Array, + out: cutlass.Array, + epoch: cutlass.Array, + smem_a: cutlass.Array, + smem_b: cutlass.Array, + bars: cutlass.Array, + tmem_holder: cutlass.Array, + s_slot: cutlass.Array, + s_winq: cutlass.Array, + s_wink: cutlass.Array, + s_winv: cutlass.Array, + s_wq: cutlass.Array, + s_wk: cutlass.Array, + s_wv: cutlass.Array, + s_dtb: cutlass.Array, + s_onw: cutlass.Array, + s_nq: cutlass.Array, + s_nk: cutlass.Array, + s_gr: cutlass.Array, + s_q: cutlass.Array, + s_k: cutlass.Array, + s_dec: cutlass.Array, + s_beta: cutlass.Array, + s_braw: cutlass.Array, + s_v: cutlass.Array, + s_o: cutlass.Array, + s_ss: cutlass.Array, + s_rs: cutlass.Array, + s_vp: cutlass.Array, + s_bp: cutlass.Array, + s_st: cutlass.Array, + n_req, + ssm_stride, + conv_stride, + lower_bound: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + eps: cutlass.Constexpr[float], + USE_PDL: cutlass.Constexpr[bool], + lbx, +): + """One CTA of a head cluster (logical clusters 26-31): local head h = cluster - 26, V rows [32 q, 32 q + 32) + (q = cluster rank) of every request, fed by the stream clusters' Lamport buffers of this launch.""" + tidx, _, _ = cute.arch.thread_idx() + bx = lbx + lane = tidx % 32 + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + vq = cute.arch.block_idx_in_cluster() + h = bx // cutlass.Int32(CLUSTER) - cutlass.Int32(STREAM_CLUSTERS) + v0 = vq * V_CTA + ch0 = h * HD + w_full = bars.subview(0) + acc_done = bars.subview(1) + ss_ready = bars.subview(2) + + tma_ptr_w = tma_wfb.get_ptr() + if warp == 0: + prims.prefetch_tensormap(tma_ptr_w) + if prims.elect_sync(): + prims.mbarrier_init(w_full, 1) + prims.mbarrier_init(acc_done, 1) + elif warp == 2: + prims.tcgen05_alloc(tmem_holder, FB_TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + elif warp == 3: + if prims.elect_sync(): + prims.mbarrier_init(ss_ready, 1) + prims.mbarrier_arrive_expect_tx(ss_ready, CLUSTER * 2 * NR * 4) + prims.fence_mbarrier_init() + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + if warp == 0: + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(w_full, HD * HD * 2) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a, tma_ptr_w, (cutlass.Int32(0), ch0, cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0)), + w_full, + ) # fmt: skip + + # ---- Before the grid dependency: the slots, the constants, then (asynchronous copies) every request's conv + # windows and its state rows. + if tidx < n_req: + s_slot.store(slots.load(idx=tidx), idx=tidx) + a_raw = a_log.load(idx=h) + if tidx < HD: + wq4 = w_q.load(idx=(ch0 + tidx) * CONV_W, vector_size=CONV_W, alignment=16) + wk4 = w_k.load(idx=(ch0 + tidx) * CONV_W, vector_size=CONV_W, alignment=16) + for w in cutlass.range_constexpr(CONV_W): + s_wq.store(cutlass.Float32(wq4[w]), idx=w * HD + tidx) + s_wk.store(cutlass.Float32(wk4[w]), idx=w * HD + tidx) + s_dtb.store(dt_bias.load(idx=ch0 + tidx), idx=tidx) + elif tidx < HD + V_CTA: + cv_w = tidx - HD + wv4 = w_v.load(idx=(ch0 + v0 + cv_w) * CONV_W, vector_size=CONV_W, alignment=16) + for w in cutlass.range_constexpr(CONV_W): + s_wv.store(cutlass.Float32(wv4[w]), idx=w * V_CTA + cv_w) + s_onw.store(onorm_w.load(idx=v0 + cv_w), idx=cv_w) + for i_o in cutlass.range_constexpr(NR * V_CTA // THREADS): + s_o.store(cutlass.Float32(0.0), idx=tidx + i_o * THREADS) + prims.barrier_cta_sync(0) + for it_w in cutlass.range_constexpr((NR * REQ_CHUNKS + THREADS - 1) // THREADS): + q_w = tidx + it_w * THREADS + r_w = q_w // REQ_CHUNKS + u_w = q_w % REQ_CHUNKS + if r_w < n_req: + conv_w = conv.subview(s_slot.load(idx=r_w) * conv_stride) + if u_w < QK_CHUNKS: + prims.cp_async_shared_global( + s_winq.subview(r_w * WIN_QK + u_w * 8), + conv_w.subview(ch0 * WIN + u_w * 8), + 16, + CP_CG, + ) + elif u_w < 2 * QK_CHUNKS: + u_k = u_w - QK_CHUNKS + prims.cp_async_shared_global( + s_wink.subview(r_w * WIN_QK + u_k * 8), + conv_w.subview((HK + ch0) * WIN + u_k * 8), + 16, + CP_CG, + ) + else: + u_v = u_w - 2 * QK_CHUNKS + prims.cp_async_shared_global( + s_winv.subview(r_w * WIN_V + u_v * 8), + conv_w.subview((2 * HK + ch0 + v0) * WIN + u_v * 8), + 16, + CP_CG, + ) + prims.cp_async_commit_group() + for r_s in cutlass.range_constexpr(NR): + if r_s < n_req: + st_src = ssm.subview(s_slot.load(idx=r_s) * ssm_stride + (ch0 + v0) * HD) + for j_s in cutlass.range_constexpr(ST_CHUNKS // THREADS): + c_s = (tidx + j_s * THREADS) * 4 + prims.cp_async_shared_global( + s_st.subview(r_s * ST_ROWS + c_s), st_src.subview(c_s), 16, CP_CG + ) + prims.cp_async_commit_group() + exp_a = cute.math.exp(a_raw, fastmath=True) + # The windows (the first group) have landed: every CTA of the head has read the head's q / k windows once all + # four arrive here (release; the acquiring wait follows phase 1), so each may then rewrite its quarter. + prims.cp_async_wait_group(1) + prims.barrier_cta_sync(0) + prims.barrier_cluster_arrive() + + if cutlass.const_expr(USE_PDL): + prims.griddepcontrol(prims.GridDepAction.WAIT) + e = epoch.load(idx=bx) + buf = e % cutlass.Int32(BUFFERS) + if cutlass.const_expr(USE_PDL): + if tidx == 0: + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + + # ---- Phase 1 (q, k, f_a): thread i polls 16-byte chunk i of the q | k rows, threads < 128 also chunk i of the + # f_a rows; q and k as fp32, f_a into the MMA's B operand (128-byte swizzle). + st_qk = tidx // HD + t_qk = (tidx % HD) // 16 + j_qk = tidx % 16 + a0 = cutlass.Int32(SENTINEL) + a1 = cutlass.Int32(SENTINEL) + a2 = cutlass.Int32(SENTINEL) + a3 = cutlass.Int32(SENTINEL) + qk_idx = ((buf * NR + t_qk) * P1_ROWS + st_qk * HK + ch0 + j_qk * 8) // 2 + while _sent16(a0) | _sent16(a1) | _sent16(a2) | _sent16(a3): + vqk = prims.load_ext( + p1w.subview(qk_idx), dtype=cutlass.Int32, count=4, order="relaxed", scope="gpu" + ) + a0 = cutlass.Int32(vqk[0]) + a1 = cutlass.Int32(vqk[1]) + a2 = cutlass.Int32(vqk[2]) + a3 = cutlass.Int32(vqk[3]) + if st_qk == 0: + _store8(s_nq, (a0, a1, a2, a3), t_qk * HD + j_qk * 8) + else: + _store8(s_nk, (a0, a1, a2, a3), t_qk * HD + j_qk * 8) + if tidx < HD: + t_fa = tidx // 16 + j_fa = tidx % 16 + f0 = cutlass.Int32(SENTINEL) + f1 = cutlass.Int32(SENTINEL) + f2 = cutlass.Int32(SENTINEL) + f3 = cutlass.Int32(SENTINEL) + fa_idx = ((buf * NR + t_fa) * P1_ROWS + 2 * HK + j_fa * 8) // 2 + while _sent16(f0) | _sent16(f1) | _sent16(f2) | _sent16(f3): + vfa = prims.load_ext( + p1w.subview(fa_idx), dtype=cutlass.Int32, count=4, order="relaxed", scope="gpu" + ) + f0 = cutlass.Int32(vfa[0]) + f1 = cutlass.Int32(vfa[1]) + f2 = cutlass.Int32(vfa[2]) + f3 = cutlass.Int32(vfa[3]) + # Row t_fa, 16-byte unit u of 64-column chunk c: byte 1024 c + 128 t_fa + 16 (u ^ t_fa). + sw_word = ((j_fa // 8) * 1024 + t_fa * 128 + ((j_fa % 8) ^ t_fa) * 16) // 4 + smem_b.store((f0, f1, f2, f3), idx=sw_word, alignment=16) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.barrier_cta_sync(0) + prims.barrier_cluster_wait() + + # ---- f_b on the tensor cores (warp 2 issues, warps 4-7 read TMEM), rounded to bf16 as the unfused GEMM's output. + if warp == 2: + idesc = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cutlass.BFloat16, + b_dtype=cutlass.BFloat16, + n_dim=MMA_N, + m_dim=HD, + ) + desc_a_base = prims.Tcgen05SmemDesc.build( + start_address=smem_a, leading_byte_offset=16, stride_byte_offset=1024, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + desc_b_base = prims.Tcgen05SmemDesc.build( + start_address=smem_b, leading_byte_offset=16, stride_byte_offset=1024, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + tmem_acc = cutlass.inttoptr(tmem_holder.load(), 6, cutlass.Int32) + while not cute.arch.mbarrier_test_wait(w_full.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for kb in cutlass.range_constexpr(2 * (BOX_K // MMA_K)): + box = kb // (BOX_K // MMA_K) + within = kb % (BOX_K // MMA_K) + if prims.elect_sync(): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, prims.CTAGroup.CTA_1, tmem_acc, + desc_a_base + (box * ((HD * BOX_K * 2) >> 4) + within * K_STEP_U), + desc_b_base + (box * ((MMA_N * BOX_K * 2) >> 4) + within * K_STEP_U), + idesc, kb != 0, + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + if warp >= 4: + while not cute.arch.mbarrier_test_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", cutlass.inttoptr(tmem_holder.load(), 6, cutlass.Float32), num=MMA_N + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + c_fb = (warp - 4) * 32 + lane + for t_fb in cutlass.range_constexpr(NR): + s_gr.store(_bf16(cutlass.Float32(acc[t_fb])), idx=t_fb * HD + c_fb) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.barrier_cta_sync(1, thread_count=128) + if warp == 4: + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_holder.load(), 6, cutlass.Int32), FB_TMEM_COLS + ) + + # ---- q and k of request w (warp w): conv4 over the window and the new input, SiLU, then the L2 norm with + # kda_decode's sum order (per 32-channel group a butterfly, then (g0 + g2) + (g1 + g3)); q also scaled. + if warp < n_req: + tk = warp + pq = [cutlass.Float32(0.0)] * LANE_KEYS + pk = [cutlass.Float32(0.0)] * LANE_KEYS + for i in cutlass.range_constexpr(LANE_KEYS): + c = i * 32 + lane + cq = cutlass.Float32(0.0) + ck = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(WIN): + cq += _ld_bf16(s_winq.subview(tk * WIN_QK + c * WIN + w).data_ptr()) * s_wq.load( + idx=w * HD + c + ) + ck += _ld_bf16(s_wink.subview(tk * WIN_QK + c * WIN + w).data_ptr()) * s_wk.load( + idx=w * HD + c + ) + cq += s_nq.load(idx=tk * HD + c) * s_wq.load(idx=WIN * HD + c) + ck += s_nk.load(idx=tk * HD + c) * s_wk.load(idx=WIN * HD + c) + pq[i] = cq * ( + cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.math.exp(-cq, fastmath=True)) + ) + pk[i] = ck * ( + cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.math.exp(-ck, fastmath=True)) + ) + gq = [_butterfly(pq[i] * pq[i]) for i in range(LANE_KEYS)] + gk = [_butterfly(pk[i] * pk[i]) for i in range(LANE_KEYS)] + rq = ( + cute.math.rsqrt( + (gq[0] + gq[2]) + (gq[1] + gq[3]) + cutlass.Float32(1e-6), fastmath=True + ) + * scale + ) + rk = cute.math.rsqrt( + (gk[0] + gk[2]) + (gk[1] + gk[3]) + cutlass.Float32(1e-6), fastmath=True + ) + for i in cutlass.range_constexpr(LANE_KEYS): + s_q.store(pq[i] * rq, idx=tk * HD + i * 32 + lane) + s_k.store(pk[i] * rk, idx=tk * HD + i * 32 + lane) + # The q / k windows of this CTA's quarter: the two newer raw inputs, then the token's (all four CTAs of the head + # read the old ones before the cluster wait above). + for it_c in cutlass.range_constexpr((NR * CW_ITEMS + THREADS - 1) // THREADS): + q_c = tidx + it_c * THREADS + r_c = q_c // CW_ITEMS + u_c = q_c % CW_ITEMS + if r_c < n_req: + sec_c = u_c // (V_CTA * WIN) + ch_c = vq * V_CTA + (u_c % (V_CTA * WIN)) // WIN + w_c = u_c % WIN + dst_c = ( + conv.subview( + s_slot.load(idx=r_c) * conv_stride + (sec_c * HK + ch0 + ch_c) * WIN + w_c + ) + .data_ptr() + .toint() + ) + # Single 16-bit stores: a channel's entries are only 2-byte aligned. + if w_c < WIN - 1: + if sec_c == 0: + _st_b16( + dst_c, + _bf16_bits( + _ld_bf16(s_winq.subview(r_c * WIN_QK + ch_c * WIN + w_c + 1).data_ptr()) + ), + ) + else: + _st_b16( + dst_c, + _bf16_bits( + _ld_bf16(s_wink.subview(r_c * WIN_QK + ch_c * WIN + w_c + 1).data_ptr()) + ), + ) + else: + if sec_c == 0: + _st_b16(dst_c, _bf16_bits(s_nq.load(idx=r_c * HD + ch_c))) + else: + _st_b16(dst_c, _bf16_bits(s_nk.load(idx=r_c * HD + ch_c))) + prims.barrier_cta_sync(0) + + # ---- Phase 2 (v, b): the two K-half partials of this CTA's 32 v rows and of this head's b; the projection is + # their bf16-rounded sum. + if tidx < 2 * NR * (V_CTA // 4): + half_v = tidx // (NR * (V_CTA // 4)) + t_v = (tidx // (V_CTA // 4)) % NR + j_v = tidx % (V_CTA // 4) + _poll_partials( + part, + ((buf * 3 + 0) * 2 + half_v) * (8 * PART_ROWS) + t_v * PART_ROWS + ch0 + v0 + j_v * 4, + s_vp, + (half_v * NR + t_v) * V_CTA + j_v * 4, + ) + elif tidx < 2 * NR * (V_CTA // 4) + 2 * NR: + ib = tidx - 2 * NR * (V_CTA // 4) + half_b = ib // NR + t_b = ib % NR + wb = cutlass.Int32(SENTINEL) + while _sent32(wb): + wb = prims.load_ext( + part.subview(((buf * 3 + 2) * 2 + half_b) * (8 * PART_ROWS) + t_b * PART_ROWS + h), + dtype=cutlass.Int32, + order="relaxed", + scope="gpu", + ) + s_bp.store(wb.bitcast(cutlass.Float32), idx=half_b * NR + t_b) + prims.barrier_cta_sync(0) + t_cv = tidx // V_CTA + c_cv = tidx % V_CTA + nv = _bf16(s_vp.load(idx=t_cv * V_CTA + c_cv) + s_vp.load(idx=(NR + t_cv) * V_CTA + c_cv)) + if tidx < NR: + s_braw.store(_bf16(s_bp.load(idx=tidx) + s_bp.load(idx=NR + tidx)), idx=tidx) + if t_cv < n_req: + # v (one (request, channel) per thread), then its window: the two newer raw inputs and the token's. + cv = cutlass.Float32(0.0) + wv_old = [ + _ld_bf16(s_winv.subview(t_cv * WIN_V + c_cv * WIN + w).data_ptr()) for w in range(WIN) + ] + for w in cutlass.range_constexpr(WIN): + cv += wv_old[w] * s_wv.load(idx=w * V_CTA + c_cv) + cv += nv * s_wv.load(idx=WIN * V_CTA + c_cv) + s_v.store( + cv + * (cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.math.exp(-cv, fastmath=True))), + idx=t_cv * V_CTA + c_cv, + ) + dst_v = ( + conv.subview(s_slot.load(idx=t_cv) * conv_stride + (2 * HK + ch0 + v0 + c_cv) * WIN) + .data_ptr() + .toint() + ) + _st_b16(dst_v, _bf16_bits(wv_old[1])) + _st_b16(dst_v + cutlass.Int64(2), _bf16_bits(wv_old[2])) + _st_b16(dst_v + cutlass.Int64(4), _bf16_bits(nv)) + prims.barrier_cta_sync(0) + if tidx < n_req: + s_beta.store( + cutlass.Float32(1.0) + / (cutlass.Float32(1.0) + cute.math.exp(-s_braw.load(idx=tidx), fastmath=True)), + idx=tidx, + ) + # The decay of request w's keys (warp w; lane l keys 4 l .. 4 l + 3, the update's layout). + if warp < n_req: + for i in cutlass.range_constexpr(LANE_KEYS): + c = lane * LANE_KEYS + i + xg = exp_a * (s_gr.load(idx=warp * HD + c) + s_dtb.load(idx=c)) + sig = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.math.exp(-xg, fastmath=True)) + s_dec.store(cute.math.exp(lower_bound * sig, fastmath=True), idx=warp * HD + c) + prims.cp_async_wait_group(0) + prims.barrier_cta_sync(0) + + # ---- The update, request by request: warp w rows w, w + 8 and w + 16, w + 24 of this CTA's 32, lane l keys + # 4 l .. 4 l + 3 (kda_decode's layout and sum order); the rows back to the pool, the outputs to s_o. + for t in cutlass.range(n_req, unroll=1): + q4 = s_q.load(idx=t * HD + lane * LANE_KEYS, vector_size=4, alignment=16) + k4 = s_k.load(idx=t * HD + lane * LANE_KEYS, vector_size=4, alignment=16) + d4 = s_dec.load(idx=t * HD + lane * LANE_KEYS, vector_size=4, alignment=16) + beta_t = s_beta.load(idx=t) + st_dst = ssm.subview(s_slot.load(idx=t) * ssm_stride + (ch0 + v0) * HD) + for pr in cutlass.range_constexpr(2): + ra = warp + 16 * pr + rb = ra + 8 + sa4 = s_st.load( + idx=t * ST_ROWS + ra * HD + lane * LANE_KEYS, vector_size=4, alignment=16 + ) + sb4 = s_st.load( + idx=t * ST_ROWS + rb * HD + lane * LANE_KEYS, vector_size=4, alignment=16 + ) + sa = [cutlass.Float32(sa4[i]) * cutlass.Float32(d4[i]) for i in range(LANE_KEYS)] + sb = [cutlass.Float32(sb4[i]) * cutlass.Float32(d4[i]) for i in range(LANE_KEYS)] + ska = sa[0] * cutlass.Float32(k4[0]) + skb = sb[0] * cutlass.Float32(k4[0]) + for i in cutlass.range_constexpr(1, LANE_KEYS): + ska = ska + sa[i] * cutlass.Float32(k4[i]) + skb = skb + sb[i] * cutlass.Float32(k4[i]) + for offset in [16, 8, 4, 2, 1]: + ska = ska + cute.arch.shuffle_sync_bfly( + ska, offset=offset, mask=-1, mask_and_clamp=31 + ) + skb = skb + cute.arch.shuffle_sync_bfly( + skb, offset=offset, mask=-1, mask_and_clamp=31 + ) + res_a = (s_v.load(idx=t * V_CTA + ra) - ska) * beta_t + res_b = (s_v.load(idx=t * V_CTA + rb) - skb) * beta_t + for i in cutlass.range_constexpr(LANE_KEYS): + sa[i] = sa[i] + cutlass.Float32(k4[i]) * res_a + sb[i] = sb[i] + cutlass.Float32(k4[i]) * res_b + st_dst.store((sa[0], sa[1], sa[2], sa[3]), idx=ra * HD + lane * LANE_KEYS, alignment=16) + st_dst.store((sb[0], sb[1], sb[2], sb[3]), idx=rb * HD + lane * LANE_KEYS, alignment=16) + sqa = sa[0] * cutlass.Float32(q4[0]) + sqb = sb[0] * cutlass.Float32(q4[0]) + for i in cutlass.range_constexpr(1, LANE_KEYS): + sqa = sqa + sa[i] * cutlass.Float32(q4[i]) + sqb = sqb + sb[i] * cutlass.Float32(q4[i]) + for offset in [16, 8, 4, 2, 1]: + sqa = sqa + cute.arch.shuffle_sync_bfly( + sqa, offset=offset, mask=-1, mask_and_clamp=31 + ) + sqb = sqb + cute.arch.shuffle_sync_bfly( + sqb, offset=offset, mask=-1, mask_and_clamp=31 + ) + if lane == 0: + s_o.store(sqa, idx=t * V_CTA + ra) + s_o.store(sqb, idx=t * V_CTA + rb) + + # ---- Phase 3 (the output gate), then the gated RMSNorm over V: every CTA sends its two 16-row sums of squares + # per row to all four CTAs of the head. Thread (row t, V row v) reads its gate's two K-half partials first, so + # their round trip overlaps the norm exchange. + t_out = tidx // V_CTA + v_y = tidx % V_CTA + og_at = (buf * 3 + 1) * 2 * (8 * PART_ROWS) + t_out * PART_ROWS + ch0 + v0 + v_y + og_h0 = cutlass.Int32( + prims.load_ext(part.subview(og_at), dtype=cutlass.Int32, order="relaxed", scope="gpu") + ) + og_h1 = cutlass.Int32(prims.load_ext(part.subview(og_at + 8 * PART_ROWS), dtype=cutlass.Int32, order="relaxed", + scope="gpu")) # fmt: skip + prims.barrier_cta_sync(0) + x_o = s_o.load(idx=tidx) + ss = x_o * x_o + for off_ss in [8, 4, 2, 1]: + ss = ss + cute.arch.shuffle_sync_bfly(ss, offset=off_ss, mask=-1, mask_and_clamp=31) + if lane % 16 == 0: + src_slot = vq * 2 + lane // 16 + for r in cutlass.range_constexpr(CLUSTER): + _st_async_f32( + _mapa_u32(s_ss.subview(src_slot * NR + t_out).data_ptr(), r), + ss, + _mapa_u32(ss_ready.data_ptr(), r), + ) + # Completed by the head's st.async: acquire at cluster scope. + while not _test_wait_cluster(ss_ready.data_ptr(), 0): + pass + while _sent32(og_h0) | _sent32(og_h1): + og_h0 = cutlass.Int32( + prims.load_ext(part.subview(og_at), dtype=cutlass.Int32, order="relaxed", scope="gpu") + ) + og_h1 = cutlass.Int32(prims.load_ext(part.subview(og_at + 8 * PART_ROWS), dtype=cutlass.Int32, + order="relaxed", scope="gpu")) # fmt: skip + if tidx < NR: + total = cutlass.Float32(0.0) + for r in cutlass.range_constexpr(2 * CLUSTER): + total = total + s_ss.load(idx=r * NR + tidx) + s_rs.store(cute.math.rsqrt(total / HD + eps), idx=tidx) + prims.barrier_cta_sync(0) + if t_out < n_req: + z = _bf16(og_h0.bitcast(cutlass.Float32) + og_h1.bitcast(cutlass.Float32)) + gate = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.math.exp(-z, fastmath=True)) + y = s_o.load(idx=tidx) * s_rs.load(idx=t_out) * s_onw.load(idx=v_y) * gate + out.store(cutlass.BFloat16(y), idx=(t_out * (HK // HD) + h) * HD + v0 + v_y) + + if tidx == 0: + epoch.store((e + cutlass.Int32(1)) % cutlass.Int32(BUFFERS), idx=bx) + + +@cute.kernel +def k3_kda_decode_fused_kernel( + tma_w: cutlass.GridConstant[ + cuda.TensorMap + ], # W [3208, 7168] bf16, 5-D, box 64 cols x 64 rows x 2 chunks + tma_x: cutlass.GridConstant[ + cuda.TensorMap + ], # x [R, 7168] bf16, box 64 cols x 8 rows (rows >= R read as 0) + tma_wfb: cutlass.GridConstant[ + cuda.TensorMap + ], # W_fb [768, 128] bf16, a head's 128 rows by one call + p1: cutlass.Array, # int16 bits of bf16 [3][8][1664] (Lamport, written by the stream role) + p1w: cutlass.Array, # the same memory as int32 words (polled by the head role) + part: cutlass.Array, # int32 bits of fp32 [3][3][2][8][768] (Lamport) + epoch: cutlass.Array, # int32 [128]: each CTA's buffer index (launches completed mod 3) + w_q: cutlass.Array, # fp32 [768, 4] + w_k: cutlass.Array, + w_v: cutlass.Array, + a_log: cutlass.Array, # fp32 [6] + dt_bias: cutlass.Array, # fp32 [768] + onorm_w: cutlass.Array, # fp32 [128] + conv: cutlass.Array, # bf16 [slots][2304][3], slot stride conv_stride elements + ssm: cutlass.Array, # fp32 [slots][6][128][128], slot stride ssm_stride elements + slots: cutlass.Array, # int32 [R] + out: cutlass.Array, # bf16 [R][6][128] + n_req: cutlass.Int32, + ssm_stride: cutlass.Int64, + conv_stride: cutlass.Int64, + lower_bound: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + eps: cutlass.Constexpr[float], + USE_PDL: cutlass.Constexpr[bool], +): + # The stream CTAs (clusters 0-25) and the head CTAs (26-31) never share a cluster, so their large arrays alias one + # raw buffer (k3_kda_attn's carve). + raw = SmemAllocator().allocate(FUSED_RAW_BYTES, byte_alignment=1024) + stream_carve = _SmemCarver(raw, FUSED_RAW_BYTES) + ring_w = stream_carve.view(cutlass.BFloat16, FUSED_STAGES * BOX_ELEMS, 1024) + ring_x = stream_carve.view(cutlass.BFloat16, FUSED_STAGES * X_ELEMS, 1024) + mbox = stream_carve.view(cutlass.Float32, 3 * CLUSTER * 32 * 4, 16) + head_carve = _SmemCarver(raw, FUSED_RAW_BYTES) + smem_a = head_carve.view(cutlass.BFloat16, HD * HD, 1024) + smem_b = head_carve.view(cutlass.Int32, MMA_N * HD // 2, 1024) + s_st = head_carve.view(cutlass.Float32, NR * ST_ROWS) + s_winq = head_carve.view(cutlass.BFloat16, NR * WIN_QK) + s_wink = head_carve.view(cutlass.BFloat16, NR * WIN_QK) + s_winv = head_carve.view(cutlass.BFloat16, NR * WIN_V) + s_slot = head_carve.view(cutlass.Int32, NR) + s_wq = head_carve.view(cutlass.Float32, CONV_W * HD) + s_wk = head_carve.view(cutlass.Float32, CONV_W * HD) + s_wv = head_carve.view(cutlass.Float32, CONV_W * V_CTA) + s_dtb = head_carve.view(cutlass.Float32, HD) + s_onw = head_carve.view(cutlass.Float32, V_CTA) + s_nq = head_carve.view(cutlass.Float32, NR * HD) + s_nk = head_carve.view(cutlass.Float32, NR * HD) + s_gr = head_carve.view(cutlass.Float32, NR * HD) + s_q = head_carve.view(cutlass.Float32, NR * HD) + s_k = head_carve.view(cutlass.Float32, NR * HD) + s_dec = head_carve.view(cutlass.Float32, NR * HD) + s_beta = head_carve.view(cutlass.Float32, NR) + s_braw = head_carve.view(cutlass.Float32, NR) + s_v = head_carve.view(cutlass.Float32, NR * V_CTA) + s_o = head_carve.view(cutlass.Float32, NR * V_CTA) + s_ss = head_carve.view(cutlass.Float32, 2 * CLUSTER * NR) + s_rs = head_carve.view(cutlass.Float32, NR) + s_vp = head_carve.view(cutlass.Float32, 2 * NR * V_CTA) + s_bp = head_carve.view(cutlass.Float32, 2 * NR) + # Barriers and the TMEM holder stay static (both roles initialize their own). + full = cutlass.Array(cutlass.Int64, FUSED_STAGES, space=cutlass.AddressSpace.smem, alignment=8) + empty = cutlass.Array(cutlass.Int64, FUSED_STAGES, space=cutlass.AddressSpace.smem, alignment=8) + acc_done = cutlass.Array(cutlass.Int64, 3, space=cutlass.AddressSpace.smem, alignment=8) + mbox_bar = cutlass.Array(cutlass.Int64, 3, space=cutlass.AddressSpace.smem, alignment=8) + tmem_holder = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + bars = cutlass.Array(cutlass.Int64, 3, space=cutlass.AddressSpace.smem, alignment=8) + + # Launch order = logical order: the f_a / b stream clusters, the q/k/v/og stream clusters, then the head clusters. + # Beside a predecessor that still holds SMs (the residual-update sandwich before this layer), the CTAs that launch + # first are the stream CTAs, which have the most bytes to move; the head CTAs' prologue is short. Stream CTAs 0-103, + # head CTAs 104-127. + lbx = cutlass.Int32(cute.arch.block_idx()[0]) + if lbx < cutlass.Int32(STREAM_CLUSTERS * CLUSTER): + _stream_role(tma_w, tma_x, p1, part, epoch, ring_w, ring_x, full, empty, acc_done, mbox_bar, mbox, tmem_holder, + USE_PDL, FUSED_STAGES, lbx) # fmt: skip + else: + _decode_head_role(tma_wfb, p1w, part, w_q, w_k, w_v, a_log, dt_bias, onorm_w, conv, ssm, slots, out, epoch, + smem_a, smem_b, bars, tmem_holder, s_slot, s_winq, s_wink, s_winv, s_wq, s_wk, s_wv, s_dtb, + s_onw, s_nq, s_nk, s_gr, s_q, s_k, s_dec, s_beta, s_braw, s_v, s_o, s_ss, s_rs, s_vp, + s_bp, s_st, n_req, ssm_stride, conv_stride, lower_bound, scale, eps, USE_PDL, + lbx) # fmt: skip + + +@cute.jit +def k3_kda_decode( + w: cute.Tensor, # bf16 [3208, 7168] + x: cute.Tensor, # bf16 [R, 7168] + w_fb: cute.Tensor, # bf16 [768, 128] + p1: cute.Tensor, # int16 [3 * 8 * 1664] + p1w: cute.Tensor, # the same memory as int32 + part: cute.Tensor, # int32 [3 * 3 * 2 * 8 * 768] + epoch: cute.Tensor, # int32 [128] + w_q: cute.Tensor, + w_k: cute.Tensor, + w_v: cute.Tensor, + a_log: cute.Tensor, + dt_bias: cute.Tensor, + onorm_w: cute.Tensor, + conv: cute.Tensor, + ssm: cute.Tensor, + slots: cute.Tensor, + out: cute.Tensor, + n_req: cutlass.Int32, + ssm_stride: cutlass.Int64, + conv_stride: cutlass.Int64, + lower_bound: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + eps: cutlass.Constexpr[float], + USE_PDL: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + """One launch of the fused KDA projection + plain decode of ``n_req`` requests of one token.""" + tma_w = cuda.create_tensor_map_tiled( + global_address=w.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[BOX_K, PROJ_ROWS, K_IN // BOX_K, 1, 1], + global_strides=[(K_IN * 2) // 16, (BOX_K * 2) // 16, (PROJ_ROWS * K_IN * 2) // 16, + (PROJ_ROWS * K_IN * 2) // 16], + box_dims=[BOX_K, TILE, BOX_CH, 1, 1], + swizzle=cuda.TensorMapSwizzle.s128b, + ) # fmt: skip + tma_x = cuda.create_tensor_map_tiled( + global_address=x.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[K_IN, n_req], + global_strides=[(K_IN * 2) // 16], + box_dims=[BOX_K, MMA_N], + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_wfb = cuda.create_tensor_map_tiled( + global_address=w_fb.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[BOX_K, HK, HD // BOX_K, 1, 1], + global_strides=[(HD * 2) // 16, (BOX_K * 2) // 16, (HK * HD * 2) // 16, (HK * HD * 2) // 16], + box_dims=[BOX_K, HD, HD // BOX_K, 1, 1], + swizzle=cuda.TensorMapSwizzle.s128b, + ) # fmt: skip + k3_kda_decode_fused_kernel( + tma_w, tma_x, tma_wfb, p1, p1w, part, epoch, w_q, w_k, w_v, a_log, dt_bias, onorm_w, conv, ssm, slots, out, + n_req, ssm_stride, conv_stride, lower_bound, scale, eps, USE_PDL, + ).launch( + grid=((STREAM_CLUSTERS + HEAD_CLUSTERS) * CLUSTER, 1, 1), block=(THREADS, 1, 1), cluster=(CLUSTER, 1, 1), + stream=stream, use_pdl=USE_PDL, + # One CTA per SM (the shared-memory carve); without it ptxas may pick an occupancy-driven register target. + min_blocks_per_mp=1, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_attn/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_attn/op.py new file mode 100644 index 000000000000..a1f703c0df03 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_attn/op.py @@ -0,0 +1,356 @@ +# 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::k3_kda_qkvg``: Kimi K3's fused KDA projection (TP16 slice, W [3208, 7168]) as a three-phase stream. + +Outputs (Lamport buffers, see the kernel module): ``p1`` int16 [3, 8, 1664] holds the bf16 bits of the q, k and f_a +columns; ``part`` int32 [3, 3, 2, 8, 768] the fp32 bits of the two K-half partial sums of v (region 0), og (1) and b +(2, first 8 rows), whose bf16-rounded sum is the projection. A launch writes buffer e = ``epoch[cta]`` and resets buffer +(e + 1) % 3 to the sentinel; ``epoch`` int32 [104] holds each CTA's launch count mod 3. The buffers must persist across +launches and start as all-ones (``p1``, ``part``) and zeros (``epoch``). The kernel compiles on the first call for each +configuration, which must happen outside CUDA-graph capture. +""" + +from __future__ import annotations + +import os +import threading +from typing import Dict + +import torch + +K_IN = 7168 +PROJ_ROWS = 3208 +P1_NUMEL = 3 * 8 * 1664 +PART_NUMEL = 3 * 3 * 2 * 8 * 768 +CTAS = 104 +FUSED_CTAS = 128 +NT = 8 # one request's verify tokens (golden + 7 drafts) + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + + +def _arg(t: torch.Tensor, align: int = 16): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +def _index_arg(t: torch.Tensor): + """An int32 index tensor the kernels read element by element (``slots``, ``pending``), declared at its element's + alignment: the mixer passes slices such as ``state_indices[num_prefills:]``, which start on any 4-byte boundary.""" + return _arg(t, t.element_size()) + + +def use_pdl() -> bool: + return os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +def make_buffers(device: torch.device, ctas: int = CTAS): + """Fresh (p1, part, epoch) for :func:`k3_kda_qkvg` (``ctas`` = 104) or :func:`k3_kda_attn` (``FUSED_CTAS``).""" + p1 = torch.full((P1_NUMEL,), -1, dtype=torch.int16, device=device) + part = torch.full((PART_NUMEL,), -1, dtype=torch.int32, device=device) + epoch = torch.zeros(ctas, dtype=torch.int32, device=device) + return p1, part, epoch + + +@torch.library.custom_op( + "trtllm::k3_kda_qkvg", mutates_args=("p1", "part", "epoch"), device_types="cuda" +) +def k3_kda_qkvg( + x: torch.Tensor, + w: torch.Tensor, + p1: torch.Tensor, + part: torch.Tensor, + epoch: torch.Tensor, +) -> None: + """x bf16 [T <= 8, 7168], w bf16 [3208, 7168] (both contiguous); writes p1 and part, advances epoch (see the + module doc).""" + import cuda.bindings.driver as cuda_driver + + from . import k3_kda_attn_kernel as kernel + + tokens = x.shape[0] + if ( + x.dtype != torch.bfloat16 + or w.dtype != torch.bfloat16 + or tuple(w.shape) != (PROJ_ROWS, K_IN) + or x.dim() != 2 + or x.shape[1] != K_IN + or not 1 <= tokens <= 8 + or not x.is_contiguous() + or not w.is_contiguous() + or p1.numel() != P1_NUMEL + or p1.dtype != torch.int16 + or part.numel() != PART_NUMEL + or part.dtype != torch.int32 + or epoch.numel() != CTAS + or epoch.dtype != torch.int32 + ): + raise ValueError( + f"k3_kda_qkvg: unsupported call x {tuple(x.shape)} {x.dtype}, w {tuple(w.shape)} {w.dtype}, " + f"p1 {p1.numel()} {p1.dtype}, part {part.numel()} {part.dtype}, epoch {epoch.numel()} {epoch.dtype}" + ) + args = (_arg(w), _arg(x), _arg(p1.view(-1)), _arg(part.view(-1)), _arg(epoch.view(-1))) + stream = cuda_driver.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + pdl = use_pdl() + key = (pdl,) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "trtllm::k3_kda_qkvg must run once per configuration outside CUDA-graph capture first " + "(it compiles its kernel on the first call)." + ) + import cutlass.cute as cute + + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile(kernel.k3_kda_qkvg, *args, tokens, pdl, stream) + fn(*args, tokens, stream) + + +@k3_kda_qkvg.register_fake +def _(x, w, p1, part, epoch): + return None + + +@torch.library.custom_op( + "trtllm::k3_kda_attn", + mutates_args=("cs_q", "cs_k", "cs_v", "ssm", "state_tok", "p1", "part", "epoch"), + device_types="cuda", +) +def k3_kda_attn( + x: torch.Tensor, + w: torch.Tensor, + w_fb: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + a_log: torch.Tensor, + dt_bias: torch.Tensor, + onorm_w: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + ssm: torch.Tensor, + state_tok: torch.Tensor, + slots: torch.Tensor, + pending: torch.Tensor, + p1: torch.Tensor, + part: torch.Tensor, + epoch: torch.Tensor, + num_spec: int, + lower_bound: float, + scale: float, + eps: float, +) -> torch.Tensor: + """Stage B2: the fused projection of ``x`` (bf16 [8, 7168], one request's golden token and 7 drafts) and the KDA + verify of ``trtllm::k3_kda_verify`` on it, in one launch; returns the gated-norm core output bf16 [8, 6, 128]. + Weights, pools and state contract as ``k3_kda_verify`` (``slots`` int32 [1]); ``p1``, ``part`` and ``epoch`` + from :func:`make_buffers` with ``FUSED_CTAS``, persistent across launches. The head CTAs read the slot's pools + before the grid-dependency wait, so a launch must not follow another launch on the same pools directly in the + stream: a kernel that waits (or a non-PDL one) in between, as the model's other layers are.""" + import cuda.bindings.driver as cuda_driver + + from ..k3_kda_verify.op import _flat, _slots_view + from . import k3_kda_attn_kernel as kernel + + if ( + tuple(x.shape) != (NT, K_IN) + or x.dtype != torch.bfloat16 + or not x.is_contiguous() + or tuple(w.shape) != (PROJ_ROWS, K_IN) + or w.dtype != torch.bfloat16 + or not w.is_contiguous() + or tuple(w_fb.shape) != (768, 128) + or w_fb.dtype != torch.bfloat16 + or not w_fb.is_contiguous() + or num_spec != NT - 1 + or tuple(ssm.shape[1:]) != (6, 128, 128) + or tuple(state_tok.shape[1:]) != (num_spec, 6, 128, 128) + or ssm.stride()[1:] != (128 * 128, 128, 1) + or not state_tok.is_contiguous() + or cs_q.shape[-1] != 3 + num_spec + or slots.numel() != 1 + or slots.dtype != torch.int32 + or pending.dtype != torch.int32 + or p1.numel() != P1_NUMEL + or part.numel() != PART_NUMEL + or epoch.numel() != FUSED_CTAS + or ssm.dtype != torch.float32 + or state_tok.dtype != torch.float32 + or any( + t.dtype != torch.float32 + for t in (w_q, w_k, w_v, a_log, dt_bias, onorm_w, cs_q, cs_k, cs_v) + ) + ): + raise ValueError( + f"k3_kda_attn: unsupported call x {tuple(x.shape)} {x.dtype}, w {tuple(w.shape)}, " + f"w_fb {tuple(w_fb.shape)}, ssm {tuple(ssm.shape)}, state_tok {tuple(state_tok.shape)}, " + f"cs_q {tuple(cs_q.shape)}, " + f"num_spec {num_spec}, slots {tuple(slots.shape)} {slots.dtype}, epoch {epoch.numel()}" + ) + out = torch.empty(NT, 6, 128, dtype=torch.bfloat16, device=x.device) + args = ( + _arg(w), _arg(x), _arg(w_fb), _arg(p1.view(-1)), _arg(p1.view(torch.int32)), _arg(part.view(-1)), + _arg(epoch.view(-1)), _arg(_flat(w_q)), _arg(_flat(w_k)), _arg(_flat(w_v)), _arg(_flat(a_log)), + _arg(_flat(dt_bias)), _arg(_flat(onorm_w)), _arg(_flat(cs_q)), _arg(_flat(cs_k)), _arg(_flat(cs_v)), + _arg(_slots_view(ssm)), _arg(_slots_view(state_tok)), _index_arg(slots), _index_arg(pending), + _arg(_flat(out)), + ) # fmt: skip + stream = cuda_driver.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + pdl = use_pdl() + key = ("attn", float(lower_bound), float(scale), float(eps), pdl) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "trtllm::k3_kda_attn must run once per configuration outside CUDA-graph capture first " + "(it compiles its kernel on the first call)." + ) + import cutlass.cute as cute + + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kernel.k3_kda_attn, *args, ssm.stride(0), float(lower_bound), float(scale), float(eps), pdl, + stream, + ) # fmt: skip + fn(*args, ssm.stride(0), stream) + return out + + +@k3_kda_attn.register_fake +def _(x, w, w_fb, w_q, w_k, w_v, a_log, dt_bias, onorm_w, cs_q, cs_k, cs_v, ssm, state_tok, slots, pending, p1, part, + epoch, num_spec, lower_bound, scale, eps): # fmt: skip + return x.new_empty((NT, 6, 128), dtype=torch.bfloat16) + + +@torch.library.custom_op( + "trtllm::k3_kda_decode_attn", + mutates_args=("conv", "ssm", "p1", "part", "epoch"), + device_types="cuda", +) +def k3_kda_decode_attn( + x: torch.Tensor, + w: torch.Tensor, + w_fb: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + a_log: torch.Tensor, + dt_bias: torch.Tensor, + onorm_w: torch.Tensor, + conv: torch.Tensor, + ssm: torch.Tensor, + slots: torch.Tensor, + p1: torch.Tensor, + part: torch.Tensor, + epoch: torch.Tensor, + lower_bound: float, + scale: float, + eps: float, +) -> torch.Tensor: + """The fused projection of ``x`` (bf16 [R <= 8, 7168], one token of each of R requests) and the KDA plain decode + of ``trtllm::kda_decode`` on it, in one launch; returns the gated-norm core output bf16 [R, 6, 128]. + + ``conv`` bf16 [slots, 2304, 3] (q | k | v channels, the last three raw inputs; each slot dense, the slot stride a + multiple of 8 elements) and ``ssm`` fp32 [slots, 6, 128, 128] (each slot dense, the slot stride a multiple of 4) + are updated in place at ``slots`` int32 [R]; ``w_q / w_k / w_v`` fp32 [768, 4]; ``a_log`` fp32 [6]; + ``dt_bias`` fp32 [768]; ``onorm_w`` fp32 [128]; ``p1``, ``part`` and ``epoch`` from :func:`make_buffers` with + ``FUSED_CTAS``, persistent across launches. The head CTAs read the slots' pools before the grid-dependency wait, + so a launch must not follow another launch on the same pools directly in the stream: a kernel that waits (or a + non-PDL one) in between, as the model's other layers are.""" + import cuda.bindings.driver as cuda_driver + + from ..k3_kda_verify.op import _flat, _slots_view + from . import k3_kda_decode_kernel as kernel + + n_req = x.shape[0] if x.dim() == 2 else 0 + if ( + x.dim() != 2 + or not 1 <= n_req <= NT + or x.shape[1] != K_IN + or x.dtype != torch.bfloat16 + or not x.is_contiguous() + or tuple(w.shape) != (PROJ_ROWS, K_IN) + or w.dtype != torch.bfloat16 + or not w.is_contiguous() + or tuple(w_fb.shape) != (768, 128) + or w_fb.dtype != torch.bfloat16 + or not w_fb.is_contiguous() + or conv.dtype != torch.bfloat16 + or conv.dim() != 3 + or tuple(conv.shape[1:]) != (3 * 768, 3) + or conv.stride()[1:] != (3, 1) + or conv.stride(0) % 8 != 0 + or conv.data_ptr() % 16 != 0 + or ssm.dtype != torch.float32 + or tuple(ssm.shape[1:]) != (6, 128, 128) + or ssm.stride()[1:] != (128 * 128, 128, 1) + or ssm.stride(0) % 4 != 0 + or ssm.data_ptr() % 16 != 0 + or slots.numel() != n_req + or slots.dtype != torch.int32 + or not slots.is_contiguous() + or p1.numel() != P1_NUMEL + or part.numel() != PART_NUMEL + or epoch.numel() != FUSED_CTAS + or any(t.dtype != torch.float32 for t in (w_q, w_k, w_v, a_log, dt_bias, onorm_w)) + ): + raise ValueError( + f"k3_kda_decode_attn: unsupported call x {tuple(x.shape)} {x.dtype}, w {tuple(w.shape)}, " + f"w_fb {tuple(w_fb.shape)}, conv {tuple(conv.shape)} {conv.dtype} strides {conv.stride()}, " + f"ssm {tuple(ssm.shape)} {ssm.dtype} strides {ssm.stride()}, slots {tuple(slots.shape)} {slots.dtype}, " + f"epoch {epoch.numel()}" + ) + out = torch.empty(n_req, 6, 128, dtype=torch.bfloat16, device=x.device) + args = ( + _arg(w), _arg(x), _arg(w_fb), _arg(p1.view(-1)), _arg(p1.view(torch.int32)), _arg(part.view(-1)), + _arg(epoch.view(-1)), _arg(_flat(w_q)), _arg(_flat(w_k)), _arg(_flat(w_v)), _arg(_flat(a_log)), + _arg(_flat(dt_bias)), _arg(_flat(onorm_w)), _arg(_slots_view(conv)), _arg(_slots_view(ssm)), + _index_arg(slots), _arg(_flat(out)), + ) # fmt: skip + runtime = (n_req, ssm.stride(0), conv.stride(0)) + stream = cuda_driver.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + pdl = use_pdl() + key = ("decode", float(lower_bound), float(scale), float(eps), pdl) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "trtllm::k3_kda_decode_attn must run once per configuration outside CUDA-graph capture first " + "(it compiles its kernel on the first call)." + ) + import cutlass.cute as cute + + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kernel.k3_kda_decode, *args, *runtime, float(lower_bound), float(scale), float(eps), pdl, + stream, + ) # fmt: skip + fn(*args, *runtime, stream) + return out + + +@k3_kda_decode_attn.register_fake +def _(x, w, w_fb, w_q, w_k, w_v, a_log, dt_bias, onorm_w, conv, ssm, slots, p1, part, epoch, lower_bound, scale, + eps): # fmt: skip + return x.new_empty((x.shape[0], 6, 128), dtype=torch.bfloat16) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/__init__.py new file mode 100644 index 000000000000..2892fd223c30 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/__init__.py @@ -0,0 +1,19 @@ +# 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 speculative verify as one CTM kernel (``trtllm::k3_kda_verify``). + +Importing :mod:`.op` registers the torch op; nothing is imported eagerly here so that +the CuTe DSL dependency stays optional for every other model. +""" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/k3_kda_verify_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/k3_kda_verify_kernel.py new file mode 100644 index 000000000000..ad95b425e4ee --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/k3_kda_verify_kernel.py @@ -0,0 +1,827 @@ +# 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 speculative verify as one kernel: the forget-gate up-projection (f_b), the per-token pre-compute and +the delta-rule recurrence over the 1 + NUM_SPEC verify tokens, with a state committed after every verify token so +the next round starts from the accepted one instead of replaying the accepted drafts. + +Grid (H, N, 8), cluster (1, 1, 8), 256 threads. CTA (h, n, s) owns V rows [16 s, 16 s + 16) of local head h for +request n; warp w owns rows 16 s + 2 w and 16 s + 2 w + 1, and lane l keys k = 32 i + l (i < 4) of each. + +Before ``griddepcontrol.wait`` (written by the previous verify step, or constant): + the slot and P, the drafts the sampler accepted last round; the starting state into registers (P == 0: the pool + state, committed after the last golden token; else ``state_tok[slot, P - 1]``); the raw conv inputs at positions + -3..-1 before the first new token (conv-cache columns P..P+2) for q and k (128 channels) and v (this CTA's 16); + the conv weights, dt_bias, A_log and the output-norm weight; the head's W_fb^T block by TMA. +After the wait (this step's fused projection rows [q | k | v | onorm gate | f_a | b]): the new tokens' raw q, k, v, + f_a, b and output gate. +f_b: g[t, c] = bf16(sum_i f_a[t, i] W_fb[128 h + c, i]), rounded as the unfused bf16 GEMV output is. +Pre-compute: warp t = verify token t: q, k = l2norm(silu(conv4(u))) (q also scaled), the lower-bound gate, beta and + this CTA's 16 v channels. +Recurrence: the arithmetic of ``kda_mtp_decode``'s V-split path over the verify tokens, unrolled; the state after + token 0 (the golden token) is committed to the pool, the state after token t >= 1 to + ``state_tok[slot, t - 1]``. +Epilogue: the gated RMSNorm of the outputs over V (per-token sums of squares stored into every peer's shared + memory, one cluster barrier), bf16 output rows; CTA 0 rewrites the q/k conv cache (the window at the + golden token and the raw inputs of the drafts), every CTA the same for its v channels. + +With FOLD_FB False the gate comes from ``g_ext`` (the unfused f_b output); the kernel is then bit-exact against +``kda_mtp_decode`` fed the same gate and replaying the same accepted drafts. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +K = 128 # key head dim; the value head dim V is the same +V = 128 +CONV_W = 4 +THREADS = 256 +WARPS = THREADS // 32 +V_SPLIT = 8 +V_CTA = V // V_SPLIT # V rows per CTA +ROWS = V_CTA // WARPS # V rows per warp +VEC = K // 32 # keys per lane +BF16_BYTES = 2 +CHUNK = 8 # bf16 per 16-byte load +SMEM = cutlass.AddressSpace.smem +RMEM = cutlass.AddressSpace.rmem + +assert ROWS == 2, "the recurrence below processes one pair of V rows per warp" + +# f_b on the tensor cores: D[c, t] = sum_i W_fb[128 h + c, i] f_a[t, i], M = 128 channels, N = 8 tokens, K = 128. +CTA_M = 128 +MMA_N = 8 +MMA_K = 16 +TMA_K_BOX = 64 +TMA_COPY_ITERS = K // TMA_K_BOX +K_BLOCKS_PER_HALF = TMA_K_BOX // MMA_K +LEADING = 16 +STRIDE = 8 * TMA_K_BOX * BF16_BYTES +A_HALF_ELEMS = CTA_M * TMA_K_BOX +B_HALF_ELEMS = MMA_N * TMA_K_BOX +STEP = (MMA_K * BF16_BYTES) >> 4 +A_BOX = A_HALF_ELEMS >> 3 +B_BOX = B_HALF_ELEMS >> 3 +TMEM_COLS = 32 + + +@dsl_user_op +def _mapa_u32(smem_ptr, peer, *, loc=None, ip=None): + """The shared::cluster address of this CTA's shared-memory location in cluster CTA ``peer``.""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [smem_ptr.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(peer).ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _test_wait_cluster(mbar, parity, *, loc=None, ip=None): + """mbarrier.test_wait.parity.acquire.cluster on a barrier of this CTA (``mbar``, its shared-memory pointer): whether + phase ``parity`` has completed, acquiring at cluster scope. For barriers that other CTAs' st.async complete (their + complete_tx releases at cluster scope; a CTA-scope acquire does not synchronize with it).""" + done = cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tmbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" + "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + return done != cutlass.Int32(0) + + +@dsl_user_op +def _try_wait_cluster(mbar, parity, *, loc=None, ip=None): + """mbarrier.try_wait.parity.acquire.cluster on a barrier of this CTA (``mbar``, its shared-memory pointer): as + ``_test_wait_cluster``, with try_wait's bounded suspend.""" + done = cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tmbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" + "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + return done != cutlass.Int32(0) + + +@dsl_user_op +def _st_async_f32(dst, value, mbar, *, loc=None, ip=None): + """st.async of one fp32 to a shared::cluster address, completing ``mbar`` (a shared::cluster address) by 4 bytes.""" + _llvm.inline_asm( + None, + [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Float32(value).ir_value(loc=loc, ip=ip), + cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 [$0], $1, [$2];", "r,f,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +def _lo(word): + return (word << cutlass.Int32(16)).bitcast(cutlass.Float32) + + +def _hi(word): + return (word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + + +def _bf16(x): + return cutlass.Float32(cutlass.BFloat16(x)) + + +def _butterfly(value): + for offset in [16, 8, 4, 2, 1]: + value += cute.arch.shuffle_sync_bfly(value, offset=offset, mask=-1, mask_and_clamp=31) + return value + + +def _store8(dst, words, idx): + """The 8 bf16 of four int32 words, as fp32, into dst[idx : idx + 8].""" + dst.store( + (_lo(cutlass.Int32(words[0])), _hi(cutlass.Int32(words[0])), _lo(cutlass.Int32(words[1])), + _hi(cutlass.Int32(words[1]))), + idx=idx, alignment=16, + ) # fmt: skip + dst.store( + (_lo(cutlass.Int32(words[2])), _hi(cutlass.Int32(words[2])), _lo(cutlass.Int32(words[3])), + _hi(cutlass.Int32(words[3]))), + idx=idx + 4, alignment=16, + ) # fmt: skip + + +def _qk_pre(tk, lane, s_uq, s_wq, s_uk, s_wk, s_q, s_k, scale): + """q and k of verify token ``tk`` from the conv history in shared memory: conv4 + SiLU, then the L2 norm (q also + scaled), lane l keys 32 i + l (the arithmetic of kda_mtp_decode's pre-compute warps). Traced inline: its Python + loops unroll.""" + pq = [cutlass.Float32(0.0)] * VEC + for i in range(VEC): + c = i * 32 + lane + conv = cutlass.Float32(0.0) + for w in range(CONV_W - 1): + conv += s_uq.load(idx=(tk + w) * K + c) * s_wq.load(idx=w * K + c) + conv += s_uq.load(idx=(tk + CONV_W - 1) * K + c) * s_wq.load(idx=(CONV_W - 1) * K + c) + e = cute.math.exp(-conv, fastmath=True) + pq[i] = conv * cute.arch.rcp_approx(cutlass.Float32(1.0) + e) + sum_q = cutlass.Float32(0.0) + for i in range(VEC): + sum_q += pq[i] * pq[i] + sum_q = _butterfly(sum_q) + rnorm_q = cute.math.rsqrt(sum_q + 1e-06, fastmath=True) * scale + for i in range(VEC): + s_q.store(pq[i] * rnorm_q, idx=tk * K + i * 32 + lane) + pk = [cutlass.Float32(0.0)] * VEC + for i in range(VEC): + c = i * 32 + lane + conv = s_uk.load(idx=tk * K + c) * s_wk.load(idx=c) + for w in range(1, CONV_W - 1): + conv += s_uk.load(idx=(tk + w) * K + c) * s_wk.load(idx=w * K + c) + conv += s_uk.load(idx=(tk + CONV_W - 1) * K + c) * s_wk.load(idx=(CONV_W - 1) * K + c) + pk[i] = conv * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-conv, fastmath=True) + ) + sum_k = cutlass.Float32(0.0) + for i in range(VEC): + sum_k += pk[i] * pk[i] + sum_k = _butterfly(sum_k) + rnorm_k = cute.math.rsqrt(sum_k + 1e-06, fastmath=True) + for i in range(VEC): + s_k.store(pk[i] * rnorm_k, idx=tk * K + i * 32 + lane) + + +@cute.kernel +def k3_kda_verify_kernel( + tma_wfb: cutlass.GridConstant[ + cuda.TensorMap + ], # W_fb [H * K, K] bf16, the head's 128 rows by one call + tma_fa: cutlass.GridConstant[ + cuda.TensorMap + ], # the projection rows' f_a columns, box 8 tokens x 64 + proj: cutlass.Array, # int32 words of the fused projection rows, bf16 [T, 2 * proj_words] + g_ext: cutlass.Array, # int32 words of bf16 [T, H * K], the unfused f_b output (FOLD_FB False only) + w_q: cutlass.Array, # fp32 [H * K, CONV_W] + w_k: cutlass.Array, + w_v: cutlass.Array, + a_log: cutlass.Array, # fp32 [H] + dt_bias: cutlass.Array, # fp32 [H * K] + onorm_w: cutlass.Array, # fp32 [V] + cs_q: cutlass.Array, # fp32 [pool][S][H * K]: raw conv inputs, column s = position s - 2 from the golden token + cs_k: cutlass.Array, + cs_v: cutlass.Array, + ssm: cutlass.Array, # fp32 [pool][H][V][K] at slot stride ssm_stride: the state after the last golden token + state_tok: cutlass.Array, # fp32 [pool][NUM_SPEC][H][V][K]: the state after draft t + 1 of the last round + slots: cutlass.Array, # int32 [N] + pending: cutlass.Array, # int32 [pool]: drafts the sampler accepted last round + out: cutlass.Array, # bf16 [T][H][V] + proj_words: cutlass.Int32, + ssm_stride: cutlass.Int64, # fp32 elements between the pool's slots (the manager interleaves conv states) + H: cutlass.Constexpr[int], + NUM_SPEC: cutlass.Constexpr[int], + lower_bound: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + eps: cutlass.Constexpr[float], + FOLD_FB: cutlass.Constexpr[bool], + USE_PDL: cutlass.Constexpr[bool], +): + NT = NUM_SPEC + 1 # verify tokens per request + S = CONV_W - 1 + NUM_SPEC # conv-cache columns + HK = H * K + ROWS_U = ( + CONV_W - 1 + NT + ) # raw conv inputs by position: rows 0..2 before token 0, then the tokens + # The fused projection row [q | k | v | onorm gate | f_a | b | pad], in int32 words (bf16 pairs). + Q_W = 0 + K_W = HK // 2 + V_W = HK + OG_W = 3 * HK // 2 + FA_W = 2 * HK + B_COL = 4 * HK + K # bf16 column + assert NT <= WARPS and NT % 2 == 0 + + tidx, _, _ = cute.arch.thread_idx() + lane = tidx % 32 + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + h, n, vs = cute.arch.block_idx() + v0 = vs * V_CTA + ch0 = h * K + row0 = n * NT + + smem_a = cutlass.Array( + cutlass.BFloat16, CTA_M * K, space=SMEM, alignment=1024 + ) # W_fb rows, 128B-swizzled + smem_b = cutlass.Array( + cutlass.BFloat16, MMA_N * K, space=SMEM, alignment=1024 + ) # f_a, 128B-swizzled + w_full = cutlass.Array(cutlass.Int64, 1, space=SMEM, alignment=8) + x_full = cutlass.Array(cutlass.Int64, 1, space=SMEM, alignment=8) + acc_done = cutlass.Array(cutlass.Int64, 1, space=SMEM, alignment=8) + ss_ready = cutlass.Array( + cutlass.Int64, 1, space=SMEM, alignment=8 + ) # the peers' norm partials have landed + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=SMEM) + s_uq = cutlass.Array(cutlass.Float32, ROWS_U * K, space=SMEM, alignment=16) + s_uk = cutlass.Array(cutlass.Float32, ROWS_U * K, space=SMEM, alignment=16) + s_uv = cutlass.Array(cutlass.Float32, ROWS_U * V_CTA, space=SMEM, alignment=16) + s_wq = cutlass.Array(cutlass.Float32, CONV_W * K, space=SMEM, alignment=16) # [tap][channel] + s_wk = cutlass.Array(cutlass.Float32, CONV_W * K, space=SMEM, alignment=16) + s_wv = cutlass.Array(cutlass.Float32, CONV_W * V_CTA, space=SMEM, alignment=16) + s_dtb = cutlass.Array(cutlass.Float32, K, space=SMEM, alignment=16) + s_onw = cutlass.Array(cutlass.Float32, V_CTA, space=SMEM, alignment=16) + s_gr = cutlass.Array( + cutlass.Float32, NT * K, space=SMEM, alignment=16 + ) # f_b output, bf16-rounded + s_braw = cutlass.Array(cutlass.Float32, NT, space=SMEM, alignment=16) + s_og = cutlass.Array(cutlass.Float32, NT * V_CTA, space=SMEM, alignment=16) + s_q = cutlass.Array(cutlass.Float32, NT * K, space=SMEM, alignment=16) + s_k = cutlass.Array(cutlass.Float32, NT * K, space=SMEM, alignment=16) + s_dec = cutlass.Array(cutlass.Float32, NT * K, space=SMEM, alignment=16) # exp(gate) + s_kd = cutlass.Array(cutlass.Float32, NT * K, space=SMEM, alignment=16) # exp(gate) * k + s_bk = cutlass.Array(cutlass.Float32, NT * K, space=SMEM, alignment=16) # beta * k + s_beta = cutlass.Array(cutlass.Float32, NT, space=SMEM, alignment=16) + s_v = cutlass.Array(cutlass.Float32, NT * V_CTA, space=SMEM, alignment=16) + s_o = cutlass.Array( + cutlass.Float32, NT * V_CTA, space=SMEM, alignment=16 + ) # bf16-rounded raw outputs + s_ss = cutlass.Array( + cutlass.Float32, V_SPLIT * NT, space=SMEM, alignment=16 + ) # [source CTA][token] + s_rs = cutlass.Array(cutlass.Float32, NT, space=SMEM, alignment=16) + r_st = cutlass.Array(cutlass.Float32, ROWS * VEC, space=RMEM) + + # ---- Prologue: nothing here depends on the predecessor grid. + if cutlass.const_expr(USE_PDL): + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + tma_ptr_w = tma_wfb.get_ptr() + tma_ptr_x = tma_fa.get_ptr() + if cutlass.const_expr(FOLD_FB): + if warp == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_x) + if prims.elect_sync(): + prims.mbarrier_init(w_full, 1) + prims.mbarrier_init(x_full, 1) + prims.mbarrier_init(acc_done, 1) + if warp == 2: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + if warp == 3: + if prims.elect_sync(): + prims.mbarrier_init(ss_ready, 1) + prims.mbarrier_arrive_expect_tx(ss_ready, V_SPLIT * NT * 4) + prims.fence_mbarrier_init() + # Cluster formation: the peers' shared memory is addressable from here on. + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + if cutlass.const_expr(FOLD_FB): + if warp == 0: + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(w_full, CTA_M * K * BF16_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a, + tma_ptr_w, + (cutlass.Int32(0), ch0, cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0)), + w_full, + ) + + slot = slots.load(idx=n) + pend = pending.load(idx=slot) + pend = cutlass.Int32( + cutlass.select_(pend > cutlass.Int32(NUM_SPEC), cutlass.Int32(NUM_SPEC), pend) + ) + exp_a = cute.math.exp(a_log.load(idx=h), fastmath=True) + row_a = v0 + warp * ROWS + # The warp owns both rows and lane l keys 32 i + l (the arithmetic of kda_mtp_decode). + # The drafts' records, in each slot's per-token state region in place of their full states: the row innovations + # vn [NUM_SPEC][H][V], then beta * k and the decay [NUM_SPEC][H][K]. The state after the accepted drafts is the + # pool's (the golden token's) with each accepted draft's update replayed, S = fma(decay, S, vn * (beta * k)): the + # recurrence's own arithmetic, so it is bit-identical to the one the drafts reached. k3_kda_attn uses the same + # records. + CT_VN = 0 + CT_WB = NUM_SPEC * H * V + CT_WD = CT_WB + NUM_SPEC * H * K + # The slot's pool state and per-token region from their first element: the slot offset in 64 bits, once. + pool = ssm.subview(cutlass.Int64(slot) * ssm_stride) + tok = state_tok.subview(cutlass.Int64(slot) * (NUM_SPEC * H * V * K)) + st_base = (h * V + row_a) * K + for r in cutlass.range_constexpr(ROWS): + for i in cutlass.range_constexpr(VEC): + r_st.store(pool.load(idx=st_base + r * K + i * 32 + lane), idx=r * VEC + i) + for t_acc in cutlass.range(pend, unroll=1): + rec_acc = (t_acc * H + h) * V + vns_acc = [tok.load(idx=rec_acc + CT_VN + row_a + r) for r in range(ROWS)] + wbs_acc = [tok.load(idx=CT_WB + (t_acc * H + h) * K + i * 32 + lane) for i in range(VEC)] + wds_acc = [tok.load(idx=CT_WD + (t_acc * H + h) * K + i * 32 + lane) for i in range(VEC)] + sts_acc = [r_st.load(idx=j) for j in range(ROWS * VEC)] + for r in cutlass.range_constexpr(ROWS): + for _pi in cutlass.range_constexpr(VEC // 2): + _p = _pi * 2 + vb0_acc, vb1_acc = cute.arch.mul_packed_f32x2( + (vns_acc[r], vns_acc[r]), (wbs_acc[_p], wbs_acc[_p + 1]) + ) + sts_acc[r * VEC + _p], sts_acc[r * VEC + _p + 1] = cute.arch.fma_packed_f32x2( + src_a=(wds_acc[_p], wds_acc[_p + 1]), src_b=(sts_acc[r * VEC + _p], sts_acc[r * VEC + _p + 1]), + src_c=(vb0_acc, vb1_acc), + ) # fmt: skip + for j in cutlass.range_constexpr(ROWS * VEC): + r_st.store(sts_acc[j], idx=j) + + # Conv weights, transposed to [tap][channel]: q by threads 0..127, k by 128..255. + if tidx < K: + wq4 = w_q.load(idx=(ch0 + tidx) * CONV_W, vector_size=CONV_W, alignment=16) + for w in cutlass.range_constexpr(CONV_W): + s_wq.store(cutlass.Float32(wq4[w]), idx=w * K + tidx) + else: + ck_w = tidx - K + wk4 = w_k.load(idx=(ch0 + ck_w) * CONV_W, vector_size=CONV_W, alignment=16) + for w in cutlass.range_constexpr(CONV_W): + s_wk.store(cutlass.Float32(wk4[w]), idx=w * K + ck_w) + # Conv history (positions -3..-1 = columns P..P+2), the v weights, the output-norm weight, dt_bias. + QH = 3 * K // 4 # float4 loads per q (or k) history + VH = 3 * V_CTA // 4 + if tidx < QH: + m_q = tidx // (K // 4) + c4_q = (tidx % (K // 4)) * 4 + s_uq.store(cs_q.load(idx=(slot * S + pend + m_q) * HK + ch0 + c4_q, vector_size=4, alignment=16), + idx=m_q * K + c4_q, alignment=16) # fmt: skip + elif tidx < 2 * QH: + m_k = (tidx - QH) // (K // 4) + c4_k = ((tidx - QH) % (K // 4)) * 4 + s_uk.store(cs_k.load(idx=(slot * S + pend + m_k) * HK + ch0 + c4_k, vector_size=4, alignment=16), + idx=m_k * K + c4_k, alignment=16) # fmt: skip + elif tidx < 2 * QH + VH: + m_v = (tidx - 2 * QH) // (V_CTA // 4) + c4_v = ((tidx - 2 * QH) % (V_CTA // 4)) * 4 + s_uv.store(cs_v.load(idx=(slot * S + pend + m_v) * HK + ch0 + v0 + c4_v, vector_size=4, alignment=16), + idx=m_v * V_CTA + c4_v, alignment=16) # fmt: skip + elif tidx < 2 * QH + VH + V_CTA: + cv_w = tidx - (2 * QH + VH) + wv4 = w_v.load(idx=(ch0 + v0 + cv_w) * CONV_W, vector_size=CONV_W, alignment=16) + for w in cutlass.range_constexpr(CONV_W): + s_wv.store(cutlass.Float32(wv4[w]), idx=w * V_CTA + cv_w) + elif tidx < 2 * QH + VH + 2 * V_CTA: + cv_n = tidx - (2 * QH + VH + V_CTA) + s_onw.store(onorm_w.load(idx=v0 + cv_n), idx=cv_n) + if tidx < K // 4: + s_dtb.store( + dt_bias.load(idx=ch0 + tidx * 4, vector_size=4, alignment=16), + idx=tidx * 4, + alignment=16, + ) + + # ---- This step's projection rows. + if cutlass.const_expr(USE_PDL): + prims.griddepcontrol(prims.GridDepAction.WAIT) + if cutlass.const_expr(FOLD_FB): + if warp == 1: + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(x_full, MMA_N * K * BF16_BYTES) + for half in cutlass.range_constexpr(TMA_COPY_ITERS): + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(half * B_HALF_ELEMS), tma_ptr_x, + (cutlass.Int32(FA_W * 2 + half * TMA_K_BOX), row0), x_full, + ) # fmt: skip + CH = K // CHUNK # 16-byte chunks per 128-channel row + for it in cutlass.range_constexpr((2 * NT * CH + THREADS - 1) // THREADS): + item = tidx + it * THREADS + if item < 2 * NT * CH: + stream_qk = item // (NT * CH) + t_qk = (item // CH) % NT + j_qk = item % CH + words_qk = proj.load( + idx=(row0 + t_qk) * proj_words + Q_W + ch0 // 2 + stream_qk * (K_W - Q_W) + j_qk * (CHUNK // 2), + vector_size=4, alignment=16, + ) # fmt: skip + if stream_qk == 0: + _store8(s_uq, words_qk, (CONV_W - 1 + t_qk) * K + j_qk * CHUNK) + else: + _store8(s_uk, words_qk, (CONV_W - 1 + t_qk) * K + j_qk * CHUNK) + VCH = V_CTA // CHUNK # 16-byte chunks per 16-channel v (or gate) slice + n_v = NT * VCH + n_g = 0 if FOLD_FB else NT * CH + for it in cutlass.range_constexpr((2 * n_v + NT + n_g + THREADS - 1) // THREADS): + item = tidx + it * THREADS + if item < n_v: + t_v = item // VCH + j_v = item % VCH + words_v = proj.load(idx=(row0 + t_v) * proj_words + V_W + (ch0 + v0) // 2 + j_v * (CHUNK // 2), + vector_size=4, alignment=16) # fmt: skip + _store8(s_uv, words_v, (CONV_W - 1 + t_v) * V_CTA + j_v * CHUNK) + elif item < 2 * n_v: + t_og = (item - n_v) // VCH + j_og = (item - n_v) % VCH + words_og = proj.load(idx=(row0 + t_og) * proj_words + OG_W + (ch0 + v0) // 2 + j_og * (CHUNK // 2), + vector_size=4, alignment=16) # fmt: skip + _store8(s_og, words_og, t_og * V_CTA + j_og * CHUNK) + elif item < 2 * n_v + NT: + t_b = item - 2 * n_v + word_b = proj.load(idx=(row0 + t_b) * proj_words + (B_COL + h) // 2) + b_in = _lo(word_b) + if (B_COL + h) % 2 == 1: + b_in = _hi(word_b) + s_braw.store(b_in, idx=t_b) + elif item < 2 * n_v + NT + n_g: + gi = item - (2 * n_v + NT) + t_g = gi // CH + j_g = gi % CH + words_g = g_ext.load(idx=(row0 + t_g) * (HK // 2) + ch0 // 2 + j_g * (CHUNK // 2), vector_size=4, + alignment=16) # fmt: skip + _store8(s_gr, words_g, t_g * K + j_g * CHUNK) + prims.barrier_cta_sync(0) + + # ---- f_b on the tensor cores (warp 2 issues, warps 4-7 read TMEM) beside the q/k pre-compute (warps 0-3). + if cutlass.const_expr(FOLD_FB): + if warp == 2: + idesc = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cutlass.BFloat16, + b_dtype=cutlass.BFloat16, + n_dim=MMA_N, + m_dim=CTA_M, + ) + desc_a_base = prims.Tcgen05SmemDesc.build( + start_address=smem_a, leading_byte_offset=LEADING, stride_byte_offset=STRIDE, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + desc_b_base = prims.Tcgen05SmemDesc.build( + start_address=smem_b, leading_byte_offset=LEADING, stride_byte_offset=STRIDE, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + tmem_acc = cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32) + while not cute.arch.mbarrier_try_wait(w_full.data_ptr(), 0): + pass + while not cute.arch.mbarrier_try_wait(x_full.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for kb in cutlass.range_constexpr(TMA_COPY_ITERS * K_BLOCKS_PER_HALF): + box = kb // K_BLOCKS_PER_HALF + within = kb % K_BLOCKS_PER_HALF + if prims.elect_sync(): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, prims.CTAGroup.CTA_1, tmem_acc, + desc_a_base + (box * A_BOX + within * STEP), desc_b_base + (box * B_BOX + within * STEP), + idesc, kb != 0, + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + if warp < 4: + # q and k of tokens 2w and 2w + 1 (the arithmetic of kda_mtp_decode's pre-compute warps); NT not a multiple of 4 + # rounds the tokens per warp up and the warps past NT idle. + qk_per_warp = (NT + 3) // 4 + for sub in cutlass.range_constexpr(qk_per_warp): + tk = warp * qk_per_warp + sub + if cutlass.const_expr(NT % 4 == 0): + _qk_pre(tk, lane, s_uq, s_wq, s_uk, s_wk, s_q, s_k, scale) + else: + if tk < NT: + _qk_pre(tk, lane, s_uq, s_wq, s_uk, s_wk, s_q, s_k, scale) + else: + if cutlass.const_expr(FOLD_FB): + while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), num=MMA_N + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + c_fb = (warp - 4) * 32 + lane + for t_fb in cutlass.range_constexpr(NT): + s_gr.store(_bf16(cutlass.Float32(acc[t_fb])), idx=t_fb * K + c_fb) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.barrier_cta_sync(1, thread_count=128) + if warp == 4: + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32), TMEM_COLS + ) + # v (this CTA's 16 channels) and beta, one (token, channel) per thread of warps 4-7. + item_v = tidx - 128 + if item_v < NT * V_CTA: + t_cv = item_v // V_CTA + c_cv = item_v % V_CTA + conv_v = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(CONV_W): + conv_v += s_uv.load(idx=(t_cv + w) * V_CTA + c_cv) * s_wv.load(idx=w * V_CTA + c_cv) + conv_v = conv_v * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-conv_v, fastmath=True) + ) + s_v.store(conv_v, idx=t_cv * V_CTA + c_cv) + if item_v < NT: + b_pre = s_braw.load(idx=item_v) + s_beta.store( + cute.arch.rcp_approx(cutlass.Float32(1.0) + cute.math.exp(-b_pre, fastmath=True)), + idx=item_v, + ) + prims.barrier_cta_sync(0) + + # ---- The recurrence's per-token operands: warp t, token t: decay = exp(gate), decay * k, beta * k. + if warp < NT: + tg = warp + r_beta_g = s_beta.load(idx=tg) + for i_pair in cutlass.range_constexpr(VEC // 2): + dks = [] + for i in (i_pair * 2, i_pair * 2 + 1): + c = i * 32 + lane + g_raw = s_gr.load(idx=tg * K + c) + s_dtb.load(idx=c) + x = exp_a * g_raw + sig = cute.arch.rcp_approx(cutlass.Float32(1.0) + cute.math.exp(-x, fastmath=True)) + dks.append( + (cute.math.exp(lower_bound * sig, fastmath=True), s_k.load(idx=tg * K + c), c) + ) + (d0, k0v, c0), (d1, k1v, c1) = dks + bk0, bk1 = cute.arch.mul_packed_f32x2((r_beta_g, r_beta_g), (k0v, k1v)) + kd0, kd1 = cute.arch.mul_packed_f32x2((d0, d1), (k0v, k1v)) + s_dec.store(d0, idx=tg * K + c0) + s_dec.store(d1, idx=tg * K + c1) + s_kd.store(kd0, idx=tg * K + c0) + s_kd.store(kd1, idx=tg * K + c1) + s_bk.store(bk0, idx=tg * K + c0) + s_bk.store(bk1, idx=tg * K + c1) + prims.barrier_cta_sync(0) + + # ---- Recurrence over the verify tokens (register-resident state). + st = [r_st.load(idx=j) for j in range(ROWS * VEC)] + outs = [] + for tr in cutlass.range_constexpr(NT): + r_v_val = s_v.load(idx=tr * V_CTA + warp * ROWS + lane % ROWS) + r_q = [cutlass.Float32(0.0)] * VEC + r_k = [cutlass.Float32(0.0)] * VEC + r_decay = [cutlass.Float32(0.0)] * VEC + r_bk = [cutlass.Float32(0.0)] * VEC + for i in cutlass.range_constexpr(VEC): + ki = i * 32 + lane + r_q[i] = s_q.load(idx=tr * K + ki) + r_k[i] = s_kd.load(idx=tr * K + ki) + r_bk[i] = s_bk.load(idx=tr * K + ki) + r_decay[i] = s_dec.load(idx=tr * K + ki) + ra = 0 + rb = 1 + r_va = cute.arch.shuffle_sync(r_v_val, ra) + r_vb = cute.arch.shuffle_sync(r_v_val, rb) + shk_a1 = cutlass.Float32(0.0) + shk_a2 = cutlass.Float32(0.0) + shk_b1 = cutlass.Float32(0.0) + shk_b2 = cutlass.Float32(0.0) + for _pi in cutlass.range_constexpr(VEC // 2): + _p = _pi * 2 + shk_a1, shk_a2 = cute.arch.fma_packed_f32x2( + src_a=(st[ra * VEC + _p], st[ra * VEC + _p + 1]), + src_b=(r_k[_p], r_k[_p + 1]), + src_c=(shk_a1, shk_a2), + ) + shk_b1, shk_b2 = cute.arch.fma_packed_f32x2( + src_a=(st[rb * VEC + _p], st[rb * VEC + _p + 1]), + src_b=(r_k[_p], r_k[_p + 1]), + src_c=(shk_b1, shk_b2), + ) + shk_a = shk_a1 + shk_a2 + shk_b = shk_b1 + shk_b2 + for offset in [16, 8, 4, 2, 1]: + shk_a += cute.arch.shuffle_sync_bfly(shk_a, offset=offset, mask=-1, mask_and_clamp=31) + shk_b += cute.arch.shuffle_sync_bfly(shk_b, offset=offset, mask=-1, mask_and_clamp=31) + vn_a = r_va - shk_a + vn_b = r_vb - shk_b + shq_a1 = cutlass.Float32(0.0) + shq_a2 = cutlass.Float32(0.0) + shq_b1 = cutlass.Float32(0.0) + shq_b2 = cutlass.Float32(0.0) + for _pi in cutlass.range_constexpr(VEC // 2): + _p = _pi * 2 + vnbk_a0, vnbk_a1 = cute.arch.mul_packed_f32x2((vn_a, vn_a), (r_bk[_p], r_bk[_p + 1])) + vnbk_b0, vnbk_b1 = cute.arch.mul_packed_f32x2((vn_b, vn_b), (r_bk[_p], r_bk[_p + 1])) + st[ra * VEC + _p], st[ra * VEC + _p + 1] = cute.arch.fma_packed_f32x2( + src_a=(r_decay[_p], r_decay[_p + 1]), + src_b=(st[ra * VEC + _p], st[ra * VEC + _p + 1]), + src_c=(vnbk_a0, vnbk_a1), + ) + st[rb * VEC + _p], st[rb * VEC + _p + 1] = cute.arch.fma_packed_f32x2( + src_a=(r_decay[_p], r_decay[_p + 1]), + src_b=(st[rb * VEC + _p], st[rb * VEC + _p + 1]), + src_c=(vnbk_b0, vnbk_b1), + ) + shq_a1, shq_a2 = cute.arch.fma_packed_f32x2( + src_a=(st[ra * VEC + _p], st[ra * VEC + _p + 1]), + src_b=(r_q[_p], r_q[_p + 1]), + src_c=(shq_a1, shq_a2), + ) + shq_b1, shq_b2 = cute.arch.fma_packed_f32x2( + src_a=(st[rb * VEC + _p], st[rb * VEC + _p + 1]), + src_b=(r_q[_p], r_q[_p + 1]), + src_c=(shq_b1, shq_b2), + ) + # The q-dot feeds only the output, not the next token: its reduction waits until after the loop. + outs.append([shq_a1 + shq_a2, shq_b1 + shq_b2]) + # The golden token's state to the pool; a draft's record: the warp's row innovations (lane r, row r). + if cutlass.const_expr(tr == 0): + for r in cutlass.range_constexpr(ROWS): + for i in cutlass.range_constexpr(VEC): + pool.store(st[r * VEC + i], idx=st_base + r * K + i * 32 + lane) + else: + if lane < ROWS: + tok.store( + cutlass.Float32(cutlass.select_(lane == 0, vn_a, vn_b)), + idx=CT_VN + ((tr - 1) * H + h) * V + row_a + lane, + ) + + for offset in [16, 8, 4, 2, 1]: + for tr in cutlass.range_constexpr(NT): + for r in cutlass.range_constexpr(ROWS): + outs[tr][r] += cute.arch.shuffle_sync_bfly( + outs[tr][r], offset=offset, mask=-1, mask_and_clamp=31 + ) + if lane == 0: + for tr in cutlass.range_constexpr(NT): + s_o.store(_bf16(outs[tr][0]), idx=tr * V_CTA + warp * ROWS + 0) + s_o.store(_bf16(outs[tr][1]), idx=tr * V_CTA + warp * ROWS + 1) + + # ---- Output gated RMSNorm over V: per-token partial sums of squares to every CTA of the cluster. + prims.barrier_cta_sync(0) + if tidx < NT * V_CTA: + t_out = tidx // V_CTA + x_o = s_o.load(idx=tidx) + ss = x_o * x_o + for off_ss in [V_CTA >> d for d in range(1, V_CTA.bit_length()) if V_CTA >> d > 0]: + ss = ss + cute.arch.shuffle_sync_bfly(ss, offset=off_ss, mask=-1, mask_and_clamp=31) + if tidx % V_CTA == 0: + # st.async into every CTA of the cluster (this one too), completing on its mailbox barrier: no cluster + # barrier, so nothing waits for this CTA's state stores to drain. + for r in cutlass.range_constexpr(V_SPLIT): + _st_async_f32( + _mapa_u32(s_ss.subview(vs * NT + t_out).data_ptr(), r), + ss, + _mapa_u32(ss_ready.data_ptr(), r), + ) + # Every peer's partials are in this CTA's shared memory once its mailbox completes. Each peer sends only after + # its recurrence, i.e. after it consumed its reads of the q/k conv history, so CTA 0's rewrite below cannot race + # them; and no CTA exits before its own mailbox has received every peer's partials. + while not _try_wait_cluster(ss_ready.data_ptr(), 0): + pass + if tidx < NT: + total = cutlass.Float32(0.0) + for r in cutlass.range_constexpr(V_SPLIT): + total = total + s_ss.load(idx=r * NT + tidx) + s_rs.store(cute.math.rsqrt(total / V + eps), idx=tidx) + prims.barrier_cta_sync(0) + if tidx < NT * V_CTA: + t_y = tidx // V_CTA + v_y = tidx % V_CTA + z = s_og.load(idx=tidx) + gate = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.math.exp(-z, fastmath=True)) + y = s_o.load(idx=tidx) * s_rs.load(idx=t_y) * s_onw.load(idx=v_y) * gate + out.store(cutlass.BFloat16(y), idx=((row0 + t_y) * H + h) * V + v0 + v_y) + + # ---- The drafts' per-key records (beta * k and the decay), one CTA per head, after every peer's norm partials + # have arrived: each peer sends them after its recurrence, i.e. after its prologue read the last round's records. + if vs == 0: + for it_r in cutlass.range_constexpr((NUM_SPEC * K + THREADS - 1) // THREADS): + item_r = tidx + it_r * THREADS + if item_r < NUM_SPEC * K: + t_r = item_r // K + key_r = item_r % K + tok.store( + s_bk.load(idx=(t_r + 1) * K + key_r), + idx=CT_WB + (t_r * H + h) * K + key_r, + ) + tok.store( + s_dec.load(idx=(t_r + 1) * K + key_r), + idx=CT_WD + (t_r * H + h) * K + key_r, + ) + + # ---- Conv caches for the next round: column s = position s - 2 = row s + 1 of the position-indexed inputs. + if vs == 0: + if tidx < K: + for s in cutlass.range_constexpr(S): + cs_q.store(s_uq.load(idx=(s + 1) * K + tidx), idx=(slot * S + s) * HK + ch0 + tidx) + else: + ck_c = tidx - K + for s in cutlass.range_constexpr(S): + cs_k.store(s_uk.load(idx=(s + 1) * K + ck_c), idx=(slot * S + s) * HK + ch0 + ck_c) + for it in cutlass.range_constexpr((S * V_CTA + THREADS - 1) // THREADS): + item = tidx + it * THREADS + if item < S * V_CTA: + s_col = item // V_CTA + v_col = item % V_CTA + cs_v.store( + s_uv.load(idx=(s_col + 1) * V_CTA + v_col), + idx=(slot * S + s_col) * HK + ch0 + v0 + v_col, + ) + + +@cute.jit +def k3_kda_verify( + w_fb: cute.Tensor, # bf16 [H * K, K]: the f_b weight (out, in) + proj: cute.Tensor, # int32 words of the fused projection [T, proj_words] + g_ext: cute.Tensor, # int32 words of the unfused f_b output [T, H * K / 2] (FOLD_FB False) + w_q: cute.Tensor, + w_k: cute.Tensor, + w_v: cute.Tensor, + a_log: cute.Tensor, + dt_bias: cute.Tensor, + onorm_w: cute.Tensor, + cs_q: cute.Tensor, + cs_k: cute.Tensor, + cs_v: cute.Tensor, + ssm: cute.Tensor, + state_tok: cute.Tensor, + slots: cute.Tensor, + pending: cute.Tensor, + out: cute.Tensor, + proj_words: cutlass.Int32, + n_req: cutlass.Int32, + ssm_stride: cutlass.Int64, + H: cutlass.Constexpr[int], + NUM_SPEC: cutlass.Constexpr[int], + lower_bound: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + eps: cutlass.Constexpr[float], + FOLD_FB: cutlass.Constexpr[bool], + USE_PDL: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + """One launch of the verify kernel for ``n_req`` requests of 1 + NUM_SPEC tokens each.""" + # W_fb [H K, K] as five TMA dimensions (64-element column chunk, row, chunk index, 1, 1): one call lands the + # head's 128 rows, both 128-byte-swizzled halves; strides in 16-byte units. + tma_wfb = cuda.create_tensor_map_tiled( + global_address=w_fb.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[TMA_K_BOX, H * K, K // TMA_K_BOX, 1, 1], + global_strides=[(K * BF16_BYTES) // 16, (TMA_K_BOX * BF16_BYTES) // 16, (H * K * K * BF16_BYTES) // 16, + (H * K * K * BF16_BYTES) // 16], + box_dims=[TMA_K_BOX, CTA_M, TMA_COPY_ITERS, 1, 1], + swizzle=cuda.TensorMapSwizzle.s128b, + ) # fmt: skip + # The projection rows as bf16 [T, 2 proj_words]: an 8-token x 64-column box of f_a per call. + tma_fa = cuda.create_tensor_map_tiled( + global_address=proj.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[proj_words * 2, n_req * (NUM_SPEC + 1)], + global_strides=[(proj_words * 4) // 16], + box_dims=[TMA_K_BOX, MMA_N], + swizzle=cuda.TensorMapSwizzle.s128b, + ) + k3_kda_verify_kernel( + tma_wfb, tma_fa, proj, g_ext, w_q, w_k, w_v, a_log, dt_bias, onorm_w, cs_q, cs_k, cs_v, ssm, state_tok, slots, + pending, out, proj_words, ssm_stride, H, NUM_SPEC, lower_bound, scale, eps, FOLD_FB, USE_PDL, + ).launch( + grid=(H, n_req, V_SPLIT), block=(THREADS, 1, 1), cluster=(1, 1, V_SPLIT), stream=stream, use_pdl=USE_PDL, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/op.py new file mode 100644 index 000000000000..7e181d958fec --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/op.py @@ -0,0 +1,205 @@ +# 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::k3_kda_verify``: Kimi K3's KDA speculative verify for 1 + num_spec tokens per request, from the fused +projection rows to the gated-norm core output, with a committed state per verify token. + +State contract (differs from ``trtllm::kda_mtp_decode``'s replay caches): after a call the pool ``ssm`` holds the +state after each request's golden token and ``state_tok[slot, t - 1]`` the state after its draft t; the conv caches +hold the raw inputs at positions -2..num_spec around the golden token. The next call reads the state after the +drafts the sampler accepted (``pending[slot]``) directly, so nothing is replayed. The kernel compiles on the first +call for each configuration, which must happen outside CUDA-graph capture. +""" + +from __future__ import annotations + +import os +import threading +from typing import Dict, Optional + +import torch + +K = 128 +CONV_W = 4 + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + + +def _arg(t: torch.Tensor, align: int = 16): + from cutlass.cute.runtime import from_dlpack + + # detach(): DLPack refuses tensors that require grad (weights are parameters). + return from_dlpack(t.detach(), assumed_align=align).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +def _index_arg(t: torch.Tensor): + """An int32 index tensor the kernel reads element by element (``slots``, ``pending``), declared at its element's + alignment: the mixer passes slices such as ``state_indices[num_prefills:]``, which start on any 4-byte boundary.""" + return _arg(t, t.element_size()) + + +def _flat(t: torch.Tensor) -> torch.Tensor: + """A 1-D view of a dense tensor's storage (a view with its last two dims transposed is fine), never a copy.""" + if t.is_contiguous(): + return t.view(-1) + swapped = t.transpose(-1, -2) + if swapped.is_contiguous(): + return swapped.reshape(-1) + raise ValueError(f"expected a dense tensor, got shape {tuple(t.shape)} strides {t.stride()}") + + +def _slots_view(t: torch.Tensor) -> torch.Tensor: + """A [slots, slot elements] view of a pool whose slots are dense but may be strided (the Mamba cache manager + coalesces each slot's per-layer states), never a copy. The kernels address it flat from the first slot at 64-bit + offsets; the view keeps every extent within the DSL's 32-bit sizes however far the last slot lies.""" + return t.view(t.shape[0], -1) + + +def _words(t: torch.Tensor) -> torch.Tensor: + return _flat(t).view(torch.int32) + + +def use_pdl() -> bool: + return os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +@torch.library.custom_op( + "trtllm::k3_kda_verify", + mutates_args=("cs_q", "cs_k", "cs_v", "ssm", "state_tok"), + device_types="cuda", +) +def k3_kda_verify( + proj: torch.Tensor, + w_fb: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + a_log: torch.Tensor, + dt_bias: torch.Tensor, + onorm_w: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + ssm: torch.Tensor, + state_tok: torch.Tensor, + slots: torch.Tensor, + pending: torch.Tensor, + num_spec: int, + lower_bound: float, + scale: float, + eps: float, + g_ext: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Gated-norm KDA core output [T, H, V] bf16 for T = N * (1 + num_spec) verify tokens. + + ``proj`` bf16 [T, 4 H K + K + H + pad]: the fused projection rows [q | k | v | onorm gate | f_a | b | pad]; + ``w_fb`` bf16 [H K, K], the f_b weight (out, in); ``w_q/w_k/w_v`` fp32 [H K, 4]; ``a_log`` fp32 [H]; + ``dt_bias`` fp32 [H K]; ``onorm_w`` fp32 [V]; ``cs_*`` fp32 [pool, H K, 3 + num_spec] (dim-contiguous); + ``ssm`` fp32 [pool, H, V, K] (each slot dense, slots at any stride); ``state_tok`` fp32 + [pool, num_spec, H, V, K]; ``slots`` int32 [N]; ``pending`` int32 [pool]. With ``g_ext`` (bf16 [T, H K], the + unfused f_b output) the gate is read instead of computed from f_a.""" + import cuda.bindings.driver as cuda_driver + + from . import k3_kda_verify_kernel as kernel + + n_req = slots.shape[0] + num_heads = w_fb.shape[0] // K + tokens = n_req * (1 + num_spec) + v_dim = onorm_w.shape[0] + if ( + proj.dtype != torch.bfloat16 + or proj.dim() != 2 + or proj.shape[0] != tokens + or proj.shape[1] % 8 != 0 + or proj.shape[1] < 4 * num_heads * K + K + num_heads + or not proj.is_contiguous() + or tuple(w_fb.shape) != (num_heads * K, K) + or w_fb.dtype != torch.bfloat16 + or not w_fb.is_contiguous() + or v_dim != K + or tuple(ssm.shape[1:]) != (num_heads, K, K) + or tuple(state_tok.shape[1:]) != (num_spec, num_heads, K, K) + or ssm.stride()[1:] != (K * K, K, 1) + or ssm.dtype != torch.float32 + or not state_tok.is_contiguous() + or state_tok.dtype != torch.float32 + or cs_q.shape[-1] != CONV_W - 1 + num_spec + or slots.dtype != torch.int32 + or pending.dtype != torch.int32 + or any( + t.dtype != torch.float32 + for t in (w_q, w_k, w_v, a_log, dt_bias, onorm_w, cs_q, cs_k, cs_v) + ) + ): + raise ValueError( + f"k3_kda_verify: unsupported call (w_q/w_k/w_v/a_log/dt_bias/onorm_w/cs_* must be fp32: " + f"{[str(t.dtype) for t in (w_q, w_k, w_v, a_log, dt_bias, onorm_w, cs_q, cs_k, cs_v)]}) " + f"proj {tuple(proj.shape)} {proj.dtype} contiguous {proj.is_contiguous()}, " + f"w_fb {tuple(w_fb.shape)} {w_fb.dtype} contiguous {w_fb.is_contiguous()}, onorm_w {tuple(onorm_w.shape)}, " + f"ssm {tuple(ssm.shape)} {ssm.dtype} strides {ssm.stride()}, state_tok {tuple(state_tok.shape)} " + f"{state_tok.dtype} contiguous {state_tok.is_contiguous()}, cs_q {tuple(cs_q.shape)}, num_spec {num_spec}, " + f"slots {tuple(slots.shape)} {slots.dtype}, pending {pending.dtype} (tokens expected {tokens})" + ) + fold_fb = g_ext is None + out = torch.empty(tokens, num_heads, v_dim, dtype=torch.bfloat16, device=proj.device) + args = ( + _arg(w_fb), + _arg(_words(proj)), + _arg(_words(proj if fold_fb else g_ext)), + _arg(_flat(w_q)), + _arg(_flat(w_k)), + _arg(_flat(w_v)), + _arg(_flat(a_log)), + _arg(_flat(dt_bias)), + _arg(_flat(onorm_w)), + _arg(_flat(cs_q)), + _arg(_flat(cs_k)), + _arg(_flat(cs_v)), + _arg(_slots_view(ssm)), + _arg(_slots_view(state_tok)), + _index_arg(slots), + _index_arg(pending), + _arg(_flat(out)), + ) + runtime = (proj.shape[1] // 2, n_req, ssm.stride(0)) + stream = cuda_driver.CUstream(torch.cuda.current_stream(proj.device).cuda_stream) + pdl = use_pdl() + key = (num_heads, num_spec, float(lower_bound), float(scale), float(eps), fold_fb, pdl) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "trtllm::k3_kda_verify must run once per configuration outside CUDA-graph capture first " + "(it compiles its kernel on the first call)." + ) + import cutlass.cute as cute + + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kernel.k3_kda_verify, *args, *runtime, num_heads, num_spec, float(lower_bound), float(scale), + float(eps), fold_fb, pdl, stream, + ) # fmt: skip + fn(*args, *runtime, stream) + return out + + +@k3_kda_verify.register_fake +def _(proj, w_fb, w_q, w_k, w_v, a_log, dt_bias, onorm_w, cs_q, cs_k, cs_v, ssm, state_tok, slots, pending, + num_spec, lower_bound, scale, eps, g_ext=None): # fmt: skip + return proj.new_empty( + (proj.shape[0], w_fb.shape[0] // K, onorm_w.shape[0]), dtype=torch.bfloat16 + ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_mla/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_mla/__init__.py new file mode 100644 index 000000000000..76bebda7a90b --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_mla/__init__.py @@ -0,0 +1,19 @@ +# 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 MLA decode CTM kernels (``trtllm::k3_mla_q``). + +Importing :mod:`.op` registers the torch ops; nothing is imported eagerly here so that +the CuTe DSL dependency stays optional for every other model. +""" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_mla/k3_mla_attn_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_mla/k3_mla_attn_kernel.py new file mode 100644 index 000000000000..9d179d489649 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_mla/k3_mla_attn_kernel.py @@ -0,0 +1,1187 @@ +# 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 MLA decode attention -- CTM (prims/cute) kernel: R requests of T <= 8 query tokens, 6 heads per rank +# ============================================================================= +# +# o[i, t, h] = softmax_k( scale * q[i, t, h] . kv_i[k] ) @ kv_i[k, :512] k <= L_i - T + t (bottom-right causal) +# q = fused_q [M = R T, 6, 576], request-major (request i's tokens are rows i T .. i T + T - 1); per cluster 48 +# query rows r = 8 h + t, head-major (a head's tokens are adjacent rows) +# kv_i = request i's pages of the paged latent cache, 64 rows x 576 (512 latent | 64 rope), L_i rows +# o = [M, 6, 512] bf16 +# +# One cluster of 16 CTAs per (request, group of 6 heads): grid (16 x groups, R), one group at TP16 (a rank holding more +# heads gets one cluster per group, each reading the request's whole cache). A cluster's work does not depend on R; the +# request only selects its page-table row, length, query / output rows and workspace slots. CTA c takes the 128-row KV +# tiles c, c + 16, ... (two pages each). +# Per tile: +# S = Q K^T A = Q [64 rows (48 real), 576] from shared memory, B = the tile's K [128 rows, 576] (chunk-major: +# page p's rows of chunk j at [j][64 p..]), M 64, N 128 (Q read once per chunk), into TMEM (an M = 64 +# accumulator: row 16 w + l in lane 32 w + l) +# softmax each row's 128 values over the 4 lanes of a quad (tcgen05.ld 16x256b): masked max, exp2, sum, +# P = bf16(p) into a 128B-swizzled K-major tile, fenced to the async proxy +# O = P V A = P [64, 128], B = V = the tile's latent columns read MN-major (N 256 per MMA), M 64 +# partial (m, l) and O / l (a convex combination of V rows: fp16 keeps 2^-11) of the tile's 48 rows into the +# CTA's global workspace slot; a CTA's later tiles (L > 2048) fold into its slot: +# O / l' = (O / l) (l a / l') + O_k (b / l'), l' = l a + l_k b, a, b = exp2(m - m'), exp2(m_k - m') +# Right after its last softmax a CTA st.async's its (m, l) per row into every CTA's mailbox; warps 1 and 3 turn the +# mailbox into the merge weights exp2(m_s - M) / sum_s exp2(m_s - M) l_s while O is computed and drained. Then a +# cluster barrier (release / acquire), and CTA c merges latent columns [32 c, 32 c + 32) of all rows over the slots +# (flash-decoding combine in fp32, one round of 16-byte loads) and stores o in bf16. The KV pages below L - T are +# TMA'd before griddepcontrol.wait (earlier steps wrote them); Q (3 chunk groups, so S starts on the first) and the +# last T rows' pages after it. O = P V runs in two N = 256 halves, the first drained while the second computes. +# Q is TMA'd as 8 token rows from the request's first token: rows t >= T (the next request's tokens, or zeros past M) +# are computed and written to the workspace slot like the live rows (so the merge's 16-byte loads read 32-byte sectors +# written whole; it weights these rows 0), and never stored. +# +# fuse_vb: the merge is split by (head, 4 tokens) over CTAs 0-11 instead of by columns, and each of those CTAs applies +# v_b to its 4 merged rows: W_vb[h] (128 x 512) TMA'd into the dead KV buffer after the CTA's last P V, o rounded to +# bf16 (the unfused path's attention output) into the dead P buffer as the B operand, y^T = W_vb[h] o^T (M 128, N 8, +# K 512) in TMEM, y [M, heads * 128] bf16 stored (the o_proj input). apply_gate: y is multiplied by the output gate's +# sigmoid (bf16, from the fused projection's gate columns) with the unfused path's rounding, bf16(bf16(y) * s). +# +# no_cluster (R x groups > CLUSTER_WAVE: more 16-CTA clusters than co-reside, which would run a second wave): the grid +# launches without a cluster, CTA c = block x mod 16, and the cluster's mailbox and barriers go through the workspace's +# tail instead. At its last tile a CTA stores its live rows' (m, l) into its slot of an fp32 exchange and arrives on its +# (request, head group)'s first counter; warps 1 and 3 wait for the launch's 16 arrivals and read the slots' (m, l) +# from there. After the drain every CTA arrives on the second counter, which CTAs 0-11 wait for instead of the cluster +# barrier. Counters only grow: a CTA reads its counter after the grid wait and before it arrives, so the launch's 16 +# arrivals end at (count & ~15) + 16. The merge sums the same values in the same order, so the outputs do not depend on +# the mode. All 16 R groups CTAs must co-reside (1 per SM), as their waits are spins. +# +# Warps: 0 TMA, 2 TMEM allocation (512 columns) + MMA, 4-7 softmax / partial writer, 1 and 3 merge weights, all 8 in +# the merge. +# ============================================================================= +"""CTM kernel: Kimi K3 MLA decode attention over the paged latent cache, split-KV over a 16-CTA cluster.""" + +from __future__ import annotations + +import math + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +HEADS = 6 +MAX_TOKENS = 8 # query tokens per request (T) +MAX_REQUESTS = 8 # requests per launch (R): the workspace holds R x groups x 16 slots +ROWS = HEADS * MAX_TOKENS # 48 query rows +# A softmax thread's rows r and r + 8 lie in one 16-row block, so they are both real or both past ROWS. +assert ROWS % 16 == 0 +QK = 576 +LATENT = 512 +PAGE = 64 +TILE = 128 # KV rows per tile = 2 pages +PAGES_PER_TILE = TILE // PAGE +CHUNK = 64 # 128-byte swizzle width in bf16 +QK_CHUNKS = QK // CHUNK # 9 +LAT_CHUNKS = LATENT // CHUNK # 8 +CLUSTER = 16 +MMA_M = 64 +MMA_K = 16 +THREADS = 256 +EPI_THREADS = 128 +TMEM_COLS = 512 +ELEM_BYTES = 2 +LAT_PER_CTA = LATENT // CLUSTER # 32 merged columns per CTA +LOG2E = 1.4426950408889634 +NEG_INF = float("-inf") + +# Shared-memory tiles, bf16 elements. +Q_CHUNK_ELEMS = ROWS * CHUNK # 3072 (6 KB): Q chunk j = rows 0..47 of columns [64 j, 64 j + 64) +Q_ELEMS = QK_CHUNKS * Q_CHUNK_ELEMS +KV_CHUNK_ELEMS = ( + TILE * CHUNK +) # 8192 (16 KB): both pages' rows of one 64-column chunk (chunk-major tile) +KV_PAGE_OFF = PAGE * CHUNK # page p's 64 rows start 8 KB into each chunk block +KV_ELEMS = QK_CHUNKS * KV_CHUNK_ELEMS +P_HALF_ELEMS = MMA_M * CHUNK # 4096: 64 rows x 64 KV columns +P_ELEMS = PAGES_PER_TILE * P_HALF_ELEMS +Q_OVERREAD_ELEMS = (MMA_M - ROWS) * CHUNK # the M = 64 MMA reads 16 rows past the last Q chunk + +# Descriptor offsets in 16-byte units. +LEADING = 16 +SBO = 8 * CHUNK * ELEM_BYTES # 1024 B between 8-row swizzle atoms +STEP_K = (MMA_K * ELEM_BYTES) >> 4 # K-major: 32 B per K16 step +STEP_MN = (2 * SBO) >> 4 # MN-major: 2 x SBO per K16 step (16 KV rows) +Q_CHUNK_U = (Q_CHUNK_ELEMS * ELEM_BYTES) >> 4 +KV_CHUNK_U = (KV_CHUNK_ELEMS * ELEM_BYTES) >> 4 +P_HALF_U = (P_HALF_ELEMS * ELEM_BYTES) >> 4 + +Q_GROUPS = 3 # Q TMA'd as 3 groups of 3 chunks, each on its own barrier +Q_GROUP_CHUNKS = QK_CHUNKS // Q_GROUPS +O_HALVES = LATENT // 256 +MERGE_HALF = CLUSTER // 2 # slots per thread of a merge pair +V_DIM = 128 # v_head_dim +VB_CTAS = ( + HEADS * 2 +) # fuse_vb: CTA c < 12 merges head c // 2, tokens 4 (c % 2) .. + 4, and applies v_b to them +VB_TOKENS = MAX_TOKENS // 2 +VB_ELEMS = V_DIM * LATENT + +io_dtype = cutlass.BFloat16 +ws_dtype = cutlass.Float16 # the partials O / l (convex combinations of V rows, |x| <= max |V|) +# Workspace: slot (request x groups + head group) x 16 + CTA. Slot layout, fragment-major: [64 groups of 8 columns][64 +# rows (the M = 64 accumulator)][8], so one warp's 16x256b fragment of a group (8 rows x 4 lanes x 2 values) is 128 +# contiguous bytes, and a row's 8 columns are 16. +WS_GROUP_ELEMS = MMA_M * 8 +WS_SLOT_ELEMS = (LATENT // 8) * WS_GROUP_ELEMS +# 16-CTA clusters of this kernel that co-reside on GB200 (cuOccupancyMaxActiveClusters at its shared memory: the eighth +# GPC holds 12 free SMs); a launch of more clusters takes the no_cluster mode. +CLUSTER_WAVE = 7 +# no_cluster's tail of the workspace, after the R x groups x 16 partial slots: an fp32 [slot][row][m, l] exchange, then +# two int32 counters per (request, head group), CTR_STRIDE words apart (a 128-byte line each). +CTR_STRIDE = 32 + + +def ws_sync_elems(groups: int) -> int: + """fp16 elements of the workspace's no_cluster tail (zeroed once: the counters start at 0).""" + exchange = MAX_REQUESTS * groups * CLUSTER * ROWS * 2 * 4 + counters = MAX_REQUESTS * groups * 2 * CTR_STRIDE * 4 + return (exchange + counters) // 2 + + +@dsl_user_op +def _mapa_u32(smem_ptr, peer, *, loc=None, ip=None): + """The shared::cluster address of this CTA's shared-memory location in cluster CTA ``peer``.""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [smem_ptr.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(peer).ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_async_v2(dst, a, b, mbar, *, loc=None, ip=None): + """st.async of two fp32 to a shared::cluster address, completing ``mbar`` (shared::cluster) by 8 bytes.""" + _llvm.inline_asm( + None, + [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Float32(a).ir_value(loc=loc, ip=ip), + cutlass.Float32(b).ir_value(loc=loc, ip=ip), cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.f32 [$0], {$1, $2}, [$3];", "r,f,f,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _red_add_release(addr, value, *, loc=None, ip=None): + """red.release.gpu.global.add.s32 on a global address (an arrival: no value back).""" + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip)], + "red.release.gpu.global.add.s32 [$0], $1;", "l,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _ld_acquire(addr, *, loc=None, ip=None): + """ld.acquire.gpu.global.s32 of a global address.""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.acquire.gpu.global.s32 $0, [$1];", "=r,l", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _try_wait_cluster(mbar, parity, *, loc=None, ip=None): + """mbarrier.try_wait.parity.acquire.cluster on a barrier of this CTA (``mbar``, its shared-memory pointer): whether + phase ``parity`` has completed, acquiring at cluster scope. The barrier is completed by other CTAs' st.async, whose + complete_tx releases at cluster scope.""" + done = cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tmbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" + "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + return done != cutlass.Int32(0) + + +@cute.kernel +def k3_mla_attn_kernel( + tma_q: cutlass.GridConstant[ + cuda.TensorMap + ], # fused_q [M, heads, 576]: box 64 cols x 8 tokens x 6 heads x 3 chunks + tma_kv: cutlass.GridConstant[ + cuda.TensorMap + ], # pool as [pages * 64 rows, 576]: box 64 cols x 64 rows x 9 chunks + tma_vb: cutlass.GridConstant[ + cuda.TensorMap + ], # fuse_vb: v_b_proj as [heads * 128 rows, 512]: box 64 x 128 x 8 + page_table: cutlass.Array, # int32: request i's pages (>= ceil(L_i / 64)) at [i * pt_stride, ...) + seq_len: cutlass.Array, # int32 [R]: L_i, request i's KV rows including its T new ones + ws_o: cutlass.Array, # fp16 [R, head groups, 16 CTAs, 64 column groups, 64 rows, 8] per-CTA partial O / l + out: cutlass.Array, # bf16 [M * heads * 512] (fuse_vb: [M * heads * 128]) + gate: cutlass.Array, # apply_gate: bf16 [M, gate_ld], sigmoid(g) of head h at columns gate_col0 + 128 h + tokens: cutlass.Int32, # T, the query tokens of each request (M = R T) + pt_stride: cutlass.Int32, # page-table elements between consecutive requests' rows + scale_log2: cutlass.Float32, # softmax scale * log2(e) + page_offset: cutlass.Int32, # added to every page-table entry (the layer's slot in an interleaved pool) + gate_col0: cutlass.Int32, + gate_ld: cutlass.Int32, + total_heads: cutlass.Constexpr[int], # heads of the rank: one cluster per 6 + fuse_vb: cutlass.Constexpr[bool], + apply_gate: cutlass.Constexpr[bool], + no_cluster: cutlass.Constexpr[bool], +): + tx, _, _ = cute.arch.thread_idx() + warp_id = cute.arch.warp_idx() + if cutlass.const_expr(not no_cluster): + cta = cute.arch.block_idx_in_cluster() + bid, req, _ = cute.arch.block_idx() + if cutlass.const_expr(no_cluster): + cta = bid % cutlass.Int32(CLUSTER) + hg = bid // cutlass.Int32(CLUSTER) # head group: heads 6 hg .. 6 hg + 5 + ws_slot0 = (req * cutlass.Int32(total_heads // HEADS) + hg) * cutlass.Int32(CLUSTER) + my_slot = ws_slot0 + cta + tok0 = req * tokens # the request's first row of q, out and gate + pt0 = req * pt_stride # the request's page-table row + ptr_q = tma_q.get_ptr() + ptr_kv = tma_kv.get_ptr() + ptr_vb = tma_vb.get_ptr() + + # Read before the grid dependency: the host writes the lengths and page table before the step. + kv_len = seq_len.load(idx=req) + n_tiles = (kv_len + cutlass.Int32(TILE - 1)) // cutlass.Int32(TILE) + n_pages = (kv_len + cutlass.Int32(PAGE - 1)) // cutlass.Int32(PAGE) + my_count = cutlass.Int32(0) + if cta < n_tiles: + my_count = (n_tiles - cta + cutlass.Int32(CLUSTER - 1)) // cutlass.Int32(CLUSTER) + q_rows = tokens * cutlass.Int32(HEADS) # live rows: 8 h + t with t < T + # Pages holding a row >= L - T are written this step (the cache append): they load after the wait. + fresh_page = (kv_len - tokens) // cutlass.Int32(PAGE) + # Slots (CTAs) that had a tile: min(n_tiles, 16). + n_valid = cutlass.Int32( + cutlass.select_(n_tiles < cutlass.Int32(CLUSTER), n_tiles, cutlass.Int32(CLUSTER)) + ) + + smem_kv = cutlass.Array(io_dtype, KV_ELEMS, space=cutlass.AddressSpace.smem, alignment=1024) + smem_q = cutlass.Array(io_dtype, Q_ELEMS, space=cutlass.AddressSpace.smem, alignment=1024) + # Right after Q so the M = 64 MMA's read past the last Q chunk stays in shared memory. + smem_p = cutlass.Array(io_dtype, P_ELEMS, space=cutlass.AddressSpace.smem, alignment=1024) + kv_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + q_full = cutlass.Array(cutlass.Int64, Q_GROUPS, space=cutlass.AddressSpace.smem, alignment=8) + s_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + p_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + o_full = cutlass.Array(cutlass.Int64, O_HALVES, space=cutlass.AddressSpace.smem, alignment=8) + o_drained = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + ml_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + vb_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + vb_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + # Mailbox: [source CTA][row] (m, l) of every CTA with a tile (st.async, completed by bytes); warps 1 and 3 then + # overwrite m with the merge weight exp2(m_s - M) / L. + ml_mail = cutlass.Array( + cutlass.Float32, CLUSTER * ROWS * 2, space=cutlass.AddressSpace.smem, alignment=16 + ) + if cutlass.const_expr(no_cluster): + # The workspace's tail: the (m, l) exchange [slot][row][2] and the counters; this CTA's targets in sync_target. + groups = total_heads // HEADS + tail = MAX_REQUESTS * groups * CLUSTER * WS_SLOT_ELEMS + ml_glob = cutlass.Array( + ws_o.data_ptr(tail), shape=(MAX_REQUESTS * groups * CLUSTER * ROWS * 2,), dtype=cutlass.Float32, + bounds_check=False, addrspace=cutlass.AddressSpace.gmem.value, alignment=16, + ) # fmt: skip + ctrs = cutlass.Array( + ws_o.data_ptr(tail + MAX_REQUESTS * groups * CLUSTER * ROWS * 4), + shape=(MAX_REQUESTS * groups * 2 * CTR_STRIDE,), dtype=cutlass.Int32, bounds_check=False, + addrspace=cutlass.AddressSpace.gmem.value, alignment=16, + ) # fmt: skip + ctr_ml = (req * cutlass.Int32(groups) + hg) * cutlass.Int32(2 * CTR_STRIDE) + ctr_o = ctr_ml + cutlass.Int32(CTR_STRIDE) + sync_target = cutlass.Array(cutlass.Int32, 2, space=cutlass.AddressSpace.smem, alignment=8) + + if warp_id == 0: + prims.prefetch_tensormap(ptr_q) + prims.prefetch_tensormap(ptr_kv) + if cutlass.const_expr(fuse_vb): + prims.prefetch_tensormap(ptr_vb) + if prims.elect_sync(): + prims.mbarrier_init(kv_full, 1) + for i in cutlass.range_constexpr(Q_GROUPS): + prims.mbarrier_init(q_full.subview(i), 1) + prims.mbarrier_init(s_full, 1) + prims.mbarrier_init(p_full, EPI_THREADS) + for i in cutlass.range_constexpr(O_HALVES): + prims.mbarrier_init(o_full.subview(i), 1) + prims.mbarrier_init(o_drained, EPI_THREADS) + prims.mbarrier_init(ml_full, 1) + prims.mbarrier_init(vb_full, 1) + prims.mbarrier_init(vb_done, 1) + if cutlass.const_expr(not no_cluster): + # Every CTA with a tile sends (m, l) of the M * 6 live rows (8 bytes each). + prims.mbarrier_arrive_expect_tx(ml_full, n_valid * q_rows * cutlass.Int32(8)) + if warp_id == 2: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + if cutlass.const_expr(not no_cluster): + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + tmem_base = tmem_ptr_i32.load() + + if warp_id == 0: + # ===================================================================== + # TMA: each tile's two pages (the pages written earlier before the grid + # wait), Q after it; the next tile once the previous one's O is done. + # ===================================================================== + if prims.elect_sync(): + for k in range(my_count): + tile = cta + k * cutlass.Int32(CLUSTER) + if k > 0: + # The KV buffer is free once the previous tile's P V is done (its last half's commit). + while not cute.arch.mbarrier_try_wait( + o_full.subview(O_HALVES - 1).data_ptr(), + (k - cutlass.Int32(1)) & cutlass.Int32(1), + ): + pass + prims.mbarrier_arrive_expect_tx(kv_full, KV_ELEMS * ELEM_BYTES) + for p in cutlass.range_constexpr(PAGES_PER_TILE): + page_idx = tile * cutlass.Int32(PAGES_PER_TILE) + cutlass.Int32(p) + page_c = cutlass.Int32( + cutlass.select_(page_idx < n_pages, page_idx, n_pages - cutlass.Int32(1)) + ) + row0 = (page_table.load(idx=pt0 + page_c) + page_offset) * cutlass.Int32(PAGE) + if k == 0: + if page_c >= fresh_page: + prims.griddepcontrol(prims.GridDepAction.WAIT) + # Chunk-major tile: page p's rows of chunk j at [j][64 p .. 64 p + 64), so one MMA covers both + # pages. + for j in cutlass.range_constexpr(QK_CHUNKS): + prims.cp_async_bulk_tensor_shared_cta_global( + smem_kv.subview(j * KV_CHUNK_ELEMS + p * KV_PAGE_OFF), + ptr_kv, + (cutlass.Int32(0), row0, cutlass.Int32(j)), + kv_full, + ) + if k == 0: + prims.griddepcontrol(prims.GridDepAction.WAIT) + for g in cutlass.range_constexpr(Q_GROUPS): + prims.mbarrier_arrive_expect_tx( + q_full.subview(g), Q_GROUP_CHUNKS * Q_CHUNK_ELEMS * ELEM_BYTES + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_q.subview(g * Q_GROUP_CHUNKS * Q_CHUNK_ELEMS), + ptr_q, + ( + cutlass.Int32(0), + tok0, + hg * cutlass.Int32(HEADS), + cutlass.Int32(g * Q_GROUP_CHUNKS), + ), + q_full.subview(g), + ) + # Dependents (v_b / o_proj) wait for this whole grid before reading its output. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if cutlass.const_expr(no_cluster): + if tx == cutlass.Int32(0): + prims.griddepcontrol(prims.GridDepAction.WAIT) + sync_target.store(ctrs.load(idx=ctr_o, is_volatile=True), idx=1) + if cutlass.const_expr(fuse_vb): + if cta < cutlass.Int32(VB_CTAS): + # W_vb[h] into the KV buffer once the CTA's last P V has read it. + if prims.elect_sync(): + if my_count > cutlass.Int32(0): + while not cute.arch.mbarrier_try_wait( + o_full.subview(O_HALVES - 1).data_ptr(), + (my_count - cutlass.Int32(1)) & cutlass.Int32(1), + ): + pass + prims.mbarrier_arrive_expect_tx(vb_full, VB_ELEMS * ELEM_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_kv, + ptr_vb, + ( + cutlass.Int32(0), + (hg * cutlass.Int32(HEADS) + cta // cutlass.Int32(2)) + * cutlass.Int32(V_DIM), + cutlass.Int32(0), + ), + vb_full, + ) + elif warp_id == 2: + # ===================================================================== + # MMA: S = Q K^T (per page N 64), then O = P V (N 256 x 2), per tile. + # ===================================================================== + idesc_s = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, a_dtype=io_dtype, b_dtype=io_dtype, n_dim=TILE, m_dim=MMA_M + ) + idesc_o = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=io_dtype, + b_dtype=io_dtype, + n_dim=256, + m_dim=MMA_M, + b_major=1, + ) + swz = prims.Tcgen05SmemSwizzle.SWIZZLE_128B + desc_q = prims.Tcgen05SmemDesc.build( + start_address=smem_q, leading_byte_offset=LEADING, stride_byte_offset=SBO, layout=swz + ) + desc_k = prims.Tcgen05SmemDesc.build( + start_address=smem_kv, leading_byte_offset=LEADING, stride_byte_offset=SBO, layout=swz + ) + desc_p = prims.Tcgen05SmemDesc.build( + start_address=smem_p, leading_byte_offset=LEADING, stride_byte_offset=SBO, layout=swz + ) + # V read MN-major: N = latent across 64-column chunks KV_CHUNK_ELEMS apart (leading byte offset). + desc_v = prims.Tcgen05SmemDesc.build( + start_address=smem_kv, + leading_byte_offset=KV_CHUNK_ELEMS * ELEM_BYTES, + stride_byte_offset=SBO, + layout=swz, + ) + tmem_s = cutlass.inttoptr(tmem_base, 6, cutlass.Int32) + for k in range(my_count): + phase = k & cutlass.Int32(1) + if k > 0: + # S overwrites O's first columns: the previous tile's O has been read out. + while not cute.arch.mbarrier_try_wait( + o_drained.data_ptr(), phase ^ cutlass.Int32(1) + ): + pass + while not cute.arch.mbarrier_try_wait(kv_full.data_ptr(), phase): + pass + # S per Q chunk group (each group's barrier completes once; later tiles pass it at once). + for g in cutlass.range_constexpr(Q_GROUPS): + while not cute.arch.mbarrier_try_wait(q_full.subview(g).data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for jj in cutlass.range_constexpr(Q_GROUP_CHUNKS): + j = g * Q_GROUP_CHUNKS + jj + for kk in cutlass.range_constexpr(CHUNK // MMA_K): + if prims.elect_sync(): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, prims.CTAGroup.CTA_1, tmem_s, + desc_q + (j * Q_CHUNK_U + kk * STEP_K), + desc_k + (j * KV_CHUNK_U + kk * STEP_K), + idesc_s, not (j == 0 and kk == 0), + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_commit(s_full) + while not cute.arch.mbarrier_try_wait(p_full.data_ptr(), phase): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + # O in two N = 256 halves, each committed on its own barrier (the first is drained during the second). + for nh in cutlass.range_constexpr(O_HALVES): + for kk in cutlass.range_constexpr(TILE // MMA_K): + if prims.elect_sync(): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_base + cutlass.Int32(nh * 256), 6, cutlass.Int32), + desc_p + ((kk // (CHUNK // MMA_K)) * P_HALF_U + (kk % (CHUNK // MMA_K)) * STEP_K), + desc_v + (nh * 4 * KV_CHUNK_U + kk * STEP_MN), + idesc_o, kk != 0, + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_commit(o_full.subview(nh)) + elif warp_id >= 4: + # ===================================================================== + # Softmax and the partial (m, l, O) -> this CTA's workspace slot. + # 16x256b: lane l of warp w holds rows 16 w + l / 4 and + 8, columns + # 8 g + 2 (l % 4) + {0, 1} of every 8-column group g. + # ===================================================================== + lane = tx % 32 + w = warp_id - 4 + quad = lane % cutlass.Int32(4) + r0 = w * cutlass.Int32(16) + lane // cutlass.Int32(4) + r1 = r0 + cutlass.Int32(8) + # Last visible KV row of each query row (token t = r % 8): L - T + t. + lim0 = kv_len - tokens + r0 % cutlass.Int32(MAX_TOKENS) + lim1 = kv_len - tokens + r1 % cutlass.Int32(MAX_TOKENS) + # The slot takes every real row's partial (rows t >= T too), the mailboxes the live rows' (m, l). + real0 = r0 < cutlass.Int32(ROWS) + real1 = r1 < cutlass.Int32(ROWS) + live0 = real0 & (r0 % cutlass.Int32(MAX_TOKENS) < tokens) + live1 = real1 & (r1 % cutlass.Int32(MAX_TOKENS) < tokens) + # The CTA's running (m, l) per row over its tiles (the slot's O is relative to m). + m_run0 = cutlass.Float32(NEG_INF) + m_run1 = cutlass.Float32(NEG_INF) + l_run0 = cutlass.Float32(0.0) + l_run1 = cutlass.Float32(0.0) + for k in range(my_count): + phase = k & cutlass.Int32(1) + tile = cta + k * cutlass.Int32(CLUSTER) + kv0 = tile * cutlass.Int32(TILE) + while not cute.arch.mbarrier_try_wait(s_full.data_ptr(), phase): + pass + if cutlass.const_expr(no_cluster): + count_ml = ctrs.load(idx=ctr_ml, is_volatile=True) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + s = prims.tcgen05_ld( + "16x256b", cutlass.inttoptr(tmem_base, 6, cutlass.Float32), num=TILE // 8 + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + # Masked, scaled scores and their row maxima (the quad's 4 lanes share the two rows). + m0 = cutlass.Float32(NEG_INF) + m1 = cutlass.Float32(NEG_INF) + v0 = [] + v1 = [] + for g in cutlass.range_constexpr(TILE // 8): + for e in cutlass.range_constexpr(2): + col = kv0 + cutlass.Int32(8 * g + e) + cutlass.Int32(2) * quad + a = cutlass.Float32(s[4 * g + e]) * scale_log2 + b = cutlass.Float32(s[4 * g + 2 + e]) * scale_log2 + a = cutlass.Float32(cutlass.select_(col <= lim0, a, cutlass.Float32(NEG_INF))) + b = cutlass.Float32(cutlass.select_(col <= lim1, b, cutlass.Float32(NEG_INF))) + v0.append(a) + v1.append(b) + m0 = cute.arch.fmax(m0, a) + m1 = cute.arch.fmax(m1, b) + for offset in (1, 2): + m0 = cute.arch.fmax(m0, cute.arch.shuffle_sync_bfly(m0, offset=offset)) + m1 = cute.arch.fmax(m1, cute.arch.shuffle_sync_bfly(m1, offset=offset)) + # A row with nothing visible in this tile: m = -inf, p = 0 (the merge weights it by 0). + base0 = cutlass.Float32( + cutlass.select_(m0 == cutlass.Float32(NEG_INF), cutlass.Float32(0.0), m0) + ) + base1 = cutlass.Float32( + cutlass.select_(m1 == cutlass.Float32(NEG_INF), cutlass.Float32(0.0), m1) + ) + l0 = cutlass.Float32(0.0) + l1 = cutlass.Float32(0.0) + for g in cutlass.range_constexpr(TILE // 8): + pv0 = [] + pv1 = [] + for e in cutlass.range_constexpr(2): + pa = cute.math.exp2(v0[2 * g + e] - base0, fastmath=True) + pb = cute.math.exp2(v1[2 * g + e] - base1, fastmath=True) + l0 = l0 + pa + l1 = l1 + pb + pv0.append(pa) + pv1.append(pb) + # P [row][kv] bf16, K-major 128B swizzle: half g // 8 (= page), chunk g % 8, elements 2 quad, +1. + half = g // 8 + chunk = g % 8 + pair0 = cutlass.Vector.from_elements( + (pv0[0].to(io_dtype), pv0[1].to(io_dtype)), io_dtype + ) + pair1 = cutlass.Vector.from_elements( + (pv1[0].to(io_dtype), pv1[1].to(io_dtype)), io_dtype + ) + off0 = ( + cutlass.Int32(half * P_HALF_ELEMS) + + r0 * cutlass.Int32(CHUNK) + + ((cutlass.Int32(chunk) ^ (r0 % cutlass.Int32(8))) * cutlass.Int32(8)) + + cutlass.Int32(2) * quad + ) + off1 = ( + cutlass.Int32(half * P_HALF_ELEMS) + + r1 * cutlass.Int32(CHUNK) + + ((cutlass.Int32(chunk) ^ (r1 % cutlass.Int32(8))) * cutlass.Int32(8)) + + cutlass.Int32(2) * quad + ) + smem_p.store(pair0, idx=off0, vector_size=2, alignment=4) + smem_p.store(pair1, idx=off1, vector_size=2, alignment=4) + for offset in (1, 2): + l0 = l0 + cute.arch.shuffle_sync_bfly(l0, offset=offset) + l1 = l1 + cute.arch.shuffle_sync_bfly(l1, offset=offset) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.fence_proxy(prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) + prims.mbarrier_arrive(p_full) + # Fold scales (k > 0; the slot's m is finite then: an earlier tile of a CTA is fully visible). + m_new0 = cute.arch.fmax(m_run0, m0) + m_new1 = cute.arch.fmax(m_run1, m1) + fa0 = cute.math.exp2(m_run0 - m_new0, fastmath=True) + fa1 = cute.math.exp2(m_run1 - m_new1, fastmath=True) + fb0 = cute.math.exp2(m0 - m_new0, fastmath=True) + fb1 = cute.math.exp2(m1 - m_new1, fastmath=True) + first = k == cutlass.Int32(0) + l_new0 = cutlass.Float32(cutlass.select_(first, l0, l_run0 * fa0 + l0 * fb0)) + l_new1 = cutlass.Float32(cutlass.select_(first, l1, l_run1 * fa1 + l1 * fb1)) + # The slot holds O / l (rows with nothing visible: 0): scales of the old slot (ca) and of this tile (cb). + inv0 = cutlass.Float32( + cutlass.select_( + l_new0 > cutlass.Float32(0.0), + cutlass.Float32(1.0) / l_new0, + cutlass.Float32(0.0), + ) + ) + inv1 = cutlass.Float32( + cutlass.select_( + l_new1 > cutlass.Float32(0.0), + cutlass.Float32(1.0) / l_new1, + cutlass.Float32(0.0), + ) + ) + ca0 = cutlass.Float32(cutlass.select_(first, cutlass.Float32(0.0), l_run0 * fa0 * inv0)) + ca1 = cutlass.Float32(cutlass.select_(first, cutlass.Float32(0.0), l_run1 * fa1 * inv1)) + cb0 = cutlass.Float32(cutlass.select_(first, inv0, fb0 * inv0)) + cb1 = cutlass.Float32(cutlass.select_(first, inv1, fb1 * inv1)) + l_run0 = l_new0 + l_run1 = l_new1 + m_run0 = cutlass.Float32(cutlass.select_(first, m0, m_new0)) + m_run1 = cutlass.Float32(cutlass.select_(first, m1, m_new1)) + if k == my_count - cutlass.Int32(1): + if cutlass.const_expr(no_cluster): + # The CTA's final (m, l) of its live rows into its slot of the exchange, then its arrival. + if quad == cutlass.Int32(0): + if live0: + ml_glob.store( + cutlass.Vector.from_elements((m_run0, l_run0), cutlass.Float32), + idx=(my_slot * cutlass.Int32(ROWS) + r0) * cutlass.Int32(2), vector_size=2, + alignment=8, + ) # fmt: skip + if live1: + ml_glob.store( + cutlass.Vector.from_elements((m_run1, l_run1), cutlass.Float32), + idx=(my_slot * cutlass.Int32(ROWS) + r1) * cutlass.Int32(2), vector_size=2, + alignment=8, + ) # fmt: skip + prims.barrier_cta_sync(1, thread_count=EPI_THREADS) + if tx == cutlass.Int32(4 * 32): + _red_add_release(ctrs.data_ptr(ctr_ml).toint(), 1) + sync_target.store( + (count_ml & cutlass.Int32(-16)) + cutlass.Int32(16), idx=0 + ) + prims.mbarrier_arrive(ml_full) + else: + # The CTA's final (m, l) of its live rows into slot `cta` of every CTA's mailbox (this one too). + if quad == cutlass.Int32(0): + for dst in cutlass.range_constexpr(CLUSTER): + mb = _mapa_u32(ml_full.data_ptr(), dst) + if live0: + _st_async_v2( + _mapa_u32( + ml_mail.data_ptr((cta * cutlass.Int32(ROWS) + r0) * cutlass.Int32(2)), dst + ), + m_run0, l_run0, mb, + ) # fmt: skip + if live1: + _st_async_v2( + _mapa_u32( + ml_mail.data_ptr((cta * cutlass.Int32(ROWS) + r1) * cutlass.Int32(2)), dst + ), + m_run1, l_run1, mb, + ) # fmt: skip + base_o0 = ( + my_slot * cutlass.Int32(WS_SLOT_ELEMS) + + r0 * cutlass.Int32(8) + + cutlass.Int32(2) * quad + ) + base_o1 = ( + my_slot * cutlass.Int32(WS_SLOT_ELEMS) + + r1 * cutlass.Int32(8) + + cutlass.Int32(2) * quad + ) + # O / l -> ws_o[slot, row, :] in fp16 (folded into it for k > 0), 128 columns per load, each N = 256 half + # as soon as its MMAs are done. + for part in cutlass.range_constexpr(LATENT // 128): + if cutlass.const_expr(part % 2 == 0): + while not cute.arch.mbarrier_try_wait( + o_full.subview(part // 2).data_ptr(), phase + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + o = prims.tcgen05_ld( + "16x256b", + cutlass.inttoptr(tmem_base + cutlass.Int32(part * 128), 6, cutlass.Float32), + num=16, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if first: + for g in cutlass.range_constexpr(16): + col = (part * 16 + g) * WS_GROUP_ELEMS + if real0: + ws_o.store( + cutlass.Vector.from_elements( + ((cutlass.Float32(o[4 * g]) * cb0).to(ws_dtype), + (cutlass.Float32(o[4 * g + 1]) * cb0).to(ws_dtype)), + ws_dtype, + ), + idx=base_o0 + cutlass.Int32(col), vector_size=2, alignment=4, + ) # fmt: skip + if real1: + ws_o.store( + cutlass.Vector.from_elements( + ((cutlass.Float32(o[4 * g + 2]) * cb1).to(ws_dtype), + (cutlass.Float32(o[4 * g + 3]) * cb1).to(ws_dtype)), + ws_dtype, + ), + idx=base_o1 + cutlass.Int32(col), vector_size=2, alignment=4, + ) # fmt: skip + elif real0: + # This thread wrote these slot entries itself (program order): read, rescale, add, write back. + # Rows past ROWS are never written, so they are not read either; a thread's two rows are both + # real or both past ROWS (ROWS % 16 == 0), so warp 7 skips this whole branch. + old0 = [] + old1 = [] + for g in cutlass.range_constexpr(16): + col = (part * 16 + g) * WS_GROUP_ELEMS + old0.append( + ws_o.load(idx=base_o0 + cutlass.Int32(col), vector_size=2, alignment=4) + ) + old1.append( + ws_o.load(idx=base_o1 + cutlass.Int32(col), vector_size=2, alignment=4) + ) + for g in cutlass.range_constexpr(16): + col = (part * 16 + g) * WS_GROUP_ELEMS + ws_o.store( + cutlass.Vector.from_elements( + ( + (cutlass.Float32(old0[g][0]) * ca0 + + cutlass.Float32(o[4 * g]) * cb0).to(ws_dtype), + (cutlass.Float32(old0[g][1]) * ca0 + + cutlass.Float32(o[4 * g + 1]) * cb0).to(ws_dtype), + ), + ws_dtype, + ), + idx=base_o0 + cutlass.Int32(col), vector_size=2, alignment=4, + ) # fmt: skip + ws_o.store( + cutlass.Vector.from_elements( + ( + (cutlass.Float32(old1[g][0]) * ca1 + + cutlass.Float32(o[4 * g + 2]) * cb1).to(ws_dtype), + (cutlass.Float32(old1[g][1]) * ca1 + + cutlass.Float32(o[4 * g + 3]) * cb1).to(ws_dtype), + ), + ws_dtype, + ), + idx=base_o1 + cutlass.Int32(col), vector_size=2, alignment=4, + ) # fmt: skip + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(o_drained) + if cutlass.const_expr(no_cluster): + if my_count == cutlass.Int32(0): + if tx == cutlass.Int32(4 * 32): + prims.griddepcontrol(prims.GridDepAction.WAIT) + count_ml0 = ctrs.load(idx=ctr_ml, is_volatile=True) + _red_add_release(ctrs.data_ptr(ctr_ml).toint(), 1) + sync_target.store((count_ml0 & cutlass.Int32(-16)) + cutlass.Int32(16), idx=0) + prims.mbarrier_arrive(ml_full) + else: + # ===================================================================== + # Warps 1 and 3: merge weights from the (m, l) mailbox while O runs. + # ===================================================================== + t1 = (warp_id // cutlass.Int32(2)) * cutlass.Int32(32) + tx % cutlass.Int32(32) + if cutlass.const_expr(no_cluster): + # This CTA's own arrival; the other CTAs' (m, l) are acquired through the counter. + while not cute.arch.mbarrier_try_wait(ml_full.data_ptr(), 0): + pass + target = sync_target.load(idx=0) + while _ld_acquire(ctrs.data_ptr(ctr_ml).toint()) - target < cutlass.Int32(0): + pass + else: + # Completed by the cluster's st.async: acquire at cluster scope. + while not _try_wait_cluster(ml_full.data_ptr(), 0): + pass + # A quad of lanes per row (16 rows per pass, 3 passes), lane j of the quad over slots j, j + 4, j + 8, j + 12: + # M = max_s m_s, weight_s = exp2(m_s - M) l_s / sum_s exp2(m_s - M) l_s (the slots hold O / l), over m_s. + sub = t1 % cutlass.Int32(4) + for ps in cutlass.range_constexpr(ROWS // 16): + row = cutlass.Int32(ps * 16) + t1 // cutlass.Int32(4) + live = (row < cutlass.Int32(ROWS)) & (row % cutlass.Int32(MAX_TOKENS) < tokens) + row_c = cutlass.Int32(cutlass.select_(live, row, cutlass.Int32(0))) + ms = [] + ls = [] + for i in cutlass.range_constexpr(CLUSTER // 4): + sl = sub + cutlass.Int32(4 * i) + ok = sl < n_valid + e = ( + cutlass.Int32(cutlass.select_(ok, sl, cutlass.Int32(0))) * cutlass.Int32(ROWS) + + row_c + ) * cutlass.Int32(2) + # Only a live lane's own slot entries are read as m: the m words are overwritten with the + # weights below by the lane that owns them. A lane past n_valid, or of a dead row, reads (and + # drops) an l word instead, which nothing writes here. + e_m = cutlass.Int32(cutlass.select_(ok & live, e, e + cutlass.Int32(1))) + if cutlass.const_expr(no_cluster): + eg = ws_slot0 * cutlass.Int32(ROWS * 2) + e + ms.append( + cutlass.Float32( + cutlass.select_( + ok, + ml_glob.load(idx=ws_slot0 * cutlass.Int32(ROWS * 2) + e_m), + cutlass.Float32(NEG_INF), + ) + ) + ) + ls.append( + cutlass.Float32( + cutlass.select_( + ok, ml_glob.load(idx=eg + cutlass.Int32(1)), cutlass.Float32(0.0) + ) + ) + ) + else: + ms.append( + cutlass.Float32( + cutlass.select_(ok, ml_mail.load(idx=e_m), cutlass.Float32(NEG_INF)) + ) + ) + ls.append( + cutlass.Float32( + cutlass.select_( + ok, ml_mail.load(idx=e + cutlass.Int32(1)), cutlass.Float32(0.0) + ) + ) + ) + mx = ms[0] + for i in cutlass.range_constexpr(1, CLUSTER // 4): + mx = cute.arch.fmax(mx, ms[i]) + for offset in (1, 2): + mx = cute.arch.fmax(mx, cute.arch.shuffle_sync_bfly(mx, offset=offset)) + bs = [cute.math.exp2(ms[i] - mx, fastmath=True) for i in range(CLUSTER // 4)] + den = bs[0] * ls[0] + for i in cutlass.range_constexpr(1, CLUSTER // 4): + den = den + bs[i] * ls[i] + for offset in (1, 2): + den = den + cute.arch.shuffle_sync_bfly(den, offset=offset) + inv = cutlass.Float32(1.0) / den + if live: + for i in cutlass.range_constexpr(CLUSTER // 4): + sl = sub + cutlass.Int32(4 * i) + if sl < n_valid: + ml_mail.store( + bs[i] * ls[i] * inv, + idx=(sl * cutlass.Int32(ROWS) + row) * cutlass.Int32(2), + ) + + # ========================================================================= + # Every tile's partial is in the workspace (release / acquire over the + # cluster, all threads; no_cluster: the second counter); CTA c merges + # latent columns [32 c, 32 c + 32). + # ========================================================================= + if cutlass.const_expr(no_cluster): + if cutlass.const_expr(fuse_vb): + merging = VB_CTAS + else: + merging = CLUSTER + prims.barrier_cta_sync(0) + if tx == cutlass.Int32(0): + _red_add_release(ctrs.data_ptr(ctr_o).toint(), 1) + if cta < cutlass.Int32(merging): + target_o = (sync_target.load(idx=1) & cutlass.Int32(-16)) + cutlass.Int32(16) + while _ld_acquire(ctrs.data_ptr(ctr_o).toint()) - target_o < cutlass.Int32(0): + pass + prims.barrier_cta_sync(0) + else: + prims.barrier_cluster_arrive() + prims.barrier_cluster_wait() + if cutlass.const_expr(fuse_vb): + if cta < cutlass.Int32(VB_CTAS): + # Rows 8 h + t of tokens t = 4 th + j (j < 4; adjacent rows), all 512 columns: thread -> row j = tx % 4, + # columns 8 (tx / 4) .. + 8, so a warp reads 8 column groups x 4 adjacent 16-byte rows; the slots in two + # batches of 8 (16 loads of 16 bytes in flight per batch). + h_loc = cta // cutlass.Int32(2) + j_row = tx % cutlass.Int32(VB_TOKENS) + t_tok = (cta % cutlass.Int32(2)) * cutlass.Int32(VB_TOKENS) + j_row + live = t_tok < tokens + # A token t >= T reads its own row too (weight 0); its column of the v_b product is never stored. + row = h_loc * cutlass.Int32(MAX_TOKENS) + t_tok + c8 = (tx // cutlass.Int32(VB_TOKENS)) * cutlass.Int32(8) + acc = [cutlass.Float32(0.0)] * 8 + for half in cutlass.range_constexpr(2): + vals = [] + wts = [] + for j in cutlass.range_constexpr(MERGE_HALF): + sl = cutlass.Int32(half * MERGE_HALF + j) + ok = sl < n_valid + sl_c = cutlass.Int32(cutlass.select_(ok, sl, cutlass.Int32(0))) + base = ( + (ws_slot0 + sl_c) * cutlass.Int32(WS_SLOT_ELEMS) + + (c8 // cutlass.Int32(8)) * cutlass.Int32(WS_GROUP_ELEMS) + + row * cutlass.Int32(8) + ) + vals.append(ws_o.load(idx=base, vector_size=8, alignment=16)) + wts.append( + cutlass.Float32( + cutlass.select_( + ok & live, + ml_mail.load( + idx=(sl_c * cutlass.Int32(ROWS) + row) * cutlass.Int32(2) + ), + cutlass.Float32(0.0), + ) + ) + ) + for j in cutlass.range_constexpr(MERGE_HALF): + for e in cutlass.range_constexpr(8): + acc[e] = acc[e] + wts[j] * cutlass.Float32(vals[j][e]) + # bf16 o (the unfused path's attention output) -> the v_b B tile in the dead P buffer (token j, columns + # c8 .. c8 + 8), K-major 128B swizzle: chunk c8 / 64, 16-byte vector ((c8 % 64) / 8) ^ j. + smem_p.store( + cutlass.Vector.from_elements(tuple(v.to(io_dtype) for v in acc), io_dtype), + idx=(c8 // cutlass.Int32(CHUNK)) * cutlass.Int32(8 * CHUNK) + j_row * cutlass.Int32(CHUNK) + + (((c8 % cutlass.Int32(CHUNK)) // cutlass.Int32(8)) ^ j_row) * cutlass.Int32(8), + vector_size=8, + alignment=16, + ) # fmt: skip + prims.fence_proxy(prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) + prims.barrier_cta_sync(0) + if cta < cutlass.Int32(VB_CTAS): + if warp_id == 2: + # y^T [128 v, 8 tokens] = W_vb[h] (M 128, K 512 from the KV buffer) x o^T (the B tile), fp32 in TMEM. + idesc_vb = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, a_dtype=io_dtype, b_dtype=io_dtype, n_dim=8, m_dim=128 + ) + swz = prims.Tcgen05SmemSwizzle.SWIZZLE_128B + desc_w = prims.Tcgen05SmemDesc.build( + start_address=smem_kv, + leading_byte_offset=LEADING, + stride_byte_offset=SBO, + layout=swz, + ) + desc_o = prims.Tcgen05SmemDesc.build( + start_address=smem_p, + leading_byte_offset=LEADING, + stride_byte_offset=SBO, + layout=swz, + ) + while not cute.arch.mbarrier_try_wait(vb_full.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for jc in cutlass.range_constexpr(LAT_CHUNKS): + for kk in cutlass.range_constexpr(CHUNK // MMA_K): + if prims.elect_sync(): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_base, 6, cutlass.Int32), + desc_w + (jc * ((V_DIM * CHUNK * ELEM_BYTES) >> 4) + kk * STEP_K), + desc_o + (jc * ((8 * CHUNK * ELEM_BYTES) >> 4) + kk * STEP_K), + idesc_vb, not (jc == 0 and kk == 0), + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_commit(vb_done) + elif warp_id >= 4: + # y[t, h, v]: lane v = 32 (warp - 4) + lane of the accumulator, tokens in its 8 columns. + v_idx = (warp_id - cutlass.Int32(4)) * cutlass.Int32(32) + tx % cutlass.Int32(32) + h_glob = hg * cutlass.Int32(HEADS) + cta // cutlass.Int32(2) + # The gate of this thread's tokens, loaded before the v_b wait so its latency hides under the MMA (a + # token past T reads the request's first row and is not used). + gate_pre = [] + if cutlass.const_expr(apply_gate): + for j in cutlass.range_constexpr(VB_TOKENS): + t_pre = (cta % cutlass.Int32(2)) * cutlass.Int32(VB_TOKENS) + cutlass.Int32( + j + ) + t_row = tok0 + cutlass.Int32( + cutlass.select_(t_pre < tokens, t_pre, cutlass.Int32(0)) + ) + gate_pre.append( + cutlass.Float32( + gate.load( + idx=t_row * gate_ld + + gate_col0 + + h_glob * cutlass.Int32(V_DIM) + + v_idx + ) + ) + ) + while not cute.arch.mbarrier_try_wait(vb_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + yv = prims.tcgen05_ld( + "32x32b", cutlass.inttoptr(tmem_base, 6, cutlass.Float32), num=8 + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + for j in cutlass.range_constexpr(VB_TOKENS): + t_out = (cta % cutlass.Int32(2)) * cutlass.Int32(VB_TOKENS) + cutlass.Int32(j) + if t_out < tokens: + y_out = cutlass.Float32(yv[j]).to(io_dtype) + if cutlass.const_expr(apply_gate): + y_out = (cutlass.Float32(y_out) * gate_pre[j]).to(io_dtype) + out.store( + y_out, + idx=((tok0 + t_out) * cutlass.Int32(total_heads) + h_glob) + * cutlass.Int32(V_DIM) + + v_idx, + ) + # Every TMEM reader of this CTA has waited for its loads. + prims.barrier_cta_sync(0) + if warp_id == 4: + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_base, 6, cutlass.Int32), TMEM_COLS) + else: + if warp_id == 4: + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_base, 6, cutlass.Int32), TMEM_COLS) + # 4 groups of 8 columns x 48 rows = 192 items (item = 48 group + row, so consecutive pairs read consecutive + # 16-byte rows of a group); thread pair (2 p, 2 p + 1) takes items p and p + 128 (< 192), the even thread over + # slots [0, 8), the odd one over [8, 16): all 16 of a thread's workspace loads (8 fp16 each) are in flight + # together, then the pair's halves are added by one shuffle. Slots >= n_valid load slot 0 with weight 0, rows + # of tokens t >= T their own row with weight 0 (never stored). + hs = tx % cutlass.Int32(2) + pair = tx // cutlass.Int32(2) + vals = [] + wts = [] + for i in cutlass.range_constexpr(2): + item = pair + cutlass.Int32(128 * i) + item_c = cutlass.Int32(cutlass.select_(item < cutlass.Int32(ROWS * 4), item, pair)) + row = item_c % cutlass.Int32(ROWS) + live = (row % cutlass.Int32(MAX_TOKENS) < tokens) & (item < cutlass.Int32(ROWS * 4)) + grp = cta * cutlass.Int32(LAT_PER_CTA // 8) + item_c // cutlass.Int32(ROWS) + for j in cutlass.range_constexpr(MERGE_HALF): + sl = hs * cutlass.Int32(MERGE_HALF) + cutlass.Int32(j) + ok = sl < n_valid + sl_c = cutlass.Int32(cutlass.select_(ok, sl, cutlass.Int32(0))) + vals.append( + ws_o.load( + idx=(ws_slot0 + sl_c) * cutlass.Int32(WS_SLOT_ELEMS) + + grp * cutlass.Int32(WS_GROUP_ELEMS) + + row * cutlass.Int32(8), + vector_size=8, + alignment=16, + ) + ) + wts.append( + cutlass.Float32( + cutlass.select_( + ok & live, + ml_mail.load(idx=(sl_c * cutlass.Int32(ROWS) + row) * cutlass.Int32(2)), + cutlass.Float32(0.0), + ) + ) + ) + for i in cutlass.range_constexpr(2): + item = pair + cutlass.Int32(128 * i) + item_c = cutlass.Int32(cutlass.select_(item < cutlass.Int32(ROWS * 4), item, pair)) + row = item_c % cutlass.Int32(ROWS) + grp = cta * cutlass.Int32(LAT_PER_CTA // 8) + item_c // cutlass.Int32(ROWS) + acc = [] + for e in cutlass.range_constexpr(8): + a = cutlass.Float32(0.0) + for j in cutlass.range_constexpr(MERGE_HALF): + a = a + wts[i * MERGE_HALF + j] * cutlass.Float32(vals[i * MERGE_HALF + j][e]) + acc.append(a + cute.arch.shuffle_sync_bfly(a, offset=1)) + out_row = ( + (tok0 + row % cutlass.Int32(MAX_TOKENS)) * cutlass.Int32(total_heads) + + hg * cutlass.Int32(HEADS) + + row // cutlass.Int32(MAX_TOKENS) + ) + if hs == cutlass.Int32(0): + if (row % cutlass.Int32(MAX_TOKENS) < tokens) & (item < cutlass.Int32(ROWS * 4)): + out.store( + cutlass.Vector.from_elements(tuple(v.to(io_dtype) for v in acc), io_dtype), + idx=out_row * cutlass.Int32(LATENT) + grp * cutlass.Int32(8), + vector_size=8, + alignment=16, + ) + + +def _q_map(q, num_tokens, total_heads): + """fused_q [M, heads, 576] as (64-column chunk, token, head, chunk index): one call at (token i T, head 6 g, chunk + 3 j) lands the group's [3][6 heads x 8 tokens][64] (row 8 h + t holds token i T + t; tokens >= M zero).""" + return cuda.create_tensor_map_tiled( + global_address=q.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[CHUNK, num_tokens, total_heads, QK_CHUNKS], + global_strides=[ + (total_heads * QK * ELEM_BYTES) // 16, + (QK * ELEM_BYTES) // 16, + (CHUNK * ELEM_BYTES) // 16, + ], + box_dims=[CHUNK, MAX_TOKENS, HEADS, Q_GROUP_CHUNKS], + swizzle=cuda.TensorMapSwizzle.s128b, + ) + + +def _vb_map(w_vb, rows): + """v_b_proj [heads * 128, 512] as (64-column chunk, row, chunk index): one call lands a head's [8][128][64].""" + return cuda.create_tensor_map_tiled( + global_address=w_vb.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[CHUNK, rows, LAT_CHUNKS], + global_strides=[(LATENT * ELEM_BYTES) // 16, (CHUNK * ELEM_BYTES) // 16], + box_dims=[CHUNK, V_DIM, LAT_CHUNKS], + swizzle=cuda.TensorMapSwizzle.s128b, + ) + + +def _kv_map(pool, total_rows, row_stride): + """The latent pool as [rows, 576] (row stride `row_stride` elements): one call lands a page's 64-column chunk.""" + return cuda.create_tensor_map_tiled( + global_address=pool.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[CHUNK, total_rows, QK_CHUNKS], + global_strides=[(row_stride * ELEM_BYTES) // 16, (CHUNK * ELEM_BYTES) // 16], + box_dims=[CHUNK, PAGE, 1], + swizzle=cuda.TensorMapSwizzle.s128b, + ) + + +@cute.jit +def k3_mla_attn( + q: cute.Tensor, # [M * heads * 576] bf16, M = R T + pool: cute.Tensor, # the latent pool, flat bf16: row i of page p at (p * 64 + i) * row_stride + page_table: cute.Tensor, # int32, request i's pages at [i * pt_stride, ...) + seq_len: cute.Tensor, # int32 [R] + ws_o: cute.Tensor, # fp16 [MAX_REQUESTS * heads / 6 * 16 * WS_SLOT_ELEMS] + out: cute.Tensor, # bf16 [M * heads * 512] (fuse_vb: [M * heads * 128]) + w_vb: cute.Tensor, # fuse_vb: v_b_proj [heads * 128 * 512] bf16 (else any bf16 tensor, unused) + gate: cute.Tensor, # apply_gate: bf16 [M * gate_ld] (else any bf16 tensor, unused) + tokens: cutlass.Int32, # T <= 8 + num_requests: cutlass.Int32, # R <= MAX_REQUESTS + pt_stride: cutlass.Int32, + scale_log2: cutlass.Float32, + total_rows: cutlass.Int32, + page_offset: cutlass.Int32, + gate_col0: cutlass.Int32, + gate_ld: cutlass.Int32, + row_stride: cutlass.Constexpr[int], + total_heads: cutlass.Constexpr[int], + fuse_vb: cutlass.Constexpr[bool], + apply_gate: cutlass.Constexpr[bool], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, + no_cluster: cutlass.Constexpr[bool] = False, +) -> None: + tma_q = _q_map(q, tokens * num_requests, total_heads) + tma_kv = _kv_map(pool, total_rows, row_stride) + tma_vb = _vb_map(w_vb, total_heads * V_DIM) if cutlass.const_expr(fuse_vb) else tma_kv + k3_mla_attn_kernel( + tma_q, + tma_kv, + tma_vb, + page_table, + seq_len, + ws_o, + out, + gate, + tokens, + pt_stride, + scale_log2, + page_offset, + gate_col0, + gate_ld, + total_heads, + fuse_vb, + apply_gate, + no_cluster, + ).launch( # fmt: skip + grid=(CLUSTER * (total_heads // HEADS), num_requests, 1), + block=(THREADS, 1, 1), + cluster=None if no_cluster else (CLUSTER, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) + + +def softmax_scale_log2( + qk_nope_head_dim: int, qk_rope_head_dim: int, q_scaling: float = 1.0 +) -> float: + return LOG2E / (math.sqrt(qk_nope_head_dim + qk_rope_head_dim) * q_scaling) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_mla/k3_mla_q_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_mla/k3_mla_q_kernel.py new file mode 100644 index 000000000000..e39ea633654d --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_mla/k3_mla_q_kernel.py @@ -0,0 +1,1099 @@ +# 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 MLA decode query path -- CTM (prims/cute) kernel, M <= 64 tokens, any number of heads per rank (6 at TP16) +# ============================================================================= +# +# q_n[t] = bf16(q_a[t] * rsqrt(mean(q_a[t]^2) + eps) * w_qa) (q_a_layernorm, 1536) +# q[t, h] = bf16(q_n[t] @ W_qb[192 h : 192 h + 192]^T) (q_b_proj: 128 nope | 64 pe) +# fused_q[t, h] = [ bf16(q[t, h, :128] @ W_kb[h]^T) (512) | q[t, h, 128:] (64) ] (k_b absorb; K3 is NoPE) +# +# The grid's y index is a chunk of N tokens (rows N c .. N c + N - 1, the last chunk possibly shorter; N, the MMA's +# N, is the compile-time `mma_n`: 8, or 32 for the wider steps, whose chunks of 8 would need more co-resident clusters +# than fit): each chunk is the same N-token problem below, on its own rows, the weights read again per chunk (from L2 +# after the first). Every token's sums keep the same order at either N, so a 32-token chunk's rows are bit-identical +# to the four 8-token chunks of the same rows. Everything below describes one chunk of 8 tokens; at N 32 every +# per-token step runs for four groups of 8 tokens (token t of group g is row 8 g + t of the chunk). +# +# One cluster of 6 CTAs per head. Rank r holds (TMA'd EVICT_FIRST before griddepcontrol.wait) the head's W_qb +# nope rows (128) and pe rows (64) over k-tiles {2 r, 2 r + 1}, and for r < 4 the k_b rows W_kb[h][128 r:128 r+128]. +# After the wait the epilogue warps normalize the q_a rows (the RMS over all 1536 columns) and write this rank's 256 +# columns, bf16, into a resident 128B-swizzled B tile (the layout TMA would produce), fence it to the async proxy and +# release the MMA warp (fuse_rmsnorm_qkv_rope's resident-B prolog). cluster_rms: each rank reads only its 256 columns +# and st.async's its per-token partial sums of squares to the 6 ranks, which add them in rank order. The MMA warp +# runs the nope (M 128) and pe (M 64) partial products. Split-K reduce over the cluster: nope rows [32 w, 32 w + 32) +# (TMEM lanes of epilogue warp w) belong to rank w, pe rows [32 j, 32 j + 32) (lanes 0-15 of warps 2 j, 2 j + 1: an +# M = 64 accumulator puts row 16 w + l in lane 32 w + l) to rank 4 + j; non-owners store their fp32 partials into +# slot [rank] of the owner's mailbox through DSMEM and arrive (release, cluster scope); owners add the 6 partials in +# rank order and round once. Every wait on a barrier other CTAs complete acquires at cluster scope. +# pe owners store fused_q[..., 512:]. Each nope owner writes its bf16 32 x 8 slice, as 16-byte chunks, into the k_b B +# tile of ranks 0-3 and arrives on their barrier; those ranks fence the DSMEM writes to the async proxy, run the k_b +# MMA (M 128, N 8, K 128) and store fused_q[t, h, 128 r : 128 r + 128]. +# single_hop (8-token chunks only): instead, every rank st.async's its whole nope partial (128 rows x 8 tokens fp32) +# into the mailbox of each k_b rank (completing that rank's barrier by bytes); a k_b rank adds the 6 partials in rank +# order itself (the same fp32 sums, so the same bits), writes its B tile and runs the k_b MMA: one DSMEM hop instead +# of two. +# +# kv_mode (the KV half, warp 1 of CTA j < the chunk's tokens: token t = N c + j): token t's latent row +# ag[t, 1536:2048] RMS-normalized (flashinfer's order, like q_a) and its rope columns ag[t, 2048:2112] (K3 is NoPE: +# copied) form the 576-column cache row, stored (1) into the paged latent pool: the M tokens are R requests of T each +# (request-major), token t is token u = t - i T of request i = t / T, at position p = L_i - T + u (row +# (page_table[i][p / 64] + page_offset) * 64 + p % 64, 64-bit element index; not stored when p < 0), or (2) into +# kv_out [M, 576] (check builds). Page table and lengths are host-filled: read before the grid wait. The attention +# after this kernel reads the new rows after its own grid wait. +# +# Warps: 0 weight TMA (+ early dependent trigger), 1 KV half (kv_mode), 2 TMEM allocation + MMA, 4-7 norm prolog / +# reduce / epilogue, 3 idle. +# ============================================================================= +"""CTM kernel for Kimi K3 MLA's decode query path: q_a RMSNorm, q_b projection, k_b absorption -> fused_q.""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +HEADS = 6 # per rank at TP16 (the kernel takes the rank's head count as a compile-time argument) +NOPE = 128 +PE = 64 +QK = NOPE + PE # 192 q_b rows per head +LATENT = 512 # kv_lora_rank +Q_LORA = 1536 # q_lora_rank = q_b's K +FUSED = LATENT + PE # 576 +PAGE = 64 # latent cache rows per page +CLUSTER = 6 # CTAs per head +K_TILES = Q_LORA // 128 # 12 +MY_TILES = K_TILES // CLUSTER # 2 q_b k-tiles per rank +KB_RANKS = LATENT // 128 # 4 ranks hold a 128-row k_b tile +MMA_N = 8 # tokens per chunk (the MMA's N), and the group of tokens every per-token step works on +WIDE_CHUNKS = ( + 16, + 32, +) # the other chunk sizes (cluster_rms, two-hop): 64 tokens are 12 clusters in chunks of 32 +CTA_K = 128 +MMA_K = 16 +TMA_K_BOX = 64 +TMA_COPY_ITERS = CTA_K // TMA_K_BOX +K_BLOCKS_PER_HALF = TMA_K_BOX // MMA_K +THREADS = 256 +EPI_THREADS = 128 +TMEM_NOPE = 0 # the nope accumulator's TMEM column; pe and k_b follow at N and 2 N (tmem_cols) +ELEM_BYTES = 2 +VEC = 8 +EVICT_FIRST = 0x12F0000000000000 + +LEADING = 16 +STRIDE = 8 * TMA_K_BOX * ELEM_BYTES # 1024 B between 8-row swizzle atoms +STEP = (MMA_K * ELEM_BYTES) >> 4 +HALF_128 = ( + 128 * TMA_K_BOX * ELEM_BYTES +) >> 4 # one 64-column half of a 128-row tile, 16-byte units +HALF_64 = (64 * TMA_K_BOX * ELEM_BYTES) >> 4 +TILE_128 = 2 * HALF_128 +TILE_64 = 2 * HALF_64 +NOPE_ELEMS = 128 * CTA_K # per k-tile +PE_ELEMS = 64 * CTA_K + +KB_MAIL_BYTES = ( + CLUSTER * 128 * MMA_N * 4 +) # the k_b rank's single-hop mailbox: 6 partials of 128 rows x 8 tokens + +io_dtype = cutlass.BFloat16 + + +def tmem_cols(n: int) -> int: + """TMEM columns of an n-token chunk: the nope, pe and k_b accumulators (n columns each), a power of 2 >= 32.""" + cols = 32 + while cols < 3 * n: + cols *= 2 + return cols + + +@dsl_user_op +def _mapa_u32(smem_ptr, peer, *, loc=None, ip=None): + """The shared::cluster address of this CTA's shared-memory location in cluster CTA ``peer``.""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [smem_ptr.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(peer).ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_async_v4(dst, a, b, c, d, mbar, *, loc=None, ip=None): + """st.async of four fp32 to a shared::cluster address, completing ``mbar`` (shared::cluster) by 16 bytes.""" + _llvm.inline_asm( + None, + [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Float32(a).ir_value(loc=loc, ip=ip), + cutlass.Float32(b).ir_value(loc=loc, ip=ip), cutlass.Float32(c).ir_value(loc=loc, ip=ip), + cutlass.Float32(d).ir_value(loc=loc, ip=ip), cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.f32 [$0], {$1, $2, $3, $4}, [$5];", + "r,f,f,f,f,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _st_async_f32(dst, value, mbar, *, loc=None, ip=None): + """st.async of one fp32 to a shared::cluster address, completing ``mbar`` (shared::cluster) by 4 bytes.""" + _llvm.inline_asm( + None, + [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Float32(value).ir_value(loc=loc, ip=ip), + cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 [$0], $1, [$2];", "r,f,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _st_async_v4_b32(dst, w0, w1, w2, w3, mbar, *, loc=None, ip=None): + """st.async of 16 bytes (four 32-bit words) to a shared::cluster address, completing ``mbar`` by 16 bytes.""" + _llvm.inline_asm( + None, + [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Int32(w0).ir_value(loc=loc, ip=ip), + cutlass.Int32(w1).ir_value(loc=loc, ip=ip), cutlass.Int32(w2).ir_value(loc=loc, ip=ip), + cutlass.Int32(w3).ir_value(loc=loc, ip=ip), cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], + "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, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _try_wait_cluster(mbar, parity, *, loc=None, ip=None): + """mbarrier.try_wait.parity.acquire.cluster on a barrier of this CTA (``mbar``, its shared-memory pointer): whether + phase ``parity`` has completed, acquiring at cluster scope. The barrier is completed by other CTAs' st.async, whose + complete_tx releases at cluster scope.""" + done = cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tmbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" + "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + return done != cutlass.Int32(0) + + +def _swz(t, k, n=MMA_N): + """Element index of the 8-column vector holding (row t < n, column k < 128) in an n x 128 bf16 K-major tile, + 128B swizzle (two 64-column halves of 8-row swizzle atoms); add k % 8 for the element itself.""" + half = k // cutlass.Int32(TMA_K_BOX) + chunk = (k % cutlass.Int32(TMA_K_BOX)) // cutlass.Int32(VEC) + phase = t if n == MMA_N else t % cutlass.Int32(8) + return ( + half * cutlass.Int32(n * TMA_K_BOX) + + t * cutlass.Int32(TMA_K_BOX) + + (chunk ^ phase) * cutlass.Int32(VEC) + ) + + +def _plus(base, offset: int): + """base + offset; a zero offset traces nothing, so 8-token chunks keep exactly their per-token instructions.""" + return base if offset == 0 else base + cutlass.Int32(offset) + + +@cute.kernel +def k3_mla_q_kernel( + tma_nope: cutlass.GridConstant[cuda.TensorMap], # W_qb [1152, 1536], box 128 rows + tma_pe: cutlass.GridConstant[cuda.TensorMap], # W_qb [1152, 1536], box 64 rows + tma_kb: cutlass.GridConstant[cuda.TensorMap], # W_kb [6 * 512, 128], box 128 rows + ag: cutlass.Array, # [M, ag_cols] bf16: q_a in columns [0, 1536) + w_qa: cutlass.Array, # [1536] bf16, q_a_layernorm weight + fused_q: cutlass.Array, # [M, heads, 576] bf16, out + w_kv: cutlass.Array, # [512] bf16, kv_a_layernorm weight (kv_mode) + kv_pool: cutlass.Array, # the latent pool, flat bf16: row i of page p at (p * 64 + i) * row_stride (kv_mode 1) + page_table: cutlass.Array, # int32, request i's pages at [i * pt_stride, ...) (kv_mode 1) + seq_len: cutlass.Array, # int32 [R], request i's KV length including its T tokens of this step (kv_mode 1) + kv_out: cutlass.Array, # [M, 576] bf16 (kv_mode 2) + total_tokens: cutlass.Int32, # M = R T + tokens: cutlass.Int32, # T (kv_mode 1) + pt_stride: cutlass.Int32, # page-table elements between consecutive requests' rows (kv_mode 1) + eps: cutlass.Float32, + kv_eps: cutlass.Float32, + page_offset: cutlass.Int32, # added to every page-table entry (the layer's slot in an interleaved pool) + ag_cols: cutlass.Constexpr[int], + total_heads: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + single_hop: cutlass.Constexpr[bool], + kv_mode: cutlass.Constexpr[int], # 0 off, 1 into the pool, 2 into kv_out + row_stride: cutlass.Constexpr[int], + cluster_rms: cutlass.Constexpr[ + bool + ], # the q_a RMS from the cluster's partial sums (each rank reads its 256 columns) + mma_n: cutlass.Constexpr[ + int + ], # tokens per chunk: MMA_N, or one of WIDE_CHUNKS (cluster_rms, two-hop) +): + tx, _, _ = cute.arch.thread_idx() + bx, chunk, _ = cute.arch.block_idx() + warp_id = cute.arch.warp_idx() + rank = cute.arch.block_idx_in_cluster() + head = bx // cutlass.Int32(CLUSTER) + groups = mma_n // MMA_N # per-token steps run for each group of 8 of the chunk's tokens + half_n = ( + mma_n * TMA_K_BOX * ELEM_BYTES + ) >> 4 # one 64-column half of the B tile, 16-byte units + b_elems = mma_n * CTA_K + col_pe, col_kb = ( + mma_n, + 2 * mma_n, + ) # TMEM columns of the pe and k_b accumulators (nope at column 0) + # This chunk's tokens: rows tok0 .. tok0 + num_tokens - 1 of ag and fused_q. + tok0 = chunk * cutlass.Int32(mma_n) + num_tokens = cutlass.Int32( + cutlass.select_( + total_tokens - tok0 < cutlass.Int32(mma_n), total_tokens - tok0, cutlass.Int32(mma_n) + ) + ) + ptr_nope = tma_nope.get_ptr() + ptr_pe = tma_pe.get_ptr() + ptr_kb = tma_kb.get_ptr() + + # Same allocation order in every CTA: mapa() addresses the peers' mailbox, B2 tile and barriers. + smem_wn = cutlass.Array( + io_dtype, MY_TILES * NOPE_ELEMS, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_wp = cutlass.Array( + io_dtype, MY_TILES * PE_ELEMS, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_wk = cutlass.Array(io_dtype, NOPE_ELEMS, space=cutlass.AddressSpace.smem, alignment=1024) + smem_b = cutlass.Array( + io_dtype, MY_TILES * b_elems, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_b2 = cutlass.Array(io_dtype, b_elems, space=cutlass.AddressSpace.smem, alignment=1024) + stage_n = cutlass.Array(io_dtype, mma_n * 32, space=cutlass.AddressSpace.smem, alignment=16) + stage_o = cutlass.Array(io_dtype, mma_n * 128, space=cutlass.AddressSpace.smem, alignment=16) + # [source rank][token][row within the owned 32] fp32 partials (the own rank's slot stays unused). + mailbox = cutlass.Array( + cutlass.Float32, CLUSTER * 32 * mma_n, space=cutlass.AddressSpace.smem, alignment=16 + ) + norm_part = cutlass.Array( + cutlass.Float32, 4 * mma_n, space=cutlass.AddressSpace.smem, alignment=16 + ) + rrms = cutlass.Array(cutlass.Float32, mma_n, space=cutlass.AddressSpace.smem, alignment=16) + # cluster_rms: [source rank][token] partial sums of squares over the source's 256 q_a columns. + norm_mail = cutlass.Array( + cutlass.Float32, CLUSTER * mma_n, space=cutlass.AddressSpace.smem, alignment=16 + ) + norm_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + w_full = cutlass.Array( + cutlass.Int64, 2 * MY_TILES + 1, space=cutlass.AddressSpace.smem, alignment=8 + ) + b_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + acc_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + mail_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + qn_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + acc2_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + # single_hop: [source rank][nope row][token] fp32 partials, completed by bytes; the B tile's local readiness. + # (Not in 32-token builds, which are two-hop: 96 KB.) + kb_mail = None + if cutlass.const_expr(mma_n == MMA_N): + kb_mail = cutlass.Array( + cutlass.Float32, CLUSTER * 128 * MMA_N, space=cutlass.AddressSpace.smem, alignment=16 + ) + kb_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + b2_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + + if warp_id == 0: + prims.prefetch_tensormap(ptr_nope) + prims.prefetch_tensormap(ptr_pe) + prims.prefetch_tensormap(ptr_kb) + if prims.elect_sync(): + for i in cutlass.range_constexpr(2 * MY_TILES + 1): + prims.mbarrier_init(w_full.subview(i), 1) + prims.mbarrier_init(b_ready, EPI_THREADS) + prims.mbarrier_init(acc_done, 1) + # Both two-hop mailboxes complete by bytes (st.async), armed here once. + prims.mbarrier_init(mail_full, 1) + prims.mbarrier_init(qn_full, 1) + prims.mbarrier_init(acc2_done, 1) + prims.mbarrier_init(kb_full, 1) + prims.mbarrier_init(b2_ready, EPI_THREADS) + if cutlass.const_expr(single_hop): + if rank < cutlass.Int32(KB_RANKS): + prims.mbarrier_arrive_expect_tx(kb_full, KB_MAIL_BYTES) + prims.mbarrier_arrive_expect_tx(mail_full, (CLUSTER - 1) * mma_n * 32 * 4) + if cutlass.const_expr(cluster_rms): + prims.mbarrier_init(norm_full, 1) + prims.mbarrier_arrive_expect_tx(norm_full, CLUSTER * mma_n * 4) + if rank < cutlass.Int32(KB_RANKS): + # 4 owners x the chunk's tokens x 32 columns, bf16. + prims.mbarrier_arrive_expect_tx(qn_full, KB_RANKS * mma_n * 32 * ELEM_BYTES) + if warp_id == 2: + prims.tcgen05_alloc(tmem_ptr_i32, tmem_cols(mma_n)) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + # Cluster formation: the peers' shared memory and barriers are addressable from here on. + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + tmem_base = tmem_ptr_i32.load() + + if warp_id == 0: + # ===================================================================== + # Weights, ahead of the grid dependency: q_b nope and pe rows over this + # rank's two k-tiles, and (ranks 0-3) the k_b tile. + # ===================================================================== + if prims.elect_sync(): + row_nope = head * cutlass.Int32(QK) + for i in cutlass.range_constexpr(MY_TILES): + k = rank * cutlass.Int32(MY_TILES) + cutlass.Int32(i) + prims.mbarrier_arrive_expect_tx(w_full.subview(i), NOPE_ELEMS * ELEM_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_wn.subview(i * NOPE_ELEMS), + ptr_nope, + ( + cutlass.Int32(0), + row_nope, + k * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), + cutlass.Int32(0), + ), + w_full.subview(i), + l2_cache_hint=EVICT_FIRST, + ) + prims.mbarrier_arrive_expect_tx(w_full.subview(MY_TILES + i), PE_ELEMS * ELEM_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_wp.subview(i * PE_ELEMS), + ptr_pe, + ( + cutlass.Int32(0), + row_nope + cutlass.Int32(NOPE), + k * cutlass.Int32(TMA_COPY_ITERS), + cutlass.Int32(0), + cutlass.Int32(0), + ), + w_full.subview(MY_TILES + i), + l2_cache_hint=EVICT_FIRST, + ) + if rank < cutlass.Int32(KB_RANKS): + prims.mbarrier_arrive_expect_tx( + w_full.subview(2 * MY_TILES), NOPE_ELEMS * ELEM_BYTES + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_wk, + ptr_kb, + ( + cutlass.Int32(0), + head * cutlass.Int32(LATENT) + rank * cutlass.Int32(128), + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + ), + w_full.subview(2 * MY_TILES), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(trigger_early): + # Dependents may launch now; they wait for this whole grid before reading fused_q. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + elif warp_id == 1: + if cutlass.const_expr(kv_mode != 0): + # ===================================================================== + # KV half: CTA j < the chunk's tokens, token tok0 + j's cache row (see + # kv_mode above). + # ===================================================================== + if bx < num_tokens: + tok = tok0 + bx + lane = tx % 32 + # Lane l: latent vectors l and l + 32 (8 columns each) and, for l < 8, rope vector l. + wk = [] + for j in cutlass.range_constexpr(2): + wk.append( + w_kv.load( + idx=(lane + cutlass.Int32(32 * j)) * cutlass.Int32(VEC), + vector_size=VEC, + alignment=16, + ) + ) + dst = cutlass.Int64(tok) * cutlass.Int64(FUSED) + keep = cutlass.Boolean(True) + if cutlass.const_expr(kv_mode == 1): + req = tok // tokens + pos = seq_len.load(idx=req) - tokens + (tok - req * tokens) + keep = pos >= cutlass.Int32(0) + pos_c = cutlass.Int32(cutlass.select_(keep, pos, cutlass.Int32(0))) + page = ( + page_table.load(idx=req * pt_stride + pos_c // cutlass.Int32(PAGE)) + + page_offset + ) + dst = ( + cutlass.Int64(page) * cutlass.Int64(PAGE) + + cutlass.Int64(pos_c % cutlass.Int32(PAGE)) + ) * cutlass.Int64(row_stride) + prims.griddepcontrol(prims.GridDepAction.WAIT) + src = tok * cutlass.Int32(ag_cols) + cutlass.Int32(Q_LORA) + xs = [] + for j in cutlass.range_constexpr(2): + xs.append( + ag.load( + idx=src + (lane + cutlass.Int32(32 * j)) * cutlass.Int32(VEC), + vector_size=VEC, + alignment=16, + ) + ) + pe_v = cutlass.Int32( + cutlass.select_(lane < cutlass.Int32(PE // VEC), lane, cutlass.Int32(0)) + ) + pe = ag.load( + idx=src + cutlass.Int32(LATENT) + pe_v * cutlass.Int32(VEC), + vector_size=VEC, + alignment=16, + ) + ss = cutlass.Float32(0.0) + for j in cutlass.range_constexpr(2): + for e in cutlass.range_constexpr(VEC): + xf = cutlass.Float32(xs[j][e]) + ss = ss + xf * xf + for offset in (16, 8, 4, 2, 1): + ss = ss + cute.arch.shuffle_sync_bfly(ss, offset=offset) + r = cute.math.rsqrt(ss * cutlass.Float32(1.0 / LATENT) + kv_eps, fastmath=True) + dst_arr = kv_pool if cutlass.const_expr(kv_mode == 1) else kv_out + if keep: + for j in cutlass.range_constexpr(2): + outs = [] + for e in cutlass.range_constexpr(VEC): + # flashinfer's RMSNorm order: (x * rrms) * w in fp32, one bf16 rounding. + outs.append( + (cutlass.Float32(xs[j][e]) * r * cutlass.Float32(wk[j][e])).to( + io_dtype + ) + ) + dst_arr.store( + cutlass.Vector.from_elements(tuple(outs), io_dtype), + idx=dst + + cutlass.Int64((lane + cutlass.Int32(32 * j)) * cutlass.Int32(VEC)), + vector_size=VEC, + alignment=16, + ) + if lane < cutlass.Int32(PE // VEC): + dst_arr.store( + pe, + idx=dst + + cutlass.Int64(cutlass.Int32(LATENT) + lane * cutlass.Int32(VEC)), + vector_size=VEC, + alignment=16, + ) + elif warp_id == 2: + # ===================================================================== + # MMA: q_b partials (nope M 128, pe M 64) once B is normalized; then, + # on ranks 0-3, the k_b product once q_nope has arrived. + # ===================================================================== + idesc128 = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, a_dtype=io_dtype, b_dtype=io_dtype, n_dim=mma_n, m_dim=128 + ) + idesc64 = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, a_dtype=io_dtype, b_dtype=io_dtype, n_dim=mma_n, m_dim=64 + ) + swz = prims.Tcgen05SmemSwizzle.SWIZZLE_128B + desc_wn = prims.Tcgen05SmemDesc.build( + start_address=smem_wn, + leading_byte_offset=LEADING, + stride_byte_offset=STRIDE, + layout=swz, + ) + desc_wp = prims.Tcgen05SmemDesc.build( + start_address=smem_wp, + leading_byte_offset=LEADING, + stride_byte_offset=STRIDE, + layout=swz, + ) + desc_wk = prims.Tcgen05SmemDesc.build( + start_address=smem_wk, + leading_byte_offset=LEADING, + stride_byte_offset=STRIDE, + layout=swz, + ) + desc_b = prims.Tcgen05SmemDesc.build( + start_address=smem_b, leading_byte_offset=LEADING, stride_byte_offset=STRIDE, layout=swz + ) + desc_b2 = prims.Tcgen05SmemDesc.build( + start_address=smem_b2, + leading_byte_offset=LEADING, + stride_byte_offset=STRIDE, + layout=swz, + ) + tmem_nope = cutlass.inttoptr(tmem_base + cutlass.Int32(TMEM_NOPE), 6, cutlass.Int32) + tmem_pe = cutlass.inttoptr(tmem_base + cutlass.Int32(col_pe), 6, cutlass.Int32) + tmem_kb = cutlass.inttoptr(tmem_base + cutlass.Int32(col_kb), 6, cutlass.Int32) + while not cute.arch.mbarrier_try_wait(b_ready.data_ptr(), 0): + pass + for i in cutlass.range_constexpr(MY_TILES): + while not cute.arch.mbarrier_try_wait(w_full.subview(i).data_ptr(), 0): + pass + while not cute.arch.mbarrier_try_wait(w_full.subview(MY_TILES + i).data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for kb in cutlass.range_constexpr(TMA_COPY_ITERS * K_BLOCKS_PER_HALF): + box = kb // K_BLOCKS_PER_HALF + within = kb % K_BLOCKS_PER_HALF + d_b = desc_b + (i * 2 * half_n + box * half_n + within * STEP) + acc = not (i == 0 and kb == 0) + if prims.elect_sync(): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, prims.CTAGroup.CTA_1, tmem_nope, + desc_wn + (i * TILE_128 + box * HALF_128 + within * STEP), d_b, idesc128, acc, + ) # fmt: skip + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, prims.CTAGroup.CTA_1, tmem_pe, + desc_wp + (i * TILE_64 + box * HALF_64 + within * STEP), d_b, idesc64, acc, + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + if rank < cutlass.Int32(KB_RANKS): + while not cute.arch.mbarrier_try_wait(w_full.subview(2 * MY_TILES).data_ptr(), 0): + pass + if cutlass.const_expr(single_hop): + # This rank's epilogue wrote the B tile (generic stores, each writer fenced to the async proxy). + while not cute.arch.mbarrier_try_wait(b2_ready.data_ptr(), 0): + pass + else: + # q_nope arrives from the 4 owners' st.async stores, then feeds the tensor core (async proxy). + while not _try_wait_cluster(qn_full.data_ptr(), 0): + pass + prims.fence_proxy(prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for kb in cutlass.range_constexpr(TMA_COPY_ITERS * K_BLOCKS_PER_HALF): + box = kb // K_BLOCKS_PER_HALF + within = kb % K_BLOCKS_PER_HALF + if prims.elect_sync(): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, prims.CTAGroup.CTA_1, tmem_kb, + desc_wk + (box * HALF_128 + within * STEP), desc_b2 + (box * half_n + within * STEP), + idesc128, kb != 0, + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_commit(acc2_done) + elif warp_id >= 4: + tid = tx - cutlass.Int32(EPI_THREADS) + lane = tx % 32 + w = warp_id - 4 + # ===================================================================== + # Norm prolog: the RMS of every token's q_a row (1536 columns), then this + # rank's 256 columns normalized into the resident B tile. + # ===================================================================== + last_row = num_tokens - cutlass.Int32(1) + if cutlass.const_expr(cluster_rms): + # Thread tid: token t = tid / 16 of each group, this rank's vectors c and c + 16 (c = tid % 16; columns + # 256 rank + 8 c). Each rank reads only its own 256 columns; the per-token partial sums of squares go to + # every rank of the cluster (st.async into slot [rank][t], completing norm_full by bytes), which adds the + # 6 in rank order: the same total, hence the same rrms, on every rank. + cr_t_me = tid // cutlass.Int32(16) + cr_c_me = tid % cutlass.Int32(16) + cr_wk = [] + for cr_j in cutlass.range_constexpr(2): + cr_wk.append( + w_qa.load( + idx=(rank * cutlass.Int32(32) + cr_c_me + cutlass.Int32(16 * cr_j)) + * cutlass.Int32(VEC), + vector_size=VEC, + alignment=16, + ) + ) + prims.griddepcontrol(prims.GridDepAction.WAIT) + cr_t = [] + cr_live_t = [] + cr_row_t = [] + for cr_g in cutlass.range_constexpr(groups): + cr_t.append(_plus(cr_t_me, MMA_N * cr_g)) + cr_live_t.append(cr_t[cr_g] < num_tokens) + cr_row_t.append( + cutlass.Int32(cutlass.select_(cr_live_t[cr_g], cr_t[cr_g], last_row)) + ) + cr_xk = [] + cr_part = [] + for cr_g in cutlass.range_constexpr(groups): + cr_part_g = cutlass.Float32(0.0) + for cr_j in cutlass.range_constexpr(2): + cr_xv = ag.load( + idx=(tok0 + cr_row_t[cr_g]) * cutlass.Int32(ag_cols) + + (rank * cutlass.Int32(32) + cr_c_me + cutlass.Int32(16 * cr_j)) + * cutlass.Int32(VEC), + vector_size=VEC, + alignment=16, + ) + cr_xk.append(cr_xv) + for cr_e in cutlass.range_constexpr(VEC): + cr_xf = cutlass.Float32( + cutlass.select_( + cr_live_t[cr_g], cutlass.Float32(cr_xv[cr_e]), cutlass.Float32(0.0) + ) + ) + cr_part_g = cr_part_g + cr_xf * cr_xf + cr_part.append(cr_part_g) + for cr_g in cutlass.range_constexpr(groups): + for cr_off in (8, 4, 2, 1): + cr_part[cr_g] = cr_part[cr_g] + cute.arch.shuffle_sync_bfly( + cr_part[cr_g], offset=cr_off + ) + if cr_c_me == cutlass.Int32(0): + for cr_g in cutlass.range_constexpr(groups): + for cr_dst in cutlass.range_constexpr(CLUSTER): + _st_async_f32( + _mapa_u32(norm_mail.data_ptr(rank * cutlass.Int32(mma_n) + cr_t[cr_g]), cr_dst), + cr_part[cr_g], _mapa_u32(norm_full.data_ptr(), cr_dst), + ) # fmt: skip + while not _try_wait_cluster(norm_full.data_ptr(), 0): + pass + for cr_g in cutlass.range_constexpr(groups): + cr_total = norm_mail.load(idx=cr_t[cr_g]) + for cr_src in cutlass.range_constexpr(1, CLUSTER): + cr_total = cr_total + norm_mail.load( + idx=cutlass.Int32(cr_src * mma_n) + cr_t[cr_g] + ) + cr_r = cutlass.Float32( + cutlass.select_( + cr_live_t[cr_g], + cute.math.rsqrt( + cr_total * cutlass.Float32(1.0 / Q_LORA) + eps, fastmath=True + ), + cutlass.Float32(0.0), + ) + ) + for cr_j in cutlass.range_constexpr(2): + cr_outs = [] + for cr_e in cutlass.range_constexpr(VEC): + # flashinfer's RMSNorm order: (x * rrms) * w in fp32, one bf16 rounding. + cr_outs.append( + ( + cutlass.Float32(cr_xk[2 * cr_g + cr_j][cr_e]) + * cr_r + * cutlass.Float32(cr_wk[cr_j][cr_e]) + ).to(io_dtype) + ) + cr_k_local = (cr_c_me + cutlass.Int32(16 * cr_j)) * cutlass.Int32(VEC) + smem_b.store( + cutlass.Vector.from_elements(tuple(cr_outs), io_dtype), + idx=(cr_k_local // cutlass.Int32(CTA_K)) * cutlass.Int32(b_elems) + + _swz(cr_t[cr_g], cr_k_local % cutlass.Int32(CTA_K), mma_n), + vector_size=VEC, + alignment=16, + ) + else: + # RMS pass: thread tid loads 16-byte vector v = tid + 128 j (j = 1 only for tid < 64) of every token's + # 1536-column q_a row. This rank's own 32 vectors (v = 32 rank + lane) sit in warp rank % 4, pass rank // 4: + # that warp normalizes them from registers, so q_a is read once. + owner = w == rank % cutlass.Int32(4) + own_j = rank // cutlass.Int32(4) + # The norm weight of the owner lane's 8 columns is a weight: loaded ahead of the grid dependency. + wv = w_qa.load( + idx=(rank * cutlass.Int32(32) + lane) * cutlass.Int32(VEC), + vector_size=VEC, + alignment=16, + ) + prims.griddepcontrol(prims.GridDepAction.WAIT) + # Rows past M and vectors past the row are loaded clamped and count zero (no loads out of bounds, no + # values escaping runtime branches). + sums = [cutlass.Float32(0.0)] * MMA_N + xs = [] + for t in cutlass.range_constexpr(MMA_N): + row_t = cutlass.Int32( + cutlass.select_(cutlass.Int32(t) < num_tokens, cutlass.Int32(t), last_row) + ) + for j in cutlass.range_constexpr((Q_LORA // VEC + EPI_THREADS - 1) // EPI_THREADS): + v = tid + cutlass.Int32(j * EPI_THREADS) + live = (cutlass.Int32(t) < num_tokens) & (v < cutlass.Int32(Q_LORA // VEC)) + v_c = cutlass.Int32( + cutlass.select_(v < cutlass.Int32(Q_LORA // VEC), v, cutlass.Int32(0)) + ) + xv = ag.load( + idx=(tok0 + row_t) * cutlass.Int32(ag_cols) + v_c * cutlass.Int32(VEC), + vector_size=VEC, + alignment=16, + ) + xs.append(xv) + for e in cutlass.range_constexpr(VEC): + xf = cutlass.Float32( + cutlass.select_(live, cutlass.Float32(xv[e]), cutlass.Float32(0.0)) + ) + sums[t] = sums[t] + xf * xf + for t in cutlass.range_constexpr(MMA_N): + for offset in (16, 8, 4, 2, 1): + sums[t] = sums[t] + cute.arch.shuffle_sync_bfly(sums[t], offset=offset) + if lane == 0: + for t in cutlass.range_constexpr(MMA_N): + norm_part.store(sums[t], idx=w * cutlass.Int32(MMA_N) + cutlass.Int32(t)) + prims.barrier_cta_sync(1, thread_count=EPI_THREADS) + if tid < cutlass.Int32(MMA_N): + total = norm_part.load(idx=tid) + for ww in cutlass.range_constexpr(1, 4): + total = total + norm_part.load(idx=cutlass.Int32(ww * MMA_N) + tid) + rrms.store( + cute.math.rsqrt(total * cutlass.Float32(1.0 / Q_LORA) + eps, fastmath=True), + idx=tid, + ) + prims.barrier_cta_sync(1, thread_count=EPI_THREADS) + # The owner warp: 8 tokens x its lane's 8 columns (CTA column 8 lane) into the B tile (rows >= M are zero). + n_pass = (Q_LORA // VEC + EPI_THREADS - 1) // EPI_THREADS + if owner: + for jo in cutlass.range_constexpr(n_pass): + if own_j == cutlass.Int32(jo): + for t in cutlass.range_constexpr(MMA_N): + xv = xs[t * n_pass + jo] + r = cutlass.Float32( + cutlass.select_( + cutlass.Int32(t) < num_tokens, + rrms.load(idx=cutlass.Int32(t)), + cutlass.Float32(0.0), + ) + ) + outs = [] + for e in cutlass.range_constexpr(VEC): + # flashinfer's RMSNorm order: (x * rrms) * w in fp32, one bf16 rounding. + outs.append( + (cutlass.Float32(xv[e]) * r * cutlass.Float32(wv[e])).to( + io_dtype + ) + ) + k_local = lane * cutlass.Int32(VEC) + smem_b.store( + cutlass.Vector.from_elements(tuple(outs), io_dtype), + idx=(k_local // cutlass.Int32(CTA_K)) * cutlass.Int32(b_elems) + + _swz(cutlass.Int32(t), k_local % cutlass.Int32(CTA_K)), + vector_size=VEC, + alignment=16, + ) + prims.fence_proxy(prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) + prims.mbarrier_arrive(b_ready) + + # ===================================================================== + # q_b partials -> row owners (DSMEM), rank-order sum, bf16. + # ===================================================================== + while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc_n = prims.tcgen05_ld( + "32x32b", + cutlass.inttoptr(tmem_base + cutlass.Int32(TMEM_NOPE), 6, cutlass.Float32), + num=mma_n, + ) + acc_p = prims.tcgen05_ld( + "32x32b", + cutlass.inttoptr(tmem_base + cutlass.Int32(col_pe), 6, cutlass.Float32), + num=mma_n, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + # nope rows 32 w + lane -> rank w; pe rows 16 w + lane (lane < 16) -> rank 4 + w // 2, local row + # 16 (w % 2) + lane. + pe_owner = cutlass.Int32(KB_RANKS) + w // cutlass.Int32(2) + pe_row = (w % cutlass.Int32(2)) * cutlass.Int32(16) + lane + if cutlass.const_expr(single_hop): + # Nope row 32 w + lane, 8 tokens -> slot [rank] of every k_b rank's mailbox (this rank's too). + base_kb = (rank * cutlass.Int32(128) + w * cutlass.Int32(32) + lane) * cutlass.Int32( + MMA_N + ) + for dst in cutlass.range_constexpr(KB_RANKS): + mbar = _mapa_u32(kb_full.data_ptr(), dst) + for h4 in cutlass.range_constexpr(MMA_N // 4): + _st_async_v4( + _mapa_u32(kb_mail.data_ptr(base_kb + cutlass.Int32(4 * h4)), dst), + acc_n[4 * h4], acc_n[4 * h4 + 1], acc_n[4 * h4 + 2], acc_n[4 * h4 + 3], mbar, + ) # fmt: skip + else: + if w != rank: + # st.async into the owner's mailbox [source][token][row]: a warp's 32 lanes fill 128 contiguous bytes + # per token; the owner's barrier completes when all 5 sources' bytes have landed. + mb_n = _mapa_u32(mail_full.data_ptr(), w) + for t in cutlass.range_constexpr(mma_n): + _st_async_f32( + _mapa_u32(mailbox.data_ptr((rank * cutlass.Int32(mma_n) + cutlass.Int32(t)) * cutlass.Int32(32) + + lane), w), + acc_n[t], mb_n, + ) # fmt: skip + if lane < cutlass.Int32(16): + if pe_owner != rank: + mb_p = _mapa_u32(mail_full.data_ptr(), pe_owner) + for t in cutlass.range_constexpr(mma_n): + _st_async_f32( + _mapa_u32(mailbox.data_ptr((rank * cutlass.Int32(mma_n) + cutlass.Int32(t)) * cutlass.Int32(32) + + pe_row), pe_owner), + acc_p[t], mb_p, + ) # fmt: skip + if cutlass.const_expr(single_hop): + if rank < cutlass.Int32(KB_RANKS): + # Nope row tid: the 6 partials in rank order (the owner's sum of the two-hop path), bf16, into the + # k_b B tile (token t, column tid), fenced to the async proxy for the MMA warp. + while not _try_wait_cluster(kb_full.data_ptr(), 0): + pass + parts = [] + for src in cutlass.range_constexpr(CLUSTER): + for h4 in cutlass.range_constexpr(MMA_N // 4): + parts.append( + kb_mail.load( + idx=(cutlass.Int32(src * 128) + tid) * cutlass.Int32(MMA_N) + + cutlass.Int32(4 * h4), + vector_size=4, + alignment=16, + ) + ) + for t in cutlass.range_constexpr(MMA_N): + total = cutlass.Float32(0.0) + for src in cutlass.range_constexpr(CLUSTER): + total = total + cutlass.Float32(parts[src * (MMA_N // 4) + t // 4][t % 4]) + smem_b2.store( + total.to(io_dtype), + idx=_swz(cutlass.Int32(t), tid) + tid % cutlass.Int32(VEC), + ) + prims.fence_proxy(prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) + prims.mbarrier_arrive(b2_ready) + elif cutlass.const_expr(mma_n == MMA_N): + if w == rank: + # Rank w < 4 owns nope rows 32 w + lane: sum the 6 partials, round, stage for the k_b broadcast. + while not _try_wait_cluster(mail_full.data_ptr(), 0): + pass + for t in cutlass.range_constexpr(MMA_N): + total = cutlass.Float32(0.0) + for src in cutlass.range_constexpr(CLUSTER): + part = mailbox.load(idx=cutlass.Int32((src * MMA_N + t) * 32) + lane) + total = total + cutlass.Float32( + cutlass.select_( + rank == cutlass.Int32(src), cutlass.Float32(acc_n[t]), part + ) + ) + stage_n.store(total.to(io_dtype), idx=cutlass.Int32(t * 32) + lane) + prims.bar_warp_sync(0xFFFFFFFF) + # Lane l: token l // 4, 8 of the owner's 32 columns (16 bytes) -> st.async into the k_b B tile of ranks + # 0-3 (this one too), completing their q_nope barrier. + ct = lane // cutlass.Int32(4) + c4 = lane % cutlass.Int32(4) + stage_w = cutlass.Array( + cutlass.inttoptr(stage_n.data_ptr().toint(), 3, cutlass.Int32), + shape=MMA_N * 32 // 2, + ) + words = stage_w.load( + idx=(ct * cutlass.Int32(32) + c4 * cutlass.Int32(VEC)) // cutlass.Int32(2), + vector_size=4, + alignment=16, + ) + dst_idx = _swz(ct, rank * cutlass.Int32(32) + c4 * cutlass.Int32(VEC)) + for dst in cutlass.range_constexpr(KB_RANKS): + _st_async_v4_b32( + _mapa_u32(smem_b2.data_ptr(dst_idx), dst), words[0], words[1], words[2], words[3], + _mapa_u32(qn_full.data_ptr(), dst), + ) # fmt: skip + else: + if rank < cutlass.Int32(KB_RANKS): + # Rank r < 4 owns nope rows 32 r + lane. Its warp r puts its own partials into the mailbox's slot [r]; + # then epilogue warp v sums the 6 partials of tokens 8 v .. 8 v + 7 in rank order, rounds, stages and + # sends them as above (each group's sums and transfers are an 8-token chunk's). + if w == rank: + for t in cutlass.range_constexpr(mma_n): + mailbox.store( + cutlass.Float32(acc_n[t]), + idx=(rank * cutlass.Int32(mma_n) + cutlass.Int32(t)) * cutlass.Int32(32) + + lane, + ) + prims.barrier_cta_sync(1, thread_count=EPI_THREADS) + while not _try_wait_cluster(mail_full.data_ptr(), 0): + pass + if w < cutlass.Int32(groups): + for t in cutlass.range_constexpr(MMA_N): + g_tok = w * cutlass.Int32(MMA_N) + cutlass.Int32(t) + g_total = cutlass.Float32(0.0) + for src in cutlass.range_constexpr(CLUSTER): + g_total = g_total + cutlass.Float32( + mailbox.load( + idx=(cutlass.Int32(src * mma_n) + g_tok) * cutlass.Int32(32) + + lane + ) + ) + stage_n.store(g_total.to(io_dtype), idx=g_tok * cutlass.Int32(32) + lane) + prims.bar_warp_sync(0xFFFFFFFF) + g_ct = w * cutlass.Int32(MMA_N) + lane // cutlass.Int32(4) + g_c4 = lane % cutlass.Int32(4) + g_stage = cutlass.Array( + cutlass.inttoptr(stage_n.data_ptr().toint(), 3, cutlass.Int32), + shape=mma_n * 32 // 2, + ) + g_words = g_stage.load( + idx=(g_ct * cutlass.Int32(32) + g_c4 * cutlass.Int32(VEC)) + // cutlass.Int32(2), + vector_size=4, + alignment=16, + ) + g_dst = _swz(g_ct, rank * cutlass.Int32(32) + g_c4 * cutlass.Int32(VEC), mma_n) + for dst in cutlass.range_constexpr(KB_RANKS): + _st_async_v4_b32( + _mapa_u32(smem_b2.data_ptr(g_dst), dst), g_words[0], g_words[1], g_words[2], g_words[3], + _mapa_u32(qn_full.data_ptr(), dst), + ) # fmt: skip + if rank >= cutlass.Int32(KB_RANKS): + if pe_owner == rank: + if lane < cutlass.Int32(16): + # Rank 4 + j owns pe rows 32 j + pe_row: sum, round, store fused_q[t, head, 512 + row]. + while not _try_wait_cluster(mail_full.data_ptr(), 0): + pass + p = (rank - cutlass.Int32(KB_RANKS)) * cutlass.Int32(32) + pe_row + for t in cutlass.range_constexpr(mma_n): + total = cutlass.Float32(0.0) + for src in cutlass.range_constexpr(CLUSTER): + part = mailbox.load(idx=cutlass.Int32((src * mma_n + t) * 32) + pe_row) + total = total + cutlass.Float32( + cutlass.select_( + rank == cutlass.Int32(src), cutlass.Float32(acc_p[t]), part + ) + ) + if cutlass.Int32(t) < num_tokens: + fused_q.store( + total.to(io_dtype), + idx=(tok0 + cutlass.Int32(t)) * cutlass.Int32(total_heads * FUSED) + + head * cutlass.Int32(FUSED) + + cutlass.Int32(LATENT) + + p, + ) + # ===================================================================== + # k_b product (ranks 0-3): fused_q[t, head, 128 rank + i]. + # ===================================================================== + if rank < cutlass.Int32(KB_RANKS): + while not cute.arch.mbarrier_try_wait(acc2_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc_k = prims.tcgen05_ld( + "32x32b", + cutlass.inttoptr(tmem_base + cutlass.Int32(col_kb), 6, cutlass.Float32), + num=mma_n, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + row = w * cutlass.Int32(32) + lane + for t in cutlass.range_constexpr(mma_n): + stage_o.store( + cutlass.Float32(acc_k[t]).to(io_dtype), idx=cutlass.Int32(t * 128) + row + ) + # Every TMEM reader has waited for its loads; the staging tile is complete. + prims.barrier_cta_sync(1, thread_count=EPI_THREADS) + if warp_id == 4: + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_base, 6, cutlass.Int32), tmem_cols(mma_n)) + if rank < cutlass.Int32(KB_RANKS): + o_tok = tid // cutlass.Int32(16) + oc = tid % cutlass.Int32(16) + for g in cutlass.range_constexpr(groups): + o_tok_g = _plus(o_tok, MMA_N * g) + if o_tok_g < num_tokens: + out_vec = stage_o.load( + idx=o_tok_g * cutlass.Int32(128) + oc * cutlass.Int32(VEC), + vector_size=VEC, + alignment=16, + ) + fused_q.store( + out_vec, + idx=(tok0 + o_tok_g) * cutlass.Int32(total_heads * FUSED) + + head * cutlass.Int32(FUSED) + + rank * cutlass.Int32(128) + + oc * cutlass.Int32(VEC), + vector_size=VEC, + alignment=16, + ) + + +def _weight_map(w, n_rows, k_in, box_rows): + """W [n_rows, k_in] as five TMA dimensions (64-column chunk, row, chunk index, 1, 1): one call per 128-column + k-tile lands both 128-byte-swizzled halves of `box_rows` rows.""" + return cuda.create_tensor_map_tiled( + global_address=w.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[TMA_K_BOX, n_rows, k_in // TMA_K_BOX, 1, 1], + global_strides=[ + (k_in * ELEM_BYTES) // 16, + (TMA_K_BOX * ELEM_BYTES) // 16, + (n_rows * k_in * ELEM_BYTES) // 16, + (n_rows * k_in * ELEM_BYTES) // 16, + ], + box_dims=[TMA_K_BOX, box_rows, TMA_COPY_ITERS, 1, 1], + swizzle=cuda.TensorMapSwizzle.s128b, + ) + + +@cute.jit +def k3_mla_q( + w_qb: cute.Tensor, # [heads * 192, 1536] bf16 + w_kb: cute.Tensor, # [heads * 512, 128] bf16 (k_b_proj_trans rows) + ag: cute.Tensor, # [M * ag_cols] bf16: q_a in columns [0, 1536) of each row + w_qa: cute.Tensor, # [1536] bf16 + fused_q: cute.Tensor, # [M * heads * 576] bf16 + w_kv: cute.Tensor, # [512] bf16 (kv_mode) + kv_pool: cute.Tensor, # flat bf16 latent pool (kv_mode 1) + page_table: cute.Tensor, # int32, request i's pages at [i * pt_stride, ...) (kv_mode 1) + seq_len: cute.Tensor, # int32 [R] (kv_mode 1) + kv_out: cute.Tensor, # [M * 576] bf16 (kv_mode 2) + num_tokens: cutlass.Int32, # M + tokens: cutlass.Int32, # T = M / R (kv_mode 1) + pt_stride: cutlass.Int32, # (kv_mode 1) + eps: cutlass.Float32, + kv_eps: cutlass.Float32, + page_offset: cutlass.Int32, + ag_cols: cutlass.Constexpr[int], + total_heads: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + single_hop: cutlass.Constexpr[bool], + kv_mode: cutlass.Constexpr[int], + row_stride: cutlass.Constexpr[int], + cluster_rms: cutlass.Constexpr[bool], + use_pdl: cutlass.Constexpr[bool], + mma_n: cutlass.Constexpr[ + int + ], # tokens per chunk: MMA_N, or one of WIDE_CHUNKS with cluster_rms, two-hop + stream: cuda_driver.CUstream, +) -> None: + if cutlass.const_expr( + mma_n != MMA_N and (mma_n not in WIDE_CHUNKS or single_hop or not cluster_rms) + ): + raise ValueError( + f"k3_mla_q: chunks of {mma_n} tokens need one of {WIDE_CHUNKS}, cluster_rms and two-hop" + ) + tma_nope = _weight_map(w_qb, total_heads * QK, Q_LORA, 128) + tma_pe = _weight_map(w_qb, total_heads * QK, Q_LORA, 64) + tma_kb = _weight_map(w_kb, total_heads * LATENT, NOPE, 128) + k3_mla_q_kernel( + tma_nope, + tma_pe, + tma_kb, + ag, + w_qa, + fused_q, + w_kv, + kv_pool, + page_table, + seq_len, + kv_out, + num_tokens, + tokens, + pt_stride, + eps, + kv_eps, + page_offset, + ag_cols, + total_heads, + trigger_early, + single_hop, + kv_mode, + row_stride, + cluster_rms, + mma_n, + ).launch( # fmt: skip + grid=(total_heads * CLUSTER, (num_tokens + mma_n - 1) // mma_n, 1), + block=(THREADS, 1, 1), + cluster=(CLUSTER, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_mla/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_mla/op.py new file mode 100644 index 000000000000..0c3269ba5118 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_mla/op.py @@ -0,0 +1,612 @@ +# 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. +"""Torch ops of the Kimi K3 MLA decode CTM kernels. + +A decode step is M = R x T tokens: R <= 8 requests of T <= 8 tokens each, request-major (the rows of request i are +i T .. i T + T - 1). Request i's KV pages are row i of ``page_table`` (int32 [R, W] with unit column stride and any row +stride, e.g. a slice of the attention metadata's kv_cache_block_offsets; [W] for one request) and its KV length, +including its T new tokens, is ``seq_len[i]`` (int32 [R]). + +``trtllm::k3_mla_q``: the decode query path (M <= 64; 6 heads per rank at TP16, 24 at TP4): q_a RMSNorm, q_b projection +and k_b absorption in one launch, producing the attention's ``fused_q`` [M, heads * 576] = [q_nope @ W_kb^T | q_pe] per +head; ``trtllm::k3_mla_qkv`` also stores the KV half (kv_a RMSNorm, rope columns) into the paged latent cache. +``trtllm::k3_mla_attn`` and its ``_out`` / ``_vb_out`` forms: the attention of every request over its pages, over a +caller-owned workspace (:func:`make_attn_workspace`). Compiled on the first call for its shapes, which must happen +outside CUDA-graph capture. +""" + +from __future__ import annotations + +import os +import threading +from typing import Dict, Optional + +import torch + +MAX_TOKENS = 64 # tokens of a k3_mla_q / k3_mla_qkv call +MAX_REQUESTS = 8 +MAX_REQUEST_TOKENS = 8 # T +# k3_mla_q's tokens per chunk (the MMA's N) by call size, as (largest call, chunk): a chunk is a cluster of 6 CTAs per +# head, 23 such clusters are resident at once on GB200, and a wider chunk takes longer per CTA. 8 up to 24 tokens (18 +# clusters), 16 up to 48 (18), 32 up to 64 (12). +CHUNK_TOKENS = ((24, 8), (48, 16), (64, 32)) + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + + +def _arg(t: torch.Tensor, align: int = 16): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +def _use_pdl() -> bool: + return os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +_arg_dummies: Dict[tuple, torch.Tensor] = {} + + +def _dummy(device: torch.device, dtype: torch.dtype) -> torch.Tensor: + """A tensor argument the build does not touch (the kv arguments when the KV half is off).""" + key = (device.index, dtype) + t = _arg_dummies.get(key) + if t is None: + t = _arg_dummies[key] = torch.zeros(8, dtype=dtype, device=device) + return t + + +def _pool_base(pool: torch.Tensor, row_stride: int) -> torch.Tensor: + """The pool's first page as the kernels' pointer carrier: they index the pool with 64-bit offsets (or through a + tensor map sized separately), and a flat view of a multi-GB pool would not fit a 32-bit dynamic shape.""" + return pool.view(-1)[: 64 * row_stride] + + +def _requests(page_table: torch.Tensor, seq_len: torch.Tensor, num_tokens: int): + """``(rows, row_stride, R, T)`` of a step of ``num_tokens`` tokens over R = ``seq_len.numel()`` requests: the + page-table rows as one flat int32 view (request i's row at ``i * row_stride``), or None when the arguments do not + describe R requests of T = num_tokens / R tokens (see the module docstring).""" + num_requests = seq_len.numel() + if not ( + page_table.dtype == seq_len.dtype == torch.int32 + and page_table.device == seq_len.device + and seq_len.dim() == 1 + and seq_len.is_contiguous() + and 0 < num_requests <= num_tokens + and num_tokens % num_requests == 0 + and page_table.dim() in (1, 2) + and page_table.stride(-1) == 1 + and page_table.shape[-1] > 0 + ): + return None + if page_table.dim() == 1: + return (page_table, 0, 1, num_tokens) if num_requests == 1 else None + width = page_table.shape[1] + row_stride = page_table.stride(0) if num_requests > 1 else width + if page_table.shape[0] != num_requests or row_stride < width: + return None + rows = page_table.as_strided(((num_requests - 1) * row_stride + width,), (1,)) + return rows, row_stride, num_requests, num_tokens // num_requests + + +def supports_kv(ag: torch.Tensor, w_kv: torch.Tensor, pool: torch.Tensor, row_stride: int, page_table: torch.Tensor, + seq_len: torch.Tensor) -> bool: # fmt: skip + """Whether ``k3_mla_qkv`` can store the KV half of the ``ag.shape[0]`` tokens: the latent (512) and rope (64) + columns after q_a in ``ag``, a dense bf16 pool with rows of ``row_stride`` elements, int32 page-table rows and + lengths of R requests (see the module docstring).""" + from . import k3_mla_q_kernel as kernel + + return ( + ag.shape[1] >= kernel.Q_LORA + kernel.FUSED + and tuple(w_kv.shape) == (kernel.LATENT,) + and w_kv.dtype == torch.bfloat16 + and w_kv.is_contiguous() + and pool.dtype == torch.bfloat16 + and pool.is_contiguous() + and row_stride >= kernel.FUSED + and row_stride % 8 == 0 + and pool.numel() >= 64 * row_stride + and _requests(page_table, seq_len, ag.shape[0]) is not None + ) + + +def supports_q( + ag: torch.Tensor, w_qa: torch.Tensor, w_qb: torch.Tensor, w_kb: torch.Tensor +) -> bool: + """Whether ``k3_mla_q`` runs: q_a in the first 1536 columns of dense bf16 rows, the TP16 per-rank shapes.""" + from . import k3_mla_q_kernel as kernel + + return ( + ag.is_cuda + and ag.dtype == w_qa.dtype == w_qb.dtype == w_kb.dtype == torch.bfloat16 + and ag.dim() == 2 + and 0 < ag.shape[0] <= MAX_TOKENS + and ag.shape[1] >= kernel.Q_LORA + and ag.shape[1] % 8 == 0 + and ag.is_contiguous() + and tuple(w_qa.shape) == (kernel.Q_LORA,) + and w_kb.dim() == 3 + and tuple(w_kb.shape[1:]) == (kernel.LATENT, kernel.NOPE) + and 0 < w_kb.shape[0] * kernel.CLUSTER <= 148 + and tuple(w_qb.shape) == (w_kb.shape[0] * kernel.QK, kernel.Q_LORA) + and w_qa.is_contiguous() + and w_qb.is_contiguous() + and w_kb.is_contiguous() + ) + + +def _launch_q( + ag, w_qa, eps, w_qb, w_kb, trigger_early=True, single_hop=False, kv=None, cluster_rms=True +) -> torch.Tensor: + """``kv``: None (query path only) or dict(w, eps) plus either pool, row_stride, page_table, page_offset, seq_len + (the KV half into the paged pool) or out (into a dense [M, 576] tensor). ``cluster_rms``: each rank of a head's + cluster reads only its 256 q_a columns and the ranks exchange per-token partial sums of squares (st.async; False: + every CTA reads all 1536 columns). ``single_hop``: the single-hop k_b reduce (same bits, slower on GB200). + The call runs in chunks of CHUNK_TOKENS' size for its token count when cluster_rms and two-hop (and, with the KV + half, when the heads launch a CTA per token of a chunk), else in chunks of 8: every token's bits are the same.""" + import cuda.bindings.driver as cuda_driver + + if not supports_q(ag, w_qa, w_qb, w_kb): + raise ValueError( + f"k3_mla_q: unsupported call ag {tuple(ag.shape)} {ag.dtype}, w_qa {tuple(w_qa.shape)}, " + f"w_qb {tuple(w_qb.shape)}, w_kb {tuple(w_kb.shape)}" + ) + from . import k3_mla_q_kernel as kernel + + num_tokens, ag_cols = ag.shape + heads = w_kb.shape[0] + mma_n = next(n for limit, n in CHUNK_TOKENS if num_tokens <= limit) + if single_hop or not cluster_rms or (kv is not None and heads * kernel.CLUSTER < mma_n): + mma_n = kernel.MMA_N + out = torch.empty(num_tokens, heads * kernel.FUSED, dtype=torch.bfloat16, device=ag.device) + bf16_dummy, i32_dummy = _dummy(ag.device, torch.bfloat16), _dummy(ag.device, torch.int32) + kv_mode, row_stride, kv_eps, page_offset = 0, kernel.FUSED, 0.0, 0 + tokens, pt_stride = num_tokens, 0 + w_kv, kv_pool, page_table, seq_len, kv_out = ( + bf16_dummy, + bf16_dummy, + i32_dummy, + i32_dummy, + bf16_dummy, + ) + if kv is not None: + w_kv, kv_eps = kv["w"], float(kv["eps"]) + if heads * kernel.CLUSTER < kernel.MMA_N: + raise ValueError( + f"k3_mla_q: the KV half takes {kernel.MMA_N} CTAs per 8 tokens, {heads} heads launch " + f"{heads * kernel.CLUSTER}" + ) + if "out" in kv: + kv_mode, kv_out = 2, kv["out"] + if not ( + kv_out.is_contiguous() + and kv_out.dtype == torch.bfloat16 + and kv_out.numel() == num_tokens * kernel.FUSED + ): + raise ValueError( + f"k3_mla_q: kv out {tuple(kv_out.shape)} {kv_out.dtype} is not a dense [M, 576] bf16" + ) + else: + kv_mode, row_stride, page_offset = 1, int(kv["row_stride"]), int(kv["page_offset"]) + page_table, seq_len = kv["page_table"], kv["seq_len"] + if not supports_kv(ag, w_kv, kv["pool"], row_stride, page_table, seq_len): + raise ValueError( + f"k3_mla_q: unsupported KV half: ag {tuple(ag.shape)}, w_kv {tuple(w_kv.shape)} {w_kv.dtype}, pool " + f"{kv['pool'].dtype} rows of {row_stride}, page table {tuple(page_table.shape)} " + f"{page_table.dtype} strides {page_table.stride()}, lengths {tuple(seq_len.shape)} {seq_len.dtype}" + ) + page_table, pt_stride, _, tokens = _requests(page_table, seq_len, num_tokens) + kv_pool = _pool_base(kv["pool"], row_stride) + args = ( + _arg(w_qb), + _arg(w_kb.view(heads * kernel.LATENT, kernel.NOPE)), + _arg(ag.view(-1)), + _arg(w_qa), + _arg(out.view(-1)), + _arg(w_kv), + _arg(kv_pool), + # int32 page-table rows and lengths: views into the metadata buffers, read by scalar loads. + _arg(page_table, 4), + _arg(seq_len, 4), + _arg(kv_out.view(-1)), + ) + stream = cuda_driver.CUstream(torch.cuda.current_stream(ag.device).cuda_stream) + use_pdl = _use_pdl() + key = ( + "k3_mla_q", + ag_cols, + heads, + trigger_early, + single_hop, + kv_mode, + row_stride, + cluster_rms, + use_pdl, + mma_n, + ) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "trtllm::k3_mla_q must run once outside CUDA-graph capture first (it compiles its kernel)." + ) + import cutlass.cute as cute + + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kernel.k3_mla_q, *args, num_tokens, tokens, pt_stride, float(eps), kv_eps, page_offset, ag_cols, + heads, trigger_early, single_hop, kv_mode, row_stride, cluster_rms, use_pdl, mma_n, stream, + ) # fmt: skip + fn(*args, num_tokens, tokens, pt_stride, float(eps), kv_eps, page_offset, stream) + return out + + +@torch.library.custom_op("trtllm::k3_mla_q", mutates_args=()) +def k3_mla_q( + ag: torch.Tensor, + w_qa: torch.Tensor, + eps: float, + w_qb: torch.Tensor, + w_kb: torch.Tensor, + trigger_early: bool = True, +) -> torch.Tensor: + """``fused_q`` [M <= 64, heads * 576] bf16 from ``ag`` [M, C >= 1536] (q_a = ag[:, :1536]): per head, + ``[bf16(q_nope @ w_kb[h]^T) | q_pe]`` with ``q = bf16(rmsnorm(q_a) @ w_qb^T)``; each bf16 rounding of the + unfused RMSNorm -> q_b -> bmm chain is kept. Every 8-token chunk of rows computes as an M <= 8 call would.""" + return _launch_q(ag, w_qa, eps, w_qb, w_kb, trigger_early) + + +@k3_mla_q.register_fake +def _(ag, w_qa, eps, w_qb, w_kb, trigger_early=True): + return ag.new_empty((ag.shape[0], w_kb.shape[0] * 576), dtype=torch.bfloat16) + + +@torch.library.custom_op("trtllm::k3_mla_qkv", mutates_args=("pool",)) +def k3_mla_qkv( + ag: torch.Tensor, + w_qa: torch.Tensor, + eps: float, + w_qb: torch.Tensor, + w_kb: torch.Tensor, + w_kv: torch.Tensor, + kv_eps: float, + pool: torch.Tensor, + row_stride: int, + page_table: torch.Tensor, + page_offset: int, + seq_len: torch.Tensor, + trigger_early: bool = True, +) -> torch.Tensor: + """``k3_mla_q`` plus the KV half in the same launch: the cache row ``[bf16(rmsnorm(ag[t, 1536:2048]) * w_kv) | + ag[t, 2048:2112]]`` (K3 is NoPE) of token t = i T + u (token u of request i, see the module docstring) stored into + the paged latent ``pool`` at position ``pos = seq_len[i] - T + u``, row ``(page_table[i][pos // 64] + + page_offset) * 64 + pos % 64`` of ``row_stride`` elements (``seq_len[i] >= T``; a token with ``pos < 0`` is not + stored). The page table and lengths are read before the grid dependency wait (written before the graph runs).""" + kv = dict(w=w_kv, eps=kv_eps, pool=pool, row_stride=row_stride, page_table=page_table, page_offset=page_offset, + seq_len=seq_len) # fmt: skip + return _launch_q(ag, w_qa, eps, w_qb, w_kb, trigger_early, kv=kv) + + +@k3_mla_qkv.register_fake +def _( + ag, + w_qa, + eps, + w_qb, + w_kb, + w_kv, + kv_eps, + pool, + row_stride, + page_table, + page_offset, + seq_len, + trigger_early=True, +): + return ag.new_empty((ag.shape[0], w_kb.shape[0] * 576), dtype=torch.bfloat16) + + +@torch.library.custom_op("trtllm::k3_mla_qkv_out", mutates_args=("kv_out",)) +def k3_mla_qkv_out( + ag: torch.Tensor, + w_qa: torch.Tensor, + eps: float, + w_qb: torch.Tensor, + w_kb: torch.Tensor, + w_kv: torch.Tensor, + kv_eps: float, + kv_out: torch.Tensor, + trigger_early: bool = True, +) -> torch.Tensor: + """``k3_mla_qkv`` with the cache rows stored densely into ``kv_out`` [M, 576] instead of the pool (checks).""" + return _launch_q( + ag, w_qa, eps, w_qb, w_kb, trigger_early, kv=dict(w=w_kv, eps=kv_eps, out=kv_out) + ) + + +@k3_mla_qkv_out.register_fake +def _(ag, w_qa, eps, w_qb, w_kb, w_kv, kv_eps, kv_out, trigger_early=True): + return ag.new_empty((ag.shape[0], w_kb.shape[0] * 576), dtype=torch.bfloat16) + + +# --------------------------------------------------------------------------------------------------------------- +# trtllm::k3_mla_attn: decode attention over the paged latent cache (R <= 8 requests of T <= 8 tokens, one cluster of +# 16 CTAs per request and 6 heads) +# --------------------------------------------------------------------------------------------------------------- +def attn_workspace_elems(groups: int) -> int: + """fp16 elements of a ``k3_mla_attn`` workspace for calls of ``groups`` head groups (heads / 6): the per-CTA + partials of MAX_REQUESTS x groups x 16 slots, then the no_cluster mode's (m, l) exchange and arrival counters.""" + from . import k3_mla_attn_kernel as kernel + + return ( + kernel.MAX_REQUESTS * groups * kernel.CLUSTER * kernel.WS_SLOT_ELEMS + + kernel.ws_sync_elems(groups) + ) + + +def make_attn_workspace(device: torch.device, groups: int) -> torch.Tensor: + """A new workspace for the ``k3_mla_attn`` calls of ``groups`` head groups on ``device``: fp16 + [attn_workspace_elems(groups)], the partial slots uninitialized (a call reads only the words it wrote) and the + no_cluster tail zeroed (the arrival counters start at 0). It allocates, so it refuses to run under CUDA-graph + capture; the zeroing is ordered on the device's current stream.""" + from . import k3_mla_attn_kernel as kernel + + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("k3_mla_attn workspaces allocate: make them outside CUDA-graph capture.") + if groups < 1: + raise ValueError(f"k3_mla_attn workspace: {groups} head groups") + ws = torch.empty(attn_workspace_elems(groups), dtype=torch.float16, device=device) + ws[kernel.MAX_REQUESTS * groups * kernel.CLUSTER * kernel.WS_SLOT_ELEMS :].zero_() + return ws + + +def _check_attn_workspace(workspace: torch.Tensor, device: torch.device, groups: int) -> None: + elems = attn_workspace_elems(groups) + if not ( + workspace.dtype == torch.float16 + and workspace.dim() == 1 + and workspace.is_contiguous() + and workspace.numel() == elems + and workspace.device == device + ): + raise ValueError( + f"k3_mla_attn: workspace {tuple(workspace.shape)} {workspace.dtype} on {workspace.device} is not one for " + f"{groups} head group(s) on {device}: fp16 [{elems}] (make_attn_workspace)" + ) + + +def supports_attn( + q: torch.Tensor, + pool: torch.Tensor, + row_stride: int, + page_table: torch.Tensor, + seq_len: torch.Tensor, +) -> bool: + """Whether ``k3_mla_attn`` takes the call: ``q`` [M = R T, heads * 576] bf16 with heads a multiple of 6, R <= 8 + requests of T <= 8 tokens (page-table rows and lengths as in the module docstring), a dense bf16 pool with rows of + ``row_stride`` elements.""" + from . import k3_mla_attn_kernel as kernel + + if not ( + q.is_cuda + and q.dtype == pool.dtype == torch.bfloat16 + and q.dim() == 2 + and q.shape[1] % (kernel.HEADS * kernel.QK) == 0 + and q.is_contiguous() + and pool.is_contiguous() + and row_stride >= kernel.QK + and row_stride % 8 == 0 + ): + return False + requests = _requests(page_table, seq_len, q.shape[0]) + return ( + requests is not None + and requests[2] <= kernel.MAX_REQUESTS + and requests[3] <= kernel.MAX_TOKENS + ) + + +def _launch_attn( + q, + pool, + row_stride, + page_table, + seq_len, + softmax_scale, + workspace, + out=None, + page_offset=0, + w_vb=None, + gate=None, + gate_col0=0, +): + import cuda.bindings.driver as cuda_driver + + from . import k3_mla_attn_kernel as kernel + + if not supports_attn(q, pool, row_stride, page_table, seq_len): + raise ValueError( + f"k3_mla_attn: unsupported call q {tuple(q.shape)} {q.dtype}, row_stride {row_stride}, page table " + f"{tuple(page_table.shape)} {page_table.dtype} strides {page_table.stride()}, lengths " + f"{tuple(seq_len.shape)} {seq_len.dtype}" + ) + num_tokens = q.shape[0] + page_rows, pt_stride, num_requests, tokens = _requests(page_table, seq_len, num_tokens) + total_heads = q.shape[1] // kernel.QK + fuse_vb = w_vb is not None + if fuse_vb and not ( + w_vb.dtype == torch.bfloat16 + and w_vb.is_contiguous() + and tuple(w_vb.shape) == (total_heads, kernel.V_DIM, kernel.LATENT) + ): + raise ValueError( + f"k3_mla_attn: v_b weight {tuple(w_vb.shape)} {w_vb.dtype} is not a dense [heads, 128, 512] bf16" + ) + width = kernel.V_DIM if fuse_vb else kernel.LATENT + apply_gate = gate is not None + if apply_gate and not ( + fuse_vb + and gate.dtype == torch.bfloat16 + and gate.dim() == 2 + and gate.shape[0] == num_tokens + and gate.stride(1) == 1 + and 0 <= gate_col0 + and gate_col0 + total_heads * kernel.V_DIM <= gate.shape[1] + ): + raise ValueError( + f"k3_mla_attn: gate {tuple(gate.shape)} {gate.dtype} col0 {gate_col0} does not fit the v_b output" + ) + groups = total_heads // kernel.HEADS + _check_attn_workspace(workspace, q.device, groups) + # More 16-CTA clusters than co-reside would run a second wave: launch without a cluster instead when every CTA fits + # on the SMs at once (one per SM; its waits are spins). + clusters = num_requests * groups + no_cluster = ( + clusters > kernel.CLUSTER_WAVE + and clusters * kernel.CLUSTER + <= torch.cuda.get_device_properties(q.device).multi_processor_count + ) + if out is None: + out = torch.empty(num_tokens, total_heads * width, dtype=torch.bfloat16, device=q.device) + elif not ( + out.is_contiguous() + and out.dtype == torch.bfloat16 + and out.numel() == num_tokens * total_heads * width + ): + raise ValueError( + f"k3_mla_attn: output {tuple(out.shape)} {out.dtype} is not a dense [M, heads * {width}] bf16" + ) + total_rows = pool.numel() // row_stride + gate_flat = gate.as_strided((gate.numel(),), (1,)) if apply_gate else q.view(-1) + gate_ld = gate.stride(0) if apply_gate else 0 + args = (_arg(q.view(-1)), _arg(_pool_base(pool, row_stride)), _arg(page_rows, 4), _arg(seq_len, 4), + _arg(workspace), _arg(out.view(-1)), _arg((w_vb if fuse_vb else q).view(-1)), _arg(gate_flat)) # fmt: skip + stream = cuda_driver.CUstream(torch.cuda.current_stream(q.device).cuda_stream) + use_pdl = _use_pdl() + scale_log2 = float(softmax_scale) * kernel.LOG2E + key = ("k3_mla_attn", row_stride, total_heads, fuse_vb, apply_gate, use_pdl, no_cluster) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "trtllm::k3_mla_attn must run once outside CUDA-graph capture first (it compiles its kernel)." + ) + import cutlass.cute as cute + + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kernel.k3_mla_attn, *args, tokens, num_requests, pt_stride, scale_log2, total_rows, + int(page_offset), int(gate_col0), int(gate_ld), row_stride, total_heads, fuse_vb, apply_gate, + use_pdl, stream, no_cluster, + ) # fmt: skip + fn( + *args, + tokens, + num_requests, + pt_stride, + scale_log2, + total_rows, + int(page_offset), + int(gate_col0), + int(gate_ld), + stream, + ) + return out + + +@torch.library.custom_op("trtllm::k3_mla_attn", mutates_args=("workspace",)) +def k3_mla_attn( + q: torch.Tensor, + pool: torch.Tensor, + row_stride: int, + page_table: torch.Tensor, + seq_len: torch.Tensor, + softmax_scale: float, + workspace: torch.Tensor, +) -> torch.Tensor: + """MLA decode attention of R <= 8 requests of T <= 8 tokens: ``q`` [M = R T, heads * 576] (``fused_q``, heads a + multiple of 6, request-major) against the paged latent cache ``pool`` (flat bf16; row i of page p at ``(p * 64 + + i) * row_stride``, 512 latent then 64 rope columns), request i's pages ``page_table[i]`` and length ``seq_len[i]`` + = L_i (rows including its T new ones; see the module docstring), causal bottom-right (token t of request i sees + rows <= L_i - T + t). Returns ``[M, heads * 512]`` bf16. The page table and lengths are read before the grid + dependency wait (they must be written before the CUDA graph runs); q and the pages holding rows >= L_i - T after + it. Request i's rows are computed as the R = 1 call on its own rows, pages and length would compute them. + ``workspace``: from :func:`make_attn_workspace` for this device and heads / 6 head groups. A call writes and reads + its partials there, all after its grid dependency wait, and in the no_cluster mode (more than CLUSTER_WAVE + clusters, all of whose CTAs fit on the SMs) adds 16 to the arrival counters of its requests' head groups; calls + on one workspace must run one at a time.""" + return _launch_attn(q, pool, row_stride, page_table, seq_len, softmax_scale, workspace) + + +@torch.library.custom_op("trtllm::k3_mla_attn_out", mutates_args=("out", "workspace")) +def k3_mla_attn_out( + q: torch.Tensor, + pool: torch.Tensor, + row_stride: int, + page_table: torch.Tensor, + page_offset: int, + seq_len: torch.Tensor, + softmax_scale: float, + out: torch.Tensor, + workspace: torch.Tensor, +) -> None: + """``k3_mla_attn`` into ``out`` (dense [M, heads * 512] bf16), with ``page_offset`` added to every page-table + entry (the layer's slot in a layer-interleaved pool).""" + _launch_attn( + q, + pool, + row_stride, + page_table, + seq_len, + softmax_scale, + workspace, + out=out, + page_offset=page_offset, + ) + + +@torch.library.custom_op("trtllm::k3_mla_attn_vb_out", mutates_args=("out", "workspace")) +def k3_mla_attn_vb_out( + q: torch.Tensor, + pool: torch.Tensor, + row_stride: int, + page_table: torch.Tensor, + page_offset: int, + seq_len: torch.Tensor, + softmax_scale: float, + w_vb: torch.Tensor, + out: torch.Tensor, + workspace: torch.Tensor, + gate: Optional[torch.Tensor] = None, + gate_col0: int = 0, +) -> None: + """``k3_mla_attn`` with v_b applied in the same launch: ``out`` [M, heads * 128] = per head + ``bf16(bf16(o) @ w_vb[h]^T)`` for the attention output o, ``w_vb`` = v_b_proj [heads, 128, 512] bf16. With ``gate`` + (bf16 [M, C], sigmoid of head h's gate at columns ``gate_col0 + 128 h``) the output is ``bf16(y * s)``, the + unfused output gate.""" + _launch_attn( + q, pool, row_stride, page_table, seq_len, softmax_scale, workspace, out=out, page_offset=page_offset, + w_vb=w_vb, gate=gate, gate_col0=gate_col0, + ) # fmt: skip + + +@k3_mla_attn.register_fake +def _(q, pool, row_stride, page_table, seq_len, softmax_scale, workspace): + return q.new_empty((q.shape[0], q.shape[1] // 576 * 512), dtype=torch.bfloat16) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index f830d3c23191..2028050aa034 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -3048,6 +3048,14 @@ def _create_kv_cache_manager( if is_kda_mtp_verify_available(): kda_extra_kwargs["kda_replay_num_spec"] = ( spec_config.tokens_per_gen_step - 1) + # A model whose KDA verify starts from the state after the + # accepted drafts (instead of replaying them) asks for the + # state after every draft with `kda_token_states`. + backbone = getattr(getattr(model_engine, "model", None), + "model", None) + if (issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2) + and getattr(backbone, "kda_token_states", False)): + kda_extra_kwargs["kda_token_states"] = True if is_glm5_next: # The manager places an indexer buffer on every attention layer. from ..attention.backends.sparse.glm_kpool import \ diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index ad417618950b..6b60514f8dfa 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -330,6 +330,14 @@ def use_kda_replay_update(self) -> bool: """Whether KDA fused verification owns per-slot replay caches.""" return getattr(self, "_use_kda_replay_update", False) + @property + def keeps_kda_token_states(self) -> bool: + """Whether the KDA replay caches come with the state after every + draft of the last verify round (``kda_state_tok``), so that a verify + kernel can start the next round from the state after the accepted + drafts instead of replaying them.""" + return False + @abstractmethod def get_conv_states(self, layer_idx: int) -> torch.Tensor: """Return conv states for specific layer. @@ -438,6 +446,14 @@ class SpeculativeState(State): kda_qkg_cache: torch.Tensor | None = None kda_v_cache: torch.Tensor | None = None kda_beta_cache: torch.Tensor | None = None + # Optional: the state after every draft of the last verify round, + # [slots, num_spec, H, V, K] fp32 (the SSM pool holds the state + # after the round's first, non-draft token). A verify kernel starts + # the next round from entry n - 1 when the round accepted n > 0 + # drafts (prev_num_accepted_tokens) instead of replaying them from + # the caches above; a slot reset for a new request has n = 0, so a + # previous owner's entries are never read. + kda_state_tok: torch.Tensor | None = None # Replay path: compact double-buffered cache # prev_num_accepted_tokens: # accepted tokens (always >= 1 if drafting). @@ -3024,6 +3040,7 @@ def __init__( conv_state_layout: Literal["x_b_c", "q_k_v"] = "x_b_c", kda_replay_num_spec: Optional[int] = None, qwen4_exp_ple_cache_params: Optional["Qwen4ExpPLECacheParams"] = None, + kda_token_states: bool = False, **kwargs, ) -> None: if conv_state_layout not in ("x_b_c", "q_k_v"): @@ -3040,6 +3057,9 @@ def __init__( self._kda_replay_num_spec = kda_replay_num_spec self._use_kda_replay_update = kda_replay_num_spec is not None + # Only with the KDA replay caches: also keep the state after every + # verify token (kda_state_tok). + self._kda_token_states = kda_token_states and self._use_kda_replay_update if self._use_kda_replay_update: if use_replay_state_update: raise ValueError( @@ -3420,6 +3440,10 @@ def get_disagg_role_layouts(self) -> Dict[DataRole, RoleLayout]: def use_gdn_cached_replay_all_layer_commit(self) -> bool: return getattr(self, "_use_gdn_cached_replay_all_layer_commit", False) + @property + def keeps_kda_token_states(self) -> bool: + return getattr(self, "kda_state_tok", None) is not None + def _setup_mtp_intermediate_states(self, spec_config, max_batch_size: int) -> None: if not self.use_kda_replay_update: @@ -3450,6 +3474,7 @@ def _allocate_pool_replay_buffers( self.kda_qkg_cache = None self.kda_v_cache = None self.kda_beta_cache = None + self.kda_state_tok = None if (not self.use_kda_replay_update or self.local_num_mamba_layers == 0): return allocated @@ -3504,8 +3529,17 @@ def allocate_dim_contiguous_conv_cache() -> torch.Tensor: self.ssm_state_shape[0]), device, ) + if self._kda_token_states: + self.kda_state_tok = torch.zeros( + (self.local_num_mamba_layers, cache_size, num_spec, + *self.ssm_state_shape), + dtype=torch.float32, + device=device, + ) + per_token = (", with per-token states" + if self.kda_state_tok is not None else "") logger.info("Mamba Cache (kda-replay) is allocated for " - f"{cache_size} state slots") + f"{cache_size} state slots{per_token}") return True def _commit_gdn_cached_replay_history_layers( @@ -3801,6 +3835,12 @@ def _max_resident_sequences(self) -> int: def _mamba_state_bytes_per_slot(self) -> int: base_bytes = self.local_num_mamba_layers * (self.ssm_bytes + self.conv_bytes) + if getattr(self, "_kda_token_states", False): + # kda_state_tok: an fp32 SSM state per draft of every slot. + base_bytes += (self.local_num_mamba_layers * + self._kda_replay_num_spec * + math.prod(self.ssm_state_shape) * + torch.float32.itemsize) local_ple_layers = sum(layer_id in self.pp_layers for layer_id in self._ple_layer_ids) if local_ple_layers == 0: @@ -3945,6 +3985,32 @@ def _minimum_live_gpu_quota(self) -> int: state_quota + attention_block_quota, ) + def _ssm_pool_at_live_floor(self, kv_cache_config: KvCacheConfig) -> bool: + """Whether the SSM pool group keeps only its min-slots floor (one slot + per resident sequence plus the reserved dummies) instead of its share + of the typical step's ratio. + + With per-token KDA states, ``kda_state_tok`` takes num_spec states per + SSM slot outside the cache quota (``_allocate_pool_replay_buffers``), + addressed by the request's SSM slot. Without block reuse no SSM slot + holds anything but a live request or a reserved dummy, so slots past + the floor are never used and would only add that memory. + """ + return (getattr(self, "_kda_token_states", False) + and self.local_num_mamba_layers > 0 + and not kv_cache_config.enable_block_reuse + and self._attention_cache_bytes_per_token() > 0) + + def _quota_filling_request_capacity(self, gpu_quota: int) -> int: + """A typical request capacity whose resident requests' attention pages + alone take ``gpu_quota``. The typical step's ratio then gives the SSM + pool group fewer slots than its floor, so the min-slots constraint + sets its size and attention gets the rest of the quota.""" + bytes_per_token = (self._max_resident_sequences() * + self._attention_cache_bytes_per_token()) + # KVCacheDesc.capacity is a 32-bit int. + return min(-(-gpu_quota // bytes_per_token), 1 << 30) + def _build_cache_config( self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: kv_cache_config = self.kv_cache_config @@ -3990,6 +4056,10 @@ def _build_cache_config( if config.initial_pool_ratio is None: typical_capacity = self._get_typical_request_capacity( kv_cache_config) + if self._ssm_pool_at_live_floor(kv_cache_config): + typical_capacity = max( + typical_capacity, + self._quota_filling_request_capacity(gpu_quota)) request_descs = self._typical_request_descs(typical_capacity, kv_cache_config) typical_step = BatchDesc(request_descs * @@ -4285,16 +4355,20 @@ def _relocate_kda_replay_slots(self, old_slots: List[int], destination_slots = torch.tensor([new for _, new in moves], dtype=torch.long, device=device) - replay_buffers = ( + replay_buffers = [ self.kda_conv_q, self.kda_conv_k, self.kda_conv_v, self.kda_qkg_cache, self.kda_v_cache, self.kda_beta_cache, - ) + ] + assert all(replay_buffer is not None + for replay_buffer in replay_buffers) + # The per-token verify states belong to the slot as well. + if self.kda_state_tok is not None: + replay_buffers.append(self.kda_state_tok) for replay_buffer in replay_buffers: - assert replay_buffer is not None replay_buffer.index_copy_( 1, destination_slots, @@ -4498,6 +4572,8 @@ def mamba_layer_cache( "kda_v_cache": self.kda_v_cache[layer_offset], "kda_beta_cache": self.kda_beta_cache[layer_offset], } + if self.kda_state_tok is not None: + spec_kwargs["kda_state_tok"] = self.kda_state_tok[layer_offset] if self.mamba_ssm_rand_seed is not None: spec_kwargs["mamba_ssm_rand_seed"] = self.mamba_ssm_rand_seed return PythonMambaCacheManager.SpeculativeState( @@ -4781,4 +4857,5 @@ def shutdown(self): self.kda_qkg_cache = None self.kda_v_cache = None self.kda_beta_cache = None + self.kda_state_tok = None super().shutdown() diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 2ca039299f9c..17dfe563471a 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -135,6 +135,35 @@ l0_b200: - unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py - unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py - unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py + # Kimi K3 decode kernels (CuTe DSL): the GEMVs, the head GEMV and the + # embedding, and the cluster-wait and tcgen05-fence lints over them. + - unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_embed.py + - unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_ctm_gemv.py + - unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_ctm_gemv_wide.py + - unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_decode_gemv.py + - unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_head_gemv.py + - unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_cluster_waits.py + - unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_tcgen05_fences.py + # The Kimi K3 targets' routing and the claims check of their stock imports. l0_b300's modeling_v2 + # entry, the only other one that collects them, is waived (nvbugs/6853741). + - unittest/_torch/modeling_v2/test_modeling_v2_routing.py -k "kimi_k3" + - unittest/_torch/modeling_v2/test_modeling_v2_claims.py::test_uncertified_generic_calls_name_every_stock_import + # modeling_v2 catalog entries with sm_100 receipts (the K3 ones skip on + # other architectures). + - unittest/_torch/modeling_v2/norm/test_modeling_v2_k3_embed_norm.py + - unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_fwd.py + - unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py + - unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_rmsnorm.py + - unittest/_torch/modeling_v2/activation/test_modeling_v2_k3_situ_mul.py + - unittest/_torch/modeling_v2/gemm/test_modeling_v2_cublas_mm.py + - unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv.py + - unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_swiglu.py + - unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_long.py + - unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_wide.py + - unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_decode_gemv.py + - unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_head_gemv.py + # The Kimi K3 target's decode GEMVs, LM head and embedding on those entries. + - unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_gemv.py # KDA runtime: host-derived prefill metadata and the bf16 state pool # round-trip. Both are single-device cases that use GPU 0 only. - unittest/_torch/modules/kimi_kda/test_kda_host_metadata.py @@ -375,6 +404,16 @@ l0_b200: - examples/visual_gen/test_minimax_h3_e2e.py::test_minimax_h3_public_visual_gen_api_smoke TIMEOUT (30) - examples/visual_gen/test_minimax_h3_e2e.py::test_minimax_h3_diffusers_lpips_and_audio_reference[t2va] TIMEOUT (30) - examples/visual_gen/test_minimax_h3_e2e.py::test_minimax_h3_diffusers_lpips_and_audio_reference[fl2va] TIMEOUT (30) + # ------------- Kimi K3 KDA / MLA decode kernels (sm_100) and their modeling_v2 entries --------------- + - unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_kda_attn.py + - unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_kda_decode_attn.py + - unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_kda_pools_past_2g.py + - unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_kda_verify.py + - unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_mla_attn.py + - unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_mla_q.py + - unittest/_torch/modeling_v2/ssm + - unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_qkv.py + - unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_attn_vb_out.py # ------------- Host perf module regression tests (6 representative scenarios) --------------- - perf/host_perf/test_module_scheduler.py::test_scheduler_production[production_gen_only_bs8] - perf/host_perf/test_module_scheduler.py::test_scheduler_production[production_mixed_32gen_4ctx] diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index bfb2c39eba38..521187488659 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -39,6 +39,8 @@ l0_cpu: # runs -m cpu_only, so this entry contributes only the 3 marked files under attention/. - unittest/_torch/attention - unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py + # Kimi K3 MLA decode view from the attention metadata: every test is cpu_only (CPU tensors). + - unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_mla_decode_view.py - unittest/_torch/compilation/test_auto_multi_stream.py - unittest/_torch/compilation/test_remove_copy_pass.py::test_remove_copy_preserves_minimax_producer_outputs - unittest/_torch/distributed @@ -48,6 +50,8 @@ l0_cpu: - unittest/_torch/peft - unittest/_torch/memory - unittest/_torch/modeling + # l0_b300 selects unittest/_torch/modeling_v2 with -m "not cpu_only", so this entry is what runs it. + - unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_step.py - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_fp8_indexer_rejects_different_qk_norm_epsilons - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_moe_reduces_only_local_terms - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_decoder_layer_sets_post_fusion_from_moe_scheduler diff --git a/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_cluster_waits.py b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_cluster_waits.py new file mode 100644 index 000000000000..f9433e11a165 --- /dev/null +++ b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_cluster_waits.py @@ -0,0 +1,97 @@ +# 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. +"""Source check of the mailbox waits in Kimi K3 kernels whose barriers other CTAs complete with st.async: the +complete_tx releases at cluster scope, so the waiting CTA must acquire at cluster scope +(``mbarrier.{test,try}_wait.parity.acquire.cluster``, the kernels' ``_test_wait_cluster`` / ``_try_wait_cluster``); +the DSL's ``mbarrier_test_wait`` / ``mbarrier_try_wait`` acquire at CTA scope, which the PTX memory model does not let +synchronize with another CTA's release. + + pytest test_k3_cluster_waits.py +""" + +import importlib.util +import re + +import pytest + +# Kernel module -> the barriers in it that other CTAs complete with st.async. +MAILBOXES = { + "tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv.k3_ctm_gemv_kernel": ("mail_full",), +} +# Kernel module -> the header of the blocks whose mailbox waits are on this CTA's own arrival and keep CTA scope (a +# kernel that arrives on its mailbox itself and acquires the other CTAs' data through a counter). +OWN_ARRIVAL = {} +CTA_WAIT = re.compile(r"mbarrier_(test|try)_wait\(") +CLUSTER_WAIT = re.compile(r"_(test|try)_wait_cluster\(") + + +def block_header(code, i): + """The header of the block that line i is in: the nearest earlier code line indented less.""" + indent = len(code[i]) - len(code[i].lstrip()) + for ln in reversed(code[:i]): + if ln.strip() and len(ln) - len(ln.lstrip()) < indent: + return ln.strip() + return "" + + +def cta_scope_mailbox_waits(lines, mailboxes, own_arrival=None): + """Line numbers of CTA-scope waits on a mailbox barrier (the barrier named on the wait's line or the next; + a name is a regular expression), except those directly in a block whose header matches ``own_arrival``.""" + code = [ln.split("#", 1)[0] for ln in lines] + found = [] + for i, ln in enumerate(code): + if CTA_WAIT.search(ln): + window = ln + (code[i + 1] if i + 1 < len(code) else "") + if any(re.search(rf"\b{name}\b", window) for name in mailboxes): + if own_arrival is None or not re.fullmatch(own_arrival, block_header(code, i)): + found.append(i + 1) + return found + + +def waited_at_cluster_scope(lines, name): + """Whether a cluster-scope wait names the barrier (on the wait's line or the next).""" + code = [ln.split("#", 1)[0] for ln in lines] + return any( + CLUSTER_WAIT.search(ln) + and re.search(rf"\b{name}\b", ln + (code[i + 1] if i + 1 < len(code) else "")) + for i, ln in enumerate(code) + ) + + +def test_checker_catches_the_pattern(): + bad = [" while not cute.arch.mbarrier_test_wait(mail_full.data_ptr(), 0):", " pass"] + good = [" while not _test_wait_cluster(mail_full.data_ptr(), 0):", " pass"] + assert cta_scope_mailbox_waits(bad, ("mail_full",)) == [1] + assert cta_scope_mailbox_waits(good, ("mail_full",)) == [] + assert waited_at_cluster_scope(good, "mail_full") + assert not waited_at_cluster_scope(bad, "mail_full") + # Only the own-arrival block keeps CTA scope; the same wait before it or in its else branch is flagged. + own = r"if cutlass\.const_expr\(no_cluster\):" + branches = ( + bad + ["if cutlass.const_expr(no_cluster):", " # own arrival"] + bad + ["else:"] + bad + ) + assert cta_scope_mailbox_waits(branches, ("mail_full",), own) == [1, 8] + + +@pytest.mark.parametrize("module", list(MAILBOXES), ids=[m.rsplit(".", 1)[1] for m in MAILBOXES]) +def test_mailbox_waits_acquire_at_cluster_scope(module): + path = importlib.util.find_spec(module).origin + with open(path) as f: + lines = f.read().split("\n") + found = cta_scope_mailbox_waits(lines, MAILBOXES[module], OWN_ARRIVAL.get(module)) + assert not found, f"{path}: CTA-scope waits on st.async mailboxes at lines {found}" + # Every listed mailbox is still waited on, so a renamed barrier or helper cannot pass unchecked. + missing = [name for name in MAILBOXES[module] if not waited_at_cluster_scope(lines, name)] + assert not missing, f"{path}: no cluster-scope wait on {missing}" diff --git a/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_ctm_gemv.py b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_ctm_gemv.py new file mode 100644 index 000000000000..972631c8d5c1 --- /dev/null +++ b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_ctm_gemv.py @@ -0,0 +1,344 @@ +# 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 CTM decode GEMVs (trtllm::k3_ctm_gemv, _gated, _swiglu, _long, _tail) and trtllm::k3_situ_mul at the Kimi K3 +TP16 per-rank shapes and the call sites' split / ring / push flags, at every M in 1..8: error against an fp64 product +of the unfused activation (torch's bf16 roundings) and against the stock path (cuBLAS F.linear; the Triton SiTU), +run-to-run identical bits, each M's rows bit-identical to the same rows of the 8-row call; split 1 the bits of +k3_decode_gemv, the long GEMV's sigmoid columns torch.sigmoid of its own plain output. 0 and 9 rows are refused.""" + +import functools + +import pytest +import torch +import torch.nn.functional as F + + +def _is_sm100() -> bool: + if not torch.cuda.is_available(): + return False + major, _ = torch.cuda.get_device_capability() + return major == 10 + + +pytestmark = pytest.mark.skipif(not _is_sm100(), reason="needs an SM100-family GPU") + +M_ALL = list(range(1, 9)) +TOL = 8e-3 # max |y - ref| / max |ref| + +# k3_ctm_gemv: (N, K, split, push). The MLA o_proj shape at splits 1 and 2, the drafter o_proj (tuned and synthetic). +PLAIN = { + "o_proj_s1": (7168, 768, 1, False), + "o_proj_s2": (7168, 768, 2, False), + "drafter_o_proj": (7168, 384, 1, True), + "drafter_o_proj_dummy": (7168, 256, 1, True), +} +# k3_ctm_gemv_long: (N, K, sig_col0, split, ring, push), as each call site passes them. +LONG = { + "mla_qkv_a_gate": ( + 2880, + 7168, + 2112, + 6, + 6, + True, + ), # [W_a; W_g], the gate columns as bf16(sigmoid) + "dense_gate_up": (4224, 7168, -1, 4, 5, False), + "dense_down": (7168, 2112, -1, 2, 6, False), # K 2112 ends in a half k-tile + "drafter_qkv": (512, 7168, -1, 8, 6, True), + "drafter_gate_up": (1792, 7168, -1, 8, 6, True), + "drafter_gate_up_dummy": (1536, 7168, -1, 8, 6, True), + "kda_qkvg": ( + 3208, + 7168, + -1, + 5, + 6, + True, + ), # KDA q/k/v/g/f_a/b: 25 whole tiles and an 8-row last one +} +# k3_ctm_gemv_swiglu: (N, K, split, push), the drafter down projection (tuned and synthetic). +SWIGLU = {"drafter_down": (7168, 896, 2, True), "drafter_down_dummy": (7168, 768, 2, True)} +AG_COLS, G_COL0, K_O = ( + 2880, + 2112, + 768, +) # MLA: [q_a 1536 | kv_a 576 | gate 768] rows of the fused projection +HIDDEN, LATENT, WIDTH, PAD, ACT = 7168, 3584, 224, 256, 384 +EPS = 1e-6 +SITU = { + "k3": (4.0, 25.0), + "plain": (1.0, None), +} # (beta, linear_beta): the K3 checkpoint's, and the defaults + + +def _ops(): + import tensorrt_llm # noqa: F401 (registers the trtllm ops) + from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import op # noqa: F401 + from tensorrt_llm._torch.cute_dsl_kernels.k3_decode_gemv import op as _op # noqa: F401 + + return torch.ops.trtllm + + +def _ctm(): + from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import op + + return op + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.int16) + + +def _rel(y: torch.Tensor, ref: torch.Tensor) -> float: + return (y.double() - ref.double()).abs().max().item() / ref.double().abs().max().item() + + +def _report(op, case, m, y, ref, stock=None, **flags): + extra = "" if stock is None else f" vs_stock={_rel(y, stock):.2e}" + marks = " ".join(f"{k}={v}" for k, v in flags.items()) + abs_err = (y.double() - ref.double()).abs().max().item() + print( + f"OPCHECK op={op} case={case} M={m} abs={abs_err:.3e} rel={_rel(y, ref):.3e}{extra} {marks}" + ) + + +@functools.lru_cache(maxsize=None) +def _weight(n: int, k: int, seed: int, scale: float = 0.02) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(n, k, generator=gen, device="cuda") * scale).bfloat16() + + +@functools.lru_cache(maxsize=None) +def _rows(k: int, seed: int, scale: float = 1.0) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(8, k, generator=gen, device="cuda") * scale).bfloat16() + + +def _checks(call, args8, m): + """The M-row call (the first M rows of each 8-row input), a rerun and the 8-row call.""" + args = [a[:m].contiguous() for a in args8] + y = call(*args) + again = call(*args) + y8 = call(*args8) + return args, y, torch.equal(_bits(y), _bits(again)), torch.equal(_bits(y), _bits(y8[:m])) + + +@pytest.mark.parametrize("m", M_ALL) +@pytest.mark.parametrize("name", list(PLAIN)) +def test_k3_ctm_gemv(name, m): + ops = _ops() + n, k, split, push = PLAIN[name] + w = _weight(n, k, 1, 0.03) + assert _ctm().supports(_rows(k, 2)[:m], w, split) + (x,), y, det, minv = _checks( + lambda x_: ops.k3_ctm_gemv(x_, w, True, split, push), [_rows(k, 2)], m + ) + ref = x.double() @ w.double().t() + decode = ops.k3_decode_gemv(x, w, True) + same_decode = torch.equal(_bits(y), _bits(decode)) + _report( + "k3_ctm_gemv", + name, + m, + y, + ref, + F.linear(x, w), + det=det, + rows_as_m8=minv, + eq_k3_decode_gemv=same_decode, + ) + assert _rel(y, ref) <= TOL and _rel(y, F.linear(x, w)) <= TOL + assert det and minv + if split == 1: + assert same_decode + + +@pytest.mark.parametrize("m", M_ALL) +@pytest.mark.parametrize("gate_sigmoid", [False, True], ids=["sigmoid_given", "sigmoid_in_kernel"]) +@pytest.mark.parametrize("split", [2, 1]) +def test_k3_ctm_gemv_gated(split, gate_sigmoid, m): + """The MLA o_proj with its output gate: (a * sigmoid(g)) @ W^T, g the gate columns of the fused projection rows + (or a * s with s = bf16(sigmoid(g)) already there, the in-model form), against the unfused torch roundings.""" + ops = _ops() + w = _weight(HIDDEN, K_O, 3, 0.03) + a8 = _rows(K_O, 4) + if gate_sigmoid: + ag8 = _rows(AG_COLS, 5, 2.0) + else: + gen = torch.Generator(device="cuda").manual_seed(6) + ag8 = torch.rand(8, AG_COLS, generator=gen, device="cuda").bfloat16() + assert _ctm().supports_gated(a8, ag8, G_COL0, w, split) + + def call(a_, ag_): + return ops.k3_ctm_gemv_gated(a_, ag_, G_COL0, w, True, split, gate_sigmoid) + + (a, ag), y, det, minv = _checks(call, [a8, ag8], m) + g = ag[:, G_COL0 : G_COL0 + K_O] + b = a * (g.sigmoid() if gate_sigmoid else g) # bf16, as the unfused model path + ref = b.double() @ w.double().t() + decode = ops.k3_decode_gemv(b, w, True) + same_decode = torch.equal(_bits(y), _bits(decode)) + _report("k3_ctm_gemv_gated", f"s{split}_{'sig' if gate_sigmoid else 'given'}", m, y, ref, F.linear(b, w), + det=det, rows_as_m8=minv, eq_unfused_k3_decode_gemv=same_decode) # fmt: skip + assert _rel(y, ref) <= TOL and _rel(y, F.linear(b, w)) <= TOL + assert det and minv + if split == 1: + assert same_decode + + +def _silu_and_mul(gu: torch.Tensor) -> torch.Tensor: + """silu_and_mul's fp32 arithmetic and one bf16 rounding.""" + k = gu.shape[1] // 2 + return (F.silu(gu[:, :k].float()) * gu[:, k:].float()).bfloat16() + + +@pytest.mark.parametrize("m", M_ALL) +@pytest.mark.parametrize("name", list(SWIGLU)) +def test_k3_ctm_gemv_swiglu(name, m): + ops = _ops() + n, k, split, push = SWIGLU[name] + w = _weight(n, k, 7) + gu8 = _rows(2 * k, 8, 2.0) + assert _ctm().supports_swiglu(gu8, w, split) + (gu,), y, det, minv = _checks( + lambda g_: ops.k3_ctm_gemv_swiglu(g_, w, True, split, push), [gu8], m + ) + act = _silu_and_mul(gu) + ref = act.double() @ w.double().t() + _report("k3_ctm_gemv_swiglu", name, m, y, ref, F.linear(act, w), det=det, rows_as_m8=minv) + assert _rel(y, ref) <= TOL and _rel(y, F.linear(act, w)) <= TOL + assert det and minv + + +@pytest.mark.parametrize("m", M_ALL) +@pytest.mark.parametrize("name", list(LONG)) +def test_k3_ctm_gemv_long(name, m): + ops = _ops() + n, k, sig, split, ring, push = LONG[name] + w = _weight(n, k, 9) + x8 = _rows(k, 10) + assert _ctm().supports_long(x8, w, split, ring) + + def call(x_, sig_col0): + return ops.k3_ctm_gemv_long(x_, w, sig_col0, split, ring, True, push) + + (x,), y, det, minv = _checks(lambda x_: call(x_, sig), [x8], m) + plain = call(x, -1) if sig >= 0 else y + ref = x.double() @ w.double().t() + cols = slice(0, sig if sig >= 0 else n) + # Columns >= sig_col0 hold bf16(sigmoid(bf16(x @ W^T))): torch.sigmoid of the op's own plain output. + sigmoid_ok = sig < 0 or ( + torch.equal(_bits(y[:, :sig]), _bits(plain[:, :sig])) + and torch.equal(_bits(y[:, sig:]), _bits(plain[:, sig:].sigmoid())) + ) + _report("k3_ctm_gemv_long", name, m, plain, ref, F.linear(x, w), det=det, rows_as_m8=minv, + sigmoid_cols=sigmoid_ok) # fmt: skip + assert _rel(plain, ref) <= TOL and _rel(plain, F.linear(x, w)) <= TOL + assert _rel(y[:, cols], ref[:, cols]) <= TOL + assert det and minv and sigmoid_ok + + +@pytest.mark.parametrize("m", M_ALL) +@pytest.mark.parametrize("rank", [0, 7, 15]) +def test_k3_ctm_gemv_tail(rank, m): + """The row-parallel MoE tail on the CTM kernel: the bits of k3_decode_gemv_tail.""" + ops = _ops() + w = _weight(HIDDEN, PAD + ACT, 11, 0.03).clone() + w[:, WIDTH:PAD] = 0 + lo = rank * WIDTH + + def call(lat_, act_): + return ops.k3_ctm_gemv_tail(lat_, act_, w, lo, WIDTH, EPS, True) + + (lat, act), y, det, minv = _checks(call, [_rows(LATENT, 12, 0.8), _rows(ACT, 13, 0.5)], m) + lat64 = lat.double() + normed = lat64 * torch.rsqrt(lat64.pow(2).mean(dim=1, keepdim=True) + EPS) + ref = ( + torch.cat([normed[:, lo : lo + WIDTH], act.double()], dim=1) + @ torch.cat([w[:, :WIDTH], w[:, PAD:]], dim=1).double().t() + ) + decode = ops.k3_decode_gemv_tail(lat, act, w, lo, WIDTH, EPS, True) + same_decode = torch.equal(_bits(y), _bits(decode)) + _report( + "k3_ctm_gemv_tail", + f"rank{rank}", + m, + y, + ref, + det=det, + rows_as_m8=minv, + eq_k3_decode_gemv_tail=same_decode, + ) + assert _rel(y, ref) <= TOL + assert det and minv and same_decode + + +def _situ_ref(gu, beta, linear_beta): + """SituAndMul's eager fp32 path (in fp64).""" + k = gu.shape[1] // 2 + g, u = gu[:, :k].double(), gu[:, k:].double() + a = beta * torch.tanh(g / beta) * torch.sigmoid(g) + if linear_beta is not None: + u = linear_beta * torch.tanh(u / linear_beta) + return a * u + + +@pytest.mark.parametrize("m", M_ALL) +@pytest.mark.parametrize("situ", list(SITU)) +def test_k3_situ_mul(situ, m): + """The dense MLP's SiTU-and-mul at TP16 (gate_up [M, 4224] -> [M, 2112]) against the Triton SituAndMul.""" + ops = _ops() + from tensorrt_llm._torch.modules.situ import SituAndMul + + beta, linear_beta = SITU[situ] + gu8 = _rows(2 * 2112, 14, 2.0) + (gu,), y, det, minv = _checks(lambda g_: ops.k3_situ_mul(g_, beta, linear_beta), [gu8], m) + stock = SituAndMul(beta=beta, linear_beta=linear_beta, use_fused_activation=True)(gu) + ref = _situ_ref(gu, beta, linear_beta) + identical = (_bits(y) == _bits(stock)).float().mean().item() + _report( + "k3_situ_mul", + situ, + m, + y, + ref, + stock, + det=det, + rows_as_m8=minv, + frac_eq_triton=f"{identical:.4f}", + ) + assert y.shape == (m, 2112) + assert _rel(y, ref) <= TOL and _rel(y, stock) <= TOL + assert det and minv + + +@pytest.mark.parametrize("m", [0, 9, 16]) +def test_token_limit(m): + op = _ctm() + _ops() + x = torch.zeros(m, 7168, dtype=torch.bfloat16, device="cuda") + for n, k, _, split, ring, _ in LONG.values(): + assert not op.supports_long( + torch.zeros(m, k, dtype=torch.bfloat16, device="cuda"), _weight(n, k, 9), split, ring + ) + with pytest.raises(ValueError): + torch.ops.trtllm.k3_ctm_gemv_long(x, _weight(2880, 7168, 9), 2112, 6, 6, True, True) + assert not op.supports( + torch.zeros(m, 768, dtype=torch.bfloat16, device="cuda"), _weight(7168, 768, 1, 0.03), 1 + ) + assert not op.supports_situ_mul(torch.zeros(m, 4224, dtype=torch.bfloat16, device="cuda")) + assert not op.supports_swiglu( + torch.zeros(m, 1792, dtype=torch.bfloat16, device="cuda"), _weight(7168, 896, 7), 2 + ) diff --git a/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_ctm_gemv_wide.py b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_ctm_gemv_wide.py new file mode 100644 index 000000000000..0158e2be99d3 --- /dev/null +++ b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_ctm_gemv_wide.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. +"""trtllm::k3_ctm_gemv_wide (up to 64 tokens in one MMA of 16 / 32 / 64 token columns) at the Kimi K3 TP16 per-rank +projection shapes, at every M in 1..64: error against an fp64 product and against cuBLAS (F.linear), run-to-run +identical bits, each M's rows bit-identical to the same rows of the 64-row call, the sigmoid columns torch.sigmoid of +the op's own plain output, and, where the split matches the call site's k3_ctm_gemv_long at most 8 tokens, the same +bits as that kernel. Both split-K transports give the same bits. 0 and 65 rows are refused.""" + +import functools + +import pytest +import torch +import torch.nn.functional as F + + +def _is_sm100() -> bool: + if not torch.cuda.is_available(): + return False + major, _ = torch.cuda.get_device_capability() + return major == 10 + + +pytestmark = pytest.mark.skipif(not _is_sm100(), reason="needs an SM100-family GPU") + +M_ALL = list(range(1, 65)) +TOL = 8e-3 # bf16 output: max |y - ref| / max |ref| +TOL_FP32 = 1e-4 # fp32 output + +# (N, K, sig_col0, out_fp32): the per-rank TP16 shapes of the projections of a wide decode step. +WIDE = { + "kda_qkvg": (3208, 7168, -1, False), # KDA q/k/v/g/f_a/b: 25 whole tiles and an 8-row last one + "mla_qkv_a_gate": (2880, 7168, 2112, False), # [W_a; W_g], the gate rows as bf16(sigmoid) + "o_proj": (7168, 768, -1, False), # KDA / MLA o_proj + "moe_head": (280, 7168, -1, True), # [latent down slice; router rows], fp32 (router logits) + "moe_tail": (7168, 640, -1, False), # [latent up slice | padding | shared down], 5 k-tiles + "shared_gate_up": (768, 7168, -1, False), + "dense_gate_up": (4224, 7168, -1, False), + "dense_down": (7168, 2112, -1, False), # K 2112 ends in a half k-tile + "drafter_qkv": (512, 7168, -1, False), + "drafter_gate_up": (1792, 7168, -1, False), + "drafter_o_proj": (7168, 384, -1, False), +} +# The k3_ctm_gemv_long call (split, ring, push) of the shapes that run it at most 8 tokens. +LONG_SITES = { + "mla_qkv_a_gate": (6, 6, True), + "dense_gate_up": (4, 5, False), + "dense_down": (2, 6, False), + "drafter_qkv": (8, 6, True), + "drafter_gate_up": (8, 6, True), +} + + +def _ops(): + import tensorrt_llm # noqa: F401 (registers the trtllm ops) + from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import op # noqa: F401 + + return torch.ops.trtllm + + +def _ctm(): + from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import op + + return op + + +def _bits(t: torch.Tensor) -> torch.Tensor: + t = t.contiguous() + return t.view(torch.int32) if t.dtype == torch.float32 else t.view(torch.int16) + + +def _rel(y: torch.Tensor, ref: torch.Tensor) -> float: + return (y.double() - ref.double()).abs().max().item() / ref.double().abs().max().item() + + +@functools.lru_cache(maxsize=None) +def _weight(n: int, k: int, seed: int) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(n, k, generator=gen, device="cuda") * 0.02).bfloat16() + + +@functools.lru_cache(maxsize=None) +def _rows(k: int, seed: int) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return torch.randn(64, k, generator=gen, device="cuda").bfloat16() + + +@functools.lru_cache(maxsize=None) +def _full_call(name: str) -> torch.Tensor: + """The 64-row call of a case (its rows are what every M's call must reproduce).""" + n, k, sig, fp32 = WIDE[name] + return _ops().k3_ctm_gemv_wide(_rows(k, 2), _weight(n, k, 1), sig, fp32) + + +@pytest.mark.parametrize("m", M_ALL) +@pytest.mark.parametrize("name", list(WIDE)) +def test_k3_ctm_gemv_wide(name, m): + ops = _ops() + n, k, sig, fp32 = WIDE[name] + w = _weight(n, k, 1) + x = _rows(k, 2)[:m].contiguous() + assert _ctm().supports_wide(x, w, sig, fp32) + y = ops.k3_ctm_gemv_wide(x, w, sig, fp32) + det = torch.equal(_bits(y), _bits(ops.k3_ctm_gemv_wide(x, w, sig, fp32))) + rows_as_m64 = torch.equal(_bits(y), _bits(_full_call(name)[:m])) + plain = ops.k3_ctm_gemv_wide(x, w, -1, fp32) if sig >= 0 else y + ref = x.double() @ w.double().t() + cols = slice(0, sig if sig >= 0 else n) + # Columns >= sig_col0 hold bf16(sigmoid(bf16(x @ W^T))): torch.sigmoid of the op's own plain output. + sigmoid_ok = sig < 0 or ( + torch.equal(_bits(y[:, :sig]), _bits(plain[:, :sig])) + and torch.equal(_bits(y[:, sig:]), _bits(plain[:, sig:].sigmoid())) + ) + eq_long = None + if name in LONG_SITES and m <= 8: + split, ring, push = LONG_SITES[name] + from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import k3_ctm_gemv_kernel as kernel + + wide_split = _ctm().wide_config( + n, k, kernel.wide_tile(m), torch.cuda.get_device_properties(0).multi_processor_count + )[0] + long_y = ops.k3_ctm_gemv_long(x, w, sig, split, ring, True, push) + eq_long = torch.equal(_bits(y), _bits(long_y)) if wide_split == split else None + stock = F.linear(x.float(), w.float()) if fp32 else F.linear(x, w) + print( + f"OPCHECK op=k3_ctm_gemv_wide case={name} M={m} dtype={str(y.dtype)[6:]} " + f"abs={(plain.double() - ref).abs().max().item():.3e} rel={_rel(plain, ref):.3e} " + f"vs_stock={_rel(plain, stock):.2e} det={det} rows_as_m64={rows_as_m64} sigmoid_cols={sigmoid_ok} " + f"eq_long={eq_long}" + ) # fmt: skip + assert y.shape == (m, n) and y.dtype == (torch.float32 if fp32 else torch.bfloat16) + tol = TOL_FP32 if fp32 else TOL + assert _rel(plain, ref) <= tol and _rel(plain, stock) <= tol + assert _rel(y[:, cols], ref[:, cols]) <= tol + assert det and rows_as_m64 and sigmoid_ok + assert eq_long is not False + + +@pytest.mark.parametrize("m", list(range(16, 65, 8))) +@pytest.mark.parametrize("name", list(LONG_SITES)) +def test_k3_ctm_gemv_wide_rows_as_long(name, m): + """A wide step of R x 8 tokens: every token's row is the bits k3_ctm_gemv_long gives that token at <= 8 tokens (the + call sites' split), so wide steps and batch 1 agree on these projections.""" + ops = _ops() + n, k, sig, _ = WIDE[name] + split, ring, push = LONG_SITES[name] + w = _weight(n, k, 1) + x = _rows(k, 2)[:m].contiguous() + y = ops.k3_ctm_gemv_wide(x, w, sig) + + def long_rows(r): + return ops.k3_ctm_gemv_long(x[r : r + 8].contiguous(), w, sig, split, ring, True, push) + + same = [torch.equal(_bits(y[r : r + 8]), _bits(long_rows(r))) for r in range(0, m, 8)] + print( + f"OPCHECK op=k3_ctm_gemv_wide case={name} M={m} rows_as_long_per_8={sum(same)}/{len(same)}" + ) + assert all(same) + + +@pytest.mark.parametrize("m", [0, 65]) +def test_token_limit(m): + _ops() + x = torch.zeros(m, 7168, dtype=torch.bfloat16, device="cuda") + w = _weight(3208, 7168, 1) + assert not _ctm().supports_wide(x, w) + with pytest.raises(ValueError): + torch.ops.trtllm.k3_ctm_gemv_wide(x, w) diff --git a/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_decode_gemv.py b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_decode_gemv.py new file mode 100644 index 000000000000..e3d12396fb67 --- /dev/null +++ b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_decode_gemv.py @@ -0,0 +1,171 @@ +# 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::k3_decode_gemv / k3_decode_gemv_tail (CuTe DSL decode GEMV, M <= 8) at the Kimi K3 TP16 per-rank shapes, +at every M in 1..8: error against an fp64 product and against the stock path (cuBLAS F.linear; RMSNorm -> slice -> +concat -> linear for the tail), for the tail also against torch with the kernel's arithmetic (fp32 accumulators, the +RMS on the latent one, one bf16 rounding), run-to-run identical bits and each M's rows bit-identical to the same rows +of the 8-row call. 0 and 9 rows are refused.""" + +import functools + +import pytest +import torch +import torch.nn.functional as F + + +def _is_sm100() -> bool: + if not torch.cuda.is_available(): + return False + major, _ = torch.cuda.get_device_capability() + return major == 10 + + +pytestmark = pytest.mark.skipif(not _is_sm100(), reason="needs an SM100-family GPU") + +M_ALL = list(range(1, 9)) +TOL = 8e-3 # max |y - ref| / max |ref| +# (N, K) per rank at TP16: the KDA o_proj (one CTA per weight tile) and the KDA input projection (split-K). +SHAPES = {"kda_o_proj": (7168, 768), "kda_qkvg": (3208, 7168)} +HIDDEN, LATENT, WIDTH, PAD, ACT = 7168, 3584, 224, 256, 384 +EPS = 1e-6 + + +def _ops(): + import tensorrt_llm # noqa: F401 (registers the trtllm ops) + from tensorrt_llm._torch.cute_dsl_kernels.k3_decode_gemv import op # noqa: F401 + + return torch.ops.trtllm + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.int16) + + +def _rel(y: torch.Tensor, ref: torch.Tensor) -> float: + return (y.double() - ref.double()).abs().max().item() / ref.double().abs().max().item() + + +def _report(op, case, m, y, ref, stock=None, **flags): + extra = "" if stock is None else f" vs_stock={_rel(y, stock):.2e}" + marks = " ".join(f"{k}={v}" for k, v in flags.items()) + abs_err = (y.double() - ref.double()).abs().max().item() + print( + f"OPCHECK op={op} case={case} M={m} abs={abs_err:.3e} rel={_rel(y, ref):.3e}{extra} {marks}" + ) + + +@functools.lru_cache(maxsize=None) +def _weight(n: int, k: int, seed: int) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(n, k, generator=gen, device="cuda") * 0.03).bfloat16() + + +@functools.lru_cache(maxsize=None) +def _rows(k: int, seed: int, scale: float = 1.0) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(8, k, generator=gen, device="cuda") * scale).bfloat16() + + +@pytest.mark.parametrize("m", M_ALL) +@pytest.mark.parametrize("name", list(SHAPES)) +def test_k3_decode_gemv(name, m): + ops = _ops() + n, k = SHAPES[name] + w = _weight(n, k, 1) + x8 = _rows(k, 2) + x = x8[:m].contiguous() + y = ops.k3_decode_gemv(x, w, True) + again = ops.k3_decode_gemv(x, w, True) + y8 = ops.k3_decode_gemv(x8, w, True) + ref = x.double() @ w.double().t() + deterministic = torch.equal(_bits(y), _bits(again)) + m_invariant = torch.equal(_bits(y), _bits(y8[:m])) + _report( + "k3_decode_gemv", name, m, y, ref, F.linear(x, w), det=deterministic, rows_as_m8=m_invariant + ) + assert y.shape == (m, n) + assert _rel(y, ref) <= TOL and _rel(y, F.linear(x, w)) <= TOL + assert deterministic and m_invariant + + +def _tail_weight() -> torch.Tensor: + w = _weight(HIDDEN, PAD + ACT, 3).clone() + w[:, WIDTH:PAD] = 0 + return w + + +def _tail_ref(latent, act, w, lo): + lat = latent.double() + normed = lat * torch.rsqrt(lat.pow(2).mean(dim=1, keepdim=True) + EPS) + x = torch.cat([normed[:, lo : lo + WIDTH], act.double()], dim=1) + return x @ torch.cat([w[:, :WIDTH], w[:, PAD:]], dim=1).double().t() + + +def _tail_stock(latent, act, w, lo): + from tensorrt_llm._torch.modules.rms_norm import RMSNorm + + norm = RMSNorm(hidden_size=LATENT, eps=EPS, dtype=torch.bfloat16).cuda() + norm.weight.data.fill_(1.0) + x = torch.cat([norm(latent)[:, lo : lo + WIDTH], act], dim=1) + return F.linear(x, torch.cat([w[:, :WIDTH], w[:, PAD:]], dim=1)) + + +def _tail_fp32(latent, act, w, lo): + """The tail with the kernel's arithmetic: fp32 accumulators of the latent slice and of the activation, the latent + one scaled by the RMS of the whole latent row, one bf16 rounding.""" + lat = latent.float() + scale = torch.rsqrt(lat.pow(2).mean(dim=1, keepdim=True) + EPS) + acc_lat = lat[:, lo : lo + WIDTH] @ w[:, :WIDTH].float().t() + return (acc_lat * scale + act.float() @ w[:, PAD:].float().t()).bfloat16() + + +@pytest.mark.parametrize("m", M_ALL) +@pytest.mark.parametrize("rank", [0, 7, 15]) +def test_k3_decode_gemv_tail(rank, m): + ops = _ops() + w = _tail_weight() + lat8, act8 = _rows(LATENT, 4, 0.8), _rows(ACT, 5, 0.5) + lat, act, lo = lat8[:m].contiguous(), act8[:m].contiguous(), rank * WIDTH + + def call(lat_, act_): + return ops.k3_decode_gemv_tail(lat_, act_, w, lo, WIDTH, EPS, True) + + y = call(lat, act) + again = call(lat, act) + y8 = call(lat8, act8) + fp32 = _tail_fp32(lat, act, w, lo) + ref = _tail_ref(lat, act, w, lo) + stock = _tail_stock(lat, act, w, lo) + deterministic = torch.equal(_bits(y), _bits(again)) + m_invariant = torch.equal(_bits(y), _bits(y8[:m])) + _report("k3_decode_gemv_tail", f"rank{rank}", m, y, ref, stock, det=deterministic, rows_as_m8=m_invariant, + vs_fp32=f"{_rel(y, fp32):.2e}") # fmt: skip + assert _rel(y, ref) <= TOL and _rel(y, stock) <= TOL and _rel(y, fp32) <= TOL + assert deterministic and m_invariant + + +@pytest.mark.parametrize("m", [0, 9, 16]) +def test_token_limit(m): + ops = _ops() + from tensorrt_llm._torch.cute_dsl_kernels.k3_decode_gemv import op + + w = _weight(*SHAPES["kda_o_proj"], 1) + x = torch.zeros(m, w.shape[1], dtype=torch.bfloat16, device="cuda") + assert not op.supports(x, w) + with pytest.raises(ValueError): + ops.k3_decode_gemv(x, w, True) + lat = torch.zeros(m, LATENT, dtype=torch.bfloat16, device="cuda") + act = torch.zeros(m, ACT, dtype=torch.bfloat16, device="cuda") + assert not op.supports_tail(lat, act, _tail_weight(), WIDTH) diff --git a/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_embed.py b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_embed.py new file mode 100644 index 000000000000..408bb5b1ef35 --- /dev/null +++ b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_embed.py @@ -0,0 +1,574 @@ +# 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::k3_embed`` (a decode step's embedding rows, one 16-byte load per thread) and ``trtllm::k3_embed_norm`` +(the rows into a snapshot-bank slot plus layer 0's input RMSNorm, one launch) at the in-model shape: Kimi K3's +replicated bf16 embedding table [163840, 7168], M = 1 .. 64 tokens (every M; each compiles its own kernel). + +The reference is what the model runs without them (KimiLinearModel._embed and layer 0's input_layernorm): +nn.Embedding's ``F.embedding`` (an index_select) for the rows, and the RMSNorm module, whose bf16 path is +``trtllm::flashinfer_rmsnorm`` (``flashinfer.norm.rmsnorm``), for the norm. Both ops claim bit-identity (bf16 bits). +k3_embed_norm reproduces flashinfer's CuTe DSL RMSNormKernel (one 128-thread CTA per row for 6144 < H <= 16384, see +k3_embed_kernel), so its checks need flashinfer (the model's gate for the op) on that kernel. The table is replicated +(no vocabulary shard), so the model never passes an id outside [0, V); the ops define one as a zero row, checked +against the masked reference. + +Checks (in the first three, every call runs twice and the reruns must be bit-identical): + +* k3_embed: int32 ids (the engine's) at every M, int64 ids at a few; ids with 0, V - 1, repeats, one id M times; +* k3_embed_norm at every M, called as the model calls it (the module's weight and eps, a bank slot as ``raw``): the + normed rows against the module on the F.embedding rows, the raw rows, the bank's other slots untouched; eps 1e-6 + and 1e-5, bank slots 0 (the model's) and 1; +* ids outside [0, V) (negative, V, the int32 extremes, int64 ids past 2^32): zero rows, normed as the module norms + them; +* every table row through both ops (64 consecutive ids per call); +* CUDA graphs of both ops at M = 1 .. 16, 17 .. 32, 33 .. 48 and 49 .. 64, captured once and replayed with the ids + rewritten in place, bit-identical to the eager ops and the reference on every replay; the same ids replayed again + give the same bits. + +Table: ``python3 test_k3_embed.py report``. Timing: ``python3 test_k3_embed.py time`` (CUDA graphs of back-to-back +calls, each on its own random ids, median us per call over 15 replays at every M: F.embedding vs k3_embed, and the +rows plus the snapshot copy plus the RMSNorm module vs k3_embed_norm). +""" + +import math +import os +import statistics +import sys + +import pytest +import torch +import torch.nn.functional as F + +VOCAB = 163840 # Kimi K3's vocabulary (a replicated nn.Embedding) +HIDDEN = 7168 +EPS = 1e-6 +EPS_ALT = 1e-5 +MAX_TOKENS = 64 # k3_embed's MAX_TOKENS +TOKENS = list(range(1, MAX_TOKENS + 1)) +INT64_TOKENS = [1, 7, 8, 16, 33, 64] +OUTSIDE_TOKENS = [1, 8, 64] +GRAPH_FAMILIES = [list(range(lo, lo + 16)) for lo in range(1, MAX_TOKENS + 1, 16)] +BANK_SLOTS = 4 # attention-residual snapshots; the model passes slot 0 as raw +SENTINEL = 7.0 +ID_KINDS = ("0, V - 1, repeats", "V - 1, 0", "random", "one id x M") +GRAPH_KINDS = ("random", "0, V - 1, repeats", "outside [0, V)", "one id x M", "V - 1, 0") +OUTSIDE_IDS = (-1, VOCAB, VOCAB + 7, -VOCAB, 2**31 - 1, -(2**31)) +OUTSIDE_IDS_64 = OUTSIDE_IDS + (2**32, 2**32 + 5, 3 - 2**32, 2**40) + + +def _sm90() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (9, 0) + + +pytestmark = pytest.mark.skipif( + not _sm90(), reason="needs SM90 or newer (griddepcontrol, programmatic dependent launch)" +) + + +def _op(): + import tensorrt_llm # noqa: F401 (registers the trtllm ops) + from tensorrt_llm._torch.cute_dsl_kernels.k3_embed import op + + return op + + +def _flashinfer_norm() -> bool: + """Whether layer 0's RMSNorm runs flashinfer's CuTe DSL RMSNormKernel: the model takes k3_embed_norm only with + flashinfer, and FLASHINFER_USE_CUDA_NORM=1 selects flashinfer's CUDA kernel, which is not the one the op + reproduces.""" + from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE + + return IS_FLASHINFER_AVAILABLE and os.environ.get("FLASHINFER_USE_CUDA_NORM", "0") != "1" + + +def _need_norm(): + if not _flashinfer_norm(): + pytest.skip("needs flashinfer's CuTe DSL RMSNorm (the kernel k3_embed_norm reproduces)") + + +# ---------------------------------------------------------------------------------------------------------------- +# Inputs and the model's path +# ---------------------------------------------------------------------------------------------------------------- + +_state = {} + + +def table() -> torch.Tensor: + """The embedding table, built once: N(0, 1) rows at log-uniform scales 1e-4 .. 1e2, row 1 all zero.""" + if "table" not in _state: + with torch.inference_mode(False): + gen = torch.Generator(device="cuda").manual_seed(20260930) + t = torch.randn(VOCAB, HIDDEN, generator=gen, device="cuda", dtype=torch.bfloat16) + scale = torch.empty(VOCAB, 1, device="cuda") + scale.uniform_(math.log(1e-4), math.log(1e2), generator=gen) + t.mul_(scale.exp_().bfloat16()) + t[1].zero_() + _state["table"] = t + return _state["table"] + + +def layer0_norm(eps): + """Layer 0's input_layernorm as KimiLinearDecoderLayer builds it (the RMSNorm module, bf16): weights around 1, + every 101st negative.""" + key = ("norm", eps) + if key not in _state: + from tensorrt_llm._torch.modules.rms_norm import RMSNorm + + with torch.inference_mode(False): + gen = torch.Generator(device="cuda").manual_seed(7) + w = 1.0 + 0.2 * torch.randn(HIDDEN, generator=gen, device="cuda") + w[::101] *= -1.0 + norm = RMSNorm(hidden_size=HIDDEN, eps=eps, dtype=torch.bfloat16, device="cuda") + norm.requires_grad_(False) + norm.weight.copy_(w) + _state[key] = norm + return _state[key] + + +def make_ids(m, kind, gen, dtype=torch.int32): + """``m`` ids in [0, V) of one of ID_KINDS.""" + ids = torch.randint(0, VOCAB, (m,), generator=gen, device="cuda") + if kind == "0, V - 1, repeats": + ids[0] = 0 + if m >= 2: + ids[-1] = VOCAB - 1 + if m >= 4: + ids[1] = VOCAB - 1 + ids[2] = 0 + if m >= 6: + ids[4] = ids[3] + elif kind == "V - 1, 0": + ids[0] = VOCAB - 1 + if m >= 2: + ids[-1] = 0 + elif kind == "one id x M": + ids = ids[:1].repeat(m) + return ids.to(dtype) + + +def make_outside_ids(m, dtype, gen): + """``m`` ids: the even positions outside [0, V) (cycling through the dtype's out-of-range values), the others + random in [0, V).""" + bad = OUTSIDE_IDS_64 if dtype == torch.int64 else OUTSIDE_IDS + ids = torch.randint(0, VOCAB, (m,), generator=gen, device="cuda") + ids[0::2] = torch.tensor([bad[i % len(bad)] for i in range((m + 1) // 2)], device="cuda") + return ids.to(dtype) + + +def rows_ref(ids): + """The model's rows: F.embedding (nn.Embedding's index_select); an id outside [0, V) gives the ops' zero row.""" + valid = (ids >= 0) & (ids < VOCAB) + if bool(valid.all()): + return F.embedding(ids, table()) + rows = F.embedding(torch.where(valid, ids, torch.zeros_like(ids)), table()) + return rows.masked_fill_(~valid[:, None], 0.0) + + +def embed(ids): + """k3_embed as KimiLinearModel._embed calls it.""" + return torch.ops.trtllm.k3_embed(ids, table()) + + +def new_bank(m): + """A snapshot bank [BANK_SLOTS, m, H] filled with SENTINEL.""" + return torch.full((BANK_SLOTS, m, HIDDEN), SENTINEL, dtype=torch.bfloat16, device="cuda") + + +def embed_norm_into(ids, norm, raw): + """k3_embed_norm as KimiLinearModel._embed_norm calls it: the norm module's weight and eps, ``raw`` a bank slot.""" + return torch.ops.trtllm.k3_embed_norm(ids, table(), norm.weight, norm.variance_epsilon, raw) + + +def embed_norm(ids, eps=EPS, slot=0): + """k3_embed_norm into ``slot`` of a fresh bank: (normed rows, bank).""" + bank = new_bank(ids.numel()) + return embed_norm_into(ids, layer0_norm(eps), bank[slot]), bank + + +def same(a, b) -> bool: + """Bit-identical bf16 tensors.""" + return a.shape == b.shape and torch.equal(a.view(torch.int16), b.view(torch.int16)) + + +def differ(a, b): + """The number of rows of ``a`` that differ from ``b`` in any bit (on the device).""" + return (a.view(torch.int16) != b.view(torch.int16)).any(dim=1).sum() + + +def check_embed(ids, kind): + got = embed(ids) + return {f"{kind}: rows": same(got, rows_ref(ids)), f"{kind}: rerun": same(embed(ids), got)} + + +def check_norm(ids, kind, eps=EPS, slot=0): + normed, bank = embed_norm(ids, eps, slot) + normed_again, bank_again = embed_norm(ids, eps, slot) + rows = rows_ref(ids) + others = [s for s in range(BANK_SLOTS) if s != slot] + return { + f"{kind}: normed": same(normed, layer0_norm(eps)(rows)), + f"{kind}: raw": same(bank[slot], rows), + f"{kind}: other slots": bool((bank[others] == SENTINEL).all()), + f"{kind}: rerun": same(normed_again, normed) and same(bank_again, bank), + } + + +def result(op_name, m, case, checks): + """A table row: ``checks`` maps each comparison to whether it held.""" + failed = [k for k, ok in checks.items() if not ok] + identical = "yes" if not failed else "no: " + ", ".join(failed) + return dict(op=op_name, m=m, case=case, identical=identical, ok=not failed) + + +# ---------------------------------------------------------------------------------------------------------------- +# Measurements +# ---------------------------------------------------------------------------------------------------------------- + + +def measure_embed(m): + """k3_embed at M tokens: int32 ids of every kind, int64 ids at INT64_TOKENS.""" + gen = torch.Generator(device="cuda").manual_seed(1000 + m) + checks = {} + for kind in ID_KINDS: + checks.update(check_embed(make_ids(m, kind, gen), kind)) + out = [result("k3_embed", m, "int32 ids: " + " / ".join(ID_KINDS), checks)] + if m in INT64_TOKENS: + ids = make_ids(m, ID_KINDS[0], gen, torch.int64) + out.append(result("k3_embed", m, f"int64 ids: {ID_KINDS[0]}", check_embed(ids, "int64"))) + return out + + +def measure_norm(m): + """k3_embed_norm at M tokens: int32 ids of every kind into slot 0 at EPS, random ids into slot 1 at EPS_ALT, + int64 ids at INT64_TOKENS.""" + gen = torch.Generator(device="cuda").manual_seed(2000 + m) + checks = {} + for kind in ID_KINDS: + checks.update(check_norm(make_ids(m, kind, gen), kind)) + case = f"int32 ids: {' / '.join(ID_KINDS)}; eps {EPS:g}, slot 0" + out = [result("k3_embed_norm", m, case, checks)] + checks = check_norm(make_ids(m, "random", gen), "random", EPS_ALT, 1) + out.append(result("k3_embed_norm", m, f"int32 ids: random; eps {EPS_ALT:g}, slot 1", checks)) + if m in INT64_TOKENS: + checks = check_norm(make_ids(m, ID_KINDS[0], gen, torch.int64), "int64") + case = f"int64 ids: {ID_KINDS[0]}; eps {EPS:g}, slot 0" + out.append(result("k3_embed_norm", m, case, checks)) + return out + + +def measure_outside(m, dtype, with_norm): + """Both ops on ids outside [0, V) at the even positions.""" + gen = torch.Generator(device="cuda").manual_seed(3000 + m) + ids = make_outside_ids(m, dtype, gen) + values = ", ".join(map(str, OUTSIDE_IDS_64 if dtype == torch.int64 else OUTSIDE_IDS)) + case = f"{str(dtype).split('.')[-1]} ids, even positions outside [0, V): {values}" + out = [result("k3_embed", m, case, check_embed(ids, "outside"))] + if with_norm: + out.append(result("k3_embed_norm", m, case, check_norm(ids, "outside"))) + return out + + +def measure_every_row(with_norm): + """Every table row through k3_embed (and k3_embed_norm), MAX_TOKENS consecutive ids per call.""" + tab = table() + norm = layer0_norm(EPS) if with_norm else None + bank = new_bank(MAX_TOKENS) + every_id = torch.arange(VOCAB, dtype=torch.int32, device="cuda") + bad = torch.zeros(3, dtype=torch.int64, device="cuda") + for lo in range(0, VOCAB, MAX_TOKENS): + ids = every_id[lo : lo + MAX_TOKENS] + rows = F.embedding(ids, tab) + bad[0] += differ(embed(ids), rows) + if with_norm: + normed = embed_norm_into(ids, norm, bank[0]) + bad[1] += differ(bank[0], rows) + bad[2] += differ(normed, norm(rows)) + bad = bad.tolist() + case = f"all {VOCAB} rows, {MAX_TOKENS} consecutive int32 ids per call" + out = [result("k3_embed", "all", case, {f"rows ({bad[0]} differ)": bad[0] == 0})] + if with_norm: + checks = {f"raw ({bad[1]} rows differ)": bad[1] == 0} + checks[f"normed ({bad[2]} rows differ)"] = bad[2] == 0 + out.append(result("k3_embed_norm", "all", f"{case}, eps {EPS:g}", checks)) + return out + + +def measure_graph(ms, with_norm): + """One CUDA graph with both ops at every M of ``ms``, captured once and replayed with the ids rewritten in place + (one GRAPH_KINDS kind per replay), each replay against the eager ops and the reference; then the last ids + replayed again against the first replay of them.""" + tab = table() + norm = layer0_norm(EPS) if with_norm else None + gen = torch.Generator(device="cuda").manual_seed(4000 + ms[0]) + ids_buf = {m: torch.zeros(m, dtype=torch.int32, device="cuda") for m in ms} + banks = {m: new_bank(m) for m in ms} + outs = {} + + def body(): + for m in ms: + outs["k3_embed", m] = torch.ops.trtllm.k3_embed(ids_buf[m], tab) + if with_norm: + outs["k3_embed_norm", m] = embed_norm_into(ids_buf[m], norm, banks[m][0]) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + body() # every M compiles on its first call, outside capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + body() + torch.cuda.synchronize() + family = f"{ms[0]}-{ms[-1]}" + out = [] + for rep, kind in enumerate(GRAPH_KINDS): + for m in ms: + if kind == "outside [0, V)": + ids_buf[m].copy_(make_outside_ids(m, torch.int32, gen)) + else: + ids_buf[m].copy_(make_ids(m, kind, gen)) + banks[m].fill_(SENTINEL) + graph.replay() + torch.cuda.synchronize() + checks = {"k3_embed": {}, "k3_embed_norm": {}} + for m in ms: + ids = ids_buf[m] + rows = rows_ref(ids) + got = outs["k3_embed", m] + checks["k3_embed"][f"M {m} = reference"] = same(got, rows) + checks["k3_embed"][f"M {m} = eager"] = same(got, embed(ids)) + if with_norm: + got = outs["k3_embed_norm", m] + eager, eager_bank = embed_norm(ids) + checks["k3_embed_norm"].update({ + f"M {m} normed = reference": same(got, norm(rows)), + f"M {m} raw = reference": same(banks[m][0], rows), + f"M {m} = eager": same(got, eager) and same(banks[m], eager_bank), + }) # fmt: skip + for name, op_checks in checks.items(): + if op_checks: + out.append(result(name, family, f"replay {rep}: {kind} ids rewritten", op_checks)) + first = {k: v.clone() for k, v in outs.items()} + first_banks = {m: b.clone() for m, b in banks.items()} + for b in banks.values(): + b.fill_(SENTINEL) + graph.replay() + torch.cuda.synchronize() + for name in ("k3_embed", "k3_embed_norm") if with_norm else ("k3_embed",): + rerun = {f"M {m}": same(outs[name, m], first[name, m]) for m in ms} + if name == "k3_embed_norm": + rerun.update({f"M {m} bank": same(banks[m], first_banks[m]) for m in ms}) + out.append(result(name, family, "the last ids replayed again", rerun)) + return out + + +# ---------------------------------------------------------------------------------------------------------------- +# Tests +# ---------------------------------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize("m", TOKENS) +def test_embed(m): + _op() + with torch.inference_mode(): + bad = [r for r in measure_embed(m) if not r["ok"]] + assert not bad, bad + + +@pytest.mark.parametrize("m", TOKENS) +def test_embed_norm(m): + _op() + _need_norm() + with torch.inference_mode(): + bad = [r for r in measure_norm(m) if not r["ok"]] + assert not bad, bad + + +@pytest.mark.parametrize("dtype", [torch.int32, torch.int64], ids=["int32", "int64"]) +@pytest.mark.parametrize("m", OUTSIDE_TOKENS) +def test_ids_outside_vocab(m, dtype): + _op() + with torch.inference_mode(): + bad = [r for r in measure_outside(m, dtype, _flashinfer_norm()) if not r["ok"]] + assert not bad, bad + + +def test_every_row(): + _op() + with torch.inference_mode(): + bad = [r for r in measure_every_row(_flashinfer_norm()) if not r["ok"]] + assert not bad, bad + + +@pytest.mark.parametrize("ms", GRAPH_FAMILIES, ids=[f"M{f[0]}-{f[-1]}" for f in GRAPH_FAMILIES]) +def test_graph_replay(ms): + _op() + with torch.inference_mode(): + bad = [r for r in measure_graph(ms, _flashinfer_norm()) if not r["ok"]] + assert not bad, bad + + +def test_supports(): + """The model's calls are supported at every M in 1 .. 64 (hidden 7168, int32 or int64 ids); M = 0 or 65, other id + dtypes and widths outside the norm's geometry are refused.""" + op = _op() + with torch.inference_mode(): + tab = table() + norm = layer0_norm(EPS) + bank = torch.empty(2, MAX_TOKENS, HIDDEN, dtype=torch.bfloat16, device="cuda") + + def ids(n, dtype=torch.int32): + return torch.zeros(n, dtype=dtype, device="cuda") + + for m in TOKENS: + raw = bank[0, :m] + assert op.supports(ids(m), tab) and op.supports(ids(m, torch.int64), tab), m + assert op.supports_norm(ids(m), tab, norm.weight, raw), m + assert op.norm_supports_hidden(HIDDEN) + assert not op.norm_supports_hidden(6144) and not op.norm_supports_hidden(HIDDEN + 512) + assert not op.supports(ids(0), tab) and not op.supports(ids(MAX_TOKENS + 1), tab) + assert not op.supports(ids(8, torch.int16), tab) + assert not op.supports_norm(ids(8), tab, norm.weight, bank[0]) # raw [64, H] for 8 ids + with pytest.raises((ValueError, RuntimeError)): + torch.ops.trtllm.k3_embed(ids(MAX_TOKENS + 1), tab) + + +# ---------------------------------------------------------------------------------------------------------------- +# Error table (python3 test_k3_embed.py report) and timing (python3 test_k3_embed.py time) +# ---------------------------------------------------------------------------------------------------------------- + + +def report() -> int: + _op() + with_norm = _flashinfer_norm() + rows = [] + with torch.inference_mode(): + print(f"{torch.cuda.get_device_name()}; table [{VOCAB}, {HIDDEN}] bf16") + for m in TOKENS: + rows += measure_embed(m) + if with_norm: + rows += measure_norm(m) + for m in OUTSIDE_TOKENS: + for dtype in (torch.int32, torch.int64): + rows += measure_outside(m, dtype, with_norm) + rows += measure_every_row(with_norm) + for ms in GRAPH_FAMILIES: + rows += measure_graph(ms, with_norm) + what = { + "k3_embed": "identical: rows = F.embedding (zero rows for ids outside [0, V)); rerun = first run; in a " + "graph replay also = the eager op", + "k3_embed_norm": "identical: normed = the RMSNorm module on the reference rows; raw (the bank slot) = the " + "reference rows; the bank's other slots untouched; rerun = first run; in a graph replay also = the eager op", + } + for name in ("k3_embed", "k3_embed_norm"): + print(f"\n## {name}\n\n{what[name]}\n") + if name == "k3_embed_norm" and not with_norm: + print( + "skipped: needs flashinfer's CuTe DSL RMSNorm (the kernel k3_embed_norm reproduces)" + ) + continue + print("| M | case | identical | result |") + print("| :-- | :-- | :-- | :-- |") + for r in (r for r in rows if r["op"] == name): + verdict = "PASS" if r["ok"] else "FAIL" + print(f"| {r['m']} | {r['case']} | {r['identical']} | {verdict} |") + ok_all = all(r["ok"] for r in rows) + print("\nALL PASS" if ok_all else "\nFAIL") + return 0 if ok_all else 1 + + +def time_graph(body, calls, replays=15): + """Per-call us of a CUDA graph of ``calls`` back-to-back ``body(i)``: median (min, max) over ``replays`` replays. + ``body(-1)`` runs first, outside capture.""" + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + body(-1) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + for i in range(calls): + body(i) + torch.cuda.synchronize() + for _ in range(3): + graph.replay() + torch.cuda.synchronize() + per_call = [] + for _ in range(replays): + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + graph.replay() + end.record() + torch.cuda.synchronize() + per_call.append(start.elapsed_time(end) * 1e3 / calls) + return statistics.median(per_call), min(per_call), max(per_call) + + +def timing() -> None: + _op() + with_norm = _flashinfer_norm() + calls = 32 + with torch.inference_mode(): + tab = table() + norm = layer0_norm(EPS) if with_norm else None + gen = torch.Generator(device="cuda").manual_seed(11) + print(f"{torch.cuda.get_device_name()}; table [{VOCAB}, {HIDDEN}] bf16; graphs of {calls} back-to-back calls, " + "each on its own random int32 ids, 15 replays: median (min-max) us per call") # fmt: skip + print("| M | F.embedding | k3_embed | F.embedding + snapshot + RMSNorm | k3_embed + snapshot + RMSNorm " + "| k3_embed_norm |") # fmt: skip + print("| --: | --: | --: | --: | --: | --: |") + for m in TOKENS: + idss = [ + torch.randint(0, VOCAB, (m,), generator=gen, device="cuda", dtype=torch.int32) + for _ in range(calls) + ] + bank = new_bank(m) + + def unfused(rows, bank=bank): + """Layer 0 without k3_embed_norm: the bank's snapshot of the rows and the input RMSNorm.""" + bank[0].copy_(rows) + return norm(rows) + + arms = [ + lambda i: F.embedding(idss[i], tab), + lambda i: torch.ops.trtllm.k3_embed(idss[i], tab), + ] + if with_norm: + arms += [ + lambda i: unfused(F.embedding(idss[i], tab)), + lambda i: unfused(torch.ops.trtllm.k3_embed(idss[i], tab)), + lambda i: embed_norm_into(idss[i], norm, bank[0]), + ] + res = [[] for _ in arms] + for rep in range(3): # alternating order + order = range(len(arms)) if rep % 2 == 0 else reversed(range(len(arms))) + for a in order: + res[a].append(time_graph(arms[a], calls)) + cells = [] + for timings in res: + meds = sorted(x[0] for x in timings) + lo, hi = min(x[1] for x in timings), max(x[2] for x in timings) + cells.append(f"{meds[1]:.2f} ({lo:.2f}-{hi:.2f})") + cells += ["n/a"] * (5 - len(cells)) + print(f"| {m} | " + " | ".join(cells) + " |", flush=True) + + +if __name__ == "__main__": + if len(sys.argv) > 1 and sys.argv[1] == "time": + timing() + elif len(sys.argv) > 1 and sys.argv[1] == "report": + sys.exit(report()) + else: + sys.exit(pytest.main([__file__, "-q", "-p", "no:cacheprovider", *sys.argv[1:]])) diff --git a/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_head_gemv.py b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_head_gemv.py new file mode 100644 index 000000000000..7e72d01c8481 --- /dev/null +++ b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_head_gemv.py @@ -0,0 +1,252 @@ +# 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::k3_head_gemv (an lm_head vocab shard, stream-K, M <= 8) at the Kimi K3 TP16 shard [10240, 7168], at every +M in 1..8 (default schedule): error against an fp64 product and cuBLAS F.linear, run-to-run identical bits, each M's +rows bit-identical to the same rows of the 8-row call, and calls of different M back to back leaving the workspace +counters at zero (the next call's bits unchanged). The caller-owned workspace: two of one shape interleaved, a CUDA +graph replayed between eager calls on the same workspace, a workspace of another shape refused, and creation refused +under capture. Stream-K weights with fewer 128 x 128 tiles than SMs, or exactly as many, complete and match an fp64 +product at every M.""" + +import functools +import os +import subprocess +import sys + +import pytest +import torch +import torch.nn.functional as F + + +def _is_sm100() -> bool: + if not torch.cuda.is_available(): + return False + major, _ = torch.cuda.get_device_capability() + return major == 10 + + +pytestmark = pytest.mark.skipif(not _is_sm100(), reason="needs an SM100-family GPU") + +M_ALL = list(range(1, 9)) +TOL = 8e-3 +VOCAB_SHARD, HIDDEN = 10240, 7168 # 163840 / 16 rows per rank +# [N, K] weights with fewer 128 x 128 tiles than an SM100 GPU has SMs (56, 108, 120 and 144 tiles). +FEW_TILES = [(128, 7168), (1152, 1536), (1280, 1536), (2304, 1024)] +# For the few-tile child; import, compile and the calls take about a minute. +FEW_TILES_TIMEOUT_S = 300 + +# The few-tile shapes ("NxK" arguments) one after another at every M: prints each call's error against an fp64 product +# once the call has completed. +_FEW_TILES_CHILD = r""" +import sys +import torch +import tensorrt_llm # noqa: F401 +from tensorrt_llm._torch.cute_dsl_kernels.k3_head_gemv import op # noqa: F401 +gen = torch.Generator(device="cuda").manual_seed(20261002) +for shape in sys.argv[1:]: + n, k = map(int, shape.split("x")) + w = (torch.randn(n, k, generator=gen, device="cuda") * 0.02).bfloat16() + x8 = torch.randn(8, k, generator=gen, device="cuda").bfloat16() + ws = op.K3HeadGemvWorkspace.create(n, k, w.device) + for m in range(1, 9): + x = x8[:m].contiguous() + y = torch.ops.trtllm.k3_head_gemv(x, w, ws.partials, ws.flags, ws.claim) + torch.cuda.synchronize() + ref = x.double() @ w.double().t() + print("REL", n, k, m, ((y.double() - ref).abs().max() / ref.abs().max()).item(), flush=True) +""" + + +def _ops(): + import tensorrt_llm # noqa: F401 + from tensorrt_llm._torch.cute_dsl_kernels.k3_head_gemv import op # noqa: F401 + + return torch.ops.trtllm + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.int16) + + +def _rel(y: torch.Tensor, ref: torch.Tensor) -> float: + return (y.double() - ref.double()).abs().max().item() / ref.double().abs().max().item() + + +@functools.lru_cache(maxsize=None) +def _inputs(): + gen = torch.Generator(device="cuda").manual_seed(20261001) + w = (torch.randn(VOCAB_SHARD, HIDDEN, generator=gen, device="cuda") * 0.02).bfloat16() + x8 = torch.randn(8, HIDDEN, generator=gen, device="cuda").bfloat16() + return w, x8 + + +@functools.lru_cache(maxsize=None) +def _workspace(): + _ops() + from tensorrt_llm._torch.cute_dsl_kernels.k3_head_gemv import op + + return op.K3HeadGemvWorkspace.create(VOCAB_SHARD, HIDDEN, torch.device("cuda")) + + +def _call(x, w, ws): + return torch.ops.trtllm.k3_head_gemv(x, w, ws.partials, ws.flags, ws.claim) + + +@pytest.mark.parametrize("m", M_ALL) +def test_k3_head_gemv(m): + _ops() + from tensorrt_llm._torch.cute_dsl_kernels.k3_head_gemv import op + + w, x8 = _inputs() + ws = _workspace() + x = x8[:m].contiguous() + assert op.supports(x, w) + y = _call(x, w, ws) + y8 = _call(x8, w, ws) + again = _call(x, w, ws) + ref = x.double() @ w.double().t() + stock = F.linear(x, w) + det = torch.equal(_bits(y), _bits(again)) + minv = torch.equal(_bits(y), _bits(y8[:m])) + abs_err = (y.double() - ref).abs().max().item() + print(f"OPCHECK op=k3_head_gemv case=lm_head_shard M={m} abs={abs_err:.3e} rel={_rel(y, ref):.3e} " + f"vs_stock={_rel(y, stock):.2e} det={det} rows_as_m8={minv}") # fmt: skip + assert y.shape == (m, VOCAB_SHARD) + assert _rel(y, ref) <= TOL and _rel(y, stock) <= TOL + assert det and minv + assert not ws.flags.any() and not ws.claim.any() + + +def test_k3_head_gemv_mixed_m_sequence(): + """M 8, 1, 5, 8, 3 back to back on one stream and one workspace: each the bits of its own call.""" + w, x8 = _inputs() + ws = _workspace() + single = {m: _call(x8[:m].contiguous(), w, ws) for m in (1, 3, 5, 8)} + seq = [_call(x8[:m].contiguous(), w, ws) for m in (8, 1, 5, 8, 3)] + for m, y in zip((8, 1, 5, 8, 3), seq): + assert torch.equal(_bits(y), _bits(single[m])) + + +def test_k3_head_gemv_two_workspaces_interleaved(): + """Two workspaces of one shape (e.g. two heads of the same shard size), their calls alternating on one stream: + every call the bits of its own single call, both workspaces' counters at zero after.""" + _ops() + from tensorrt_llm._torch.cute_dsl_kernels.k3_head_gemv import op + + w, x8 = _inputs() + w2 = (w.float() * -0.5).bfloat16() + wa, wb = _workspace(), op.K3HeadGemvWorkspace.create(VOCAB_SHARD, HIDDEN, w.device) + single = {(m, i): _call(x8[:m].contiguous(), (w, w2)[i], wa) for m in (1, 4, 8) for i in (0, 1)} + for m in (8, 1, 4, 1, 8): + for i, ws in ((0, wa), (1, wb)): + assert torch.equal( + _bits(_call(x8[:m].contiguous(), (w, w2)[i], ws)), _bits(single[(m, i)]) + ) + assert not wa.flags.any() and not wb.flags.any() + + +def test_k3_head_gemv_graph_replays_between_eager_calls(): + """A CUDA graph of M 1, 8 and 3 calls on one workspace, replayed three times with eager calls on the same + workspace between the replays: every result the bits of its eager call, the counters at zero after.""" + w, x8 = _inputs() + ws = _workspace() + xs = {m: x8[:m].contiguous() for m in (1, 3, 8)} + want = {m: _call(xs[m], w, ws) for m in xs} # also compiles before the capture + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = {m: _call(xs[m], w, ws) for m in (1, 8, 3)} + for _ in range(3): + graph.replay() + between = _call(xs[8], w, ws) + torch.cuda.synchronize() + assert all(torch.equal(_bits(out[m]), _bits(want[m])) for m in xs) + assert torch.equal(_bits(between), _bits(want[8])) + assert not ws.flags.any() + + +def test_k3_head_gemv_refuses_a_workspace_of_another_shape(): + """A workspace created for another weight shape is refused (too small, or its flag words laid out for another + tile count) instead of running the kernel over it.""" + _ops() + from tensorrt_llm._torch.cute_dsl_kernels.k3_head_gemv import op + + w, x8 = _inputs() + other = op.K3HeadGemvWorkspace.create(VOCAB_SHARD // 2, HIDDEN, w.device) + with pytest.raises(ValueError, match="workspace"): + _call(x8[:1].contiguous(), w, other) + + +def test_k3_head_gemv_workspace_not_created_under_capture(): + _ops() + from tensorrt_llm._torch.cute_dsl_kernels.k3_head_gemv import op + + graph = torch.cuda.CUDAGraph() + with pytest.raises(RuntimeError, match="capture"): + with torch.cuda.graph(graph): + op.K3HeadGemvWorkspace.create(VOCAB_SHARD, HIDDEN, torch.device("cuda")) + + +def _tiles_equal_to_sms(sms): + """An [N, K] weight with exactly ``sms`` 128 x 128 tiles, each 128-row tile split over at least 2 k-tiles.""" + k_tiles = max([d for d in range(2, 17) if sms % d == 0] or [1]) + return 128 * (sms // k_tiles), 128 * k_tiles + + +def test_k3_head_gemv_few_tiles(): + """Stream-K weights with fewer 128 x 128 tiles than SMs, and one with exactly as many: every call completes and + matches an fp64 product. A stuck call would block its process, so the shapes run in a child process with a + deadline.""" + sms = torch.cuda.get_device_properties(0).multi_processor_count + shapes = [(n, k) for n, k in FEW_TILES if (n // 128) * (k // 128) < sms] + assert shapes + shapes.append(_tiles_equal_to_sms(sms)) + # Without a compute-sanitizer target's injection variables: they would load the tool into the child too, and + # there CUDA initialization can stall. + env = { + k: v + for k, v in os.environ.items() + if not k.startswith(("NV_SANITIZER_", "NVIDIA_PROCESS_INJECTION_")) + and k != "NVTX_INJECTION64_PATH" + } + child = subprocess.Popen( + [sys.executable, "-c", _FEW_TILES_CHILD, *(f"{n}x{k}" for n, k in shapes)], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + env=env, + ) + try: + out = child.communicate(timeout=FEW_TILES_TIMEOUT_S)[0] + except subprocess.TimeoutExpired: + child.kill() + out = child.communicate()[0] + f"\nno result after {FEW_TILES_TIMEOUT_S} s" + rel = {shape: {} for shape in shapes} + for _, n, k, m, r in (ln.split() for ln in out.splitlines() if ln.startswith("REL ")): + rel[(int(n), int(k))][int(m)] = float(r) + for (n, k), by_m in rel.items(): + print(f"OPCHECK op=k3_head_gemv case=few_tiles N={n} K={k} tiles={(n // 128) * (k // 128)} sms={sms} " + f"rel={by_m}") # fmt: skip + done = all(set(by_m) == set(M_ALL) for by_m in rel.values()) + assert child.returncode == 0 and done, out[-3000:] + assert max(max(by_m.values()) for by_m in rel.values()) <= TOL, rel + + +@pytest.mark.parametrize("m", [0, 9, 16]) +def test_token_limit(m): + _ops() + from tensorrt_llm._torch.cute_dsl_kernels.k3_head_gemv import op + + w, _ = _inputs() + assert not op.supports(torch.zeros(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), w) diff --git a/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_kda_attn.py b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_kda_attn.py new file mode 100644 index 000000000000..61f0debe682c --- /dev/null +++ b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_kda_attn.py @@ -0,0 +1,322 @@ +# 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::k3_kda_attn`` (Kimi K3's fused KDA projection and speculative verify of one request's 8 tokens in one +launch, the DSpark batch-1 path) at the per-rank TP16 shape (W [3208, 7168], 6 heads, K = V = 128, conv width 4). + +The reference is the same x through ``trtllm::k3_kda_qkvg`` (the projection stream alone, the same stream clusters; +its Lamport buffers decoded into the ``[q | k | v | og | f_a | b]`` rows) and ``trtllm::k3_kda_verify`` on those rows, +on a copy of the same pools: the same arithmetic, so outputs, conv caches, pool state and the drafts' records are +compared bit for bit. A request keeps its slot for several rounds, so its pending count (the drafts the previous +round accepted) runs through 0..7, then the next request starts on another slot; slots outside the batch stay +untouched; pools dense and strided. Also: CUDA-graph replays with rewritten inputs, the launch counter near 2^31 +(buffers inside guard bands: nothing written outside them, the same bits as a run from zero), and the slot and the +pending counts given as slices of longer index tensors at any element offset. + + pytest test_k3_kda_attn.py +""" + +import pytest +import torch + +H = 6 +K = V = 128 +W = 4 +HK = H * K +PROJ = 4 * HK + K + H + 2 # 3208 +K_IN = 7168 +NUM_SPEC = 7 +NT = NUM_SPEC + 1 +LOWER_BOUND = -5.0 +EPS = 1e-5 +SCALE = K**-0.5 +POOL = 11 +SLOT_ROUNDS = 9 # pending 0, then 1, 6, 3, 0, 5, 2, 7, 4 (pending_schedule) + + +def _sm100() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10 + + +pytestmark = pytest.mark.skipif(not _sm100(), reason="needs SM100 (tcgen05, TMA, clusters)") + + +def _ops(): + import tensorrt_llm # noqa: F401 (registers the trtllm ops) + from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_attn import op # noqa: F401 + from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_verify import op as _verify # noqa: F401 + + return op + + +def make_weights(seed: int) -> dict: + g = torch.Generator(device="cuda").manual_seed(seed) + + def rnd(*s, scale=1.0): + return torch.randn(*s, generator=g, device="cuda") * scale + + return { + "w": rnd(PROJ, K_IN, scale=0.02).bfloat16(), "w_fb": (rnd(HK, K) * 0.05).bfloat16(), + "w_q": rnd(HK, W, scale=0.3), "w_k": rnd(HK, W, scale=0.3), "w_v": rnd(HK, W, scale=0.3), + "a_log": rnd(H, scale=0.5), "dt_bias": rnd(HK, scale=0.5), "onorm_w": (1 + 0.1 * rnd(V)).float(), + } # fmt: skip + + +def make_pools(seed: int, layout: str) -> dict: + """Conv caches [pool, HK, W - 1 + NUM_SPEC] (dim-contiguous), the SSM state [pool, H, V, K] (dense, or each slot + followed by the cache manager's conv bytes), the drafts' records and the pending counts.""" + g = torch.Generator(device="cuda").manual_seed(seed) + s = W - 1 + NUM_SPEC + p = {name: (torch.randn(POOL, s, HK, generator=g, device="cuda") * 0.5).transpose(1, 2) + for name in ("cs_q", "cs_k", "cs_v")} # fmt: skip + state = torch.randn(POOL, H, V, K, generator=g, device="cuda") * 0.05 + if layout == "strided": + buf = torch.zeros(POOL, state[0].numel() + 3 * HK * (W - 1), device="cuda") + view = buf[:, : state[0].numel()].view(state.shape) + view.copy_(state) + state = view + p["state"] = state + p["state_tok"] = torch.zeros(POOL, NUM_SPEC, H, V, K, device="cuda") + p["pending"] = torch.zeros(POOL, dtype=torch.int32, device="cuda") + return p + + +def clone_pools(p: dict) -> dict: + out = {} + for name, t in p.items(): + if name.startswith("cs_"): + out[name] = t.transpose(1, 2).clone().transpose(1, 2) + elif not t.is_contiguous(): + buf = torch.zeros(t.shape[0], t.stride(0), dtype=t.dtype, device=t.device) + view = buf[:, : t[0].numel()].view(t.shape) + view.copy_(t) + out[name] = view + else: + out[name] = t.clone() + return out + + +POOL_NAMES = ("cs_q", "cs_k", "cs_v", "state", "state_tok") + + +def pending_schedule(rnd: int) -> int: + return (5 * rnd + 1) % (NUM_SPEC + 1) + + +class Fused: + """k3_kda_attn on its own pools and Lamport buffers.""" + + def __init__(self, wt, pools, bufs=None): + op = _ops() + self.wt, self.p = wt, pools + self.bufs = ( + bufs if bufs is not None else op.make_buffers(torch.device("cuda"), op.FUSED_CTAS) + ) + + def __call__(self, x, slot): + wt, p = self.wt, self.p + return torch.ops.trtllm.k3_kda_attn( + x, wt["w"], wt["w_fb"], wt["w_q"], wt["w_k"], wt["w_v"], wt["a_log"], wt["dt_bias"], wt["onorm_w"], + p["cs_q"], p["cs_k"], p["cs_v"], p["state"], p["state_tok"], slot, p["pending"], *self.bufs, NUM_SPEC, + LOWER_BOUND, SCALE, EPS, + ) # fmt: skip + + +class Unfused: + """k3_kda_qkvg's rows (its Lamport buffers decoded), then k3_kda_verify on them.""" + + def __init__(self, wt, pools): + op = _ops() + self.wt, self.p = wt, pools + self.bufs = op.make_buffers(torch.device("cuda"), op.CTAS) + + def rows(self, x): + p1, part, epoch = self.bufs + buf = int(epoch[0].item()) % 3 + torch.ops.trtllm.k3_kda_qkvg(x, self.wt["w"], p1, part, epoch) + qkfa = p1.view(3, 8, 2 * HK + K)[buf].view(torch.bfloat16) + parts = part.view(3, 3, 2, 8, HK)[buf].view(torch.float32) + v, og, b = ((parts[r, 0] + parts[r, 1]).bfloat16() for r in range(3)) + rows = torch.zeros(NT, PROJ, dtype=torch.bfloat16, device="cuda") + rows[:, : 2 * HK] = qkfa[:, : 2 * HK] + rows[:, 2 * HK : 3 * HK] = v + rows[:, 3 * HK : 4 * HK] = og + rows[:, 4 * HK : 4 * HK + K] = qkfa[:, 2 * HK :] + rows[:, 4 * HK + K : 4 * HK + K + H] = b[:, :H] + return rows + + def __call__(self, x, slot): + wt, p = self.wt, self.p + return torch.ops.trtllm.k3_kda_verify( + self.rows(x), wt["w_fb"], wt["w_q"], wt["w_k"], wt["w_v"], wt["a_log"], wt["dt_bias"], wt["onorm_w"], + p["cs_q"], p["cs_k"], p["cs_v"], p["state"], p["state_tok"], slot, p["pending"], NUM_SPEC, LOWER_BOUND, + SCALE, EPS, None, + ) # fmt: skip + + +def same_pools(a: dict, b: dict) -> bool: + return all(torch.equal(a[n], b[n]) for n in POOL_NAMES) + + +def others_untouched(p: dict, before: dict, slot: int) -> bool: + others = [s for s in range(POOL) if s != slot] + return all(torch.equal(p[n][others], before[n][others]) for n in POOL_NAMES) + + +def slot_order(seed: int): + return torch.randperm(POOL, generator=torch.Generator().manual_seed(seed)).tolist() + + +@pytest.mark.parametrize("layout", ["dense", "strided"]) +def test_rounds(layout): + """Two requests, one after the other, each on its slot for SLOT_ROUNDS rounds: bits of the unfused path.""" + with torch.inference_mode(): + wt = make_weights(100) + pools = make_pools(200, layout) + fused, ref = Fused(wt, clone_pools(pools)), Unfused(wt, clone_pools(pools)) + gen = torch.Generator(device="cuda").manual_seed(400) + order = slot_order(300) + for rnd in range(2 * SLOT_ROUNDS): + s = order[rnd // SLOT_ROUNDS] + slot = torch.tensor([s], dtype=torch.int32, device="cuda") + x = torch.randn(NT, K_IN, generator=gen, device="cuda").bfloat16() + before = {n: fused.p[n].clone() for n in POOL_NAMES} + got, want = fused(x, slot), ref(x, slot) + torch.cuda.synchronize() + pend = int(fused.p["pending"][s]) + assert torch.equal(got, want), (rnd, pend) + assert bool(torch.isfinite(got.float()).all()), (rnd, pend) + assert same_pools(fused.p, ref.p), (rnd, pend) + assert others_untouched(fused.p, before, s), (rnd, pend) + for p in (fused.p, ref.p): + p["pending"][s] = pending_schedule(rnd) + + +def test_graph_replay(): + """One captured call replayed with x, the slot and the pending counts rewritten in place: the bits of eager calls on + a copy of the pools.""" + op = _ops() + with torch.inference_mode(): + wt = make_weights(110) + pools = make_pools(210, "dense") + eager, graphed = Fused(wt, clone_pools(pools)), Fused(wt, clone_pools(pools)) + gen = torch.Generator(device="cuda").manual_seed(410) + order = slot_order(310) + x_in = torch.zeros(NT, K_IN, dtype=torch.bfloat16, device="cuda") + s_in = torch.zeros(1, dtype=torch.int32, device="cuda") + # Compiles outside capture, on scratch pools and buffers. + Fused(wt, clone_pools(pools), op.make_buffers(torch.device("cuda"), op.FUSED_CTAS))( + x_in, s_in + ) + torch.cuda.synchronize() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + y = graphed(x_in, s_in) + torch.cuda.current_stream().wait_stream(stream) + for rnd in range(2 * SLOT_ROUNDS): + s = order[rnd // SLOT_ROUNDS] + slot = torch.tensor([s], dtype=torch.int32, device="cuda") + x = torch.randn(NT, K_IN, generator=gen, device="cuda").bfloat16() + x_in.copy_(x) + s_in.copy_(slot) + graph.replay() + want = eager(x, slot) + torch.cuda.synchronize() + assert torch.equal(y, want), rnd + assert same_pools(graphed.p, eager.p), rnd + for p in (graphed.p, eager.p): + p["pending"][s] = pending_schedule(rnd) + + +def _alone(path, x, slot): + """One launch, complete before the next: the head CTAs read the slot's pools before their grid-dependency wait, + so a launch right behind another on the same slot could read them mid-update (the model never runs one layer's + KDA twice in a row; its next launch on these pools is a step later).""" + out = path(x, slot) + torch.cuda.synchronize() + return out + + +def _guarded(numel: int, dtype, fill: int): + """``numel`` all-ones words with a whole set of the op's buffers of ``fill`` words on each side.""" + guard = numel + 4096 + big = torch.full((guard + numel + guard,), fill, dtype=dtype, device="cuda") + big[guard : guard + numel] = -1 + return big, big[guard : guard + numel], guard + + +def test_epoch_wrap(): + """Every CTA's counter preset to 2^31 - 2, as after 2^31 launches on a device (the buffers are shared by every KDA + layer): four launches write nothing outside the buffers, give the bits of a run from zero, and keep the counter + in 0..2.""" + op = _ops() + with torch.inference_mode(): + wt = make_weights(120) + pools = make_pools(220, "dense") + gen = torch.Generator(device="cuda").manual_seed(420) + xs = [torch.randn(NT, K_IN, generator=gen, device="cuda").bfloat16() for _ in range(4)] + slot = torch.tensor([3], dtype=torch.int32, device="cuda") + fresh = Fused(wt, clone_pools(pools)) + want = [_alone(fresh, x, slot) for x in xs] + big1, p1, g1 = _guarded(op.P1_NUMEL, torch.int16, 0x1234) + big2, part, g2 = _guarded(op.PART_NUMEL, torch.int32, 0x12345678) + epoch = torch.full((op.FUSED_CTAS,), 2**31 - 2, dtype=torch.int32, device="cuda") + wrapped = Fused(wt, clone_pools(pools), (p1, part, epoch)) + got = [_alone(wrapped, x, slot) for x in xs] + for big, guard, numel, fill in ( + (big1, g1, op.P1_NUMEL, 0x1234), + (big2, g2, op.PART_NUMEL, 0x12345678), + ): + assert bool((big[:guard] == fill).all()), "stores before the buffers" + assert bool((big[guard + numel :] == fill).all()), "stores after the buffers" + assert all(torch.equal(a, b) for a, b in zip(got, want)) + assert same_pools(wrapped.p, fresh.p) + assert bool(((epoch >= 0) & (epoch < 3)).all()), epoch.unique().tolist() + + +def _at_offset(t: torch.Tensor, offset: int) -> torch.Tensor: + """``t``'s values in a longer tensor, starting at element ``offset`` (a slice such as + ``state_indices[num_prefills:]``: its data pointer is aligned to the element only).""" + buf = torch.zeros(offset + t.numel(), dtype=t.dtype, device=t.device) + buf[offset:] = t + view = buf[offset:] + assert view.data_ptr() % 16 == offset * t.element_size() % 16 + return view + + +@pytest.mark.parametrize("offset", [0, 1, 2, 3]) +@pytest.mark.parametrize("which", ["slot", "pending"]) +def test_index_offset(which, offset): + """The slot or the pending counts as a slice of a longer index tensor that starts at element ``offset``: over two + rounds, the second on a nonzero pending count, the outputs and the pools bit for bit those of the same indices in + tensors of their own.""" + with torch.inference_mode(): + wt = make_weights(130) + pools = make_pools(230, "dense") + gen = torch.Generator(device="cuda").manual_seed(430) + s = 5 + slot = torch.tensor([s], dtype=torch.int32, device="cuda") + ref, got = Fused(wt, clone_pools(pools)), Fused(wt, clone_pools(pools)) + got_slot = _at_offset(slot, offset) if which == "slot" else slot + if which == "pending": + got.p["pending"] = _at_offset(got.p["pending"], offset) + for rnd in range(2): + x = torch.randn(NT, K_IN, generator=gen, device="cuda").bfloat16() + want, out = _alone(ref, x, slot), _alone(got, x, got_slot) + assert torch.equal(out, want), rnd + assert same_pools(got.p, ref.p), rnd + for p in (ref.p, got.p): + p["pending"][s] = pending_schedule(rnd) diff --git a/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_kda_decode_attn.py b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_kda_decode_attn.py new file mode 100644 index 000000000000..5aad370eabd3 --- /dev/null +++ b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_kda_decode_attn.py @@ -0,0 +1,523 @@ +# 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::k3_kda_decode_attn`` (Kimi K3's fused KDA projection + plain decode of R requests of one token) at the +per-rank TP16 shape (6 heads, K = V = 128, conv width 4), against ``trtllm::kda_decode``. + +The reference is the model's plain-decode path fed the very projection rows the fused kernel computes: the stream +alone (``trtllm::k3_kda_qkvg``, the same stream clusters) on the same x, its Lamport buffers decoded into the +``[q | k | v | og | f_a | b]`` rows, f_b as a bf16 ``F.linear``, then ``trtllm::kda_decode`` (native, indexed state +pool, packed conv pool updated in place). A float64 torch decode on the same rows bounds both. + +Checks over rounds of distinct slots from the same initial pools: every request's output against kda_decode and +float64 (fp32 tolerance), the state rows against both, the conv pool against kda_decode bit for bit (raw inputs), +slots outside the batch untouched; pools dense and with strided slots (the cache manager's interleaving); slots +given as a slice of a longer index tensor at any element offset; the +schedule twice bit-identical; CUDA-graph replays with rewritten inputs; the launch counter near 2^31 (nothing +written outside the buffers); launches interleaved with ``trtllm::k3_kda_attn`` on one shared buffer set, as the +model runs them. + + pytest test_k3_kda_decode_attn.py + python3 test_k3_kda_decode_attn.py report +""" + +import sys + +import pytest +import torch +import torch.nn.functional as F + +H = 6 +K = V = 128 +HK = H * K +W = 4 +PROJ = 3208 +K_IN = 7168 +LOWER_BOUND = -5.0 +EPS = 1e-5 +SCALE = K**-0.5 +POOL = 11 +ROUNDS = 4 +NUM_SPEC = 7 # trtllm::k3_kda_attn verifies one request's NUM_SPEC + 1 tokens per launch +TOL_OUT = 2e-2 +TOL_STATE = 1e-3 +CONV_PAD = 128 # extra bf16 per conv slot in the strided layout (a multiple of 8) +SSM_PAD = ( + 3 * HK * (W - 1) +) # extra fp32 per state slot in the strided layout (the manager's conv bytes) + + +def _sm100() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10 + + +pytestmark = pytest.mark.skipif(not _sm100(), reason="needs SM100 (tcgen05, TMA, clusters)") + + +def _ops(): + import tensorrt_llm # noqa: F401 (registers the trtllm ops) + from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_attn import op # noqa: F401 + + +def make_weights(seed: int) -> dict: + g = torch.Generator(device="cuda").manual_seed(seed) + + def rnd(*s, scale=1.0): + return torch.randn(*s, generator=g, device="cuda") * scale + + conv = ( + rnd(3, HK, W, scale=0.3).bfloat16().float() + ) # bf16 values: the native path reads them as bf16 + return { + "w": rnd(PROJ, K_IN, scale=0.02).bfloat16(), "w_fb": rnd(HK, K, scale=0.05).bfloat16(), + "w_q": conv[0].contiguous(), "w_k": conv[1].contiguous(), "w_v": conv[2].contiguous(), + "w_t": [conv[i].t().bfloat16().contiguous() for i in range(3)], + "a_log": rnd(H, scale=0.5), "dt_bias": rnd(HK, scale=0.5), "onorm_w": (1 + 0.1 * rnd(V)).float(), + } # fmt: skip + + +def make_pools(seed: int, layout: str) -> dict: + """conv bf16 [POOL, 3 HK, W - 1] and the state fp32 [POOL, H, V, K]; "strided": each slot padded.""" + g = torch.Generator(device="cuda").manual_seed(seed) + conv = (torch.randn(POOL, 3 * HK, W - 1, generator=g, device="cuda") * 0.5).bfloat16() + state = torch.randn(POOL, H, V, K, generator=g, device="cuda") * 0.05 + if layout == "strided": + conv_buf = torch.zeros( + POOL, 3 * HK * (W - 1) + CONV_PAD, dtype=torch.bfloat16, device="cuda" + ) + conv_view = conv_buf[:, : 3 * HK * (W - 1)].view(POOL, 3 * HK, W - 1) + conv_view.copy_(conv) + conv = conv_view + st_buf = torch.zeros(POOL, H * V * K + SSM_PAD, device="cuda") + st_view = st_buf[:, : H * V * K].view(POOL, H, V, K) + st_view.copy_(state) + state = st_view + return {"conv": conv, "state": state} + + +def clone_pools(p: dict) -> dict: + out = {} + for name, t in p.items(): + if t.is_contiguous(): + out[name] = t.clone() + else: + buf = torch.zeros(t.shape[0], t.stride(0), dtype=t.dtype, device=t.device) + view = buf[:, : t[0].numel()].view(t.shape) + view.copy_(t) + out[name] = view + return out + + +def make_slots(num_requests: int, seed: int) -> torch.Tensor: + perm = torch.randperm(POOL, generator=torch.Generator().manual_seed(seed)) + return perm[:num_requests].to(torch.int32).cuda() + + +def rel(a: torch.Tensor, b: torch.Tensor) -> float: + return ((a.float() - b.float()).abs().max() / b.float().abs().max().clamp_min(1e-6)).item() + + +class Rows: + """The projection rows the fused kernel computes for x: the stream alone, its Lamport buffers decoded.""" + + def __init__(self, wt): + from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_attn import op + + self.wt = wt + self.bufs = op.make_buffers(torch.device("cuda"), op.CTAS) + + def __call__(self, x: torch.Tensor) -> torch.Tensor: + p1, part, epoch = self.bufs + buf = int(epoch[0].item()) % 3 + torch.ops.trtllm.k3_kda_qkvg(x, self.wt["w"], p1, part, epoch) + n = x.shape[0] + qkfa = p1.view(3, 8, 2 * HK + K)[buf, :n].view(torch.bfloat16) + parts = part.view(3, 3, 2, 8, HK)[buf].view(torch.float32) + v, og, b = ((parts[r, 0, :n] + parts[r, 1, :n]).bfloat16() for r in range(3)) + rows = torch.zeros(n, PROJ, dtype=torch.bfloat16, device="cuda") + rows[:, : 2 * HK] = qkfa[:, : 2 * HK] + rows[:, 2 * HK : 3 * HK] = v + rows[:, 3 * HK : 4 * HK] = og + rows[:, 4 * HK : 4 * HK + K] = qkfa[:, 2 * HK :] + rows[:, 4 * HK + K : 4 * HK + K + H] = b[:, :H] + return rows + + +class FusedPath: + def __init__(self, wt, pools): + from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_attn import op + + self.wt, self.p = wt, pools + self.bufs = op.make_buffers(torch.device("cuda"), op.FUSED_CTAS) + + def __call__(self, x, slots): + wt, p = self.wt, self.p + return torch.ops.trtllm.k3_kda_decode_attn( + x, wt["w"], wt["w_fb"], wt["w_q"], wt["w_k"], wt["w_v"], wt["a_log"], wt["dt_bias"], wt["onorm_w"], + p["conv"], p["state"], slots, *self.bufs, LOWER_BOUND, SCALE, EPS, + ) # fmt: skip + + +class NativePath: + """The model's plain decode (``forward_decode``): f_b as a bf16 F.linear, then trtllm::kda_decode.""" + + def __init__(self, wt, pools): + self.wt, self.p = wt, pools + + def __call__(self, rows, slots): + from tensorrt_llm._torch.modules.kimi_kda._kda_decode import run_kda_decode_fusion_cuda + + wt, p = self.wt, self.p + n = rows.shape[0] + + def heads(cols): + return cols.unflatten(-1, (H, K)).unsqueeze(0) + + g = F.linear(rows[:, 4 * HK : 4 * HK + K], wt["w_fb"]) + out = torch.empty(n, 1, H, V, dtype=torch.bfloat16, device="cuda") + conv = p["conv"] + run_kda_decode_fusion_cuda( + x_q=heads(rows[:, :HK]), x_k=heads(rows[:, HK : 2 * HK]), x_v=heads(rows[:, 2 * HK : 3 * HK]), + w_q_t=wt["w_t"][0], w_k_t=wt["w_t"][1], w_v_t=wt["w_t"][2], bias_q=None, bias_k=None, bias_v=None, + cs_q=conv[:, :HK], cs_k=conv[:, HK : 2 * HK], cs_v=conv[:, 2 * HK :], A_log=wt["a_log"], + g=heads(g), dt_bias=wt["dt_bias"], beta=rows[:, 4 * HK + K : 4 * HK + K + H].unsqueeze(0), + state=p["state"], onorm_g=heads(rows[:, 3 * HK : 4 * HK]), onorm_weight=wt["onorm_w"], out=out, + ssm_state_indices=slots, scale=SCALE, onorm_eps=EPS, lower_bound=LOWER_BOUND, + use_beta_sigmoid_in_kernel=True, update_conv_cache=True, + ) # fmt: skip + return out.view(n, H, V) + + +class F64Path: + """The plain KDA decode in float64 torch on the same rows (f_b in float64, rounded to bf16 as the GEMM's + output): conv4 + SiLU, q / k L2 norm (q scaled), beta sigmoid, the lower-bound gate, + S <- S d + beta (v - (S d) k) k^T, o = S q, the gated RMSNorm.""" + + def __init__(self, wt, pools): + self.wt = wt + self.conv = pools["conv"].double().clone() + self.state = pools["state"].double().clone() + + def __call__(self, rows, slots): + wt = self.wt + r = rows.double() + n = rows.shape[0] + conv_w = torch.stack([wt["w_q"], wt["w_k"], wt["w_v"]]).double() # [3, HK, W] + g = (r[:, 4 * HK : 4 * HK + K] @ wt["w_fb"].double().t()).bfloat16().double() + outs = torch.empty(n, H, V, dtype=torch.float64, device="cuda") + for i, s in enumerate(slots.tolist()): + new = r[i, : 3 * HK].view(3, HK) + win = self.conv[s].view(3, HK, W - 1) + u = torch.cat([win, new.unsqueeze(-1)], dim=-1) # [3, HK, W], oldest first + act = (u * conv_w).sum(-1) + act = act * torch.sigmoid(act) + self.conv[s] = u[:, :, 1:].reshape(3 * HK, W - 1) + q, k, v = (act[j].view(H, K) for j in range(3)) + q = q / torch.sqrt((q * q).sum(-1, keepdim=True) + 1e-6) * SCALE + k = k / torch.sqrt((k * k).sum(-1, keepdim=True) + 1e-6) + beta = torch.sigmoid(r[i, 4 * HK + K : 4 * HK + K + H]) + xg = torch.exp(wt["a_log"].double()).unsqueeze(-1) * ( + g[i].view(H, K) + wt["dt_bias"].double().view(H, K) + ) + decay = torch.exp(LOWER_BOUND * torch.sigmoid(xg)) + st = self.state[s] * decay.unsqueeze(1) # [H, V, K] * decay per key + res = (v - (st * k.unsqueeze(1)).sum(-1)) * beta.unsqueeze(-1) # [H, V] + st = st + res.unsqueeze(-1) * k.unsqueeze(1) + self.state[s] = st + o = (st * q.unsqueeze(1)).sum(-1) # [H, V] + rms = torch.rsqrt((o * o).mean(-1, keepdim=True) + EPS) + gate = torch.sigmoid(r[i, 3 * HK : 4 * HK].view(H, V)) + outs[i] = o * rms * wt["onorm_w"].double() * gate + return outs + + +def run_schedule(num_requests, layout="dense", seed=0, rounds=ROUNDS): + _ops() + wt = make_weights(100 + seed) + pools = make_pools(200 + seed, layout) + fused = FusedPath(wt, clone_pools(pools)) + native = NativePath(wt, clone_pools(pools)) + ref = F64Path(wt, pools) + rows_of = Rows(wt) + gen = torch.Generator(device="cuda").manual_seed(400 + seed) + metrics, outs = [], [] + for rnd in range(rounds): + slots = make_slots(num_requests, 300 + seed + rnd) + x = torch.randn(num_requests, K_IN, generator=gen, device="cuda").bfloat16() + state_before = fused.p["state"].clone() + conv_before = fused.p["conv"].clone() + out_f = fused(x, slots) + rows = rows_of(x) + out_n = native(rows, slots) + out_r = ref(rows, slots) + torch.cuda.synchronize() + idx = slots.long() + others = torch.ones(POOL, dtype=torch.bool, device="cuda") + others[idx] = False + m = dict(round=rnd) + m["out_vs_native"] = rel(out_f, out_n) + m["out_vs_f64"] = rel(out_f, out_r) + m["native_vs_f64"] = rel(out_n, out_r) + m["state_vs_f64"] = max(rel(fused.p["state"][s], ref.state[s]) for s in idx.tolist()) + m["state_vs_native"] = max( + rel(fused.p["state"][s], native.p["state"][s]) for s in idx.tolist() + ) + m["conv_bits_vs_native"] = torch.equal(fused.p["conv"], native.p["conv"]) + m["others_untouched"] = torch.equal( + fused.p["state"][others], state_before[others] + ) and torch.equal(fused.p["conv"][others], conv_before[others]) + metrics.append(m) + outs.append(out_f.clone()) + return metrics, outs + + +def round_ok(m) -> bool: + return (m["out_vs_native"] <= TOL_OUT and m["out_vs_f64"] <= TOL_OUT and m["state_vs_f64"] <= TOL_STATE + and m["state_vs_native"] <= TOL_STATE and m["conv_bits_vs_native"] and m["others_untouched"]) # fmt: skip + + +@pytest.mark.parametrize("num_requests", range(1, 9)) +@pytest.mark.parametrize("layout", ["dense", "strided"]) +def test_split(num_requests, layout): + with torch.inference_mode(): + metrics, _ = run_schedule(num_requests, layout) + bad = [m for m in metrics if not round_ok(m)] + assert not bad, bad + + +@pytest.mark.parametrize("num_requests", [1, 8]) +def test_deterministic(num_requests): + with torch.inference_mode(): + _, a = run_schedule(num_requests, seed=5) + _, b = run_schedule(num_requests, seed=5) + assert all(torch.equal(x, y) for x, y in zip(a, b)) + + +@pytest.mark.parametrize("num_requests", [1, 3, 8]) +def test_graph_replay(num_requests): + """Rounds captured once in a CUDA graph (static x and slots rewritten in place before every replay) against the + same rounds run eagerly from the same pools.""" + _ops() + with torch.inference_mode(): + wt = make_weights(7) + pools = make_pools(8, "dense") + eager = FusedPath(wt, clone_pools(pools)) + graphed = FusedPath(wt, clone_pools(pools)) + gen = torch.Generator(device="cuda").manual_seed(9) + xs = [ + torch.randn(num_requests, K_IN, generator=gen, device="cuda").bfloat16() + for _ in range(ROUNDS) + ] + slots = [make_slots(num_requests, 10 + r) for r in range(ROUNDS)] + want = [eager(x, s).clone() for x, s in zip(xs, slots)] + x_in = xs[0].clone() + s_in = slots[0].clone() + graphed(x_in, s_in) # compiles; this launch's result is discarded with the pools below + graphed.p = clone_pools(pools) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + y = graphed(x_in, s_in) + torch.cuda.current_stream().wait_stream(stream) + got = [] + for x, s in zip(xs, slots): + x_in.copy_(x) + s_in.copy_(s) + graph.replay() + got.append(y.clone()) + torch.cuda.synchronize() + assert all(torch.equal(a, b) for a, b in zip(want, got)) + assert torch.equal(eager.p["state"], graphed.p["state"]) and torch.equal( + eager.p["conv"], graphed.p["conv"] + ) + + +def _alone(path, x, slots): + """One launch, complete before the next: the head CTAs read the slots' pools before their grid-dependency wait, + so a launch right behind another on the same slots could read them mid-update (the model never runs one layer's + KDA twice in a row; its next launch on these pools is a step later).""" + out = path(x, slots) + torch.cuda.synchronize() + return out + + +def test_epoch_wrap(): + """Every CTA's counter preset to 2^31 - 2, as after 2^31 launches on a device (the buffers are shared by every KDA + layer): four launches write nothing outside the buffers (each inside guard bands), give the bits of a run from + zero, and keep the counter in 0..2.""" + _ops() + from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_attn import op + + with torch.inference_mode(): + wt = make_weights(13) + pools = make_pools(14, "dense") + gen = torch.Generator(device="cuda").manual_seed(15) + xs = [torch.randn(4, K_IN, generator=gen, device="cuda").bfloat16() for _ in range(4)] + slots = make_slots(4, 16) + fresh = FusedPath(wt, clone_pools(pools)) + want = [_alone(fresh, x, slots) for x in xs] + bands = [] + for numel, dtype, fill in ( + (op.P1_NUMEL, torch.int16, 0x1234), + (op.PART_NUMEL, torch.int32, 0x12345678), + ): + guard = numel + 4096 # a whole set of buffers on each side + big = torch.full((guard + numel + guard,), fill, dtype=dtype, device="cuda") + big[guard : guard + numel] = -1 + bands.append((big, guard, numel, fill)) + epoch = torch.full((op.FUSED_CTAS,), 2**31 - 2, dtype=torch.int32, device="cuda") + wrapped = FusedPath(wt, clone_pools(pools)) + wrapped.bufs = tuple(big[guard : guard + numel] for big, guard, numel, _ in bands) + ( + epoch, + ) + got = [_alone(wrapped, x, slots) for x in xs] + for big, guard, numel, fill in bands: + assert bool((big[:guard] == fill).all()), "stores before the buffers" + assert bool((big[guard + numel :] == fill).all()), "stores after the buffers" + assert all(torch.equal(a, b) for a, b in zip(got, want)) + assert torch.equal(wrapped.p["state"], fresh.p["state"]) and torch.equal( + wrapped.p["conv"], fresh.p["conv"] + ) + assert bool(((epoch >= 0) & (epoch < 3)).all()), epoch.unique().tolist() + + +@pytest.mark.parametrize("offset", [0, 1, 2, 3]) +def test_slots_offset(offset): + """``slots`` as a slice of a longer index tensor that starts at element ``offset`` (the mixer passes + ``state_indices[num_prefills:]`` in a step with prefills, so the data pointer is only 4-byte aligned): the output + and the pools bit for bit those of the same slots in a tensor of their own.""" + _ops() + with torch.inference_mode(): + wt = make_weights(17) + pools = make_pools(18, "dense") + gen = torch.Generator(device="cuda").manual_seed(19) + x = torch.randn(3, K_IN, generator=gen, device="cuda").bfloat16() + slots = make_slots(3, 20) + idx = torch.zeros(offset + slots.numel(), dtype=torch.int32, device="cuda") + idx[offset:] = slots + view = idx[offset:] + ref = FusedPath(wt, clone_pools(pools)) + got = FusedPath(wt, clone_pools(pools)) + want = _alone(ref, x, slots) + out = _alone(got, x, view) + assert view.data_ptr() % 16 == 4 * offset % 16 + assert torch.equal(out, want) + assert torch.equal(got.p["state"], ref.p["state"]) and torch.equal(got.p["conv"], ref.p["conv"]) + + +def make_verify_pools(seed: int) -> dict: + """``trtllm::k3_kda_attn``'s pools for a layer (test_k3_kda_attn.py's dense layout): the conv caches + [POOL, HK, W - 1 + NUM_SPEC] (dim-contiguous), the SSM state, the drafts' records and the pending counts.""" + g = torch.Generator(device="cuda").manual_seed(seed) + p = {name: (torch.randn(POOL, W - 1 + NUM_SPEC, HK, generator=g, device="cuda") * 0.5).transpose(1, 2) + for name in ("cs_q", "cs_k", "cs_v")} # fmt: skip + p["state"] = torch.randn(POOL, H, V, K, generator=g, device="cuda") * 0.05 + p["state_tok"] = torch.zeros(POOL, NUM_SPEC, H, V, K, device="cuda") + p["pending"] = torch.zeros(POOL, dtype=torch.int32, device="cuda") + return p + + +def clone_verify_pools(p: dict) -> dict: + return { + name: t.transpose(1, 2).clone().transpose(1, 2) if name.startswith("cs_") else t.clone() + for name, t in p.items() + } + + +class VerifyPath: + """``trtllm::k3_kda_attn``: the same projection and the speculative verify of one request's NUM_SPEC + 1 tokens.""" + + def __init__(self, wt, pools, bufs): + self.wt, self.p, self.bufs = wt, pools, bufs + + def __call__(self, x, slot): + wt, p = self.wt, self.p + return torch.ops.trtllm.k3_kda_attn( + x, wt["w"], wt["w_fb"], wt["w_q"], wt["w_k"], wt["w_v"], wt["a_log"], wt["dt_bias"], wt["onorm_w"], + p["cs_q"], p["cs_k"], p["cs_v"], p["state"], p["state_tok"], slot, p["pending"], *self.bufs, NUM_SPEC, + LOWER_BOUND, SCALE, EPS, + ) # fmt: skip + + +def test_shared_buffers_with_verify(): + """The model keeps one (p1, part, epoch) set per device for this op and ``trtllm::k3_kda_attn`` + (kimi_kda_mixer.py): plain-decode and verify launches interleaved in stream order on one set give the bits of the + same launches on a set of their own, one at a time. Both run the same stream role (all 8 token rows published, + the same words re-armed) and advance every CTA's index once per launch, so either may follow the other.""" + _ops() + from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_attn import op + + dev = torch.device("cuda") + with torch.inference_mode(): + wt = make_weights(23) + pools = make_pools(24, "dense") + vpools = make_verify_pools(25) + gen = torch.Generator(device="cuda").manual_seed(26) + calls = [] + for i, r in enumerate((1, 4, 8, 3)): + x = torch.randn(r, K_IN, generator=gen, device="cuda").bfloat16() + calls.append(("decode", x, make_slots(r, 27 + i))) + x = torch.randn(NUM_SPEC + 1, K_IN, generator=gen, device="cuda").bfloat16() + calls.append(("verify", x, torch.tensor([i], dtype=torch.int32, device="cuda"))) + own = { + "decode": FusedPath(wt, clone_pools(pools)), + "verify": VerifyPath( + wt, clone_verify_pools(vpools), op.make_buffers(dev, op.FUSED_CTAS) + ), + } + want = [_alone(own[kind], x, s) for kind, x, s in calls] + shared_bufs = op.make_buffers(dev, op.FUSED_CTAS) + shared = { + "decode": FusedPath(wt, clone_pools(pools)), + "verify": VerifyPath(wt, clone_verify_pools(vpools), shared_bufs), + } + shared["decode"].bufs = shared_bufs + got = [shared[kind](x, s) for kind, x, s in calls] + torch.cuda.synchronize() + assert all(torch.equal(a, b) for a, b in zip(got, want)) + for kind, path in shared.items(): + assert all(torch.equal(t, own[kind].p[name]) for name, t in path.p.items()), kind + epoch = shared_bufs[2] + assert bool((epoch == epoch[0]).all()), epoch.unique().tolist() + + +def report() -> int: + """Per split and layout: the worst round's errors.""" + print(f"{torch.cuda.get_device_name()}; {ROUNDS} rounds per split; rel = max |a - b| / max |b|") + print("| R | layout | out vs kda_decode | out vs f64 | kda_decode vs f64 | state vs f64 | state vs kda_decode | " + "conv bits | others untouched | result |") # fmt: skip + print("| --: | :-- | --: | --: | --: | --: | --: | :-- | :-- | :-- |") + ok_all = True + with torch.inference_mode(): + for layout in ("dense", "strided"): + for r in range(1, 9): + metrics, _ = run_schedule(r, layout) + ok = all(round_ok(m) for m in metrics) + ok_all &= ok + worst = {k: max(m[k] for m in metrics) for k in ("out_vs_native", "out_vs_f64", "native_vs_f64", + "state_vs_f64", "state_vs_native")} # fmt: skip + conv_ok = all(m["conv_bits_vs_native"] for m in metrics) + others_ok = all(m["others_untouched"] for m in metrics) + print(f"| {r} | {layout} | {worst['out_vs_native']:.2e} | {worst['out_vs_f64']:.2e} | " + f"{worst['native_vs_f64']:.2e} | {worst['state_vs_f64']:.2e} | {worst['state_vs_native']:.2e} | " + f"{conv_ok} | {others_ok} | {'PASS' if ok else 'FAIL'} |", flush=True) # fmt: skip + print("ALL PASS" if ok_all else "FAIL") + return 0 if ok_all else 1 + + +if __name__ == "__main__": + if len(sys.argv) > 1 and sys.argv[1] == "report": + sys.exit(report()) + sys.exit(pytest.main([__file__, "-q", "-p", "no:cacheprovider", *sys.argv[1:]])) diff --git a/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_kda_pools_past_2g.py b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_kda_pools_past_2g.py new file mode 100644 index 000000000000..c9884b6c384b --- /dev/null +++ b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_kda_pools_past_2g.py @@ -0,0 +1,212 @@ +# 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 fused KDA ops (``trtllm::k3_kda_decode_attn``, ``trtllm::k3_kda_attn``, ``trtllm::k3_kda_verify``) on pools +whose last slot starts past element 2^31, at the per-rank TP16 shape (6 heads, K = V = 128, conv width 4). + +The Mamba cache manager coalesces a rank's KDA layers inside each slot, so a layer's state pool is a view at slot +stride layers x 98,304 fp32 elements (its conv pool at layers x 6,912 bf16): with enough slots the pools reach past +2^31 elements. Here each pool is the last layer's view of such a pool ([slots, layers, ...] underneath), and the +per-token verify states are dense at 7 x 98,304 elements per slot; 3,122 slots put the last slot of each past 2^31. + +Each op runs once on requests in the pools' first, middle and last slots (``k3_kda_attn``: one request, the last slot) +and once on small pools holding copies of those slots; the outputs and every slot the call updates must agree bit for +bit. About 19 GB of device memory. + + pytest test_k3_kda_pools_past_2g.py +""" + +import pytest +import torch + +H = 6 +K = V = 128 +HK = H * K +W = 4 +PROJ = 4 * HK + K + H + 2 # the fused [q | k | v | onorm gate | f_a | b | pad] row (3208 columns) +K_IN = 7168 +NUM_SPEC = 7 +NT = NUM_SPEC + 1 +LOWER_BOUND = -5.0 +EPS = 1e-5 +SCALE = K**-0.5 +SLOTS = 3122 +SSM_LAYERS = 8 +CONV_LAYERS = 100 +TWO_G = 2**31 + + +def _sm100() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10 + + +pytestmark = pytest.mark.skipif(not _sm100(), reason="needs SM100 (tcgen05, TMA, clusters)") + + +def _ops(): + import tensorrt_llm # noqa: F401 (registers the trtllm ops) + from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_attn import op + from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_verify import op as _verify # noqa: F401 + + return op + + +def make_weights(seed: int) -> dict: + g = torch.Generator(device="cuda").manual_seed(seed) + + def rnd(*s, scale=1.0): + return torch.randn(*s, generator=g, device="cuda") * scale + + return { + "w": rnd(PROJ, K_IN, scale=0.02).bfloat16(), "w_fb": rnd(HK, K, scale=0.05).bfloat16(), + "w_q": rnd(HK, W, scale=0.3), "w_k": rnd(HK, W, scale=0.3), "w_v": rnd(HK, W, scale=0.3), + "a_log": rnd(H, scale=0.5), "dt_bias": rnd(HK, scale=0.5), "onorm_w": (1 + 0.1 * rnd(V)).float(), + } # fmt: skip + + +def layer_view(layers: int, shape: tuple, dtype: torch.dtype) -> torch.Tensor: + """The last layer's view of a pool that keeps ``layers`` layers' states in each of its slots.""" + return torch.zeros(SLOTS, layers, *shape, dtype=dtype, device="cuda")[:, -1] + + +def last_slot_offset(t: torch.Tensor) -> int: + return (t.shape[0] - 1) * t.stride(0) + + +def verify_pools() -> dict: + """Conv caches [slots, HK, W - 1 + NUM_SPEC] (dim-contiguous), the SSM state, the drafts' records, the pending + counts.""" + p = { + name: torch.zeros(SLOTS, W - 1 + NUM_SPEC, HK, device="cuda").transpose(1, 2) + for name in ("cs_q", "cs_k", "cs_v") + } + p["ssm"] = layer_view(SSM_LAYERS, (H, V, K), torch.float32) + p["state_tok"] = torch.zeros(SLOTS, NUM_SPEC, H, V, K, device="cuda") + p["pending"] = torch.zeros(SLOTS, dtype=torch.int32, device="cuda") + assert last_slot_offset(p["ssm"]) >= TWO_G and last_slot_offset(p["state_tok"]) >= TWO_G + return p + + +def fill_verify_slots(p: dict, used: list, pending: list, seed: int) -> dict: + """Random contents in the used slots; returns small pools holding copies of them, in the same layouts.""" + g = torch.Generator(device="cuda").manual_seed(seed) + for i, s in enumerate(used): + for name in ("cs_q", "cs_k", "cs_v"): + p[name][s] = torch.randn(HK, W - 1 + NUM_SPEC, generator=g, device="cuda") * 0.5 + p["ssm"][s] = torch.randn(H, V, K, generator=g, device="cuda") * 0.05 + p["state_tok"][s] = torch.rand(NUM_SPEC, H, V, K, generator=g, device="cuda") * 0.5 + p["pending"][s] = pending[i] + small = { + name: torch.stack([p[name][s].t() for s in used]).transpose(1, 2) + for name in ("cs_q", "cs_k", "cs_v") + } + for name in ("ssm", "state_tok", "pending"): + small[name] = p[name][used] + return small + + +def same(a: torch.Tensor, b: torch.Tensor) -> bool: + """Bit-equal (NaN-safe).""" + bits = torch.int16 if a.element_size() == 2 else torch.int32 + return torch.equal(a.contiguous().view(bits), b.contiguous().view(bits)) + + +def check_slots(p: dict, small: dict, used: list, names: tuple) -> None: + for i, s in enumerate(used): + for name in names: + assert same(p[name][s], small[name][i]), f"slot {s} {name}" + + +def test_decode_attn(): + op = _ops() + with torch.inference_mode(): + wt = make_weights(1) + p = { + "ssm": layer_view(SSM_LAYERS, (H, V, K), torch.float32), + "conv": layer_view(CONV_LAYERS, (3 * HK, W - 1), torch.bfloat16), + } + assert last_slot_offset(p["ssm"]) >= TWO_G and last_slot_offset(p["conv"]) >= TWO_G + used = [0, SLOTS // 2, SLOTS - 1] + g = torch.Generator(device="cuda").manual_seed(2) + for s in used: + p["ssm"][s] = torch.randn(H, V, K, generator=g, device="cuda") * 0.05 + p["conv"][s] = (torch.randn(3 * HK, W - 1, generator=g, device="cuda") * 0.5).bfloat16() + small = {name: t[used] for name, t in p.items()} + x = torch.randn(len(used), K_IN, generator=g, device="cuda").bfloat16() + + def call(q: dict, slots: list) -> torch.Tensor: + bufs = op.make_buffers(torch.device("cuda"), op.FUSED_CTAS) + return torch.ops.trtllm.k3_kda_decode_attn( + x, wt["w"], wt["w_fb"], wt["w_q"], wt["w_k"], wt["w_v"], wt["a_log"], wt["dt_bias"], wt["onorm_w"], + q["conv"], q["ssm"], torch.tensor(slots, dtype=torch.int32, device="cuda"), *bufs, LOWER_BOUND, + SCALE, EPS, + ) # fmt: skip + + want = call(small, list(range(len(used)))) + got = call(p, used) + torch.cuda.synchronize() + assert same(got, want) + check_slots(p, small, used, ("ssm", "conv")) + + +def test_attn(): + op = _ops() + with torch.inference_mode(): + wt = make_weights(3) + p = verify_pools() + used = [SLOTS - 1] + small = fill_verify_slots(p, used, [5], 4) + x = torch.randn( + NT, K_IN, generator=torch.Generator(device="cuda").manual_seed(5), device="cuda" + ).bfloat16() + + def call(q: dict, slot: int) -> torch.Tensor: + bufs = op.make_buffers(torch.device("cuda"), op.FUSED_CTAS) + return torch.ops.trtllm.k3_kda_attn( + x, wt["w"], wt["w_fb"], wt["w_q"], wt["w_k"], wt["w_v"], wt["a_log"], wt["dt_bias"], wt["onorm_w"], + q["cs_q"], q["cs_k"], q["cs_v"], q["ssm"], q["state_tok"], + torch.tensor([slot], dtype=torch.int32, device="cuda"), q["pending"], *bufs, NUM_SPEC, LOWER_BOUND, + SCALE, EPS, + ) # fmt: skip + + want = call(small, 0) + got = call(p, used[0]) + torch.cuda.synchronize() + assert same(got, want) + check_slots(p, small, used, ("cs_q", "cs_k", "cs_v", "ssm", "state_tok")) + + +def test_verify(): + _ops() + with torch.inference_mode(): + wt = make_weights(6) + p = verify_pools() + used = [0, SLOTS // 2, SLOTS - 1] + small = fill_verify_slots(p, used, [0, 3, 7], 7) + g = torch.Generator(device="cuda").manual_seed(8) + proj = torch.randn(len(used) * NT, PROJ, generator=g, device="cuda").bfloat16() + + def call(q: dict, slots: list) -> torch.Tensor: + return torch.ops.trtllm.k3_kda_verify( + proj, wt["w_fb"], wt["w_q"], wt["w_k"], wt["w_v"], wt["a_log"], wt["dt_bias"], wt["onorm_w"], + q["cs_q"], q["cs_k"], q["cs_v"], q["ssm"], q["state_tok"], + torch.tensor(slots, dtype=torch.int32, device="cuda"), q["pending"], NUM_SPEC, LOWER_BOUND, SCALE, + EPS, + ) # fmt: skip + + want = call(small, list(range(len(used)))) + got = call(p, used) + torch.cuda.synchronize() + assert same(got, want) + check_slots(p, small, used, ("cs_q", "cs_k", "cs_v", "ssm", "state_tok")) diff --git a/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_kda_verify.py b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_kda_verify.py new file mode 100644 index 000000000000..d7136753192e --- /dev/null +++ b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_kda_verify.py @@ -0,0 +1,588 @@ +# 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::k3_kda_verify`` (Kimi K3's KDA speculative verify with a committed state per verify token) for R requests +of T = 1 + num_spec tokens, at the per-rank TP16 shape (6 heads, K = V = 128, conv width 4). + +Two references, run over the same schedule of rounds from the same initial pools, with distinct slots and a pending +(accepted-draft) count per request that changes every round: + +* main's path, as ``KimiDeltaAttention.forward_verify_fused`` calls it: f_b as a bf16 ``F.linear``, then + ``trtllm::kda_mtp_decode`` replaying each request's accepted drafts from its replay caches + (``cu_seqlens[n] = n T - pending[n]``, ``num_accepted_tokens = pending``), then the gated RMSNorm + (``rms_norm_gated_token_major``, sigmoid gate). The report says which ``kda_mtp_decode`` was loaded. +* a float64 torch delta rule (the "fp32 reference"; float64 so that no TF32 enters) that keeps each request's + committed token history: conv4 + SiLU, q / k L2 norm, beta sigmoid, the lower-bound gate, + S <- S d + beta (v - (S d) k) k^T, o = S q, the gated RMSNorm. + +Checks per round: the outputs of every request against both references (fp32 tolerance), the committed pool state +(after each golden token) against the history, the conv caches against main's (raw inputs: bit-exact); with the gate +taken from the unfused f_b output (``g_ext``) the pool state against main's bit for bit. The drafts' records in +``state_tok`` are checked through the next round, which starts from the accepted ones. Also: requests isolated, the +schedule twice bit-identical, CUDA-graph replays with rewritten inputs, the slots and the pending counts given as +slices of longer index tensors at any element offset. + + pytest test_k3_kda_verify.py + python3 test_k3_kda_verify.py report | time +""" + +import inspect +import statistics +import sys + +import pytest +import torch +import torch.nn.functional as F + +H = 6 +K = V = 128 +W = 4 +HK = H * K +PROJ = 4 * HK + K + H + 2 # the fused [q | k | v | onorm gate | f_a | b | pad] row (3208 columns) +LOWER_BOUND = -5.0 +EPS = 1e-5 +SCALE = K**-0.5 +POOL = 11 +ROUNDS = 6 +TOL_OUT = 2e-2 # bf16 outputs; main rounds the core to bf16 before its norm +# fp32 states against the fp32 history: approximate exp / rcp in the kernels, and the bf16 f_b gate (one bf16 ulp +# moves a decay by up to ~1 %; main's kda_mtp_decode measures 2.05e-4 at 6x8) +TOL_STATE = 1e-3 + +SPLITS = [(r, t) for t in (8, 4, 2) for r in range(1, 9)] + + +def _sm100() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10 + + +pytestmark = pytest.mark.skipif(not _sm100(), reason="needs SM100 (tcgen05, TMA, clusters)") + + +def _ops(): + import tensorrt_llm # noqa: F401 (registers the trtllm ops) + import tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops # noqa: F401 (kda_mtp_decode) + from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_verify import op # noqa: F401 (k3_kda_verify) + + +def main_kda_variant() -> str: + """'main' when the loaded kda_mtp_decode has main's signature, else 'port' (a modified one).""" + _ops() + from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_kda_mtp_ops as mtp + + return ( + "port" + if "accepted_by_slot" in inspect.signature(mtp.kda_mtp_decode_impl).parameters + else "main" + ) + + +def make_weights(seed: int) -> dict: + g = torch.Generator(device="cuda").manual_seed(seed) + + def rnd(*s, scale=1.0): + return torch.randn(*s, generator=g, device="cuda") * scale + + return { + "w_q": rnd(HK, W, scale=0.3), "w_k": rnd(HK, W, scale=0.3), "w_v": rnd(HK, W, scale=0.3), + "a_log": rnd(H, scale=0.5), "dt_bias": rnd(HK, scale=0.5), "onorm_w": (1 + 0.1 * rnd(V)).float(), + "w_fb": (rnd(HK, K) * 0.05).bfloat16(), + } # fmt: skip + + +def strided(state: torch.Tensor) -> torch.Tensor: + """The same values in the Mamba cache manager's layout: each slot's SSM state followed by its conv state bytes.""" + pool = torch.zeros(state.shape[0], state[0].numel() + 3 * HK * (W - 1), device=state.device) + view = pool[:, : state[0].numel()].view(state.shape) + view.copy_(state) + return view + + +def make_pools(seed: int, num_spec: int, layout: str) -> dict: + """The initial pools: conv caches [pool, HK, W - 1 + num_spec] (dim-contiguous), the SSM state [pool, H, V, K].""" + g = torch.Generator(device="cuda").manual_seed(seed) + s = W - 1 + num_spec + p = {name: (torch.randn(POOL, s, HK, generator=g, device="cuda") * 0.5).transpose(1, 2) + for name in ("cs_q", "cs_k", "cs_v")} # fmt: skip + p["state"] = torch.randn(POOL, H, V, K, generator=g, device="cuda") * 0.05 + if layout == "strided": + p["state"] = strided(p["state"]) + return p + + +def clone_pools(p: dict) -> dict: + out = {} + for name, t in p.items(): + if name.startswith("cs_"): + out[name] = ( + t.transpose(1, 2).clone().transpose(1, 2) + ) # a copy in the same dim-contiguous layout + elif name == "state" and not t.is_contiguous(): + out[name] = strided(t) + else: + out[name] = t.clone() + return out + + +def split_proj(proj: torch.Tensor): + t = proj.shape[0] + x_q, x_k, x_v, og = (proj[:, i * HK : (i + 1) * HK] for i in range(4)) + f_a = proj[:, 4 * HK : 4 * HK + K] + beta = proj[:, 4 * HK + K : 4 * HK + K + H] + return t, x_q, x_k, x_v, og, f_a, beta + + +class MainPath: + """main's verify (forward_verify_fused): its replay caches, a per-request pending count.""" + + def __init__(self, wt, pools, slots, num_spec): + self.wt, self.p, self.slots, self.num_spec = wt, pools, slots, num_spec + self.p["qkg_cache"] = torch.zeros(POOL, num_spec, 3, HK, device="cuda") + self.p["v_cache"] = torch.zeros(POOL, num_spec, HK, device="cuda") + self.p["beta_cache"] = torch.zeros(POOL, num_spec, H, device="cuda") + + def __call__(self, proj, pending_req): + from tensorrt_llm._torch.modules.mamba.layernorm_gated import rms_norm_gated_token_major + + t, x_q, x_k, x_v, og, f_a, beta = split_proj(proj) + n, steps = self.slots.numel(), self.num_spec + 1 + g = F.linear(f_a, self.wt["w_fb"]) + cu = torch.arange(0, (n + 1) * steps, steps, dtype=torch.int32, device="cuda") + cu[:n].sub_(pending_req) + o = torch.ops.trtllm.kda_mtp_decode( + x_q=x_q.view(1, t, H, K), x_k=x_k.view(1, t, H, K), x_v=x_v.view(1, t, H, V), w_q=self.wt["w_q"], + w_k=self.wt["w_k"], w_v=self.wt["w_v"], cs_q=self.p["cs_q"], cs_k=self.p["cs_k"], cs_v=self.p["cs_v"], + g=g.view(1, t, H, K), beta=beta.contiguous().view(1, t, H), A_log=self.wt["a_log"], + dt_bias=self.wt["dt_bias"], recurrent_state=self.p["state"], qkg_cache=self.p["qkg_cache"], + v_cache=self.p["v_cache"], beta_cache=self.p["beta_cache"], ssm_state_indices=self.slots, cu_seqlens=cu, + num_spec=self.num_spec, num_accepted_tokens=pending_req, lower_bound=LOWER_BOUND, scale=SCALE, + ) # fmt: skip + core = rms_norm_gated_token_major(o.reshape(-1, V), og.reshape(t, H, V), self.wt["onorm_w"], EPS, + gate_activation="sigmoid") # fmt: skip + return core.view(t, H, V), g + + +class K3Path: + """k3_kda_verify: a committed state per verify token, a per-slot pending count.""" + + def __init__(self, wt, pools, slots, num_spec): + self.wt, self.p, self.slots, self.num_spec = wt, pools, slots, num_spec + self.p["state_tok"] = torch.zeros(POOL, num_spec, H, V, K, device="cuda") + self.p["pending"] = torch.zeros(POOL, dtype=torch.int32, device="cuda") + + def __call__(self, proj, g_ext=None): + p = self.p + return torch.ops.trtllm.k3_kda_verify( + proj, self.wt["w_fb"], self.wt["w_q"], self.wt["w_k"], self.wt["w_v"], self.wt["a_log"], + self.wt["dt_bias"], self.wt["onorm_w"], p["cs_q"], p["cs_k"], p["cs_v"], p["state"], p["state_tok"], + self.slots, p["pending"], self.num_spec, LOWER_BOUND, SCALE, EPS, g_ext, + ) # fmt: skip + + +def conv_silu(win, raw, c, w): + """Channel group ``c`` (q, k, v) of the causal conv over the window (oldest first) and the new raw input, SiLU.""" + x = win[0][c] * w[:, 0] + win[1][c] * w[:, 1] + win[2][c] * w[:, 2] + raw[c] * w[:, 3] + return x * torch.sigmoid(x) + + +class Fp32Path: + """The delta rule over each request's committed history (raw conv inputs and the state after the last committed + token), in float64: no TF32 even where cuBLAS is told to use it.""" + + def __init__(self, wt, pools, slots, num_spec): + self.wt, self.num_spec = wt, num_spec + self.slots = slots.tolist() + # Committed raw inputs, oldest first: the conv caches' window columns 0..2 (q, k, v). + self.seq = [[torch.stack([pools[c][s, :, i].double() for c in ("cs_q", "cs_k", "cs_v")]) for i in range(3)] + for s in self.slots] # fmt: skip + self.state = [pools["state"][s].double().clone() for s in self.slots] + self.last = None + + def __call__(self, proj): + """Outputs [T, H, V] and the per-token states of every request.""" + t_total, x_q, x_k, x_v, og, f_a, beta_raw = split_proj(proj) + steps = self.num_spec + 1 + # f_b's output is bf16 in the model. + g_all = F.linear(f_a.double(), self.wt["w_fb"].double()).bfloat16().double() + wq, wk, wv = (self.wt[n].double() for n in ("w_q", "w_k", "w_v")) + exp_a = self.wt["a_log"].double().exp() + dt_bias = self.wt["dt_bias"].double().view(H, K) + onorm_w = self.wt["onorm_w"].double() + out = torch.empty(t_total, H, V, dtype=torch.float64, device="cuda") + self.last = [] + for n in range(len(self.slots)): + seq = list(self.seq[n]) + s_cur = self.state[n] + states, raws = [], [] + for t in range(steps): + row = n * steps + t + raw = torch.stack([x_q[row].double(), x_k[row].double(), x_v[row].double()]) + win = seq[-3:] + q = conv_silu(win, raw, 0, wq).view(H, K) + k = conv_silu(win, raw, 1, wk).view(H, K) + v = conv_silu(win, raw, 2, wv).view(H, V) + q = q * torch.rsqrt((q * q).sum(-1, keepdim=True) + 1e-6) * SCALE + k = k * torch.rsqrt((k * k).sum(-1, keepdim=True) + 1e-6) + beta = torch.sigmoid(beta_raw[row].double()) + gk = LOWER_BOUND * torch.sigmoid(exp_a[:, None] * (g_all[row].view(H, K) + dt_bias)) + decay = gk.exp() + sd = s_cur * decay[:, None, :] + vn = v - torch.einsum("hvk,hk->hv", sd, k) + s_cur = sd + beta[:, None, None] * vn[:, :, None] * k[:, None, :] + o = torch.einsum("hvk,hk->hv", s_cur, q) + rms = torch.rsqrt((o * o).mean(-1, keepdim=True) + EPS) + out[row] = o * rms * onorm_w * torch.sigmoid(og[row].double().view(H, V)) + states.append(s_cur) + raws.append(raw) + seq.append(raw) + self.last.append((states, raws)) + return out + + def commit(self, pending): + """The sampler accepted ``pending[n]`` drafts of the last round: the golden token and those drafts commit.""" + for n, p in enumerate(pending): + states, raws = self.last[n] + self.seq[n] = (self.seq[n] + raws[: p + 1])[-3:] + self.state[n] = states[p] + + +def make_slots(num_requests: int, seed: int) -> torch.Tensor: + perm = torch.randperm(POOL, generator=torch.Generator().manual_seed(seed)) + return perm[:num_requests].to(torch.int32).cuda() + + +def pending_schedule(num_requests: int, num_spec: int, rnd: int): + """Accepted drafts per request after round ``rnd``: every count 0..num_spec, different per request.""" + return [(3 * n + 5 * rnd + 1 + (rnd * n) % 3) % (num_spec + 1) for n in range(num_requests)] + + +def rel(a: torch.Tensor, b: torch.Tensor) -> float: + return ((a.float() - b.float()).abs().max() / b.float().abs().max().clamp_min(1e-6)).item() + + +def run_schedule(num_requests, steps, fold=True, layout="dense", seed=0, rounds=ROUNDS): + """Main's path, k3_kda_verify and the fp32 history over ``rounds`` rounds; per-round metrics.""" + _ops() + num_spec = steps - 1 + wt = make_weights(100 + seed) + pools = make_pools(200 + seed, num_spec, layout) + slots = make_slots(num_requests, 300 + seed) + main = MainPath(wt, clone_pools(pools), slots, num_spec) + k3 = K3Path(wt, clone_pools(pools), slots, num_spec) + ref = Fp32Path(wt, pools, slots, num_spec) + pending_req = torch.zeros(num_requests, dtype=torch.int32, device="cuda") + gen = torch.Generator(device="cuda").manual_seed(400 + seed) + rows, outs = [], [] + for rnd in range(rounds): + proj = torch.randn(num_requests * steps, PROJ, generator=gen, device="cuda").bfloat16() + out_main, g = main(proj, pending_req) + out_k3 = k3(proj, None if fold else g.contiguous()) + out_ref = ref(proj) + torch.cuda.synchronize() + r = dict(round=rnd, pending=pending_req.tolist()) + r["out_vs_main"] = rel(out_k3, out_main) + r["out_vs_fp32"] = rel(out_k3, out_ref) + r["main_vs_fp32"] = rel(out_main, out_ref) + # The pool state after the golden token (state_tok holds the drafts' compact records, which the next round's + # outputs check against the history). + st_err, main_st_err, st_bits = 0.0, 0.0, True + for n, s in enumerate(slots.tolist()): + states, _ = ref.last[n] + st_err = max(st_err, rel(k3.p["state"][s], states[0])) + main_st_err = max(main_st_err, rel(main.p["state"][s], states[0])) + st_bits &= torch.equal(k3.p["state"][s], main.p["state"][s]) + r["state_vs_fp32"] = st_err + r["main_state_vs_fp32"] = main_st_err + r["state_bits_vs_main"] = st_bits + r["conv_bits_vs_main"] = all( + torch.equal(k3.p[c], main.p[c]) for c in ("cs_q", "cs_k", "cs_v") + ) + rows.append(r) + outs.append(out_k3.clone()) + nxt = pending_schedule(num_requests, num_spec, rnd) + pending_req.copy_(torch.tensor(nxt, dtype=torch.int32)) + k3.p["pending"][slots.long()] = pending_req + ref.commit(nxt) + return rows, outs + + +def round_ok(r, fold: bool) -> bool: + ok = (r["out_vs_main"] <= TOL_OUT and r["out_vs_fp32"] <= TOL_OUT and r["state_vs_fp32"] <= TOL_STATE + and r["conv_bits_vs_main"]) # fmt: skip + if not fold: # the same gate as main's: the recurrence's arithmetic is kda_mtp_decode's + ok &= r["state_bits_vs_main"] + return ok + + +@pytest.mark.parametrize("num_requests,steps", SPLITS, ids=[f"{r}x{t}" for r, t in SPLITS]) +def test_split(num_requests, steps): + with torch.inference_mode(): + rows, outs = run_schedule(num_requests, steps, fold=True) + again, outs2 = run_schedule(num_requests, steps, fold=True) + for r in rows: + assert round_ok(r, True), r + assert all(torch.equal(a, b) for a, b in zip(outs, outs2)), "the schedule twice differs" + + +@pytest.mark.parametrize("num_requests,steps", [(1, 8), (3, 4), (8, 8), (8, 2)]) +@pytest.mark.parametrize("layout", ["dense", "strided"]) +def test_g_ext(num_requests, steps, layout): + """The gate from the unfused f_b output (main's bf16 F.linear): the pool state bit-exact against main's.""" + with torch.inference_mode(): + rows, _ = run_schedule(num_requests, steps, fold=False, layout=layout, seed=1) + for r in rows: + assert round_ok(r, False), r + + +@pytest.mark.parametrize("num_requests,steps", [(4, 8), (8, 4)]) +def test_strided_fold(num_requests, steps): + with torch.inference_mode(): + rows, _ = run_schedule(num_requests, steps, fold=True, layout="strided", seed=2) + for r in rows: + assert round_ok(r, True), r + + +@pytest.mark.parametrize("num_requests,steps", [(4, 8), (8, 2)]) +def test_isolation(num_requests, steps): + """A change in request 0's rows changes only request 0's outputs, pool state and per-token states.""" + _ops() + with torch.inference_mode(): + num_spec = steps - 1 + wt = make_weights(7) + pools = make_pools(8, num_spec, "dense") + slots = make_slots(num_requests, 9) + a = K3Path(wt, clone_pools(pools), slots, num_spec) + b = K3Path(wt, clone_pools(pools), slots, num_spec) + pend = torch.tensor( + pending_schedule(num_requests, num_spec, 3), dtype=torch.int32, device="cuda" + ) + for path in (a, b): + path.p["pending"][slots.long()] = pend + proj = torch.randn(num_requests * steps, PROJ, generator=torch.Generator(device="cuda").manual_seed(10), + device="cuda").bfloat16() # fmt: skip + proj_c = proj.clone() + proj_c[1, HK + 5] += ( + 1.0 # request 0, token 1, k channel 5: its outputs and states from token 1 on + ) + out_a, out_b = a(proj), b(proj_c) + torch.cuda.synchronize() + assert not torch.equal(out_a[:steps], out_b[:steps]) + assert torch.equal(out_a[steps:], out_b[steps:]) + s0 = int(slots[0]) + others = [int(s) for s in slots[1:]] + assert not torch.equal(a.p["state_tok"][s0], b.p["state_tok"][s0]) + for s in others: + assert torch.equal(a.p["state"][s], b.p["state"][s]) + assert torch.equal(a.p["state_tok"][s], b.p["state_tok"][s]) + + +@pytest.mark.parametrize("num_requests,steps", [(8, 8), (4, 2), (1, 8)]) +def test_graph_replay(num_requests, steps): + """One captured call replayed over the schedule with the projection rows and pending counts rewritten in place, + bit-identical to eager calls on a copy of the pools.""" + _ops() + with torch.inference_mode(): + num_spec = steps - 1 + wt = make_weights(11) + pools = make_pools(12, num_spec, "dense") + slots = make_slots(num_requests, 13) + graphed = K3Path(wt, clone_pools(pools), slots, num_spec) + eager = K3Path(wt, clone_pools(pools), slots, num_spec) + proj = torch.zeros(num_requests * steps, PROJ, dtype=torch.bfloat16, device="cuda") + holder = {} + snap = clone_pools( + {k: v for k, v in graphed.p.items() if k in ("cs_q", "cs_k", "cs_v", "state")} + ) + snap_tok = graphed.p["state_tok"].clone() + graphed(proj) # compiles outside capture + torch.cuda.synchronize() + for name in ("cs_q", "cs_k", "cs_v", "state"): + graphed.p[name].copy_(snap[name]) + graphed.p["state_tok"].copy_(snap_tok) + stream = torch.cuda.Stream() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + holder["out"] = graphed(proj) + gen = torch.Generator(device="cuda").manual_seed(14) + for rnd in range(ROUNDS): + proj.copy_(torch.randn(proj.shape, generator=gen, device="cuda").bfloat16()) + graph.replay() + want = eager(proj.clone()) + torch.cuda.synchronize() + assert torch.equal(holder["out"], want), f"round {rnd}" + for name in ("cs_q", "cs_k", "cs_v", "state", "state_tok"): + assert torch.equal(graphed.p[name], eager.p[name]), f"round {rnd} {name}" + nxt = torch.tensor( + pending_schedule(num_requests, num_spec, rnd), dtype=torch.int32, device="cuda" + ) + graphed.p["pending"][slots.long()] = nxt + eager.p["pending"][slots.long()] = nxt + + +def _at_offset(t: torch.Tensor, offset: int) -> torch.Tensor: + """``t``'s values in a longer tensor, starting at element ``offset`` (a slice such as + ``state_indices[num_contexts:]``: its data pointer is aligned to the element only).""" + buf = torch.zeros(offset + t.numel(), dtype=t.dtype, device=t.device) + buf[offset:] = t + view = buf[offset:] + assert view.data_ptr() % 16 == offset * t.element_size() % 16 + return view + + +@pytest.mark.parametrize("offset", [0, 1, 2, 3]) +@pytest.mark.parametrize("which", ["slots", "pending"]) +def test_index_offset(which, offset): + """The slots or the pending counts as a slice of a longer index tensor that starts at element ``offset``: over two + rounds, the second on nonzero pending counts, the outputs and the pools bit for bit those of the same indices in + tensors of their own.""" + _ops() + with torch.inference_mode(): + num_requests, steps = 3, 8 + num_spec = steps - 1 + wt = make_weights(15) + pools = make_pools(16, num_spec, "dense") + slots = make_slots(num_requests, 17) + ref = K3Path(wt, clone_pools(pools), slots, num_spec) + got = K3Path( + wt, + clone_pools(pools), + _at_offset(slots, offset) if which == "slots" else slots, + num_spec, + ) + if which == "pending": + got.p["pending"] = _at_offset(got.p["pending"], offset) + gen = torch.Generator(device="cuda").manual_seed(18) + for rnd in range(2): + proj = torch.randn(num_requests * steps, PROJ, generator=gen, device="cuda").bfloat16() + want, out = ref(proj), got(proj) + torch.cuda.synchronize() + assert torch.equal(out, want), rnd + for name in ("cs_q", "cs_k", "cs_v", "state", "state_tok"): + assert torch.equal(got.p[name], ref.p[name]), (rnd, name) + nxt = torch.tensor( + pending_schedule(num_requests, num_spec, rnd), dtype=torch.int32, device="cuda" + ) + ref.p["pending"][slots.long()] = nxt + got.p["pending"][slots.long()] = nxt + + +# ---------------------------------------------------------------------------------------------------------------- +# Report and timing (python3 test_k3_kda_verify.py report | time) +# ---------------------------------------------------------------------------------------------------------------- + + +def report() -> int: + with torch.inference_mode(): + print(f"{torch.cuda.get_device_name()}; kda_mtp_decode: {main_kda_variant()}") + print("| split | mode | layout | rounds | k3 vs main | k3 vs fp32 | main vs fp32 | state vs fp32 " + "| main state vs fp32 | state bits = main | conv bits = main | deterministic | result |") # fmt: skip + print("| :-- | :-- | :-- | --: | --: | --: | --: | --: | --: | :-- | :-- | :-- | :-- |") + ok_all = True + cases = [(r, t, True, "dense") for r, t in SPLITS] + cases += [ + (r, t, False, layout) + for r, t in [(1, 8), (3, 4), (8, 8), (8, 2)] + for layout in ("dense", "strided") + ] + cases += [(4, 8, True, "strided"), (8, 4, True, "strided")] + for r, t, fold, layout in cases: + rows, outs = run_schedule( + r, t, fold=fold, layout=layout, seed=0 if fold and layout == "dense" else 1 + ) + _, outs2 = run_schedule( + r, t, fold=fold, layout=layout, seed=0 if fold and layout == "dense" else 1 + ) + det = all(torch.equal(a, b) for a, b in zip(outs, outs2)) + ok = all(round_ok(x, fold) for x in rows) and det + ok_all &= ok + print(f"| {r}x{t} | {'fold' if fold else 'g_ext'} | {layout} | {len(rows)} | " + f"{max(x['out_vs_main'] for x in rows):.2e} | {max(x['out_vs_fp32'] for x in rows):.2e} | " + f"{max(x['main_vs_fp32'] for x in rows):.2e} | {max(x['state_vs_fp32'] for x in rows):.2e} | " + f"{max(x['main_state_vs_fp32'] for x in rows):.2e} | " + f"{all(x['state_bits_vs_main'] for x in rows)} | {all(x['conv_bits_vs_main'] for x in rows)} | " + f"{det} | {'PASS' if ok else 'FAIL'} |", flush=True) # fmt: skip + print("ALL PASS" if ok_all else "FAIL") + return 0 if ok_all else 1 + + +def time_graph(body, calls, replays=15): + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + body(0) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + for i in range(calls): + body(i) + torch.cuda.synchronize() + for _ in range(3): + graph.replay() + torch.cuda.synchronize() + per_call = [] + for _ in range(replays): + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + graph.replay() + end.record() + torch.cuda.synchronize() + per_call.append(start.elapsed_time(end) * 1e3 / calls) + return statistics.median(per_call), min(per_call), max(per_call) + + +def timing(layers: int = 48) -> None: + """Graphs of back-to-back calls over ``layers`` copies of the weights and pools (HBM-cold), steady pending.""" + _ops() + print(f"{torch.cuda.get_device_name()}; kda_mtp_decode: {main_kda_variant()}; graphs over {layers} layer copies, " + "15 replays: median (min-max) us per call") # fmt: skip + print("| split | main: f_b + kda_mtp_decode + gated norm | k3_kda_verify |") + print("| :-- | --: | --: |") + with torch.inference_mode(): + for r, t in SPLITS: + num_spec = t - 1 + slots = make_slots(r, 1) + mains, k3s = [], [] + for i in range(layers): + wt = make_weights(1000 + i) + pools = make_pools(2000 + i, num_spec, "dense") + mains.append(MainPath(wt, clone_pools(pools), slots, num_spec)) + k3s.append(K3Path(wt, pools, slots, num_spec)) + pend = torch.tensor(pending_schedule(r, num_spec, 2), dtype=torch.int32, device="cuda") + for k3 in k3s: + k3.p["pending"][slots.long()] = pend + proj = torch.randn(r * t, PROJ, device="cuda").bfloat16() + arms = [lambda i: mains[i % layers](proj, pend), lambda i: k3s[i % layers](proj)] + res = [[], []] + for rep in range(3): + for a in (0, 1) if rep % 2 == 0 else (1, 0): + res[a].append(time_graph(arms[a], layers)) + cells = [] + for a in range(2): + meds = sorted(x[0] for x in res[a]) + cells.append( + f"{meds[1]:.2f} ({min(x[1] for x in res[a]):.2f}-{max(x[2] for x in res[a]):.2f})" + ) + print(f"| {r}x{t} | " + " | ".join(cells) + " |", flush=True) + mains.clear() + k3s.clear() + torch.cuda.empty_cache() + + +if __name__ == "__main__": + if len(sys.argv) > 1 and sys.argv[1] == "report": + sys.exit(report()) + elif len(sys.argv) > 1 and sys.argv[1] == "time": + timing() + else: + sys.exit(pytest.main([__file__, "-q", "-p", "no:cacheprovider", *sys.argv[1:]])) diff --git a/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_mla_attn.py b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_mla_attn.py new file mode 100644 index 000000000000..c359a444cbc8 --- /dev/null +++ b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_mla_attn.py @@ -0,0 +1,385 @@ +# 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::k3_mla_attn / k3_mla_attn_out / k3_mla_attn_vb_out (Kimi K3 MLA decode attention) at the TP16 shape (6 +heads per rank, latent 512 + rope 64, bf16 pool of 64-row pages) for decode steps of R requests x T tokens: against a +torch float64 reference with per-request bottom-right causal masks and against the stock CuTe DSL MLA decode +(trtllm::cute_dsl_mla_decode_fp16_blackwell); request i's rows bit-identical to the one-request call on its own +rows, pages and length; reruns bit-identical.""" + +import math + +import pytest +import torch + + +def _is_sm100() -> bool: + if not torch.cuda.is_available(): + return False + major, minor = torch.cuda.get_device_capability() + return major * 10 + minor in (100, 103) + + +pytestmark = pytest.mark.skipif(not _is_sm100(), reason="needs an SM100 / SM103 GPU") + +H, LAT, ROPE, PAGE, V = 6, 512, 64, 64, 128 +DQK = LAT + ROPE +SCALE = 1.0 / math.sqrt(128 + ROPE) +GATE_COL0 = 2112 + +# Every split of up to 8 tokens and the DSpark verify steps of R requests x 8 tokens. +SPLITS = [(1, 1), (2, 1), (3, 1), (4, 1), (5, 1), (8, 1), (1, 8), (2, 4), (4, 2)] + [ + (r, 8) for r in range(2, 9) +] +SPLIT_IDS = [f"{r}x{t}" for r, t in SPLITS] +# The TP4 shape (24 heads: 4 clusters per request, the workspace slots of several head groups) on a few splits. +ATTN_CASES = [(r, t, H) for r, t in SPLITS] + [ + (r, t, 4 * H) for r, t in ((2, 4), (8, 1), (3, 8), (8, 8)) +] +# KV lengths, assigned to the requests of a step in turn: the step's rows crossing a page boundary (579, 1989), only +# the step's tokens (8, 1), one 128-row tile minus one, > 16 tiles (several per CTA of the cluster), short contexts. +LENGTHS = (1100, 64 * 9 + 3, 8, 2049, 127, 4100, 64 * 31 + 5, 300, 64, 1) + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.int16) + + +def _make_case(seed, num_requests, tokens, heads=H, row_stride=DQK, layers=1, slot=0, lens=None): + """A step of `num_requests` x `tokens`: q [M, heads, 576] (request-major), a pool of `layers` interleaved layer + slots with rows of `row_stride` elements, the page table as rows of an int32 [R, 2, W] buffer (the layout of + kv_cache_block_offsets: row stride 2 W; entries past a request's pages name a page no request owns), lengths + (`lens`, or LENGTHS in turn).""" + gen = torch.Generator(device="cuda").manual_seed(seed) + cpu_gen = torch.Generator().manual_seed(seed) + if lens is None: + lens = [max(tokens, LENGTHS[(i + seed) % len(LENGTHS)]) for i in range(num_requests)] + pages = [(n + PAGE - 1) // PAGE for n in lens] + total_pages = sum(pages) + 3 + width = max(pages) + 2 + pool = torch.randn(total_pages * layers, PAGE, row_stride, generator=gen, device="cuda") * 0.5 + pool = pool.bfloat16() + perm = (torch.randperm(total_pages, generator=cpu_gen) * layers).to(torch.int32) + offsets = torch.empty(num_requests, 2, width, dtype=torch.int32).fill_(int(perm[-1])) + start = 0 + for i, n in enumerate(pages): + offsets[i, :, :n] = perm[start : start + n] + start += n + page_table = offsets.cuda()[:, 0, :] + seq_len = torch.tensor(lens, dtype=torch.int32, device="cuda") + q = torch.randn(num_requests * tokens, heads, DQK, generator=gen, device="cuda") * 0.5 + return q.bfloat16(), pool, page_table, slot, seq_len + + +def _reference(q, pool, page_table, page_offset, seq_len, tokens): + """float64 attention of each request's tokens over its pages; token t sees rows <= L - T + t. [M, heads, 512].""" + outs = [] + for i, length in enumerate(seq_len.tolist()): + pages = page_table[i, : (length + PAGE - 1) // PAGE].long() + page_offset + kv = pool[pages].reshape(-1, pool.shape[-1])[:length, :DQK].double() + qi = q[i * tokens : (i + 1) * tokens].double() + s = torch.einsum("thd,ld->thl", qi, kv) * SCALE + limit = length - tokens + torch.arange(tokens, device=q.device) + hidden = torch.arange(length, device=q.device)[None, :] > limit[:, None] + s = s.masked_fill(hidden[:, None, :], float("-inf")) + outs.append(torch.einsum("thl,ld->thd", torch.softmax(s, dim=-1), kv[:, :LAT])) + return torch.cat(outs) + + +def _max_rel(a, b, tokens, num_requests): + """max over requests of max |a - b| / max |b| (per request, so a short context is not hidden by a long one).""" + err = 0.0 + for i in range(num_requests): + ai, bi = ( + a[i * tokens : (i + 1) * tokens].double(), + b[i * tokens : (i + 1) * tokens].double(), + ) + err = max(err, (ai - bi).abs().max().item() / max(bi.abs().max().item(), 1e-6)) + return err + + +_workspaces = {} + + +def _workspace(heads): + """This module's attention workspace for ``heads`` (heads / 6 head groups): one per group count on the current + device, made on first use and shared by every call of that shape, as a target shares one across its layers.""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_mla import op + + groups = heads // H + if groups not in _workspaces: + _workspaces[groups] = op.make_attn_workspace( + torch.device("cuda", torch.cuda.current_device()), groups + ) + return _workspaces[groups] + + +def _attn_out(q, pool, row_stride, page_table, page_offset, seq_len, workspace=None): + from tensorrt_llm._torch.cute_dsl_kernels.k3_mla import op # noqa: F401 + + m, heads = q.shape[0], q.shape[1] + out = torch.empty(m, heads * LAT, dtype=torch.bfloat16, device="cuda") + torch.ops.trtllm.k3_mla_attn_out( + q.reshape(m, -1), pool.view(-1), row_stride, page_table, page_offset, seq_len, SCALE, out, + _workspace(heads) if workspace is None else workspace, + ) # fmt: skip + return out.view(m, heads, LAT) + + +def _attn_vb( + q, pool, row_stride, page_table, page_offset, seq_len, w_vb, gate=None, workspace=None +): + from tensorrt_llm._torch.cute_dsl_kernels.k3_mla import op # noqa: F401 + + m, heads = q.shape[0], q.shape[1] + y = torch.empty(m, heads * V, dtype=torch.bfloat16, device="cuda") + torch.ops.trtllm.k3_mla_attn_vb_out( + q.reshape(m, -1), + pool.view(-1), + row_stride, + page_table, + page_offset, + seq_len, + SCALE, + w_vb, + y, + _workspace(heads) if workspace is None else workspace, + gate, + GATE_COL0, + ) + return y + + +@pytest.mark.parametrize( + "num_requests,tokens,heads", ATTN_CASES, ids=[f"{r}x{t}-h{h}" for r, t, h in ATTN_CASES] +) +def test_attn_out(num_requests, tokens, heads): + """k3_mla_attn_out against the float64 reference, request by request against the one-request call, and rerun. Odd + seeds use an interleaved pool (2 layers, the layer's slot as the page offset) with rows of 640 elements.""" + seed = 11 * num_requests + tokens + layers, slot, row_stride = (2, 1, 640) if seed % 2 else (1, 0, DQK) + q, pool, page_table, page_offset, seq_len = _make_case( + seed, num_requests, tokens, heads, row_stride, layers, slot + ) + out = _attn_out(q, pool, row_stride, page_table, page_offset, seq_len) + again = _attn_out(q, pool, row_stride, page_table, page_offset, seq_len) + rows = [slice(i * tokens, (i + 1) * tokens) for i in range(num_requests)] + alone = torch.cat( + [_attn_out(q[r], pool, row_stride, page_table[i : i + 1], page_offset, seq_len[i : i + 1]) + for i, r in enumerate(rows)] + ) # fmt: skip + ref = _reference(q, pool, page_table, page_offset, seq_len, tokens) + torch.cuda.synchronize() + assert _max_rel(out, ref, tokens, num_requests) <= 1e-2 + assert torch.equal(_bits(out), _bits(again)) + assert torch.equal(_bits(out), _bits(alone)) + + +@pytest.mark.parametrize("num_requests,tokens", SPLITS, ids=SPLIT_IDS) +def test_attn_returns(num_requests, tokens): + """k3_mla_attn (its own output; a flat page-table row when R = 1) gives k3_mla_attn_out's bits.""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_mla import op # noqa: F401 + + q, pool, page_table, _, seq_len = _make_case(5 + num_requests, num_requests, tokens) + m = q.shape[0] + table = page_table[0] if num_requests == 1 else page_table + out = torch.ops.trtllm.k3_mla_attn( + q.view(m, -1), pool.view(-1), DQK, table, seq_len, SCALE, _workspace(H) + ) + ref = _attn_out(q, pool, DQK, page_table, 0, seq_len) + assert torch.equal(_bits(out), _bits(ref.view(m, -1))) + + +@pytest.mark.parametrize("gated", [False, True], ids=["plain", "gated"]) +@pytest.mark.parametrize("num_requests,tokens", SPLITS, ids=SPLIT_IDS) +def test_attn_vb(num_requests, tokens, gated): + """k3_mla_attn_vb_out against the reference (attention output rounded to bf16, v_b in float64), request by request + against the one-request call, rerun; with the gate, bit-identical to torch's bf16 y * s.""" + seed = 7 * num_requests + tokens + q, pool, page_table, page_offset, seq_len = _make_case(seed, num_requests, tokens) + gen = torch.Generator(device="cuda").manual_seed(seed + 1) + m, heads = q.shape[0], q.shape[1] + w_vb = (torch.randn(heads, V, LAT, generator=gen, device="cuda") * 0.05).bfloat16() + ag = ( + torch.rand(m, GATE_COL0 + heads * V, generator=gen, device="cuda").bfloat16() + if gated + else None + ) + y = _attn_vb(q, pool, DQK, page_table, page_offset, seq_len, w_vb, ag) + again = _attn_vb(q, pool, DQK, page_table, page_offset, seq_len, w_vb, ag) + rows = [slice(i * tokens, (i + 1) * tokens) for i in range(num_requests)] + alone = torch.cat( + [_attn_vb(q[r], pool, DQK, page_table[i : i + 1], page_offset, seq_len[i : i + 1], w_vb, + ag[r] if gated else None) for i, r in enumerate(rows)] + ) # fmt: skip + o_ref = _reference(q, pool, page_table, page_offset, seq_len, tokens).bfloat16().double() + y_ref = torch.einsum("thc,hvc->thv", o_ref, w_vb.double()).reshape(m, heads * V) + torch.cuda.synchronize() + if gated: + plain = _attn_vb(q, pool, DQK, page_table, page_offset, seq_len, w_vb) + assert torch.equal(_bits(y), _bits(plain * ag[:, GATE_COL0:])) + y_ref = y_ref.bfloat16().double() * ag[:, GATE_COL0:].double() + assert _max_rel(y, y_ref, tokens, num_requests) <= 1e-2 + assert torch.equal(_bits(y), _bits(again)) + assert torch.equal(_bits(y), _bits(alone)) + + +@pytest.mark.parametrize("num_requests,tokens", SPLITS, ids=SPLIT_IDS) +def test_attn_vs_stock(num_requests, tokens): + """k3_mla_attn_out against the stock CuTe DSL MLA decode on the same step (batch R, seq_len_q T); both within 1e-2 + of the float64 reference.""" + import cutlass + + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLNVMlaDecodeBlackwellRunner + + q, pool, page_table, page_offset, seq_len = _make_case( + 3 * num_requests + tokens, num_requests, tokens + ) + m, heads = q.shape[0], q.shape[1] + size = CuteDSLNVMlaDecodeBlackwellRunner.get_max_padded_workspace_size( + heads, tokens, LAT, num_requests, cutlass.Float32 + ) + workspace = torch.empty(max(size, 1), dtype=torch.int8, device="cuda") + kv = pool[:, :, :DQK] + stock = torch.empty(num_requests, tokens, heads, LAT, dtype=torch.bfloat16, device="cuda") + qv = q.view(num_requests, tokens, heads, DQK) + torch.ops.trtllm.cute_dsl_mla_decode_fp16_blackwell( + qv[..., :LAT].permute(2, 3, 1, 0), qv[..., LAT:].permute(2, 3, 1, 0), kv[..., :LAT].permute(1, 2, 0), + kv[..., LAT:].permute(1, 2, 0), (page_table + page_offset).transpose(0, 1), seq_len, + stock.permute(2, 3, 1, 0), workspace, heads, tokens, PAGE, SCALE, 1.0, num_requests, None, None, + ) # fmt: skip + out = _attn_out(q, pool, DQK, page_table, page_offset, seq_len) + ref = _reference(q, pool, page_table, page_offset, seq_len, tokens) + torch.cuda.synchronize() + stock = stock.view(m, heads, LAT) + assert _max_rel(stock, ref, tokens, num_requests) <= 1e-2 + assert _max_rel(out, stock, tokens, num_requests) <= 1e-2 + + +def test_attn_rejects(): + """Calls outside R <= 8 requests of T <= 8 tokens, or page-table rows / lengths that do not match, are refused; + so is a workspace that is not one for the call's head groups and device (its words are left as they were).""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_mla import op + + q, pool, page_table, _, seq_len = _make_case(1, 2, 4) + q2, flat = q.view(8, -1), pool.view(-1) + assert op.supports_attn(q2, flat, DQK, page_table, seq_len) + assert not op.supports_attn(q2, flat, DQK, page_table[:1], seq_len) # 1 row, 2 lengths + assert not op.supports_attn(q2, flat, DQK, page_table[0], seq_len) # a flat row, 2 lengths + assert op.supports_attn(q2[:6], flat, DQK, page_table, seq_len) # 2 requests x 3 tokens + assert not op.supports_attn(q2[:7], flat, DQK, page_table, seq_len) # 7 tokens over 2 requests + q16, _, table16, _, len16 = _make_case(2, 1, 16) + assert not op.supports_attn(q16.view(16, -1), flat, DQK, table16, len16) # T = 16 + q9, _, table9, _, len9 = _make_case(3, 9, 1) + assert not op.supports_attn(q9.view(9, -1), flat, DQK, table9, len9) # R = 9 + with pytest.raises(ValueError): + _attn_out(q[:7], pool, DQK, page_table, 0, seq_len) + other = _workspace(4 * H) # four head groups' layout; the call has one + before = other.clone() + for workspace in (other, other.float(), _workspace(H)[:-8]): + with pytest.raises(ValueError, match="workspace"): + _attn_out(q, pool, DQK, page_table, 0, seq_len, workspace) + torch.cuda.synchronize() + assert torch.equal(other.view(torch.int16), before.view(torch.int16)) + + +# Both launch modes (clusters; no_cluster past CLUSTER_WAVE clusters) at the TP16 and TP4 head counts, with folds +# (3x5: 4100 rows, 8x1, 8x8) and CTAs without a tile (1x8-h24: 5 tiles). +POISON_CASES = [(1, 1, H), (3, 5, H), (8, 1, H), (8, 8, H), (1, 8, 4 * H), (2, 4, 4 * H)] + + +def _fill_workspace(ws, heads, value): + """Fill the data words of an attention workspace: the per-CTA partial slots (fp16) and the no_cluster (m, l) + exchange (fp32). The arrival counters after them keep their values.""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_mla import k3_mla_attn_kernel as kernel + + groups = heads // kernel.HEADS + slots = kernel.MAX_REQUESTS * groups * kernel.CLUSTER + partials = slots * kernel.WS_SLOT_ELEMS + exchange = slots * kernel.ROWS * 2 # fp32 words + ws[:partials].fill_(value) + ws[partials : partials + 2 * exchange].view(torch.float32).fill_(value) + + +@pytest.mark.parametrize( + "num_requests,tokens,heads", POISON_CASES, ids=[f"{r}x{t}-h{h}" for r, t, h in POISON_CASES] +) +def test_attn_workspace_poison(num_requests, tokens, heads): + """A call reads only workspace words it wrote itself: with the partial slots and the (m, l) exchange refilled + with NaN before each call, k3_mla_attn_out and the gated k3_mla_attn_vb_out give the bits of the same calls on a + zero-filled workspace.""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_mla import op + + seed = 13 * num_requests + tokens + q, pool, page_table, page_offset, seq_len = _make_case(seed, num_requests, tokens, heads) + gen = torch.Generator(device="cuda").manual_seed(seed + 1) + m = q.shape[0] + w_vb = (torch.randn(heads, V, LAT, generator=gen, device="cuda") * 0.05).bfloat16() + ag = torch.rand(m, GATE_COL0 + heads * V, generator=gen, device="cuda").bfloat16() + ws = op.make_attn_workspace(torch.device("cuda", torch.cuda.current_device()), heads // H) + outs = {} + for value in (0.0, float("nan")): + _fill_workspace(ws, heads, value) + o = _attn_out(q, pool, DQK, page_table, page_offset, seq_len, ws) + _fill_workspace(ws, heads, value) + y = _attn_vb(q, pool, DQK, page_table, page_offset, seq_len, w_vb, ag, ws) + outs[value == 0.0] = (o, y) + torch.cuda.synchronize() + for got, want in zip(outs[False], outs[True]): + assert not torch.isnan(got).any() + assert torch.equal(_bits(got), _bits(want)) + + +# no_cluster launches (more clusters than co-reside): 8 requests at 6 heads, 2 requests at 24 heads. +WRAP_CASES = [(8, 1, H), (8, 8, H), (2, 4, 4 * H)] + + +def _set_counters(ws, heads, value): + """Set every no_cluster arrival counter of an attention workspace (int32 words after the (m, l) exchange).""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_mla import k3_mla_attn_kernel as kernel + + groups = heads // kernel.HEADS + slots = kernel.MAX_REQUESTS * groups * kernel.CLUSTER + start = ( + slots * kernel.WS_SLOT_ELEMS + 2 * slots * kernel.ROWS * 2 + ) # fp16 elements before the counters + ws[start:].view(torch.int32).fill_(value) + + +@pytest.mark.parametrize( + "num_requests,tokens,heads", WRAP_CASES, ids=[f"{r}x{t}-h{h}" for r, t, h in WRAP_CASES] +) +def test_attn_counter_wrap(num_requests, tokens, heads): + """The no_cluster arrival counters only grow (16 per launch) and are compared by signed difference: calls with the + counters just below the int32 wrap (2^31 - 16, and -16 just below 0) give the bits of calls on zeroed ones.""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_mla import op + + seed = 19 * num_requests + tokens + q, pool, page_table, page_offset, seq_len = _make_case(seed, num_requests, tokens, heads) + gen = torch.Generator(device="cuda").manual_seed(seed + 1) + w_vb = (torch.randn(heads, V, LAT, generator=gen, device="cuda") * 0.05).bfloat16() + ws = op.make_attn_workspace(torch.device("cuda", torch.cuda.current_device()), heads // H) + outs = {} + for start in (0, 2**31 - 16, -16): + _set_counters(ws, heads, start) + # Three launches: each crosses the wrap point once the counters start 16 below it. + outs[start] = [ + _attn_out(q, pool, DQK, page_table, page_offset, seq_len, ws), + _attn_vb(q, pool, DQK, page_table, page_offset, seq_len, w_vb, workspace=ws), + _attn_out(q, pool, DQK, page_table, page_offset, seq_len, ws), + ] + torch.cuda.synchronize() + for start in (2**31 - 16, -16): + for got, want in zip(outs[start], outs[0]): + assert torch.equal(_bits(got), _bits(want)) diff --git a/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_mla_decode_view.py b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_mla_decode_view.py new file mode 100644 index 000000000000..8b5004947b2e --- /dev/null +++ b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_mla_decode_view.py @@ -0,0 +1,125 @@ +# 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_mla_decode_view: the Kimi K3 MLA decode kernels' per-layer inputs from the attention metadata for R generation +requests of T tokens (views of the metadata buffers, no copies), and a reason string for every step it cannot take.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.fmha.cute_dsl_mla import k3_mla_decode_view + +MAX_SEQS, MAX_PAGES, POOL_PAGES = 16, 32, 40 + + +def _attn(num_heads=6): + return SimpleNamespace( + num_heads=num_heads, kv_lora_rank=512, qk_rope_head_dim=64, qk_nope_head_dim=128, q_scaling=1.0, + layer_idx=3, get_local_layer_idx=lambda meta: 1, + ) # fmt: skip + + +def _meta(num_generations, num_contexts=0, pool=None, **overrides): + if pool is None: + pool = torch.zeros(POOL_PAGES, 1, 64, 1, 576, dtype=torch.bfloat16) + meta = SimpleNamespace( + num_contexts=num_contexts, + num_generations=num_generations, + beam_width=1, + tokens_per_block=64, + helix_position_offsets=None, + kv_cache_manager=SimpleNamespace(get_buffers=lambda layer_idx: pool), + kv_cache_block_offsets=torch.arange(2 * MAX_SEQS * 2 * MAX_PAGES, dtype=torch.int32).view( + 2, MAX_SEQS, 2, MAX_PAGES + ), + host_kv_cache_pool_mapping=torch.tensor([[0, 0], [1, 0]], dtype=torch.int32), + kv_lens_cuda_runtime=torch.arange(100, 100 + MAX_SEQS, dtype=torch.int32), + ) + for key, value in overrides.items(): + setattr(meta, key, value) + return meta + + +SPLITS = [(1, 1), (2, 1), (3, 1), (4, 1), (5, 1), (8, 1), (1, 8), (2, 4), (4, 2)] + [ + (r, 8) for r in range(2, 9) +] + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("num_requests,tokens", SPLITS, ids=[f"{r}x{t}" for r, t in SPLITS]) +def test_view_requests(num_requests, tokens): + """R x T: the R page-table rows of the layer's pool as a strided view of kv_cache_block_offsets, the R lengths as + a view of kv_lens_cuda_runtime, R and T.""" + meta = _meta(num_requests) + view = k3_mla_decode_view(_attn(), meta, num_requests * tokens) + assert isinstance(view, dict), view + table = view["page_table"] + assert tuple(table.shape) == (num_requests, MAX_PAGES) and table.stride() == (2 * MAX_PAGES, 1) + assert table.data_ptr() == meta.kv_cache_block_offsets[1, 0, 0].data_ptr() + assert torch.equal(table, meta.kv_cache_block_offsets[1, :num_requests, 0]) + assert view["seq_len"].data_ptr() == meta.kv_lens_cuda_runtime.data_ptr() + assert view["seq_len"].shape == (num_requests,) + assert (view["num_requests"], view["tokens_per_request"]) == (num_requests, tokens) + assert view["row_stride"] == 576 and view["page_offset"] == 0 + assert view["pool"].numel() == POOL_PAGES * 64 * 576 + assert view["softmax_scale"] == pytest.approx(192**-0.5) + + +@pytest.mark.cpu_only +def test_view_interleaved_pool(): + """A layer's view of a layer-interleaved pool: the whole pool with the layer's slot as the page offset.""" + layers = 3 + pools = torch.zeros(POOL_PAGES, layers, 1, 64, 1, 640, dtype=torch.bfloat16) + view = k3_mla_decode_view(_attn(), _meta(2, pool=pools[:, 2]), 16) + assert isinstance(view, dict), view + assert view["page_offset"] == 2 and view["row_stride"] == 640 + assert view["pool"].numel() == POOL_PAGES * layers * 64 * 640 + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "num_generations,num_tokens,overrides,attn_heads", + [ + (2, 16, dict(num_contexts=1), 6), # a context request + (0, 8, {}, 6), # no generation request + (9, 9, {}, 6), # more than 8 requests + (1, 16, {}, 6), # more than 8 tokens per request + (2, 7, {}, 6), # tokens not uniform over the requests + (2, 8, dict(beam_width=2), 6), + (2, 8, dict(is_spec_dec_tree=True), 6), + (2, 8, dict(tokens_per_block=32), 6), + (2, 8, dict(helix_position_offsets=torch.zeros(1)), 6), + (2, 8, dict(kv_cache_manager=None), 6), + (2, 8, {}, 8), # heads not a multiple of 6 + ], +) +def test_view_reasons(num_generations, num_tokens, overrides, attn_heads): + """Every step the kernels cannot take returns a reason string (and does not raise).""" + overrides = dict(overrides) + num_contexts = overrides.pop("num_contexts", 0) + reason = k3_mla_decode_view( + _attn(attn_heads), _meta(num_generations, num_contexts, **overrides), num_tokens + ) + assert isinstance(reason, str) and reason + + +@pytest.mark.cpu_only +def test_view_pool_reasons(): + """A quantized or differently laid out pool returns a reason string.""" + fp8 = torch.zeros(POOL_PAGES, 1, 64, 1, 576, dtype=torch.float8_e4m3fn) + assert isinstance(k3_mla_decode_view(_attn(), _meta(1, pool=fp8), 1), str) + narrow = torch.zeros(POOL_PAGES, 1, 64, 1, 512, dtype=torch.bfloat16) + assert isinstance(k3_mla_decode_view(_attn(), _meta(1, pool=narrow), 1), str) diff --git a/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_mla_q.py b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_mla_q.py new file mode 100644 index 000000000000..f3105d565e83 --- /dev/null +++ b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_mla_q.py @@ -0,0 +1,274 @@ +# 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::k3_mla_q / k3_mla_qkv / k3_mla_qkv_out (Kimi K3 MLA decode query path: q_a RMSNorm, q_b, k_b absorb, and the +KV half into the paged latent cache) at the TP16 shape (6 heads per rank, q_lora 1536, latent 512 + rope 64), M <= 64 +tokens: fused_q against the unfused chain (the model's RMSNorm, the q_b GEMM, the k_b bmm) and a reference with +the same bf16 roundings, every 8-token chunk bit-identical to the 8-token call on its rows; the KV rows of R requests +x T tokens at each request's positions in a sentinel-filled pool, nothing else written.""" + +import pytest +import torch + + +def _is_sm100() -> bool: + if not torch.cuda.is_available(): + return False + major, minor = torch.cuda.get_device_capability() + return major * 10 + minor in (100, 103) + + +pytestmark = pytest.mark.skipif(not _is_sm100(), reason="needs an SM100 / SM103 GPU") + +H, NOPE, PE, QK, LAT, QL, V = 6, 128, 64, 192, 512, 1536, 128 +DQK, PAGE = LAT + PE, 64 +EPS = KV_EPS = 1e-6 +SENTINEL = 0x7F7F # bf16 bits of the largest finite value: no cache row of the tests holds it + +SPLITS = [(1, 1), (2, 1), (3, 1), (4, 1), (5, 1), (8, 1), (1, 8), (2, 4), (4, 2)] + [ + (r, 8) for r in range(2, 9) +] +SPLITS += [(3, 5), (7, 3)] # chunks of 8 tokens that cut through requests +SPLIT_IDS = [f"{r}x{t}" for r, t in SPLITS] +# KV lengths, assigned to the requests of a step in turn: the step's rows crossing a page boundary (579, 1989), only +# the step's tokens (8), short and long contexts. +LENGTHS = (1100, 64 * 9 + 3, 8, 2049, 127, 4100, 64 * 31 + 5, 300, 64) + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.int16) + + +def _weights(seed, heads=H): + gen = torch.Generator(device="cuda").manual_seed(seed) + w_qa = (1.0 + 0.1 * torch.randn(QL, generator=gen, device="cuda")).bfloat16() + w_qb = (torch.randn(heads * QK, QL, generator=gen, device="cuda") * 0.03).bfloat16() + w_kb = (torch.randn(heads, LAT, NOPE, generator=gen, device="cuda") * 0.08).bfloat16() + w_kv = (1.0 + 0.1 * torch.randn(LAT, generator=gen, device="cuda")).bfloat16() + return w_qa, w_qb, w_kb, w_kv + + +def _ag(seed, m, heads=H): + """The fused projection's rows: [q_a 1536 | kv_a latent 512 | rope 64 | gate heads * 128].""" + gen = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(m, QL + DQK + heads * V, generator=gen, device="cuda") * 0.7).bfloat16() + + +def _rms_norm(x, w, eps): + from tensorrt_llm._torch.modules.rms_norm import RMSNorm + + norm = RMSNorm(hidden_size=x.shape[1], eps=eps, dtype=torch.bfloat16).cuda() + norm.weight.data.copy_(w) + return norm(x.contiguous()) + + +def _unfused_q(ag, w_qa, w_qb, w_kb): + """The model's unfused chain: q_a_layernorm, q_b_proj (GEMM), bmm with k_b_proj_trans, q_pe copied.""" + m, heads = ag.shape[0], w_kb.shape[0] + q = torch.matmul(_rms_norm(ag[:, :QL], w_qa, EPS), w_qb.t()).view(m, heads, QK) + q_abs = torch.bmm(q[..., :NOPE].transpose(0, 1), w_kb.transpose(1, 2)).transpose(0, 1) + return torch.cat([q_abs, q[..., NOPE:]], dim=-1).reshape(m, heads * DQK) + + +def _reference_q(ag, w_qa, w_qb, w_kb): + """The norm in fp32, the GEMMs in float64 (cuBLAS may run fp32 GEMMs in TF32), bf16 rounding where the model + rounds (norm output, q_b output, q_abs).""" + m, heads = ag.shape[0], w_kb.shape[0] + x = ag[:, :QL].float() + qn = (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + EPS) * w_qa.float()).bfloat16() + q = (qn.double() @ w_qb.double().t()).bfloat16().view(m, heads, QK) + q_abs = torch.einsum("thd,hcd->thc", q[..., :NOPE].double(), w_kb.double()).bfloat16() + return torch.cat([q_abs, q[..., NOPE:]], dim=-1).reshape(m, heads * DQK) + + +def _reference_kv(ag, w_kv): + """The cache rows in fp32, flashinfer's order ((x * rrms) * w, one bf16 rounding); the rope columns copied.""" + x = ag[:, QL : QL + LAT].float() + r = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + KV_EPS) + return torch.cat([((x * r) * w_kv.float()).bfloat16(), ag[:, QL + LAT : QL + DQK]], dim=-1) + + +def _max_rel(a, b): + return (a.float() - b.float()).abs().max().item() / b.float().abs().max().item() + + +def _q(ag, w_qa, w_qb, w_kb): + from tensorrt_llm._torch.cute_dsl_kernels.k3_mla import op # noqa: F401 + + return torch.ops.trtllm.k3_mla_q(ag, w_qa, EPS, w_qb, w_kb, True) + + +# Chunks of 8 tokens up to 24, of 16 up to 48 and of 32 up to 64 (op.CHUNK_TOKENS), with full and short last chunks. +# The TP4 shape (24 heads, 144 CTAs per chunk) at M 8, 40 and 64. +Q_CASES = [(m, H) for m in (1, 2, 3, 5, 8, 9, 15, 16, 24, 25, 31, 32, 33, 40, 48, 49, 57, 64)] + [ + (m, 4 * H) for m in (8, 40, 64) +] + + +@pytest.mark.parametrize("m,heads", Q_CASES, ids=[f"m{m}-h{h}" for m, h in Q_CASES]) +def test_q(m, heads): + """fused_q against the reference and the unfused chain (q_abs and q_pe), rerun; each 8-token chunk of rows + bit-identical to the call on those rows alone.""" + w_qa, w_qb, w_kb, _ = _weights(20261001, heads) + ag = _ag(m, m, heads) + y = _q(ag, w_qa, w_qb, w_kb) + again = _q(ag, w_qa, w_qb, w_kb) + chunks = torch.cat([_q(ag[c : c + 8].contiguous(), w_qa, w_qb, w_kb) for c in range(0, m, 8)]) + yu = _unfused_q(ag, w_qa, w_qb, w_kb) + yr = _reference_q(ag, w_qa, w_qb, w_kb) + torch.cuda.synchronize() + for part in (slice(0, LAT), slice(LAT, DQK)): + yk = y.view(m, heads, DQK)[..., part] + assert _max_rel(yk, yr.view(m, heads, DQK)[..., part]) <= 2e-2 + assert _max_rel(yk, yu.view(m, heads, DQK)[..., part]) <= 2e-2 + assert torch.equal(_bits(y), _bits(again)) + assert torch.equal(_bits(y), _bits(chunks)) + + +def _kv_case(seed, num_requests, tokens, row_stride=DQK, layers=1, lens=None): + """A sentinel-filled pool of `layers` interleaved layer slots (rows of `row_stride`), the page table as rows of an + int32 [R, 2, W] buffer (kv_cache_block_offsets' layout; entries past a request's pages name a page no request + owns), the lengths.""" + cpu_gen = torch.Generator().manual_seed(seed) + if lens is None: + lens = [max(tokens, LENGTHS[(i + seed) % len(LENGTHS)]) for i in range(num_requests)] + pages = [(n + PAGE - 1) // PAGE for n in lens] + total_pages = sum(pages) + 5 + width = max(pages) + 2 + pool = torch.full( + (total_pages * layers * PAGE * row_stride,), SENTINEL, dtype=torch.int16, device="cuda" + ) + perm = (torch.randperm(total_pages, generator=cpu_gen) * layers).to(torch.int32) + offsets = torch.empty(num_requests, 2, width, dtype=torch.int32).fill_(int(perm[-1])) + start = 0 + for i, n in enumerate(pages): + offsets[i, :, :n] = perm[start : start + n] + start += n + return ( + pool.view(torch.bfloat16), + offsets.cuda()[:, 0, :], + torch.tensor(lens, dtype=torch.int32, device="cuda"), + ) + + +def _kv_rows(page_table, page_offset, seq_len, tokens): + """(token, pool row) of every stored token: token u of request i at position L_i - T + u, if >= 0.""" + out = [] + for i, length in enumerate(seq_len.tolist()): + for u in range(tokens): + pos = length - tokens + u + if pos >= 0: + out.append( + ( + i * tokens + u, + (int(page_table[i, pos // PAGE]) + page_offset) * PAGE + pos % PAGE, + ) + ) + return out + + +def _qkv(ag, weights, pool, row_stride, page_table, page_offset, seq_len): + from tensorrt_llm._torch.cute_dsl_kernels.k3_mla import op # noqa: F401 + + w_qa, w_qb, w_kb, w_kv = weights + return torch.ops.trtllm.k3_mla_qkv( + ag, + w_qa, + EPS, + w_qb, + w_kb, + w_kv, + KV_EPS, + pool, + row_stride, + page_table, + page_offset, + seq_len, + True, + ) + + +@pytest.mark.parametrize("num_requests,tokens", SPLITS, ids=SPLIT_IDS) +def test_qkv(num_requests, tokens): + """The KV rows of every request at its positions (latent against the fp32 reference and the model's RMSNorm, rope + columns bit-exact), nothing else in the pool written, fused_q bit-identical to k3_mla_q, the dense variant's rows + bit-identical, rerun. Odd seeds: an interleaved pool (2 layers, slot 1) with rows of 640 elements.""" + seed = 13 * num_requests + tokens + layers, slot, row_stride = (2, 1, 640) if seed % 2 else (1, 0, DQK) + m = num_requests * tokens + weights = _weights(seed) + ag = _ag(seed + 1, m) + pool, page_table, seq_len = _kv_case(seed, num_requests, tokens, row_stride, layers) + y = _qkv(ag, weights, pool, row_stride, page_table, slot, seq_len) + stored = _kv_rows(page_table, slot, seq_len, tokens) + toks = torch.tensor([t for t, _ in stored], device="cuda") + rows = torch.tensor([r for _, r in stored], device="cuda") + pool_rows = pool.view(-1, row_stride) + got = pool_rows[rows, :DQK].clone() + rest = pool_rows.clone() + rest[rows] = torch.full_like(rest[rows].view(torch.int16), SENTINEL).view(torch.bfloat16) + ref = _reference_kv(ag, weights[3])[toks] + model = _rms_norm(ag[:, QL : QL + LAT], weights[3], KV_EPS)[toks] + dense = torch.empty(m, DQK, dtype=torch.bfloat16, device="cuda") + y_dense = torch.ops.trtllm.k3_mla_qkv_out( + ag, weights[0], EPS, weights[1], weights[2], weights[3], KV_EPS, dense, True + ) + y_q = _q(ag, *weights[:3]) + _qkv(ag, weights, pool, row_stride, page_table, slot, seq_len) + torch.cuda.synchronize() + assert len(stored) == m + assert _max_rel(got[:, :LAT], ref[:, :LAT]) <= 1e-2 + assert int((_bits(got[:, :LAT]) != _bits(model)).sum()) <= max(4, model.numel() // 1000) + assert torch.equal(_bits(got[:, LAT:]), _bits(ag[toks, QL + LAT : QL + DQK])) + assert bool((rest.view(torch.int16) == SENTINEL).all()) + assert torch.equal(_bits(y), _bits(y_q)) and torch.equal(_bits(y_dense), _bits(y_q)) + assert torch.equal(_bits(dense[toks]), _bits(got)) + assert torch.equal(_bits(pool_rows[rows, :DQK]), _bits(got)) + + +def test_qkv_short_length(): + """A request whose length is below T (positions < 0) has only its tokens at positions >= 0 stored.""" + num_requests, tokens = 2, 4 + weights = _weights(5) + ag = _ag(6, num_requests * tokens) + pool, page_table, seq_len = _kv_case(7, num_requests, tokens, lens=[2, 100]) + _qkv(ag, weights, pool, DQK, page_table, 0, seq_len) + stored = _kv_rows(page_table, 0, seq_len, tokens) + rows = torch.tensor([r for _, r in stored], device="cuda") + pool_rows = pool.view(-1, DQK) + ref = _reference_kv(ag, weights[3])[[t for t, _ in stored]] + rest = pool_rows.clone() + rest[rows] = torch.full_like(rest[rows].view(torch.int16), SENTINEL).view(torch.bfloat16) + torch.cuda.synchronize() + assert len(stored) == 6 + assert _max_rel(pool_rows[rows, :LAT], ref[:, :LAT]) <= 1e-2 + assert bool((rest.view(torch.int16) == SENTINEL).all()) + + +def test_q_rejects(): + """More than 64 tokens, or page-table rows / lengths that do not describe R requests of the call's tokens.""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_mla import op + + w_qa, w_qb, w_kb, w_kv = _weights(1) + assert op.supports_q(_ag(1, 64), w_qa, w_qb, w_kb) + assert not op.supports_q(_ag(1, 65), w_qa, w_qb, w_kb) + pool, page_table, seq_len = _kv_case(2, 2, 4) + ag = _ag(2, 8) + assert op.supports_kv(ag, w_kv, pool, DQK, page_table, seq_len) + assert not op.supports_kv( + ag[:7], w_kv, pool, DQK, page_table, seq_len + ) # 7 tokens over 2 requests + assert not op.supports_kv(ag, w_kv, pool, DQK, page_table[:1], seq_len) # 1 row, 2 lengths + assert not op.supports_kv(ag, w_kv, pool, DQK, page_table[0], seq_len) # a flat row, 2 lengths + assert not op.supports_kv(ag, w_kv, pool, DQK, page_table.long(), seq_len) # int64 table diff --git a/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_tcgen05_fences.py b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_tcgen05_fences.py new file mode 100644 index 000000000000..613e5d559a89 --- /dev/null +++ b/tests/unittest/_torch/cute_dsl_kernels/kimi_k3/test_k3_tcgen05_fences.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. +"""Source check of the tcgen05 thread-sync fences in Kimi K3 kernels (PTX ISA, tcgen05 memory consistency, canonical +sync patterns): a tcgen05 operation issued after an mbarrier wait follows ``tcgen05.fence::after_thread_sync``, and a +thread whose TMEM loads another thread's tcgen05 work must not overtake (an mbarrier arrive or a CTA barrier after +``tcgen05.wait::ld``) issues ``tcgen05.fence::before_thread_sync`` first. Without them ptxas may move the TMEM access +across the synchronization; no test of values can see it, so the kernels' sources are read. + + pytest test_k3_tcgen05_fences.py +""" + +import importlib.util +import re + +import pytest + +KERNELS = [ + "tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv.k3_ctm_gemv_kernel", + "tensorrt_llm._torch.cute_dsl_kernels.k3_decode_gemv.k3_decode_gemv_kernel", +] + +# Thread syncs that order other threads' work before a tcgen05 operation of this thread. Waits on a barrier only TMA +# completes (the weight / activation rings) are not: TMA writes and tcgen05 reads are both async proxy, ordered by the +# barrier's complete_tx, the CUTLASS mainloop pattern. +WAIT = re.compile( + r"mbarrier_try_wait\(|mbarrier_test_wait\(|_(try|test)_wait_cluster\(|(? SimpleNamespace: quant_config=None, sparse_attention_config=None, get_num_mamba_layers=lambda: 2, + # get_kv_cache_manager_cls reads the mapping and spec config (helix check). + mapping=None, + spec_config=None, ) def _capture_kimi_v2_manager_ctor( monkeypatch: pytest.MonkeyPatch, spec_config=None, + model_engine=None, ) -> tuple[tuple, dict]: """Route a Kimi config through _create_kv_cache_manager with an explicit V2 manager and capture the constructor arguments.""" @@ -324,7 +329,7 @@ def __init__(self, *args: object, **kwargs: object) -> None: assert get_kv_cache_manager_cls(model_config, kv_cache_config) is MambaHybridCacheManagerV2 _create_kv_cache_manager( - model_engine=None, + model_engine=model_engine, kv_cache_manager_cls=RecordingV2Manager, mapping=Mapping(world_size=1, tp_size=1, pp_size=1), kv_cache_config=kv_cache_config, @@ -380,6 +385,33 @@ def test_kimi_explicit_v2_manager_enables_kda_replay( assert kwargs["kda_replay_num_spec"] == spec_config.tokens_per_gen_step - 1 assert kwargs["conv_state_layout"] == "q_k_v" + assert "kda_token_states" not in kwargs + + +@pytest.mark.parametrize("backbone_asks", [True, False]) +def test_kimi_explicit_v2_manager_kda_token_states_follow_the_model( + monkeypatch: pytest.MonkeyPatch, + backbone_asks: bool, +) -> None: + """The factory asks the V2 manager for per-token KDA verify states only when + the model's backbone sets ``kda_token_states``.""" + monkeypatch.setattr( + "tensorrt_llm._torch.modules.kimi_kda._kda_kernels.is_kda_mtp_verify_available", + lambda: True, + ) + backbone = SimpleNamespace(kda_token_states=True) if backbone_asks else SimpleNamespace() + model_engine = SimpleNamespace( + model=SimpleNamespace(model=backbone), + _max_cuda_graph_batch_size=4, + is_multimodal=True, + ) + + _, kwargs = _capture_kimi_v2_manager_ctor( + monkeypatch, MTPDecodingConfig(max_draft_len=3), model_engine + ) + + assert kwargs["kda_replay_num_spec"] == 3 + assert kwargs.get("kda_token_states", False) is backbone_asks def test_kimi_explicit_v2_manager_uses_qkv_convolution_layout( @@ -2353,6 +2385,7 @@ def _build_v2_hybrid_with_mamba_layer( mamba_n_groups=1, mamba_ssm_cache_dtype=torch.float16, kda_replay_num_spec=None, + kda_token_states=False, ): """Construct a real MambaHybridCacheManagerV2.""" mamba_mask = [True] * num_mamba_layers + [False] * num_attention_layers @@ -2407,6 +2440,7 @@ def _build_v2_hybrid_with_mamba_layer( dtype=dtype, conv_state_layout=conv_state_layout, kda_replay_num_spec=kda_replay_num_spec, + kda_token_states=kda_token_states, ) @@ -3629,6 +3663,43 @@ def test_v2_kda_replay_allocates_logical_slot_caches(): ) assert layer_cache.intermediate_ssm is None assert layer_cache.intermediate_conv_window is None + assert layer_cache.kda_state_tok is None + assert not mgr.keeps_kda_token_states + finally: + mgr.shutdown() + + +@skip_no_cuda +@pytest.mark.parametrize("kda_replay_num_spec", [2, None]) +def test_v2_kda_token_states_allocated_with_the_replay_caches(kda_replay_num_spec): + """With the replay caches, ``kda_token_states`` adds the fp32 state after every + draft of each slot, per layer; without them it allocates nothing.""" + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=4, + num_mamba_layers=2, + spec_config=MTPDecodingConfig(max_draft_len=2), + conv_state_layout="q_k_v", + mamba_d_conv=5, + mamba_num_heads=6, + mamba_n_groups=6, + mamba_ssm_cache_dtype=torch.float32, + kda_replay_num_spec=kda_replay_num_spec, + kda_token_states=True, + ) + try: + if kda_replay_num_spec is None: + assert getattr(mgr, "kda_state_tok", None) is None + assert not mgr.keeps_kda_token_states + return + assert mgr.keeps_kda_token_states + for layer_idx in range(2): + layer_cache = mgr.mamba_layer_cache(layer_idx) + cache_size = layer_cache.temporal.shape[0] + states = layer_cache.kda_state_tok + assert states.shape == (cache_size, 2, *layer_cache.temporal.shape[1:]) + assert states.dtype is torch.float32 + assert states.data_ptr() == mgr.kda_state_tok[layer_idx].data_ptr() + assert not states.any() finally: mgr.shutdown() @@ -3676,6 +3747,78 @@ def test_v2_kda_replay_validates_configuration( ) +@pytest.mark.parametrize( + ("kda_token_states", "enable_block_reuse", "at_floor"), + [(True, False, True), (False, False, False), (True, True, False)], +) +def test_v2_kda_token_states_keep_ssm_pool_at_live_floor( + kda_token_states, enable_block_reuse, at_floor +): + """The per-token KDA states take memory per SSM slot outside the cache quota. Without block reuse the SSM pool + keeps only its live floor and attention gets the rest of the quota; otherwise (replay caches only, or block reuse) + the typical step's ratio sizes it.""" + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr._generation_kv_capacity_headroom = 1 + mgr._has_cp_helix = False + mgr.kv_cache_type = CacheTypeCpp.SELF + mgr.head_dim_per_layer = [64, 64] + mgr.pp_layers = [0, 1] + mgr._mamba_layer_mask = [True, False] + # One 2 MiB grain per state slot, so the pool's slot count follows its grains (a Kimi K3 slot holds 27 MB). + mgr.ssm_bytes = 2 << 20 + mgr.conv_bytes = 64 << 10 + mgr.max_attention_window_vec = [128, 128] + mgr.max_batch_size = 2 + mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + mgr.max_seq_len = 128 + mgr.max_num_tokens = 128 + mgr.tokens_per_block = 32 + mgr.num_local_layers = 2 + mgr.local_num_mamba_layers = 1 + mgr._num_reserved_dummy_slots = 1 + mgr.dtype = DataType.HALF + mgr.enable_swa_scratch_reuse = False + mgr.enable_stats = False + mgr.num_extra_kv_tokens = 0 + mgr.get_layer_bytes_per_token = lambda **kwargs: 8 + # Layer 1's 256-byte pages of 32 tokens (_base_attention_layer_configs; layer 0 becomes the SSM layer). + mgr._attention_cache_bytes_per_token = lambda: 8 + mgr._use_kda_replay_update = True + mgr._kda_token_states = kda_token_states + # The per-token states count in each slot's state bytes (_mamba_state_bytes_per_slot): num_spec fp32 states of + # the SSM state shape, 2 MiB each here. + mgr._kda_replay_num_spec = 2 + mgr.ssm_state_shape = [8, 256, 256] + mgr.kv_cache_config = KvCacheConfig( + avg_seq_len=64, + enable_block_reuse=enable_block_reuse, + enable_partial_reuse=False, + ) + base_config = KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=128 << 20)], + layers=_base_attention_layer_configs(2), + ) + runtime_manager = RuntimeKVCacheManager(mgr._build_cache_config(base_config)) + try: + slots = {} + for stats in runtime_manager.get_storage_statistics(): + sizes = stats.slot_sizes if hasattr(stats, "slot_sizes") else stats.slot_size + role = "ssm" if mgr.ssm_bytes in [int(s) for s in sizes] else "attention" + slots[role] = int(stats.total) + finally: + runtime_manager.shutdown() + + floor = mgr._max_resident_sequences() + mgr._num_reserved_dummy_slots + if at_floor: + # The floor's min-slots constraint is scaled by 1 / max_util_for_resume (0.97). + assert floor <= slots["ssm"] <= floor + 1 + # Every other 2 MiB grain goes to attention's 256-byte pages. + assert slots["attention"] >= ((128 << 20) // (2 << 20) - (floor + 2)) * ((2 << 20) // 256) + else: + assert slots["ssm"] > 10 * floor + + def test_mamba_cache_manager_delegates_kda_replay_capability() -> None: mgr = object.__new__(MambaCacheManager) mgr._impl = SimpleNamespace(use_kda_replay_update=True) @@ -3785,6 +3928,60 @@ def test_v2_kda_replay_relocates_live_slot_history(): mgr.shutdown() +@skip_no_cuda +def test_v2_kda_token_states_relocate_with_their_slot(): + mgr = _build_v2_hybrid_with_mamba_layer( + spec_config=MTPDecodingConfig(max_draft_len=2), + conv_state_layout="q_k_v", + mamba_n_groups=4, + mamba_ssm_cache_dtype=torch.float32, + kda_replay_num_spec=2, + kda_token_states=True, + ) + try: + states = mgr.kda_state_tok + assert states is not None + states.zero_() + states[:, 0].fill_(1.0) + states[:, 1].fill_(2.0) + source_zero, source_one = states[:, 0].clone(), states[:, 1].clone() + + mgr._relocate_kda_replay_slots([0, 1], [1, 2]) + + torch.testing.assert_close(states[:, 1], source_zero, rtol=0, atol=0) + torch.testing.assert_close(states[:, 2], source_one, rtol=0, atol=0) + finally: + mgr.shutdown() + # Released with the other replay buffers. + assert mgr.kda_state_tok is None + + +@skip_no_cuda +def test_v2_kda_token_states_count_in_the_per_slot_budget(): + """The capacity math sees the per-token states: a slot costs one more fp32 SSM state per draft and layer.""" + kwargs = dict( + num_mamba_layers=2, + spec_config=MTPDecodingConfig(max_draft_len=2), + conv_state_layout="q_k_v", + mamba_d_conv=5, + mamba_num_heads=6, + mamba_n_groups=6, + mamba_ssm_cache_dtype=torch.float32, + kda_replay_num_spec=2, + ) + plain = _build_v2_hybrid_with_mamba_layer(**kwargs) + try: + plain_bytes = plain._mamba_state_bytes_per_slot() + finally: + plain.shutdown() + with_states = _build_v2_hybrid_with_mamba_layer(kda_token_states=True, **kwargs) + try: + per_token = 2 * 2 * math.prod(with_states.ssm_state_shape) * 4 + assert with_states._mamba_state_bytes_per_slot() == plain_bytes + per_token + finally: + with_states.shutdown() + + @skip_no_cuda def test_v2_kda_state_index_setup_relocates_generation_history(): class StateCache: diff --git a/tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_k3_situ_mul.py b/tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_k3_situ_mul.py new file mode 100644 index 000000000000..41976eaef202 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_k3_situ_mul.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the k3_situ_mul catalog entry. + +The Kimi K3 dense MLP's SiTU-and-mul at TP16 (gate_up [M, 4224] -> [M, 2112]) with the checkpoint's and the default +(beta, linear_beta), at every M in 1..8: error against SituAndMul's formula evaluated in fp64, identical bits on a +repeated call, and each M's rows bit-identical to the same rows of the 8-row call. A CUDA-graph replay returns the +eager bits, calls outside the op's preconditions raise ValueError, and the wrapper refuses linear_beta=0.0. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.k3_situ_mul import k3_situ_mul + +assert torch.cuda.is_available(), "k3_situ_mul requires a CUDA device" +# The receipts are sm_100 ones; CI's other architectures skip this file. +pytestmark = pytest.mark.skipif( + torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" +) + +M_ALL = range(1, 9) +TOL = 8e-3 # max |y - ref| / max |ref| against the formula evaluated in fp64 +K = 2112 # the dense MLP's intermediate size per rank at TP16 +SITU = {"k3": (4.0, 25.0), "defaults": (1.0, None)} # (beta, linear_beta) + + +def _randn(rows: int, cols: int, seed: int, scale: float) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(rows, cols, generator=gen, device="cuda") * scale).bfloat16() + + +def _gate_up() -> torch.Tensor: + return _randn(8, 2 * K, 14, 2.0) + + +def _situ_ref(gu: torch.Tensor, beta: float, linear_beta: float | None) -> torch.Tensor: + """SituAndMul's eager fp32 formula, evaluated in fp64.""" + k = gu.shape[1] // 2 + g, u = gu[:, :k].double(), gu[:, k:].double() + a = beta * torch.tanh(g / beta) * torch.sigmoid(g) + if linear_beta is not None: + u = linear_beta * torch.tanh(u / linear_beta) + return a * u + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.int16) + + +def _rel(y: torch.Tensor, ref: torch.Tensor) -> float: + """max |y - ref| / max |ref|, in fp64.""" + return (y.double() - ref).abs().max().item() / ref.abs().max().item() + + +def test_call_site_cells() -> None: + gu8 = _gate_up() + for name, (beta, linear_beta) in SITU.items(): + ref8 = _situ_ref(gu8, beta, linear_beta) + y8 = k3_situ_mul(gu8, beta, linear_beta) + for m in M_ALL: + cell = f"{name} M={m}" + gu = gu8[:m].contiguous() + y = k3_situ_mul(gu, beta, linear_beta) + assert y.shape == (m, K) and y.dtype == torch.bfloat16, cell + rel = _rel(y, ref8[:m]) + assert rel <= TOL, f"{cell}: max |y - ref| / max |ref| = {rel:.2e} > {TOL}" + again = k3_situ_mul(gu, beta, linear_beta) + assert torch.equal(_bits(y), _bits(again)), f"{cell}: a repeated call changed bits" + assert torch.equal(_bits(y), _bits(y8[:m])), f"{cell}: rows differ from the 8-row call" + + +def test_cuda_graph_replay_matches_eager() -> None: + """Captured after an eager call per key and replayed with gu rewritten in place: the eager bits.""" + calls = [] + for name, (beta, linear_beta) in SITU.items(): + for m in (1, 8): + gu = _gate_up()[:m].clone() + k3_situ_mul(gu, beta, linear_beta) # compiles the key outside capture + calls.append((f"{name} M={m}", gu, beta, linear_beta)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outs = [k3_situ_mul(gu, beta, linear_beta) for _, gu, beta, linear_beta in calls] + for seed, (_, gu, *_) in enumerate(calls, start=100): + gu.copy_(_randn(gu.shape[0], gu.shape[1], seed, 2.0)) + graph.replay() + for (cell, gu, beta, linear_beta), y in zip(calls, outs): + eager = k3_situ_mul(gu, beta, linear_beta) + assert torch.equal(_bits(y), _bits(eager)), f"{cell}: the replay differs from an eager call" + + +def test_unsupported_calls_raise_value_error() -> None: + """The op's own check refuses these before compiling or launching anything.""" + gu8 = _gate_up() + flat = torch.zeros(8 * 2 * K + 8, dtype=torch.bfloat16, device="cuda") + cases = { + "0 rows": gu8[:0], + "9 rows": torch.cat([gu8, gu8[:1]]), + "width % 16 != 0": _randn(8, 2 * K - 8, 14, 2.0), + "gu 2 bytes past a 16-byte boundary": flat[1 : 1 + 8 * 2 * K].view(8, 2 * K), + "fp16 gu": gu8.half(), + "row-strided gu": torch.cat([gu8, gu8], dim=1)[:, : 2 * K], + "1-D gu": gu8[0], + } + for name, gu in cases.items(): + try: + k3_situ_mul(gu, 4.0, 25.0) + except ValueError: + continue + raise AssertionError(f"{name}: expected ValueError, the op accepted the call") + + +def test_zero_linear_beta_is_refused_before_dispatch() -> None: + """The op would run linear_beta=0.0 as 1.0 with no error; the wrapper refuses it instead.""" + gu = _gate_up()[:1] + for linear_beta in (0.0, -0.0): + try: + k3_situ_mul(gu, 4.0, linear_beta) + except AssertionError: + continue + raise AssertionError(f"linear_beta={linear_beta} should have been refused before dispatch") diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_attn_vb_out.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_attn_vb_out.py new file mode 100644 index 000000000000..065b142ed32c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_attn_vb_out.py @@ -0,0 +1,446 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the k3_mla_attn_vb_out catalog entry (and its k3_mla_attn / k3_mla_attn_out forms). + +The op reads the paged latent cache through addressing tensors the runtime derives from a KVCacheManager and a +prepared TrtllmAttentionMetadata, and keeps its partials and arrival counters in a caller-owned K3MlaAttnWorkspace. +The test builds that state for real: an MLA (SELFKONLY, kv_factor 1) KVCacheManager of three layers at Kimi K3's +latent width (576 = 512 + 64) and page size (64), every layer's rows written through the manager's own block ids, +a prepared TrtllmAttentionMetadata of R generation requests of T tokens, one TRTLLM MLA attention object per layer +(k3_mla_decode_view gives the op's pool, row stride, page table, layer slot and lengths) and workspaces made by +K3MlaAttnWorkspace.create. The reference reads each request's rows back through the manager's block ids, not +through the op's addressing. + +Covered, at Kimi K3's per-rank shapes (6 heads at TP16, 24 at TP4; latent 512, rope 64, v_head 128), in both +launch modes (up to CLUSTER_WAVE = 7 clusters of 16 CTAs; more take the no-cluster mode, which uses the workspace's +arrival counters): + +1. Cells: the gated output against a float64 reference (attention output rounded to bf16, v_b in float64, the gate + applied as bf16(y * s)); the plain v_b output, k3_mla_attn_out and k3_mla_attn on the same step. +2. Call sequences: layers x decode steps on one shared workspace eagerly, then the same steps captured once as a + CUDA graph and replayed with rewritten inputs and advanced lengths; two workspaces interleaved across layers. + Every output is bit-identical to the same call on a new workspace, and each workspace's counters account for + exactly its own no-cluster launches (16 per launch per request and head group). +3. Negative control: a workspace laid out for another head-group count, or of another dtype or size, is refused + with ValueError before any launch, and neither the output nor the workspace changes. +""" + +from typing import List + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_attn_vb_out import ( + k3_mla_attn, + k3_mla_attn_out, + k3_mla_attn_vb_out, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_attn_workspace import ( + K3MlaAttnWorkspace, +) +from tensorrt_llm._torch.attention.backends.fmha.cute_dsl_mla import k3_mla_decode_view +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping + + +def _is_sm100() -> bool: + if not torch.cuda.is_available(): + return False + major, minor = torch.cuda.get_device_capability() + # sm_100 exactly: the receipts' architecture (a missing receipt reads as unknown). + return (major, minor) == (10, 0) + + +pytestmark = pytest.mark.skipif(not _is_sm100(), reason="certified on SM100 only") + +H, NOPE, LAT, PE, QL, V = 6, 128, 512, 64, 1536, 128 +DQK, PAGE = LAT + PE, 64 +SCALE = 1.0 / (NOPE + PE) ** 0.5 +GATE_COL0 = QL + DQK # the gate's columns in the fused projection's rows +LAYERS = 3 +MAX_REQUESTS = 8 +# Request lengths before the first step, assigned in turn: rows crossing a page boundary (579, 1989), > 2048 rows +# (several tiles per CTA), a short context, exactly one page. +LENGTHS = (1100, 64 * 9 + 3, 2049, 127, 4100, 64 * 31 + 5, 300, 64) + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.int16) + + +def _no_cluster(num_requests: int, heads: int) -> bool: + """Whether a call of R requests takes the no-cluster mode: more 16-CTA clusters than co-reside, all of whose CTAs + fit on the SMs at once.""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_mla import k3_mla_attn_kernel as kernel + + clusters = num_requests * heads // H + sms = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + return clusters > kernel.CLUSTER_WAVE and clusters * kernel.CLUSTER <= sms + + +class _MlaCache: + """Real op state: an MLA (SELFKONLY) paged KV cache manager of LAYERS layers holding R requests that grow by T + tokens per decode step, every row they will reach written through the manager's block ids, a + TrtllmAttentionMetadata over it and one TRTLLM MLA attention object per layer for k3_mla_decode_view.""" + + def __init__(self, num_requests: int, tokens: int, steps: int, heads: int, seed: int): + self.tokens, self.heads = tokens, heads + self.request_ids = list(range(num_requests)) + self.cached = [LENGTHS[(i + seed) % len(LENGTHS)] for i in range(num_requests)] + final = [n + steps * tokens for n in self.cached] + pages = sum((n + PAGE - 1) // PAGE for n in final) + self.mgr = KVCacheManager( + KvCacheConfig(max_tokens=(pages + 16) * PAGE, enable_block_reuse=False), + CacheType.SELFKONLY, + num_layers=LAYERS, + num_kv_heads=1, + head_dim=DQK, + tokens_per_block=PAGE, + max_seq_len=((max(final) + PAGE - 1) // PAGE + 1) * PAGE, + max_batch_size=MAX_REQUESTS, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=DataType.BF16, + ) + self.mgr.add_dummy_requests(self.request_ids, token_nums=final) + self.attn = [ + create_attention( + "TRTLLM", + layer_idx=layer, + num_heads=heads, + head_dim=DQK, + num_kv_heads=1, + is_mla_enable=True, + q_lora_rank=QL, + kv_lora_rank=LAT, + qk_nope_head_dim=NOPE, + qk_rope_head_dim=PE, + v_head_dim=V, + ) + for layer in range(LAYERS) + ] + gen = torch.Generator(device="cuda").manual_seed(seed) + for layer in range(LAYERS): + pool = self.mgr.get_buffers(layer) + assert pool.dtype == torch.bfloat16 and tuple(pool.shape[1:]) == (1, PAGE, 1, DQK) + for blocks, length in zip(self._blocks(layer), final): + for p in range((length + PAGE - 1) // PAGE): + rows = torch.randn(PAGE, DQK, generator=gen, device="cuda") * 0.5 + pool[blocks[p], 0, :, 0] = rows.bfloat16() + + def _blocks(self, layer: int) -> List[List[int]]: + return self.mgr.get_batch_cache_indices(self.request_ids, layer) + + def metadata(self) -> TrtllmAttentionMetadata: + return TrtllmAttentionMetadata( + max_num_requests=MAX_REQUESTS, max_num_tokens=8192, kv_cache_manager=self.mgr + ) + + def prepare(self, md: TrtllmAttentionMetadata) -> None: + """The metadata of the next decode step: every request's cached rows plus its T new tokens.""" + md.seq_lens = torch.tensor([self.tokens] * len(self.request_ids), dtype=torch.int) + md.num_contexts = 0 + md.request_ids = self.request_ids + md.prompt_lens = [c + self.tokens for c in self.cached] + md.kv_cache_params = KVCacheParams( + use_cache=True, num_cached_tokens_per_seq=list(self.cached) + ) + md.prepare() + + def view(self, md: TrtllmAttentionMetadata, layer: int) -> dict: + view = k3_mla_decode_view(self.attn[layer], md, len(self.request_ids) * self.tokens) + assert isinstance(view, dict), view + assert view["row_stride"] == DQK and view["page_offset"] == self.mgr.layer_offsets[layer] + assert view["softmax_scale"] == pytest.approx(SCALE) + return view + + def advance(self) -> None: + self.cached = [c + self.tokens for c in self.cached] + + def reference(self, layer: int, q: torch.Tensor) -> torch.Tensor: + """float64 attention of each request's T tokens over its rows [0, L) read through the manager's block ids; + token t sees rows <= L - T + t. [M, heads, 512].""" + pool, blocks, outs = self.mgr.get_buffers(layer), self._blocks(layer), [] + t = self.tokens + for i, cached in enumerate(self.cached): + length = cached + t + pages = [blocks[i][p] for p in range((length + PAGE - 1) // PAGE)] + kv = pool[pages, 0, :, 0].reshape(-1, DQK)[:length].double() + qi = q[i * t : (i + 1) * t].view(t, self.heads, DQK).double() + s = torch.einsum("thd,ld->thl", qi, kv) * SCALE + limit = length - t + torch.arange(t, device="cuda") + hidden = torch.arange(length, device="cuda")[None, :] > limit[:, None] + s = s.masked_fill(hidden[:, None, :], float("-inf")) + outs.append(torch.einsum("thl,ld->thd", torch.softmax(s, dim=-1), kv[:, :LAT])) + return torch.cat(outs) + + def shutdown(self) -> None: + self.mgr.shutdown() + + +def _inputs(seed: int, m: int, heads: int): + """q (fused_q rows), the v_b weight and the gate (sigmoid values at GATE_COL0 + 128 h).""" + gen = torch.Generator(device="cuda").manual_seed(seed) + q = (torch.randn(m, heads * DQK, generator=gen, device="cuda") * 0.5).bfloat16() + w_vb = (torch.randn(heads, V, LAT, generator=gen, device="cuda") * 0.05).bfloat16() + gate = torch.rand(m, GATE_COL0 + heads * V, generator=gen, device="cuda").bfloat16() + return q, w_vb, gate + + +def _vb(view: dict, q, w_vb, gate, workspace: K3MlaAttnWorkspace) -> torch.Tensor: + m, heads = q.shape[0], q.shape[1] // DQK + out = torch.empty(m, heads * V, dtype=torch.bfloat16, device="cuda") + k3_mla_attn_vb_out( + q, + view["pool"], + view["row_stride"], + view["page_table"], + view["page_offset"], + view["seq_len"], + view["softmax_scale"], + w_vb, + out, + workspace, + gate, + GATE_COL0, + ) + return out + + +def _vb_reference(cache: _MlaCache, layer: int, q, w_vb, gate) -> torch.Tensor: + m, heads = q.shape[0], q.shape[1] // DQK + o = cache.reference(layer, q).bfloat16().double() + y = torch.einsum("thc,hvc->thv", o, w_vb.double()).reshape(m, heads * V) + if gate is None: + return y + return y.bfloat16().double() * gate[:, GATE_COL0:].double() + + +def _max_rel(a, b, tokens: int) -> float: + """max over requests of max |a - b| / max |b| (per request, so a short context is not hidden by a long one).""" + err = 0.0 + for i in range(a.shape[0] // tokens): + ai, bi = ( + a[i * tokens : (i + 1) * tokens].double(), + b[i * tokens : (i + 1) * tokens].double(), + ) + err = max(err, (ai - bi).abs().max().item() / max(bi.abs().max().item(), 1e-6)) + return err + + +def _counters(workspace: K3MlaAttnWorkspace) -> torch.Tensor: + """The no-cluster arrival counters, [8 requests, groups, 2] (the (m, l) exchange's and the drain's).""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_mla import k3_mla_attn_kernel as kernel + + g = workspace.groups + slots = kernel.MAX_REQUESTS * g * kernel.CLUSTER + # fp16 elements before the counters: the partial slots, then the fp32 (m, l) exchange. + start = slots * kernel.WS_SLOT_ELEMS + slots * kernel.ROWS * 4 + words = workspace.buffer[start:].view(torch.int32) + return words.view(kernel.MAX_REQUESTS, g, 2, kernel.CTR_STRIDE)[..., 0].clone() + + +def _check_counters(workspace: K3MlaAttnWorkspace, num_requests: int, launches: int) -> None: + """After `launches` no-cluster launches of R requests, every counter of those requests' head groups is + 16 x launches and every other counter is 0.""" + ctrs = _counters(workspace) + assert bool((ctrs[:num_requests] == 16 * launches).all()), ctrs[:num_requests] + assert bool((ctrs[num_requests:] == 0).all()) + + +# (R, T, heads): T 1 (decode without speculation) and 8 (a DSpark verify step); the cluster and no-cluster modes; +# the TP4 head count (4 head groups: several clusters per request). +CASES = [(3, 8, H), (8, 1, H), (2, 8, 4 * H)] +CASE_IDS = [f"{r}x{t}-h{h}" for r, t, h in CASES] +EAGER_STEPS, GRAPH_STEPS = 2, 2 + + +@pytest.mark.parametrize("num_requests,tokens,heads", CASES, ids=CASE_IDS) +def test_cells(num_requests, tokens, heads): + """One step of layer 0: the gated output within 1e-2 (per request) of the float64 reference and rerun-identical; + the plain v_b output within 1e-2 of its reference, and the gated one bit-identical to bf16(plain * s); + k3_mla_attn_out within 1e-2 of the attention reference and k3_mla_attn bit-identical to it.""" + seed = 17 * num_requests + tokens + heads + cache = _MlaCache(num_requests, tokens, 1, heads, seed) + m = num_requests * tokens + try: + md = cache.metadata() + cache.prepare(md) + view = cache.view(md, 0) + workspace = K3MlaAttnWorkspace.create(torch.device("cuda"), heads // H) + q, w_vb, gate = _inputs(seed, m, heads) + y = _vb(view, q, w_vb, gate, workspace) + again = _vb(view, q, w_vb, gate, workspace) + plain = _vb(view, q, w_vb, None, workspace) + o = torch.empty(m, heads * LAT, dtype=torch.bfloat16, device="cuda") + k3_mla_attn_out( + q, view["pool"], DQK, view["page_table"], view["page_offset"], view["seq_len"], + view["softmax_scale"], o, workspace, + ) # fmt: skip + assert view["page_offset"] == 0 # layer 0: k3_mla_attn takes no page offset + o_new = k3_mla_attn( + q, + view["pool"], + DQK, + view["page_table"], + view["seq_len"], + view["softmax_scale"], + workspace, + ) + torch.cuda.synchronize() + assert _max_rel(y, _vb_reference(cache, 0, q, w_vb, gate), tokens) <= 1e-2 + assert _max_rel(plain, _vb_reference(cache, 0, q, w_vb, None), tokens) <= 1e-2 + assert torch.equal(_bits(y), _bits(again)) + assert torch.equal(_bits(y), _bits(plain * gate[:, GATE_COL0:])) + ref_o = cache.reference(0, q).reshape(m, heads * LAT) + assert _max_rel(o, ref_o, tokens) <= 1e-2 + assert torch.equal(_bits(o_new), _bits(o)) + finally: + cache.shutdown() + + +@pytest.mark.parametrize("num_requests,tokens,heads", CASES, ids=CASE_IDS) +def test_layers_by_steps_eager_then_graph(num_requests, tokens, heads): + """LAYERS layers x EAGER_STEPS decode steps on one shared workspace, eagerly; then one step's LAYERS calls + captured as a CUDA graph on the same workspace and replayed for GRAPH_STEPS steps with rewritten q and the + metadata prepared for each step. Every output within 1e-2 of the reference and bit-identical to the same call + on a new workspace; in the no-cluster mode the shared workspace's counters end at 16 x its launches.""" + seed = 23 * num_requests + tokens + heads + groups = heads // H + cache = _MlaCache(num_requests, tokens, EAGER_STEPS + GRAPH_STEPS, heads, seed) + m = num_requests * tokens + shared = K3MlaAttnWorkspace.create(torch.device("cuda"), groups) + launches = 0 + try: + weights = [_inputs(seed + layer, m, heads)[1:] for layer in range(LAYERS)] + md = cache.metadata() + for step in range(EAGER_STEPS): + cache.prepare(md) + for layer in range(LAYERS): + q = _inputs(1000 * step + 10 * layer + seed, m, heads)[0] + view = cache.view(md, layer) + y = _vb(view, q, *weights[layer], shared) + launches += 1 + fresh = _vb(view, q, *weights[layer], K3MlaAttnWorkspace.create("cuda", groups)) + ref = _vb_reference(cache, layer, q, *weights[layer]) + torch.cuda.synchronize() + assert torch.equal(_bits(y), _bits(fresh)), (step, layer) + assert _max_rel(y, ref, tokens) <= 1e-2, (step, layer) + cache.advance() + + graph_md = cache.metadata().create_cuda_graph_metadata( + num_requests, max_draft_tokens=tokens - 1 + ) + cache.prepare(graph_md) + views = [cache.view(graph_md, layer) for layer in range(LAYERS)] + static_q = [ + torch.zeros(m, heads * DQK, dtype=torch.bfloat16, device="cuda") for _ in range(LAYERS) + ] + static_y = [ + torch.empty(m, heads * V, dtype=torch.bfloat16, device="cuda") for _ in range(LAYERS) + ] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for layer in range(LAYERS): + w_vb, gate = weights[layer] + k3_mla_attn_vb_out( + static_q[layer], views[layer]["pool"], DQK, views[layer]["page_table"], + views[layer]["page_offset"], views[layer]["seq_len"], SCALE, w_vb, static_y[layer], + shared, gate, GATE_COL0, + ) # fmt: skip + for step in range(EAGER_STEPS, EAGER_STEPS + GRAPH_STEPS): + if step > EAGER_STEPS: + cache.prepare(graph_md) + qs = [_inputs(1000 * step + 10 * layer + seed, m, heads)[0] for layer in range(LAYERS)] + for layer in range(LAYERS): + static_q[layer].copy_(qs[layer]) + graph.replay() + launches += LAYERS + for layer in range(LAYERS): + fresh = _vb(views[layer], qs[layer], *weights[layer], + K3MlaAttnWorkspace.create("cuda", groups)) # fmt: skip + ref = _vb_reference(cache, layer, qs[layer], *weights[layer]) + torch.cuda.synchronize() + assert torch.equal(_bits(static_y[layer]), _bits(fresh)), (step, layer) + assert _max_rel(static_y[layer], ref, tokens) <= 1e-2, (step, layer) + cache.advance() + torch.cuda.synchronize() + _check_counters(shared, num_requests, launches if _no_cluster(num_requests, heads) else 0) + finally: + cache.shutdown() + + +def test_two_workspaces_interleaved(): + """Two workspaces with the layers' calls alternating between them over 3 no-cluster steps (8 requests at TP16): + every output bit-identical to the same call on a new workspace, and each workspace's counters at 16 x its own + launches.""" + num_requests, tokens = 8, 1 + assert _no_cluster(num_requests, H) + cache = _MlaCache(num_requests, tokens, 3, H, 11) + pair = [K3MlaAttnWorkspace.create(torch.device("cuda"), 1) for _ in range(2)] + launches = [0, 0] + try: + weights = _inputs(12, num_requests, H)[1:] + md = cache.metadata() + for step in range(3): + cache.prepare(md) + for layer in range(LAYERS): + k = (step * LAYERS + layer) % 2 + q = _inputs(100 * step + layer, num_requests, H)[0] + view = cache.view(md, layer) + y = _vb(view, q, *weights, pair[k]) + launches[k] += 1 + fresh = _vb(view, q, *weights, K3MlaAttnWorkspace.create("cuda", 1)) + torch.cuda.synchronize() + assert torch.equal(_bits(y), _bits(fresh)), (step, layer) + cache.advance() + for workspace, n in zip(pair, launches): + _check_counters(workspace, num_requests, n) + finally: + cache.shutdown() + + +def test_workspace_of_another_shape_is_refused(): + """Negative control. A TP16 call (one head group) given a workspace laid out for four head groups, an fp32 + tensor of the same size, or a buffer 8 elements short raises ValueError before any launch: neither the output + nor the workspace changes.""" + cache = _MlaCache(2, 8, 1, H, 13) + try: + md = cache.metadata() + cache.prepare(md) + view = cache.view(md, 0) + q, w_vb, gate = _inputs(14, 16, H) + tp4 = K3MlaAttnWorkspace.create(torch.device("cuda"), 4) + right = K3MlaAttnWorkspace.create(torch.device("cuda"), 1) + out = torch.zeros(16, H * V, dtype=torch.bfloat16, device="cuda") + before = (_bits(tp4.buffer).clone(), _bits(right.buffer).clone()) + for bad in ( + tp4, + K3MlaAttnWorkspace(buffer=right.buffer.float(), groups=1), + K3MlaAttnWorkspace(buffer=right.buffer[:-8], groups=1), + ): + with pytest.raises(ValueError, match="workspace"): + k3_mla_attn_vb_out( + q, view["pool"], DQK, view["page_table"], view["page_offset"], view["seq_len"], SCALE, + w_vb, out, bad, gate, GATE_COL0, + ) # fmt: skip + torch.cuda.synchronize() + assert not bool(out.any()) + assert torch.equal(_bits(tp4.buffer), before[0]) and torch.equal( + _bits(right.buffer), before[1] + ) + finally: + cache.shutdown() + + +def test_create_refuses_capture(): + """A workspace is made eagerly: K3MlaAttnWorkspace.create raises under CUDA-graph capture.""" + graph = torch.cuda.CUDAGraph() + with pytest.raises(RuntimeError, match="capture"): + with torch.cuda.graph(graph): + K3MlaAttnWorkspace.create(torch.device("cuda"), 1) diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_qkv.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_qkv.py new file mode 100644 index 000000000000..feadc66a35e6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_qkv.py @@ -0,0 +1,367 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the k3_mla_qkv catalog entry (and its k3_mla_q / k3_mla_qkv_out forms). + +The op writes the paged latent cache through addressing tensors the runtime derives from a KVCacheManager and a +prepared TrtllmAttentionMetadata. The test builds that state for real: an MLA (SELFKONLY, kv_factor 1) KVCacheManager +of three layers at Kimi K3's latent width (576 = 512 + 64) and page size (64), a prepared TrtllmAttentionMetadata of R +generation requests of T tokens, and one TRTLLM MLA attention object per layer, whose k3_mla_decode_view gives the +op's pool, row stride, page table, layer slot (page offset) and lengths. Every pool of the manager starts filled with +a sentinel; the expected image of each layer is kept from the manager's own block ids (get_batch_cache_indices), not +from the op's addressing, and compared with the whole pool after every step. + +Covered, at Kimi K3's per-rank shapes (6 heads at TP16, 24 at TP4; q_lora 1536, latent 512, rope 64): + +1. Cells: fused_q against a reference with the model's bf16 roundings, the cache rows against an fp32 reference + (rope columns bit-exact), k3_mla_q's and k3_mla_qkv_out's results bit-identical to k3_mla_qkv's. +2. Call sequences on real cache objects: layers x decode steps eagerly, then the same steps captured once as a CUDA + graph and replayed with rewritten inputs and advanced lengths; two managers' calls interleaved. +3. Negative control: a call given another layer's slot. Nothing raises; it overwrites that layer's rows at the + step's positions and leaves its own slot unwritten, which the pool comparison sees. +""" + +from typing import Dict + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_qkv import ( + k3_mla_q, + k3_mla_qkv, + k3_mla_qkv_out, +) +from tensorrt_llm._torch.attention.backends.fmha.cute_dsl_mla import k3_mla_decode_view +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping + + +def _is_sm100() -> bool: + if not torch.cuda.is_available(): + return False + major, minor = torch.cuda.get_device_capability() + # sm_100 exactly: the receipts' architecture (a missing receipt reads as unknown). + return (major, minor) == (10, 0) + + +pytestmark = pytest.mark.skipif(not _is_sm100(), reason="certified on SM100 only") + +H, NOPE, PE, QK, LAT, QL, V = 6, 128, 64, 192, 512, 1536, 128 +DQK, PAGE = LAT + PE, 64 +EPS = KV_EPS = 1e-6 +LAYERS = 3 +MAX_REQUESTS = 8 +SENTINEL = 0x7F7F # bf16 bits of the largest finite value: no cache row of the test holds it +# Request lengths before the first step, assigned in turn: rows crossing a page boundary (579, 1989), > 2048 rows, +# a short context, exactly one page. +LENGTHS = (1100, 64 * 9 + 3, 2049, 127, 4100, 64 * 31 + 5, 300, 64) + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.int16) + + +def _weights(seed: int, heads: int): + gen = torch.Generator(device="cuda").manual_seed(seed) + w_qa = (1.0 + 0.1 * torch.randn(QL, generator=gen, device="cuda")).bfloat16() + w_qb = (torch.randn(heads * QK, QL, generator=gen, device="cuda") * 0.03).bfloat16() + w_kb = (torch.randn(heads, LAT, NOPE, generator=gen, device="cuda") * 0.08).bfloat16() + w_kv = (1.0 + 0.1 * torch.randn(LAT, generator=gen, device="cuda")).bfloat16() + return w_qa, w_qb, w_kb, w_kv + + +def _ag(seed: int, m: int, heads: int) -> torch.Tensor: + """The fused projection's rows: [q_a 1536 | kv_a latent 512 | rope 64 | gate heads * 128].""" + gen = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(m, QL + DQK + heads * V, generator=gen, device="cuda") * 0.7).bfloat16() + + +def _reference_q(ag, w_qa, w_qb, w_kb): + """The norm in fp32, the GEMMs in float64, bf16 rounding where the model rounds (norm output, q_b output, q_abs).""" + m, heads = ag.shape[0], w_kb.shape[0] + x = ag[:, :QL].float() + qn = (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + EPS) * w_qa.float()).bfloat16() + q = (qn.double() @ w_qb.double().t()).bfloat16().view(m, heads, QK) + q_abs = torch.einsum("thd,hcd->thc", q[..., :NOPE].double(), w_kb.double()).bfloat16() + return torch.cat([q_abs, q[..., NOPE:]], dim=-1).reshape(m, heads * DQK) + + +def _reference_kv(ag, w_kv): + """The cache rows in fp32 ((x * rrms) * w, one bf16 rounding); the rope columns copied.""" + x = ag[:, QL : QL + LAT].float() + r = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + KV_EPS) + return torch.cat([((x * r) * w_kv.float()).bfloat16(), ag[:, QL + LAT : QL + DQK]], dim=-1) + + +def _max_rel(a, b) -> float: + return (a.float() - b.float()).abs().max().item() / b.float().abs().max().item() + + +class _MlaCache: + """Real op state: an MLA (SELFKONLY) paged KV cache manager of LAYERS layers holding R requests that grow by T + tokens per decode step, a TrtllmAttentionMetadata over it, one TRTLLM MLA attention object per layer for + k3_mla_decode_view, and the expected image of every layer's pool.""" + + def __init__(self, num_requests: int, tokens: int, steps: int, heads: int, seed: int): + self.tokens, self.heads = tokens, heads + self.request_ids = list(range(num_requests)) + self.cached = [LENGTHS[(i + seed) % len(LENGTHS)] for i in range(num_requests)] + final = [n + steps * tokens for n in self.cached] + pages = sum((n + PAGE - 1) // PAGE for n in final) + self.mgr = KVCacheManager( + KvCacheConfig(max_tokens=(pages + 16) * PAGE, enable_block_reuse=False), + CacheType.SELFKONLY, + num_layers=LAYERS, + num_kv_heads=1, + head_dim=DQK, + tokens_per_block=PAGE, + max_seq_len=((max(final) + PAGE - 1) // PAGE + 1) * PAGE, + max_batch_size=MAX_REQUESTS, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=DataType.BF16, + ) + self.mgr.add_dummy_requests(self.request_ids, token_nums=final) + self.attn = [ + create_attention( + "TRTLLM", + layer_idx=layer, + num_heads=heads, + head_dim=DQK, + num_kv_heads=1, + is_mla_enable=True, + q_lora_rank=QL, + kv_lora_rank=LAT, + qk_nope_head_dim=NOPE, + qk_rope_head_dim=PE, + v_head_dim=V, + ) + for layer in range(LAYERS) + ] + self.image: Dict[int, torch.Tensor] = {} + for layer in range(LAYERS): + pool = self.mgr.get_buffers(layer) + assert pool.dtype == torch.bfloat16 and tuple(pool.shape[1:]) == (1, PAGE, 1, DQK) + pool.view(torch.int16).fill_(SENTINEL) + self.image[layer] = torch.full(pool.shape, SENTINEL, dtype=torch.int16, device="cuda") + self.blocks = { + layer: self.mgr.get_batch_cache_indices(self.request_ids, layer) + for layer in range(LAYERS) + } + + def metadata(self) -> TrtllmAttentionMetadata: + return TrtllmAttentionMetadata( + max_num_requests=MAX_REQUESTS, max_num_tokens=8192, kv_cache_manager=self.mgr + ) + + def prepare(self, md: TrtllmAttentionMetadata) -> None: + """The metadata of the next decode step: every request's cached rows plus its T new tokens.""" + md.seq_lens = torch.tensor([self.tokens] * len(self.request_ids), dtype=torch.int) + md.num_contexts = 0 + md.request_ids = self.request_ids + md.prompt_lens = [c + self.tokens for c in self.cached] + md.kv_cache_params = KVCacheParams( + use_cache=True, num_cached_tokens_per_seq=list(self.cached) + ) + md.prepare() + + def view(self, md: TrtllmAttentionMetadata, layer: int) -> dict: + view = k3_mla_decode_view(self.attn[layer], md, len(self.request_ids) * self.tokens) + assert isinstance(view, dict), view + assert view["row_stride"] == DQK and view["page_offset"] == self.mgr.layer_offsets[layer] + return view + + def record(self, layer: int, rows: torch.Tensor) -> None: + """Expect `rows` (the step's cache rows, request-major) at each token's position of `layer`.""" + for i in range(len(self.request_ids)): + for u in range(self.tokens): + pos = self.cached[i] + u + block = self.blocks[layer][i][pos // PAGE] + self.image[layer][block, 0, pos % PAGE, 0] = _bits(rows[i * self.tokens + u]) + + def advance(self) -> None: + self.cached = [c + self.tokens for c in self.cached] + + def pool_matches(self, layer: int) -> bool: + return torch.equal(self.mgr.get_buffers(layer).view(torch.int16), self.image[layer]) + + def shutdown(self) -> None: + self.mgr.shutdown() + + +def _call(view: dict, ag, weights) -> torch.Tensor: + w_qa, w_qb, w_kb, w_kv = weights + return k3_mla_qkv( + ag, + w_qa, + EPS, + w_qb, + w_kb, + w_kv, + KV_EPS, + view["pool"], + view["row_stride"], + view["page_table"], + view["page_offset"], + view["seq_len"], + ) + + +def _dense(ag, weights): + """fused_q and the step's cache rows from the dense form: the bits k3_mla_qkv stores.""" + w_qa, w_qb, w_kb, w_kv = weights + rows = torch.empty(ag.shape[0], DQK, dtype=torch.bfloat16, device="cuda") + y = k3_mla_qkv_out(ag, w_qa, EPS, w_qb, w_kb, w_kv, KV_EPS, rows) + return y, rows + + +def _check_cell(y, rows, ag, weights) -> None: + """fused_q and the stored rows against the references; the q-only form's fused_q bit-identical.""" + w_qa, w_qb, w_kb, w_kv = weights + m = ag.shape[0] + y_q = k3_mla_q(ag, w_qa, EPS, w_qb, w_kb) + ref_q = _reference_q(ag, w_qa, w_qb, w_kb).view(m, -1, DQK) + ref_kv = _reference_kv(ag, w_kv) + torch.cuda.synchronize() + for part in (slice(0, LAT), slice(LAT, DQK)): + assert _max_rel(y.view(m, -1, DQK)[..., part], ref_q[..., part]) <= 2e-2 + assert _max_rel(rows[:, :LAT], ref_kv[:, :LAT]) <= 1e-2 + assert torch.equal(_bits(rows[:, LAT:]), _bits(ag[:, QL + LAT : QL + DQK])) + assert torch.equal(_bits(y), _bits(y_q)) + + +# (R, T, heads): T 1 (decode without speculation) and 8 (a DSpark verify step), up to 8 requests, the TP4 head count. +CASES = [(3, 1, H), (8, 1, H), (2, 8, H), (8, 8, H), (2, 8, 4 * H)] +CASE_IDS = [f"{r}x{t}-h{h}" for r, t, h in CASES] +EAGER_STEPS, GRAPH_STEPS = 2, 2 + + +@pytest.mark.parametrize("num_requests,tokens,heads", CASES, ids=CASE_IDS) +def test_layers_by_steps_eager_then_graph(num_requests, tokens, heads): + """LAYERS layers x EAGER_STEPS decode steps called eagerly, then one step captured as a CUDA graph and replayed + for GRAPH_STEPS steps with rewritten inputs and the metadata prepared for each step: after every step each + layer's whole pool equals its expected image (the step's rows at the manager's blocks, nothing else written), + and every call's outputs pass the cell checks.""" + seed = 31 * num_requests + tokens + heads + weights = [_weights(seed + layer, heads) for layer in range(LAYERS)] + cache = _MlaCache(num_requests, tokens, EAGER_STEPS + GRAPH_STEPS, heads, seed) + m = num_requests * tokens + try: + md = cache.metadata() + for step in range(EAGER_STEPS): + cache.prepare(md) + for layer in range(LAYERS): + ag = _ag(1000 * step + 10 * layer + seed, m, heads) + y = _call(cache.view(md, layer), ag, weights[layer]) + y_dense, rows = _dense(ag, weights[layer]) + _check_cell(y, rows, ag, weights[layer]) + assert torch.equal(_bits(y), _bits(y_dense)) + cache.record(layer, rows) + for layer in range(LAYERS): + assert cache.pool_matches(layer), (step, layer) + cache.advance() + + graph_md = cache.metadata().create_cuda_graph_metadata( + num_requests, max_draft_tokens=tokens - 1 + ) + cache.prepare(graph_md) + views = [cache.view(graph_md, layer) for layer in range(LAYERS)] + static_ag = [torch.zeros(m, QL + DQK + heads * V, dtype=torch.bfloat16, device="cuda") + for _ in range(LAYERS)] # fmt: skip + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_y = [ + _call(views[layer], static_ag[layer], weights[layer]) for layer in range(LAYERS) + ] + for step in range(EAGER_STEPS, EAGER_STEPS + GRAPH_STEPS): + if step > EAGER_STEPS: + cache.prepare(graph_md) + inputs = [_ag(1000 * step + 10 * layer + seed, m, heads) for layer in range(LAYERS)] + for layer in range(LAYERS): + static_ag[layer].copy_(inputs[layer]) + graph.replay() + for layer in range(LAYERS): + y_dense, rows = _dense(inputs[layer], weights[layer]) + torch.cuda.synchronize() + assert torch.equal(_bits(static_y[layer]), _bits(y_dense)), (step, layer) + cache.record(layer, rows) + for layer in range(LAYERS): + assert cache.pool_matches(layer), (step, layer) + cache.advance() + finally: + cache.shutdown() + + +def test_two_managers_interleaved(): + """Two managers (two models' caches) with their calls alternating step by step and layer by layer: each pool + holds exactly its own rows.""" + caches = [_MlaCache(3, 8, 3, H, seed) for seed in (5, 6)] + weights = _weights(77, H) + try: + mds = [cache.metadata() for cache in caches] + for step in range(3): + for cache, md in zip(caches, mds): + cache.prepare(md) + for layer in range(LAYERS): + for k, (cache, md) in enumerate(zip(caches, mds)): + ag = _ag(100 * step + 10 * layer + k, 24, H) + _call(cache.view(md, layer), ag, weights) + cache.record(layer, _dense(ag, weights)[1]) + for cache in caches: + for layer in range(LAYERS): + assert cache.pool_matches(layer), (step, layer) + cache.advance() + finally: + for cache in caches: + cache.shutdown() + + +def test_wrong_layer_slot_overwrites_its_neighbour(): + """Negative control. Layer 1's call given layer 0's slot (page offset): nothing raises; layer 0's rows at the + step's positions now hold layer 1's values and layer 1's slot is unwritten. The op trusts the page offset, and the + pool comparison detects the misplacement.""" + cache = _MlaCache(3, 8, 1, H, 9) + weights = _weights(78, H) + try: + md = cache.metadata() + cache.prepare(md) + ag0, ag1 = _ag(1, 24, H), _ag(2, 24, H) + _call(cache.view(md, 0), ag0, weights) + cache.record(0, _dense(ag0, weights)[1]) + assert cache.pool_matches(0) and cache.pool_matches(1) + wrong = dict(cache.view(md, 1), page_offset=cache.view(md, 0)["page_offset"]) + _call(wrong, ag1, weights) + torch.cuda.synchronize() + assert not cache.pool_matches(0) + assert cache.pool_matches(1) # nothing written into layer 1's slot + cache.record(0, _dense(ag1, weights)[1]) + assert cache.pool_matches(0) # layer 0's step rows are now layer 1's + finally: + cache.shutdown() + + +def test_rejects_out_of_contract(): + """More than 64 tokens, or page-table rows / lengths that do not describe R requests of the call's tokens, raise + ValueError before any launch; the pool is left as it was.""" + cache = _MlaCache(2, 4, 1, H, 3) + weights = _weights(79, H) + try: + md = cache.metadata() + cache.prepare(md) + view = cache.view(md, 0) + with pytest.raises(ValueError): # 7 tokens over 2 requests + _call(view, _ag(4, 7, H), weights) + with pytest.raises(ValueError): # 1 page-table row, 2 lengths + _call(dict(view, page_table=view["page_table"][:1]), _ag(4, 8, H), weights) + with pytest.raises(ValueError): # an int64 page table + _call(dict(view, page_table=view["page_table"].long()), _ag(4, 8, H), weights) + with pytest.raises(ValueError): # 65 tokens + k3_mla_q(_ag(5, 65, H), weights[0], EPS, weights[1], weights[2]) + torch.cuda.synchronize() + assert all(cache.pool_matches(layer) for layer in range(LAYERS)) + finally: + cache.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv.py new file mode 100644 index 000000000000..4befd44bf4a4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the k3_ctm_gemv catalog entry. + +The Kimi K3 TP16 per-rank shapes with their call sites' split and push, at every M in 1..8: error against an fp64 +product of the same bf16 inputs, identical bits on a repeated call, and each M's rows bit-identical to the same rows +of the 8-row call. The schedule-only flags (push, trigger_early) move no bits, a CUDA-graph replay returns the +eager bits, and calls outside the op's preconditions raise ValueError. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv import k3_ctm_gemv + +assert torch.cuda.is_available(), "k3_ctm_gemv requires a CUDA device" +# The receipts are sm_100 ones; CI's other architectures skip this file. +pytestmark = pytest.mark.skipif( + torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" +) + +M_ALL = range(1, 9) +TOL = 8e-3 # max |y - ref| / max |ref| against an fp64 product of the same bf16 inputs + +# (N, K, split, push): the MLA o_proj at splits 1 and 2, the drafter o_proj (tuned and synthetic drafter). +CELLS = { + "mla_o_proj_s1": (7168, 768, 1, False), + "mla_o_proj_s2": (7168, 768, 2, False), + "drafter_o_proj": (7168, 384, 1, True), + "drafter_o_proj_synthetic": (7168, 256, 1, True), +} + + +def _randn(rows: int, cols: int, seed: int, scale: float) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(rows, cols, generator=gen, device="cuda") * scale).bfloat16() + + +def _weight(n: int, k: int) -> torch.Tensor: + return _randn(n, k, 1, 0.03) + + +def _rows(k: int) -> torch.Tensor: + return _randn(8, k, 2, 1.0) + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.int16) + + +def _rel(y: torch.Tensor, ref: torch.Tensor) -> float: + """max |y - ref| / max |ref|, in fp64.""" + return (y.double() - ref).abs().max().item() / ref.abs().max().item() + + +def test_call_site_cells() -> None: + for name, (n, k, split, push) in CELLS.items(): + w = _weight(n, k) + x8 = _rows(k) + ref8 = x8.double() @ w.double().t() + y8 = k3_ctm_gemv(x8, w, True, split, push) + for m in M_ALL: + cell = f"{name} M={m}" + x = x8[:m].contiguous() + y = k3_ctm_gemv(x, w, True, split, push) + assert y.shape == (m, n) and y.dtype == torch.bfloat16, cell + rel = _rel(y, ref8[:m]) + assert rel <= TOL, f"{cell}: max |y - ref| / max |ref| = {rel:.2e} > {TOL}" + again = k3_ctm_gemv(x, w, True, split, push) + assert torch.equal(_bits(y), _bits(again)), f"{cell}: a repeated call changed bits" + assert torch.equal(_bits(y), _bits(y8[:m])), f"{cell}: rows differ from the 8-row call" + + +def test_schedule_flags_move_no_bits() -> None: + """At the MLA o_proj with split 2, push=True and trigger_early=False return the call site's bits.""" + n, k, split = 7168, 768, 2 + w = _weight(n, k) + x8 = _rows(k) + for m in M_ALL: + x = x8[:m].contiguous() + base = _bits(k3_ctm_gemv(x, w, True, split, False)) + for trigger_early, push in ((True, True), (False, False)): + y = k3_ctm_gemv(x, w, trigger_early, split, push) + assert torch.equal(_bits(y), base), ( + f"M={m} trigger_early={trigger_early} push={push}: bits differ from the call site's" + ) + + +def test_cuda_graph_replay_matches_eager() -> None: + """Captured after an eager call per key and replayed with x rewritten in place: the eager bits.""" + calls = [] + for name in ("mla_o_proj_s1", "mla_o_proj_s2"): + n, k, split, push = CELLS[name] + w = _weight(n, k) + for m in (1, 8): + x = _rows(k)[:m].clone() + k3_ctm_gemv(x, w, True, split, push) # compiles the key outside capture + calls.append((f"{name} M={m}", x, w, split, push)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outs = [k3_ctm_gemv(x, w, True, split, push) for _, x, w, split, push in calls] + for seed, (_, x, *_) in enumerate(calls, start=100): + x.copy_(_randn(x.shape[0], x.shape[1], seed, 1.0)) + graph.replay() + for (cell, x, w, split, push), y in zip(calls, outs): + eager = k3_ctm_gemv(x, w, True, split, push) + assert torch.equal(_bits(y), _bits(eager)), f"{cell}: the replay differs from an eager call" + + +def test_unsupported_calls_raise_value_error() -> None: + """The op's own check refuses these before compiling or launching anything.""" + w = _weight(7168, 768) + x8 = _rows(768) + cases = { + "0 rows": (x8[:0], w, 1), + "9 rows": (torch.cat([x8, x8[:1]]), w, 1), + "fp16 x": (x8.half(), w, 1), + "row-strided x": (torch.cat([x8, x8], dim=1)[:, :768], w, 1), + "N % 128 != 0": (x8, w[:7104], 1), + "K % 128 != 0": (x8[:, :704].contiguous(), w[:, :704].contiguous(), 1), + "split 3": (x8, w, 3), + "7 k-tiles on one CTA": (_rows(896), _weight(7168, 896), 1), + } + for name, (x, weight, split) in cases.items(): + try: + k3_ctm_gemv(x, weight, True, split, False) + except ValueError: + continue + raise AssertionError(f"{name}: expected ValueError, the op accepted the call") diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_long.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_long.py new file mode 100644 index 000000000000..b842a80f0c5a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_long.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the k3_ctm_gemv_long catalog entry. + +The Kimi K3 TP16 per-rank shapes with their call sites' sig_col0, split, ring and push, at every M in 1..8: error of +the plain (sig_col0=-1) output against an fp64 product of the same bf16 inputs, the sigmoid columns torch.sigmoid of +that output bit for bit, identical bits on a repeated call, and each M's rows bit-identical to the same rows of the +8-row call. The schedule-only flags (ring, push, trigger_early) move no bits, a CUDA-graph replay returns the eager +bits, and calls outside the op's preconditions raise ValueError. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_long import ( + k3_ctm_gemv_long, +) + +assert torch.cuda.is_available(), "k3_ctm_gemv_long requires a CUDA device" +# The receipts are sm_100 ones; CI's other architectures skip this file. +pytestmark = pytest.mark.skipif( + torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" +) + +M_ALL = range(1, 9) +TOL = 8e-3 # max |y - ref| / max |ref| against an fp64 product of the same bf16 inputs + +# (N, K, sig_col0, split, ring, push), as each call site passes them. mla_qkv_a_gate is the fused +# [W_a; W_g] projection, its gate columns returned as bf16(sigmoid). +CELLS = { + "mla_qkv_a_gate": (2880, 7168, 2112, 6, 6, True), + "dense_gate_up": (4224, 7168, -1, 4, 5, False), + "dense_down": (7168, 2112, -1, 2, 6, False), # K 2112 ends in a half k-tile + "drafter_qkv": (512, 7168, -1, 8, 6, True), + "drafter_gate_up": (1792, 7168, -1, 8, 6, True), + "drafter_gate_up_synthetic": (1536, 7168, -1, 8, 6, True), + "kda_qkvg": (3208, 7168, -1, 5, 6, True), # 25 whole 128-row blocks and an 8-row last one +} + +# One flag changed from a call site's values; each must return the call site's bits. +FLAG_CHANGES = { + "dense_down": ({"push": True}, {"ring": 3}), + "dense_gate_up": ({"push": True}, {"trigger_early": False}), +} + + +def _randn(rows: int, cols: int, seed: int, scale: float) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(rows, cols, generator=gen, device="cuda") * scale).bfloat16() + + +def _weight(n: int, k: int) -> torch.Tensor: + return _randn(n, k, 9, 0.02) + + +def _rows(k: int) -> torch.Tensor: + return _randn(8, k, 10, 1.0) + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.int16) + + +def _rel(y: torch.Tensor, ref: torch.Tensor) -> float: + """max |y - ref| / max |ref|, in fp64.""" + return (y.double() - ref).abs().max().item() / ref.abs().max().item() + + +def test_call_site_cells() -> None: + for name, (n, k, sig, split, ring, push) in CELLS.items(): + w = _weight(n, k) + x8 = _rows(k) + ref8 = x8.double() @ w.double().t() + y8 = k3_ctm_gemv_long(x8, w, sig, split, ring, True, push) + for m in M_ALL: + cell = f"{name} M={m}" + x = x8[:m].contiguous() + y = k3_ctm_gemv_long(x, w, sig, split, ring, True, push) + assert y.shape == (m, n) and y.dtype == torch.bfloat16, cell + again = k3_ctm_gemv_long(x, w, sig, split, ring, True, push) + assert torch.equal(_bits(y), _bits(again)), f"{cell}: a repeated call changed bits" + assert torch.equal(_bits(y), _bits(y8[:m])), f"{cell}: rows differ from the 8-row call" + plain = k3_ctm_gemv_long(x, w, -1, split, ring, True, push) if sig >= 0 else y + rel = _rel(plain, ref8[:m]) + assert rel <= TOL, f"{cell}: max |y - ref| / max |ref| = {rel:.2e} > {TOL}" + if sig >= 0: + assert torch.equal(_bits(y[:, :sig]), _bits(plain[:, :sig])), ( + f"{cell}: columns before sig_col0 differ from the sig_col0=-1 call" + ) + assert torch.equal(_bits(y[:, sig:]), _bits(plain[:, sig:].sigmoid())), ( + f"{cell}: columns from sig_col0 on are not torch.sigmoid of the sig_col0=-1 call" + ) + + +def test_schedule_flags_move_no_bits() -> None: + """ring, push and trigger_early only schedule the call: one changed, the call site's bits come back.""" + for name, changes in FLAG_CHANGES.items(): + n, k, sig, split, ring, push = CELLS[name] + w = _weight(n, k) + x8 = _rows(k) + site = {"split": split, "ring": ring, "trigger_early": True, "push": push} + for m in M_ALL: + x = x8[:m].contiguous() + base = _bits(k3_ctm_gemv_long(x, w, sig, **site)) + for change in changes: + y = k3_ctm_gemv_long(x, w, sig, **{**site, **change}) + assert torch.equal(_bits(y), base), ( + f"{name} M={m} {change}: bits differ from the call site's" + ) + + +def test_cuda_graph_replay_matches_eager() -> None: + """Captured after an eager call per key and replayed with x rewritten in place: the eager bits.""" + calls = [] + for name in ("mla_qkv_a_gate", "dense_down"): + n, k, sig, split, ring, push = CELLS[name] + w = _weight(n, k) + for m in (1, 8): + x = _rows(k)[:m].clone() + k3_ctm_gemv_long(x, w, sig, split, ring, True, push) # compiles the key outside capture + calls.append((f"{name} M={m}", x, w, (sig, split, ring, True, push))) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outs = [k3_ctm_gemv_long(x, w, *flags) for _, x, w, flags in calls] + for seed, (_, x, _, _) in enumerate(calls, start=100): + x.copy_(_randn(x.shape[0], x.shape[1], seed, 1.0)) + graph.replay() + for (cell, x, w, flags), y in zip(calls, outs): + eager = k3_ctm_gemv_long(x, w, *flags) + assert torch.equal(_bits(y), _bits(eager)), f"{cell}: the replay differs from an eager call" + + +def test_unsupported_calls_raise_value_error() -> None: + """The op's own check refuses these before compiling or launching anything.""" + w = _weight(7168, 2112) + x8 = _rows(2112) + cases = { + "0 rows": (x8[:0], w, 2, 6), + "9 rows": (torch.cat([x8, x8[:1]]), w, 2, 6), + "fp16 x": (x8.half(), w, 2, 6), + "row-strided x": (torch.cat([x8, x8], dim=1)[:, :2112], w, 2, 6), + "K % 64 != 0": (x8[:, :2080].contiguous(), w[:, :2080].contiguous(), 2, 6), + "split 3": (x8, w, 3, 5), + "split 9": (x8, w, 9, 1), + "ring 0": (x8, w, 2, 0), + "ring past the rank's k-tiles": (x8, w, 8, 3), + "ring past shared memory": (x8, w, 2, 7), + } + for name, (x, weight, split, ring) in cases.items(): + try: + k3_ctm_gemv_long(x, weight, -1, split, ring, True, False) + except ValueError: + continue + raise AssertionError(f"{name}: expected ValueError, the op accepted the call") diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_swiglu.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_swiglu.py new file mode 100644 index 000000000000..e61be03e4683 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_swiglu.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the k3_ctm_gemv_swiglu catalog entry. + +The Kimi K3 TP16 per-rank drafter down projection with its call site's split and push, at every M in 1..8: error +against an fp64 product of torch's bf16 silu_and_mul activation, identical bits on a repeated call, and each M's rows +bit-identical to the same rows of the 8-row call. The schedule-only flags (push, trigger_early) move no bits, a +CUDA-graph replay returns the eager bits, and calls outside the op's preconditions raise ValueError. +""" + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_swiglu import ( + k3_ctm_gemv_swiglu, +) + +assert torch.cuda.is_available(), "k3_ctm_gemv_swiglu requires a CUDA device" +# The receipts are sm_100 ones; CI's other architectures skip this file. +pytestmark = pytest.mark.skipif( + torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" +) + +M_ALL = range(1, 9) +TOL = 8e-3 # max |y - ref| / max |ref| against an fp64 product of the bf16 activation + +# (N, K, split, push): the drafter down projection (tuned and synthetic drafter); K 896 is 7 k-tiles. +CELLS = { + "drafter_down": (7168, 896, 2, True), + "drafter_down_synthetic": (7168, 768, 2, True), +} + + +def _randn(rows: int, cols: int, seed: int, scale: float) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(rows, cols, generator=gen, device="cuda") * scale).bfloat16() + + +def _weight(n: int, k: int) -> torch.Tensor: + return _randn(n, k, 7, 0.02) + + +def _gate_up(k: int) -> torch.Tensor: + return _randn(8, 2 * k, 8, 2.0) + + +def _silu_and_mul(gu: torch.Tensor) -> torch.Tensor: + """silu_and_mul's fp32 arithmetic and one bf16 rounding (gate columns first).""" + k = gu.shape[1] // 2 + return (F.silu(gu[:, :k].float()) * gu[:, k:].float()).bfloat16() + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.int16) + + +def _rel(y: torch.Tensor, ref: torch.Tensor) -> float: + """max |y - ref| / max |ref|, in fp64.""" + return (y.double() - ref).abs().max().item() / ref.abs().max().item() + + +def test_call_site_cells() -> None: + for name, (n, k, split, push) in CELLS.items(): + w = _weight(n, k) + gu8 = _gate_up(k) + ref8 = _silu_and_mul(gu8).double() @ w.double().t() + y8 = k3_ctm_gemv_swiglu(gu8, w, True, split, push) + for m in M_ALL: + cell = f"{name} M={m}" + gu = gu8[:m].contiguous() + y = k3_ctm_gemv_swiglu(gu, w, True, split, push) + assert y.shape == (m, n) and y.dtype == torch.bfloat16, cell + rel = _rel(y, ref8[:m]) + assert rel <= TOL, f"{cell}: max |y - ref| / max |ref| = {rel:.2e} > {TOL}" + again = k3_ctm_gemv_swiglu(gu, w, True, split, push) + assert torch.equal(_bits(y), _bits(again)), f"{cell}: a repeated call changed bits" + assert torch.equal(_bits(y), _bits(y8[:m])), f"{cell}: rows differ from the 8-row call" + + +def test_schedule_flags_move_no_bits() -> None: + """At the drafter down, push=False and trigger_early=False return the call site's bits.""" + n, k, split, push = CELLS["drafter_down"] + w = _weight(n, k) + gu8 = _gate_up(k) + for m in M_ALL: + gu = gu8[:m].contiguous() + base = _bits(k3_ctm_gemv_swiglu(gu, w, True, split, push)) + for trigger_early, push_flag in ((True, not push), (False, push)): + y = k3_ctm_gemv_swiglu(gu, w, trigger_early, split, push_flag) + assert torch.equal(_bits(y), base), ( + f"M={m} trigger_early={trigger_early} push={push_flag}: bits differ from the call site's" + ) + + +def test_cuda_graph_replay_matches_eager() -> None: + """Captured after an eager call per key and replayed with gu rewritten in place: the eager bits.""" + n, k, split, push = CELLS["drafter_down"] + w = _weight(n, k) + calls = [] + for m in (1, 8): + gu = _gate_up(k)[:m].clone() + k3_ctm_gemv_swiglu(gu, w, True, split, push) # compiles the key outside capture + calls.append((m, gu)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outs = [k3_ctm_gemv_swiglu(gu, w, True, split, push) for _, gu in calls] + for seed, (m, gu) in enumerate(calls, start=100): + gu.copy_(_randn(m, 2 * k, seed, 2.0)) + graph.replay() + for (m, gu), y in zip(calls, outs): + eager = k3_ctm_gemv_swiglu(gu, w, True, split, push) + assert torch.equal(_bits(y), _bits(eager)), f"M={m}: the replay differs from an eager call" + + +def test_unsupported_calls_raise_value_error() -> None: + """The op's own check refuses these before compiling or launching anything.""" + w = _weight(7168, 896) + gu8 = _gate_up(896) + cases = { + "0 rows": (gu8[:0], w, 2), + "9 rows": (torch.cat([gu8, gu8[:1]]), w, 2), + "fp16 gu": (gu8.half(), w, 2), + "row-strided gu": (torch.cat([gu8, gu8], dim=1)[:, : 2 * 896], w, 2), + "gu width != 2 K": (_randn(8, 1664, 8, 2.0), w, 2), + "N % 128 != 0": (gu8, w[:7104], 2), + "split 3": (gu8, w, 3), + "7 k-tiles on one CTA": (gu8, w, 1), + } + for name, (gu, weight, split) in cases.items(): + try: + k3_ctm_gemv_swiglu(gu, weight, True, split, True) + except ValueError: + continue + raise AssertionError(f"{name}: expected ValueError, the op accepted the call") diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_wide.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_wide.py new file mode 100644 index 000000000000..af81feb06f7b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_wide.py @@ -0,0 +1,184 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the k3_ctm_gemv_wide catalog entry. + +The Kimi K3 TP16 per-rank shapes of the projections of a wide decode step, at every M in 1..64: error of the plain +(sig_col0=-1) output against an fp64 product of the same bf16 inputs, the sigmoid columns torch.sigmoid of that +output bit for bit, identical bits on a repeated call, and each M's rows bit-identical to the same rows of the +64-row call. Where the device's split equals a k3_ctm_gemv_long call site's, every token's row carries that op's +bits; a CUDA-graph replay returns the eager bits; calls outside the op's preconditions raise ValueError. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_long import ( + k3_ctm_gemv_long, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_wide import ( + k3_ctm_gemv_wide, +) + +assert torch.cuda.is_available(), "k3_ctm_gemv_wide requires a CUDA device" +# The receipts are sm_100 ones; CI's other architectures skip this file. +pytestmark = pytest.mark.skipif( + torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" +) + +M_ALL = range(1, 65) +TOL = 8e-3 # bf16 output: max |y - ref| / max |ref| against an fp64 product of the same bf16 inputs +TOL_FP32 = 1e-4 # fp32 output + +# (N, K, sig_col0, out_fp32): the per-rank TP16 shapes of the projections of a wide decode step. kda_qkvg is 25 +# whole 128-row blocks and an 8-row last one; mla_qkv_a_gate is [W_a; W_g] with its gate columns as bf16(sigmoid); +# moe_head is [latent down slice; router rows] in fp32 (router logits); moe_tail is 5 k-tiles; dense_down's K 2112 +# ends in a half k-tile. +CELLS = { + "kda_qkvg": (3208, 7168, -1, False), + "mla_qkv_a_gate": (2880, 7168, 2112, False), + "o_proj": (7168, 768, -1, False), + "moe_head": (280, 7168, -1, True), + "moe_tail": (7168, 640, -1, False), + "shared_gate_up": (768, 7168, -1, False), + "dense_gate_up": (4224, 7168, -1, False), + "dense_down": (7168, 2112, -1, False), + "drafter_qkv": (512, 7168, -1, False), + "drafter_gate_up": (1792, 7168, -1, False), + "drafter_o_proj": (7168, 384, -1, False), +} +# The k3_ctm_gemv_long call (split, ring, push) of the shapes that run it at up to 8 tokens. +LONG_SITES = { + "mla_qkv_a_gate": (6, 6, True), + "dense_gate_up": (4, 5, False), + "dense_down": (2, 6, False), + "drafter_qkv": (8, 6, True), + "drafter_gate_up": (8, 6, True), +} + + +def _randn(rows: int, cols: int, seed: int, scale: float) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(rows, cols, generator=gen, device="cuda") * scale).bfloat16() + + +def _weight(n: int, k: int) -> torch.Tensor: + return _randn(n, k, 1, 0.02) + + +def _rows(k: int) -> torch.Tensor: + return _randn(64, k, 2, 1.0) + + +def _bits(t: torch.Tensor) -> torch.Tensor: + t = t.contiguous() + return t.view(torch.int32) if t.dtype == torch.float32 else t.view(torch.int16) + + +def _rel(y: torch.Tensor, ref: torch.Tensor) -> float: + """max |y - ref| / max |ref|, in fp64.""" + return (y.double() - ref).abs().max().item() / ref.abs().max().item() + + +def test_projection_cells() -> None: + for name, (n, k, sig, fp32) in CELLS.items(): + w = _weight(n, k) + x64 = _rows(k) + ref64 = x64.double() @ w.double().t() + y64 = k3_ctm_gemv_wide(x64, w, sig, fp32) + dtype = torch.float32 if fp32 else torch.bfloat16 + tol = TOL_FP32 if fp32 else TOL + for m in M_ALL: + cell = f"{name} M={m}" + x = x64[:m].contiguous() + y = k3_ctm_gemv_wide(x, w, sig, fp32) + assert y.shape == (m, n) and y.dtype == dtype, cell + again = k3_ctm_gemv_wide(x, w, sig, fp32) + assert torch.equal(_bits(y), _bits(again)), f"{cell}: a repeated call changed bits" + assert torch.equal(_bits(y), _bits(y64[:m])), ( + f"{cell}: rows differ from the 64-row call" + ) + plain = k3_ctm_gemv_wide(x, w, -1, fp32) if sig >= 0 else y + rel = _rel(plain, ref64[:m]) + assert rel <= tol, f"{cell}: max |y - ref| / max |ref| = {rel:.2e} > {tol}" + if sig >= 0: + assert torch.equal(_bits(y[:, :sig]), _bits(plain[:, :sig])), ( + f"{cell}: columns before sig_col0 differ from the sig_col0=-1 call" + ) + assert torch.equal(_bits(y[:, sig:]), _bits(plain[:, sig:].sigmoid())), ( + f"{cell}: columns from sig_col0 on are not torch.sigmoid of the sig_col0=-1 call" + ) + + +def test_rows_match_long_at_the_same_split() -> None: + """A token's row is k3_ctm_gemv_long's for it (calls of up to 8 tokens) where the splits agree.""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import op as ctm_op + from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv.k3_ctm_gemv_kernel import wide_tile + + sms = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + compared = 0 + for name, (split, ring, push) in LONG_SITES.items(): + n, k, sig, _ = CELLS[name] + w = _weight(n, k) + x64 = _rows(k) + groups = torch.cat( + [ + k3_ctm_gemv_long(x64[r : r + 8], w, sig, split, ring, True, push) + for r in range(0, 64, 8) + ] + ) + for m in [*range(1, 9), *range(16, 65, 8)]: + if ctm_op.wide_config(n, k, wide_tile(m), sms)[0] != split: + continue + x = x64[:m].contiguous() + y = k3_ctm_gemv_wide(x, w, sig) + want = k3_ctm_gemv_long(x, w, sig, split, ring, True, push) if m <= 8 else groups[:m] + assert torch.equal(_bits(y), _bits(want)), ( + f"{name} M={m}: rows differ from k3_ctm_gemv_long's at split {split}" + ) + compared += 1 + assert compared > 0, f"no call site's split is this device's ({sms} SMs); nothing was compared" + + +def test_cuda_graph_replay_matches_eager() -> None: + """Captured after an eager call per key and replayed with x rewritten in place: the eager bits.""" + calls = [] + for name in ("mla_qkv_a_gate", "moe_head"): + n, k, sig, fp32 = CELLS[name] + w = _weight(n, k) + for m in (1, 64): + x = _rows(k)[:m].clone() + k3_ctm_gemv_wide(x, w, sig, fp32) # compiles the key outside capture + calls.append((f"{name} M={m}", x, w, sig, fp32)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outs = [k3_ctm_gemv_wide(x, w, sig, fp32) for _, x, w, sig, fp32 in calls] + for seed, (_, x, *_) in enumerate(calls, start=100): + x.copy_(_randn(x.shape[0], x.shape[1], seed, 1.0)) + graph.replay() + for (cell, x, w, sig, fp32), y in zip(calls, outs): + eager = k3_ctm_gemv_wide(x, w, sig, fp32) + assert torch.equal(_bits(y), _bits(eager)), f"{cell}: the replay differs from an eager call" + + +def test_unsupported_calls_raise_value_error() -> None: + """The op's own check refuses these before compiling or launching anything.""" + n, k = 3208, 7168 + w = _weight(n, k) + x8 = _rows(k)[:8] + flat = torch.zeros(8 * k + 8, dtype=torch.bfloat16, device="cuda") + cases = { + "0 rows": (x8[:0], w, -1, False), + "65 rows": (torch.cat([_rows(k), x8[:1]]), w, -1, False), + "x 2 bytes past a 16-byte boundary": (flat[1 : 1 + 8 * k].view(8, k), w, -1, False), + "sig_col0 = N": (x8, w, n, False), + "sigmoid columns with fp32 output": (x8, w, 100, True), + "N % 8 != 0": (x8, w[:3204], -1, False), + "K % 64 != 0": (x8[:, :7136].contiguous(), w[:, :7136].contiguous(), -1, False), + "fp16 x": (x8.half(), w, -1, False), + } + for name, (x, weight, sig, fp32) in cases.items(): + try: + k3_ctm_gemv_wide(x, weight, sig, fp32) + except ValueError: + continue + raise AssertionError(f"{name}: expected ValueError, the op accepted the call") diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_decode_gemv.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_decode_gemv.py new file mode 100644 index 000000000000..b21f4c21c4ce --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_decode_gemv.py @@ -0,0 +1,123 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the k3_decode_gemv catalog entry. + +Both kernels at Kimi K3's TP16 per-rank shapes, at every M in 1..8, against an fp64 torch product (error at most 8e-3 +of max |ref|), with run-to-run bits, each M's rows equal to the same rows of the 8-row call, trigger_early False equal +to True, a CUDA graph replayed with its input rewritten, and the refused calls. +""" + +import functools + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_decode_gemv import k3_decode_gemv + +assert torch.cuda.is_available(), "k3_decode_gemv requires a CUDA device" +# The receipts are sm_100 ones; CI's other architectures skip this file. +pytestmark = pytest.mark.skipif( + torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" +) + +TOL = 8e-3 # max |y - ref| / max |ref|, ref an fp64 product +M_ALL = range(1, 9) +# (N, K) per rank at TP16. The KDA o_proj: short K, one CTA per 128-row weight tile with the whole K resident. +SHORT_K = (7168, 768) +# The KDA input projection: split K over clusters of 4 CTAs; N is not a multiple of 128 (the last tile has 8 rows). +SPLIT_K = (3208, 7168) + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.int16) + + +def _same(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.shape == b.shape and torch.equal(_bits(a), _bits(b)) + + +def _rel(y: torch.Tensor, ref: torch.Tensor) -> float: + return (y.double() - ref).abs().max().item() / ref.abs().max().item() + + +@functools.lru_cache(maxsize=None) +def _inputs(n: int, k: int) -> tuple[torch.Tensor, torch.Tensor]: + """The weight [N, K] and 8 activation rows [8, K].""" + gen = torch.Generator(device="cuda").manual_seed(n * 100003 + k) + weight = (torch.randn(n, k, generator=gen, device="cuda") * 0.03).bfloat16() + x8 = torch.randn(8, k, generator=gen, device="cuda").bfloat16() + # The kernels read the weight before their PDL grid-dependency wait: it is written well before any call. + torch.cuda.synchronize() + return weight, x8 + + +def _check_cells(shape: tuple[int, int]) -> None: + n, k = shape + weight, x8 = _inputs(n, k) + first = {} + for trigger_early in (True, False): + y8 = k3_decode_gemv(x8, weight, trigger_early) + for m in M_ALL: + x = x8[:m] + y = k3_decode_gemv(x, weight, trigger_early) + again = k3_decode_gemv(x, weight, trigger_early) + ref = x.double() @ weight.double().t() + assert y.shape == (m, n) and y.dtype == torch.bfloat16 and y.device == x.device + assert y.is_contiguous() + err = _rel(y, ref) + assert err <= TOL, f"{shape} M={m} trigger_early={trigger_early}: rel {err:.3e}" + assert _same(y, again), f"{shape} M={m}: a rerun differs" + assert _same(y, y8[:m]), f"{shape} M={m}: rows differ from the 8-row call" + if trigger_early: + first[m] = y + else: + assert _same(y, first[m]), f"{shape} M={m}: trigger_early changes the result" + + +def test_short_k_cells() -> None: + """[7168, 768] (one CTA per 128-row tile, the whole K resident) at every M, trigger_early True and False.""" + _check_cells(SHORT_K) + + +def test_split_k_cells() -> None: + """[3208, 7168] (clusters of 4 CTAs over K) at every M, trigger_early True and False.""" + _check_cells(SPLIT_K) + + +def test_graph_replay() -> None: + """M 1 and M 8 calls on each weight in one CUDA graph, captured after eager calls compiled them and replayed with + the activation rewritten in place: every replay bit-identical to eager calls on the same rows.""" + gen = torch.Generator(device="cuda").manual_seed(5) + weights, x_bufs = [], [] + for n, k in (SHORT_K, SPLIT_K): + weight, x8 = _inputs(n, k) + weights.append(weight) + x_bufs.append(x8.clone()) + for m in (1, 8): + k3_decode_gemv(x_bufs[-1][:m], weight) # compiles outside the capture + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = [{m: k3_decode_gemv(x[:m], w) for m in (1, 8)} for w, x in zip(weights, x_bufs)] + for _ in range(3): + for x in x_bufs: + x.copy_(torch.randn(x.shape, generator=gen, device="cuda").bfloat16()) + graph.replay() + for i, (w, x) in enumerate(zip(weights, x_bufs)): + for m in (1, 8): + assert _same(out[i][m], k3_decode_gemv(x[:m], w)), f"replay of weight {i}, M={m}" + + +def test_refused_calls() -> None: + """M 0 and 9, and weights neither kernel takes, raise ValueError before launching anything.""" + weight, _ = _inputs(*SHORT_K) + for m in (0, 9): + x = torch.zeros(m, SHORT_K[1], dtype=torch.bfloat16, device="cuda") + with pytest.raises(ValueError, match="unsupported call"): + k3_decode_gemv(x, weight) + # [128, 1152]: 9 k-tiles, more than short K holds and not a multiple of 4 for split K. + # [200, 768]: short K needs whole 128-row tiles, split K more than 6 k-tiles. + for n, k in ((128, 1152), (200, 768)): + x = torch.zeros(1, k, dtype=torch.bfloat16, device="cuda") + w = torch.zeros(n, k, dtype=torch.bfloat16, device="cuda") + with pytest.raises(ValueError, match="unsupported call"): + k3_decode_gemv(x, w) diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_head_gemv.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_head_gemv.py new file mode 100644 index 000000000000..e3e82157fdfc --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_head_gemv.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the k3_head_gemv catalog entry. + +Kimi K3's TP16 per-rank LM-head vocabulary shard [10240, 7168] on the stream-K schedule, at every M in 1..8, against an +fp64 torch product (error at most 8e-3 of max |ref|), with run-to-run bits and each M's rows equal to the same rows of +the 8-row call. The caller-owned workspace is driven through real call sequences: M dipping and growing back on one +workspace, two workspaces interleaved, two weights on one workspace, and a CUDA graph replayed between eager calls on +the same workspace, every call against the same call on a fresh workspace and every flag / count word back at zero +after. Negative controls: a workspace of another shape is refused, and so is creating one under capture. +""" + +import functools + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_head_gemv import ( + K3HeadGemvWorkspace, + k3_head_gemv, +) + +assert torch.cuda.is_available(), "k3_head_gemv requires a CUDA device" +# The receipts are sm_100 ones; CI's other architectures skip this file. +pytestmark = pytest.mark.skipif( + torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" +) + +TOL = 8e-3 # max |y - ref| / max |ref|, ref an fp64 product +M_ALL = range(1, 9) +VOCAB_SHARD, HIDDEN = 10240, 7168 # Kimi K3's 163840-row LM head over 16 ranks +TILES = VOCAB_SHARD // 128 + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.int16) + + +def _same(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.shape == b.shape and torch.equal(_bits(a), _bits(b)) + + +def _rel(y: torch.Tensor, ref: torch.Tensor) -> float: + return (y.double() - ref).abs().max().item() / ref.abs().max().item() + + +@functools.lru_cache(maxsize=None) +def _inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Two shard weights (e.g. the target's LM head and DSpark's draft head) and 8 activation rows.""" + gen = torch.Generator(device="cuda").manual_seed(20261001) + weight = (torch.randn(VOCAB_SHARD, HIDDEN, generator=gen, device="cuda") * 0.02).bfloat16() + weight2 = (torch.randn(VOCAB_SHARD, HIDDEN, generator=gen, device="cuda") * 0.02).bfloat16() + x8 = torch.randn(8, HIDDEN, generator=gen, device="cuda").bfloat16() + # The kernel reads the weight before its PDL grid-dependency wait: it is written well before any call. + torch.cuda.synchronize() + return weight, weight2, x8 + + +def _workspace() -> K3HeadGemvWorkspace: + return K3HeadGemvWorkspace.create(VOCAB_SHARD, HIDDEN, torch.device("cuda")) + + +def _at_rest(ws: K3HeadGemvWorkspace) -> bool: + """Whether every flag / count word of the workspace is back at zero.""" + return not bool(ws.flags.any()) and not bool(ws.claim.any()) + + +def test_certified_cells() -> None: + """Every M on one workspace: shape, dtype, error, rerun and 8-row-call bits, the words at zero after each call; + keep_tiles 40 and 80 bit-identical to 0.""" + weight, _, x8 = _inputs() + ws = _workspace() + assert ws.schedule == "streamk" and ws.chunk_tiles == 0 + assert ws.partials.dtype == torch.float32 and ws.flags.dtype == torch.int32 + assert ws.claim.dtype == torch.int32 and ws.claim.numel() == 1 + # One flag word and one [128 x 8] fp32 partial slot per (tile, piece). + assert ws.flags.numel() % TILES == 0 and ws.partials.numel() == ws.flags.numel() * 128 * 8 + assert _at_rest(ws) + y8 = k3_head_gemv(x8, weight, ws) + for m in M_ALL: + x = x8[:m] + y = k3_head_gemv(x, weight, ws) + again = k3_head_gemv(x, weight, ws) + ref = x.double() @ weight.double().t() + assert y.shape == (m, VOCAB_SHARD) and y.dtype == torch.bfloat16 and y.device == x.device + assert y.is_contiguous() + err = _rel(y, ref) + assert err <= TOL, f"M={m}: rel {err:.3e}" + assert _same(y, again), f"M={m}: a rerun differs" + assert _same(y, y8[:m]), f"M={m}: rows differ from the 8-row call" + assert _at_rest(ws), f"M={m}: words left raised" + for m in (1, 8): + for keep in (TILES // 2, TILES): + y = k3_head_gemv(x8[:m], weight, ws, keep_tiles=keep) + assert _same(y, y8[:m]), f"M={m} keep_tiles={keep}: result differs" + assert _at_rest(ws) + + +def test_m_dipping_and_growing_on_one_workspace() -> None: + """M 8, 8, 2, 7, 8, 1, 1, 8, 3, 8 back to back on one stream and one workspace: every call bit-identical to the + same call on a fresh workspace, the words at zero after.""" + weight, _, x8 = _inputs() + seq = (8, 8, 2, 7, 8, 1, 1, 8, 3, 8) + want = {m: k3_head_gemv(x8[:m], weight, _workspace()) for m in set(seq)} + ws = _workspace() + got = [k3_head_gemv(x8[:m], weight, ws) for m in seq] + for i, (m, y) in enumerate(zip(seq, got)): + assert _same(y, want[m]), f"call {i} (M={m}) differs" + assert _at_rest(ws) + + +def test_two_workspaces_interleaved_and_two_weights_on_one() -> None: + """Two workspaces of the shard's shape, each with its own weight, their calls alternating on one stream; then both + weights alternating on one workspace: every call bit-identical to the same call on a fresh workspace, every + workspace's words at zero after.""" + weight, weight2, x8 = _inputs() + weights = (weight, weight2) + want = {} + for m in (1, 4, 8): + for i in (0, 1): + want[(m, i)] = k3_head_gemv(x8[:m], weights[i], _workspace()) + ws_a, ws_b = _workspace(), _workspace() + for m in (8, 1, 4, 1, 8): + for i, ws in ((0, ws_a), (1, ws_b)): + y = k3_head_gemv(x8[:m], weights[i], ws) + assert _same(y, want[(m, i)]), f"interleaved: M={m} weight {i}" + for m, i in ((8, 1), (8, 0), (1, 1), (4, 0), (4, 1), (1, 0)): + y = k3_head_gemv(x8[:m], weights[i], ws_a) + assert _same(y, want[(m, i)]), f"one workspace: M={m} weight {i}" + assert _at_rest(ws_a) and _at_rest(ws_b) + + +def test_graph_replays_between_eager_calls() -> None: + """A CUDA graph of M 1, 8 and 3 calls on one workspace, captured once and replayed three times on the stream of the + eager calls, with the activation rewritten in place before each replay and an eager M 8 call on the same workspace + after it: every replayed and eager result bit-identical to the same call on a fresh workspace, the words at zero + after each round.""" + weight, _, x8 = _inputs() + ws = _workspace() + x_buf = x8.clone() + for m in (1, 3, 8): + k3_head_gemv(x_buf[:m], weight, ws) # compiles outside the capture + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = {m: k3_head_gemv(x_buf[:m], weight, ws) for m in (1, 8, 3)} + gen = torch.Generator(device="cuda").manual_seed(4) + for rep in range(3): + x_buf.copy_(torch.randn(8, HIDDEN, generator=gen, device="cuda").bfloat16()) + graph.replay() + between = k3_head_gemv(x_buf, weight, ws) + fresh = _workspace() + for m in (1, 3, 8): + assert _same(out[m], k3_head_gemv(x_buf[:m], weight, fresh)), f"replay {rep}, M={m}" + assert _same(between, out[8]), f"replay {rep}: the eager call differs" + assert _at_rest(ws), f"replay {rep}: words left raised" + + +def test_refusals() -> None: + """Negative controls, none of which launches anything: a workspace created for another weight shape ([5120, 7168]: + other flag-word and partial sizes) is refused and stays at rest, and a call right after on a right workspace is + unaffected; creating a workspace under CUDA-graph capture, an unknown schedule, M 0 or 9, and N not a multiple of + 128 are refused.""" + weight, _, x8 = _inputs() + other = K3HeadGemvWorkspace.create(VOCAB_SHARD // 2, HIDDEN, torch.device("cuda")) + with pytest.raises(ValueError, match="workspace"): + k3_head_gemv(x8[:1], weight, other) + assert _at_rest(other) + ws = _workspace() + assert _same(k3_head_gemv(x8[:1], weight, ws), k3_head_gemv(x8[:1], weight, _workspace())) + graph = torch.cuda.CUDAGraph() + with pytest.raises(RuntimeError, match="capture"): + with torch.cuda.graph(graph): + _workspace() + with pytest.raises(ValueError, match="unknown schedule"): + K3HeadGemvWorkspace.create(VOCAB_SHARD, HIDDEN, torch.device("cuda"), schedule="splitk") + for m in (0, 9): + x = torch.zeros(m, HIDDEN, dtype=torch.bfloat16, device="cuda") + with pytest.raises(ValueError, match="unsupported call"): + k3_head_gemv(x, weight, ws) + odd = torch.zeros(200, HIDDEN, dtype=torch.bfloat16, device="cuda") + with pytest.raises(ValueError, match="unsupported call"): + k3_head_gemv(x8[:1], odd, ws) + assert _at_rest(ws) diff --git a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_fwd.py b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_fwd.py new file mode 100644 index 000000000000..c74c9e24671a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_fwd.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attn_res_fwd catalog entry. + +A cell is (T, N): T tokens and N = K + 1 candidates, K snapshots plus the layer residual. Every cell +checks the four outputs' shapes and dtypes, bit-identical results from two identical calls, untouched +inputs, and closeness to an fp32 torch evaluation of the contract's Semantics under the metric of +main's op test (tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py). + +The op picks its kernel from (T, N): at T == 1 the single-CTA decode kernel for N in {1, 2, 4} and the +split-K cluster kernel for N in {8, 12}; the fixed-N = 12 online variant at (T, N) == (1024, 12); the +persistent online kernel everywhere else. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.attn_res_fwd import attn_res_fwd + +assert torch.cuda.is_available(), "attn_res_fwd requires a CUDA device" +# The receipts are sm_100 ones; CI's other architectures skip this file. +pytestmark = pytest.mark.skipif( + torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" +) + +HIDDEN_SIZE = 7168 +RMS_EPS = 1e-6 +# Kimi K3's 93 layers append a snapshot every 12 layers: at most 8 snapshots, so N runs 2 ... 9. +K3_NUM_CANDIDATES = tuple(range(2, 10)) +DECODE_NUM_TOKENS = tuple(range(1, 9)) +PREFILL_NUM_TOKENS = (300, 2048) +# The thresholds of main's op test. +MIN_COSINE = 0.999 +MAX_RELATIVE_L2 = 3e-2 +INPUT_NAMES = ("layer_residual", "block_residual", "res_weight", "rms_weight") +OUTPUT_NAMES = ("output", "rsigma", "probs", "logits") + + +def _make_inputs(num_tokens: int, num_candidates: int) -> tuple[torch.Tensor, ...]: + """Inputs at the scales of main's op tests, seeded by the cell.""" + torch.manual_seed(97 * num_tokens + num_candidates) + layer_residual = ( + torch.randn(num_tokens, 1, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.05 + ) + block_residual = ( + torch.randn( + num_candidates - 1, num_tokens, 1, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda" + ) + * 0.05 + ) + res_weight = torch.randn(HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.02 + rms_weight = 1 + torch.randn(HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.02 + return layer_residual, block_residual, res_weight, rms_weight + + +def _reference( + layer_residual: torch.Tensor, + block_residual: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + rms_eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """fp32 evaluation of the contract's Semantics; snapshots first, the layer residual last.""" + values = torch.cat((block_residual, layer_residual.unsqueeze(0)), dim=0).float() + rsigma = torch.rsqrt(values.square().mean(dim=-1) + rms_eps) + score_weight = rms_weight.float() * res_weight.float() + logits = (values * rsigma.unsqueeze(-1) * score_weight).sum(dim=-1) + probs = torch.softmax(logits, dim=0) + output = (probs.unsqueeze(-1) * values).sum(dim=0).to(torch.bfloat16) + return output, rsigma, probs, logits + + +def _bits(tensor: torch.Tensor) -> torch.Tensor: + """Reinterpret as integers, so that equality means identical bits.""" + return tensor.view(torch.int16 if tensor.element_size() == 2 else torch.int32) + + +def _similarity(actual: torch.Tensor, expected: torch.Tensor) -> tuple[float, float]: + """Cosine similarity and relative L2 error, the metric of main's op tests.""" + actual_float = actual.float().flatten() + expected_float = expected.float().flatten() + cosine = torch.nn.functional.cosine_similarity(actual_float, expected_float, dim=0).item() + relative_l2 = ((actual_float - expected_float).norm() / (expected_float.norm() + 1e-12)).item() + return cosine, relative_l2 + + +def _check(num_tokens: int, num_candidates: int, rms_eps: float = RMS_EPS) -> None: + inputs = _make_inputs(num_tokens, num_candidates) + before = [tensor.clone() for tensor in inputs] + outputs = attn_res_fwd(*inputs, rms_eps) + repeat = attn_res_fwd(*inputs, rms_eps) + expected = _reference(*inputs, rms_eps) + + cell = f"T={num_tokens} N={num_candidates} rms_eps={rms_eps}" + stats_shape = (num_candidates, num_tokens, 1) + shapes = ((num_tokens, 1, HIDDEN_SIZE), stats_shape, stats_shape, stats_shape) + dtypes = (torch.bfloat16, torch.float32, torch.float32, torch.float32) + assert len(outputs) == len(OUTPUT_NAMES), f"{cell}: {len(outputs)} outputs" + for name, actual, again, reference, shape, dtype in zip( + OUTPUT_NAMES, outputs, repeat, expected, shapes, dtypes + ): + assert actual.shape == shape, f"{cell}: {name} shape {tuple(actual.shape)}" + assert actual.dtype == dtype, f"{cell}: {name} dtype {actual.dtype}" + assert actual.is_contiguous(), f"{cell}: {name} is not contiguous" + assert actual.device == inputs[0].device, f"{cell}: {name} on {actual.device}" + assert torch.equal(_bits(actual), _bits(again)), ( + f"{cell}: two identical calls disagree in {name}" + ) + cosine, relative_l2 = _similarity(actual, reference) + assert cosine > MIN_COSINE and relative_l2 < MAX_RELATIVE_L2, ( + f"{cell}: {name} cosine {cosine:.6f}, relative L2 {relative_l2:.3e}" + ) + for name, old, new in zip(INPUT_NAMES, before, inputs): + assert torch.equal(_bits(old), _bits(new)), f"{cell}: {name} was mutated" + + +def test_k3_decode_cells() -> None: + """Kimi K3's N = 2 ... 9 at T = 1 ... 8: every T == 1 kernel, then the online kernel.""" + for num_tokens in DECODE_NUM_TOKENS: + for num_candidates in K3_NUM_CANDIDATES: + _check(num_tokens, num_candidates) + + +def test_k3_prefill_cells() -> None: + """Kimi K3's N = 2 ... 9 at prefill token counts: the online kernel's CTAs loop over tokens.""" + for num_tokens in PREFILL_NUM_TOKENS: + for num_candidates in K3_NUM_CANDIDATES: + _check(num_tokens, num_candidates) + + +def test_other_dispatch_branches() -> None: + """The kernels and N extremes outside Kimi K3's range, down to the fixed-N = 12 variant.""" + for num_tokens, num_candidates in ( + (1, 1), # single-CTA kernel, no snapshot + (1, 10), # online kernel at T == 1 + (1, 11), # online kernel at T == 1 + (1, 12), # split-K kernel + (8, 1), # online kernel, a single candidate + (8, 12), # online kernel, three full chunks of four candidates + (1024, 12), # online kernel, fixed-N = 12 variant + ): + _check(num_tokens, num_candidates) + + +def test_rms_eps_enters_under_the_root() -> None: + """rms_eps is added to the mean square inside the rsqrt; 1e-2 dominates the inputs' 2.5e-3.""" + for num_tokens, num_candidates in ((1, 2), (1, 8), (1, 9), (300, 9)): + for rms_eps in (1e-5, 1e-2): + _check(num_tokens, num_candidates, rms_eps) + + +def test_cuda_graph_replay_matches_eager() -> None: + """One cell per kernel: a call captured in a CUDA graph replays the eager call's bits.""" + for num_tokens, num_candidates in ((1, 4), (1, 8), (4, 9)): + inputs = _make_inputs(num_tokens, num_candidates) + eager = attn_res_fwd(*inputs, RMS_EPS) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = attn_res_fwd(*inputs, RMS_EPS) + # One replay before the poison: under cudaMallocAsync the captured outputs are graph allocations, backed by + # memory only once the graph has run. + graph.replay() + for tensor in captured: + tensor.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + for name, expected, actual in zip(OUTPUT_NAMES, eager, captured): + assert torch.equal(_bits(actual), _bits(expected)), ( + f"T={num_tokens} N={num_candidates}: replayed {name} differs from eager" + ) + # Freed here rather than inside the next capture, where cudaMallocAsync's free would be part of it. + del captured diff --git a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py new file mode 100644 index 000000000000..d6a94e5555e4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_attn_res_rmsnorm_fwd.py @@ -0,0 +1,186 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attn_res_rmsnorm_fwd catalog entry. + +A cell is (T, N): T tokens and N = K + 1 candidates, K snapshots plus the layer residual. Every cell +checks the output's shape and dtype, bit-identical results from two identical calls, untouched inputs, +and closeness to an fp32 torch evaluation of the contract's Semantics, with both bf16 rounding +boundaries, under the metric of main's op test +(tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py). + +The kernel depends on N only: one CTA per token for N <= 4, one cluster of 8 CTAs per token for +N >= 5. With PDL enabled both release their dependents right after their own grid-dependency wait. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.attn_res_rmsnorm_fwd import ( + attn_res_rmsnorm_fwd, +) + +assert torch.cuda.is_available(), "attn_res_rmsnorm_fwd requires a CUDA device" +# The receipts are sm_100 ones; CI's other architectures skip this file. +pytestmark = pytest.mark.skipif( + torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" +) + +HIDDEN_SIZE = 7168 +RMS_EPS = 1e-6 +# Kimi K3's 93 layers append a snapshot every 12 layers: at most 8 snapshots, so N runs 2 ... 9. +K3_NUM_CANDIDATES = tuple(range(2, 10)) +DECODE_NUM_TOKENS = tuple(range(1, 9)) +# 32 is the fused path's token ceiling when KIMI_K3_ATTN_RES_TOPOLOGY is on. +LARGER_NUM_TOKENS = (32, 300) +# The thresholds of main's op test. +MIN_COSINE = 0.9999 +MAX_RELATIVE_L2 = 5e-3 +INPUT_NAMES = ("layer_residual", "block_residual", "res_weight", "rms_weight", "output_rms_weight") + + +def _make_inputs(num_tokens: int, num_candidates: int) -> tuple[torch.Tensor, ...]: + """Inputs at the scales of main's op test, seeded by the cell.""" + torch.manual_seed(97 * num_tokens + num_candidates) + layer_residual = ( + torch.randn(num_tokens, 1, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.05 + ) + block_residual = ( + torch.randn( + num_candidates - 1, num_tokens, 1, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda" + ) + * 0.05 + ) + res_weight = torch.randn(HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.02 + rms_weight = 1 + torch.randn(HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.02 + output_rms_weight = 1 + torch.randn(HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.02 + return layer_residual, block_residual, res_weight, rms_weight, output_rms_weight + + +def _reference( + layer_residual: torch.Tensor, + block_residual: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + output_rms_weight: torch.Tensor, + rms_eps: float, + output_rms_eps: float, +) -> torch.Tensor: + """fp32 evaluation of the contract's Semantics, rounding where the kernel rounds.""" + values = torch.cat((block_residual, layer_residual.unsqueeze(0)), dim=0).float() + rsigma = torch.rsqrt(values.square().mean(dim=-1, keepdim=True) + rms_eps) + logits = (values * rsigma * (rms_weight.float() * res_weight.float())).sum(dim=-1) + probs = torch.softmax(logits, dim=0) + mixed = (probs.unsqueeze(-1) * values).sum(dim=0).to(torch.bfloat16).float() + normed = mixed * torch.rsqrt(mixed.square().mean(dim=-1, keepdim=True) + output_rms_eps) + # bf16 times bf16: the exact product, rounded to bf16 once more. + return output_rms_weight * normed.to(torch.bfloat16) + + +def _bits(tensor: torch.Tensor) -> torch.Tensor: + """Reinterpret as integers, so that equality means identical bits.""" + return tensor.view(torch.int16) + + +def _similarity(actual: torch.Tensor, expected: torch.Tensor) -> tuple[float, float]: + """Cosine similarity and relative L2 error, the metric of main's op tests.""" + actual_float = actual.float().flatten() + expected_float = expected.float().flatten() + cosine = torch.nn.functional.cosine_similarity(actual_float, expected_float, dim=0).item() + relative_l2 = ((actual_float - expected_float).norm() / (expected_float.norm() + 1e-12)).item() + return cosine, relative_l2 + + +def _check( + num_tokens: int, + num_candidates: int, + rms_eps: float = RMS_EPS, + output_rms_eps: float = RMS_EPS, +) -> None: + inputs = _make_inputs(num_tokens, num_candidates) + before = [tensor.clone() for tensor in inputs] + output = attn_res_rmsnorm_fwd(*inputs, rms_eps, output_rms_eps) + repeat = attn_res_rmsnorm_fwd(*inputs, rms_eps, output_rms_eps) + expected = _reference(*inputs, rms_eps, output_rms_eps) + + cell = f"T={num_tokens} N={num_candidates} eps=({rms_eps}, {output_rms_eps})" + assert output.shape == (num_tokens, 1, HIDDEN_SIZE), f"{cell}: shape {tuple(output.shape)}" + assert output.dtype == torch.bfloat16, f"{cell}: dtype {output.dtype}" + assert output.is_contiguous(), f"{cell}: output is not contiguous" + assert output.device == inputs[0].device, f"{cell}: output on {output.device}" + assert torch.equal(_bits(output), _bits(repeat)), f"{cell}: two identical calls disagree" + cosine, relative_l2 = _similarity(output, expected) + assert cosine > MIN_COSINE and relative_l2 < MAX_RELATIVE_L2, ( + f"{cell}: cosine {cosine:.6f}, relative L2 {relative_l2:.3e}" + ) + for name, old, new in zip(INPUT_NAMES, before, inputs): + assert torch.equal(_bits(old), _bits(new)), f"{cell}: {name} was mutated" + + +def test_k3_decode_cells() -> None: + """Kimi K3's N = 2 ... 9 at T = 1 ... 8, on both kernels.""" + for num_tokens in DECODE_NUM_TOKENS: + for num_candidates in K3_NUM_CANDIDATES: + _check(num_tokens, num_candidates) + + +def test_larger_token_counts() -> None: + """Kimi K3's N = 2 ... 9 at more tokens: more CTAs or clusters, the same per-token kernel.""" + for num_tokens in LARGER_NUM_TOKENS: + for num_candidates in K3_NUM_CANDIDATES: + _check(num_tokens, num_candidates) + + +def test_other_candidate_counts() -> None: + """N = 1 (single-CTA kernel, no snapshot) and N = 10 ... 12 (split-K), outside Kimi K3's range.""" + for num_tokens in (1, 8): + for num_candidates in (1, 10, 11, 12): + _check(num_tokens, num_candidates) + + +def test_each_eps_feeds_its_own_norm() -> None: + """rms_eps scores the candidates and output_rms_eps scales the output; 1e-2 dominates either.""" + for num_tokens, num_candidates in ((1, 4), (1, 9), (8, 2)): + for rms_eps, output_rms_eps in ((1e-5, 1e-5), (1e-2, 1e-6), (1e-6, 1e-2)): + _check(num_tokens, num_candidates, rms_eps, output_rms_eps) + + +def test_cuda_graph_replay_matches_eager() -> None: + """One cell per kernel: a call captured in a CUDA graph replays the eager call's bits.""" + for num_tokens, num_candidates in ((1, 4), (1, 9), (8, 5)): + inputs = _make_inputs(num_tokens, num_candidates) + eager = attn_res_rmsnorm_fwd(*inputs, RMS_EPS, RMS_EPS) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = attn_res_rmsnorm_fwd(*inputs, RMS_EPS, RMS_EPS) + # One replay before the poison: under cudaMallocAsync the captured output is a graph allocation, backed by + # memory only once the graph has run. + graph.replay() + captured.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(_bits(captured), _bits(eager)), ( + f"T={num_tokens} N={num_candidates}: replay differs from eager" + ) + # Freed here rather than inside the next capture, where cudaMallocAsync's free would be part of it. + del captured + + +def test_chained_calls_wait_for_their_input() -> None: + """A call reading the previous call's output reads it settled, although that call triggers early. + + With PDL enabled the first call releases its dependents before it writes its output, so the second + call can start while the first still runs; its grid-dependency wait must hold its reads back. + """ + for num_tokens, num_candidates in ((1, 4), (1, 9), (8, 4), (8, 9)): + layer_residual, block_residual, res_weight, rms_weight, output_rms_weight = _make_inputs( + num_tokens, num_candidates + ) + weights = (res_weight, rms_weight, output_rms_weight, RMS_EPS, RMS_EPS) + first = attn_res_rmsnorm_fwd(layer_residual, block_residual, *weights) + chained = attn_res_rmsnorm_fwd(first, block_residual, *weights) + torch.cuda.synchronize() + settled = attn_res_rmsnorm_fwd(first.clone(), block_residual, *weights) + assert torch.equal(_bits(chained), _bits(settled)), ( + f"T={num_tokens} N={num_candidates}: the chained call read an unsettled input" + ) diff --git a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_k3_embed_norm.py b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_k3_embed_norm.py new file mode 100644 index 000000000000..89b1016e6d70 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_k3_embed_norm.py @@ -0,0 +1,211 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the k3_embed_norm catalog entry. + +Kimi K3's replicated embedding table [163840, 7168] at every token count of the K3 decode steps (N = 1..8, and +16..64 for DSpark verify steps of up to 8 requests x 8 tokens), against a native torch reference: the rows written into +`raw` bit for bit against a torch gather, the normed rows against an fp64 RMSNorm (within 8e-3 of each row's max +|ref|). Also: run-to-run bits, each row's bits independent of N, int64 ids equal to int32 ids, another eps and bank +slot, ids outside [0, V), a CUDA graph replayed with its ids rewritten, and the refused calls. +""" + +import functools +import math + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.k3_embed_norm import k3_embed_norm + +assert torch.cuda.is_available(), "k3_embed_norm requires a CUDA device" +# The receipts are sm_100 ones; CI's other architectures skip this file. +pytestmark = pytest.mark.skipif( + torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" +) + +VOCAB, HIDDEN = 163840, 7168 +EPS, EPS_ALT = 1e-6, 1e-5 +TOL = 8e-3 # per row: max |out - ref| / max |ref|, ref an fp64 RMSNorm +TOKENS = list(range(1, 9)) + list(range(16, 65, 8)) +INT64_TOKENS = (1, 8, 64) +BANK_SLOTS = 4 # the attention-residual snapshot bank; the model passes slot 0 as raw +SENTINEL = 7.0 +OUTSIDE_IDS = (-1, VOCAB, VOCAB + 7, -VOCAB, 2**31 - 1, -(2**31)) +OUTSIDE_IDS_INT64 = (2**32, 2**32 + 5, 3 - 2**32, 2**40) + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.int16) + + +def _same(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.shape == b.shape and torch.equal(_bits(a), _bits(b)) + + +@functools.lru_cache(maxsize=None) +def _table() -> tuple[torch.Tensor, int]: + """The table, built once: N(0, 1) rows at log-uniform scales 1e-4..1e2, row 1 all zero; and its smallest-scale + row's id (a row whose norm eps dominates).""" + gen = torch.Generator(device="cuda").manual_seed(20260930) + table = torch.randn(VOCAB, HIDDEN, generator=gen, device="cuda", dtype=torch.bfloat16) + scale = torch.empty(VOCAB, 1, device="cuda") + scale.uniform_(math.log(1e-4), math.log(1e2), generator=gen).exp_() + table.mul_(scale.bfloat16()) + table[1].zero_() + scale[1] = math.inf + return table, int(scale.argmin()) + + +@functools.lru_cache(maxsize=None) +def _weight() -> torch.Tensor: + """Layer 0's input-norm weight: around 1, every 101st element negative.""" + gen = torch.Generator(device="cuda").manual_seed(7) + weight = 1.0 + 0.2 * torch.randn(HIDDEN, generator=gen, device="cuda") + weight[::101] *= -1.0 + return weight.bfloat16() + + +@functools.lru_cache(maxsize=None) +def _ids64() -> torch.Tensor: + """64 int32 ids in [0, V): random, with 0, V - 1, the zero row, the smallest-scale row and repeats in front.""" + _, small = _table() + gen = torch.Generator(device="cuda").manual_seed(11) + ids = torch.randint(0, VOCAB, (64,), generator=gen, device="cuda", dtype=torch.int32) + ids[0], ids[1], ids[2], ids[4] = 0, VOCAB - 1, 1, small + ids[3] = ids[1] + ids[6] = ids[5] + torch.cuda.synchronize() + return ids + + +def _bank(n: int) -> torch.Tensor: + """A snapshot bank [BANK_SLOTS, n, H] filled with SENTINEL.""" + return torch.full((BANK_SLOTS, n, HIDDEN), SENTINEL, dtype=torch.bfloat16, device="cuda") + + +def _rows_ref(ids: torch.Tensor) -> torch.Tensor: + """The torch gather: table[ids], zero rows for ids outside [0, V).""" + table, _ = _table() + valid = (ids >= 0) & (ids < VOCAB) + rows = table.index_select(0, torch.where(valid, ids, torch.zeros_like(ids))) + return rows.masked_fill(~valid[:, None], 0.0) + + +def _norm_ref(rows: torch.Tensor, eps: float) -> torch.Tensor: + """fp64 RMSNorm: rows * rsqrt(mean(rows^2) + eps) * weight.""" + x = rows.double() + return x * torch.rsqrt(x.pow(2).mean(dim=1, keepdim=True) + eps) * _weight().double() + + +def _call(ids: torch.Tensor, eps: float = EPS, slot: int = 0) -> tuple[torch.Tensor, torch.Tensor]: + """k3_embed_norm into ``slot`` of a fresh bank: (out, bank).""" + table, _ = _table() + bank = _bank(ids.numel()) + return k3_embed_norm(ids, table, _weight(), eps, bank[slot]), bank + + +def _check(ids: torch.Tensor, eps: float = EPS, slot: int = 0) -> torch.Tensor: + """One call's checks against the reference and a rerun; returns the call's output.""" + n = ids.numel() + what = f"N={n} {ids.dtype} eps={eps:g} slot={slot}" + out, bank = _call(ids, eps, slot) + again, bank_again = _call(ids, eps, slot) + rows = _rows_ref(ids) + ref = _norm_ref(rows, eps) + assert out.shape == (n, HIDDEN) and out.dtype == torch.bfloat16 and out.device == ids.device + assert out.is_contiguous() + assert _same(bank[slot], rows), f"{what}: raw is not the gathered rows" + others = [s for s in range(BANK_SLOTS) if s != slot] + assert bool((bank[others] == SENTINEL).all()), f"{what}: wrote outside raw" + scale = ref.abs().amax(dim=1) + live = scale > 0 + err = ((out.double() - ref).abs().amax(dim=1)[live] / scale[live]).max().item() + assert err <= TOL, f"{what}: rel {err:.3e}" + assert bool((out[~live] == 0).all()), f"{what}: a zero row normed to non-zero" + assert _same(again, out) and _same(bank_again, bank), f"{what}: a rerun differs" + return out + + +def test_certified_cells() -> None: + """Every N with int32 ids, eps 1e-6, raw = slot 0; every row bit-identical to the same id's row of the 64-id + call.""" + ids64 = _ids64() + full = _check(ids64) + for n in TOKENS: + out = _check(ids64[:n]) + assert _same(out, full[:n]), f"N={n}: rows differ from the 64-id call" + + +def test_int64_ids_other_eps_and_slot() -> None: + """int64 ids at N 1, 8 and 64, bit-identical to int32 ids; eps 1e-5 into bank slot 1 at N 8.""" + ids64 = _ids64() + for n in INT64_TOKENS: + ids = ids64[:n] + assert _same(_check(ids.long()), _call(ids)[0]), f"N={n}: int64 ids differ from int32" + _check(ids64[:8], eps=EPS_ALT, slot=1) + + +def test_ids_outside_vocab() -> None: + """N 8, every out-of-range value at the even positions (int32 and int64): zero raw rows, zero normed rows, the other + rows as the reference.""" + gen = torch.Generator(device="cuda").manual_seed(3008) + for dtype, bad in ( + (torch.int32, OUTSIDE_IDS), + (torch.int64, OUTSIDE_IDS + OUTSIDE_IDS_INT64), + ): + for start in range(0, len(bad), 4): + ids = torch.randint(0, VOCAB, (8,), generator=gen, device="cuda", dtype=dtype) + values = [bad[(start + i) % len(bad)] for i in range(4)] + ids[0::2] = torch.tensor(values, dtype=dtype, device="cuda") + out = _check(ids) + assert bool((out[0::2] == 0).all()), f"{dtype} {values}: non-zero rows" + + +def test_graph_replay() -> None: + """An N = 8 call captured once after an eager call compiled it, replayed with the ids rewritten in place (random, + then some outside [0, V)): every replay bit-identical to an eager call on the same ids, raw the gathered rows.""" + table, _ = _table() + ids_buf = _ids64()[:8].clone() + bank = _bank(8) + k3_embed_norm(ids_buf, table, _weight(), EPS, bank[0]) # compiles outside the capture + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = k3_embed_norm(ids_buf, table, _weight(), EPS, bank[0]) + gen = torch.Generator(device="cuda").manual_seed(4008) + for rep in range(3): + ids = torch.randint(0, VOCAB, (8,), generator=gen, device="cuda", dtype=torch.int32) + if rep == 2: + ids[0::2] = -1 + ids_buf.copy_(ids) + bank.fill_(SENTINEL) + graph.replay() + eager, eager_bank = _call(ids_buf) + assert _same(out, eager) and _same(bank, eager_bank), f"replay {rep} differs from eager" + assert _same(bank[0], _rows_ref(ids_buf)), f"replay {rep}: raw is not the gathered rows" + + +def test_refused_calls() -> None: + """Calls outside the op's support raise ValueError before launching anything: raw stays untouched.""" + table, _ = _table() + weight = _weight() + ids8 = _ids64()[:8] + bank = _bank(8) + + def refused(ids, tab, w, raw): + with pytest.raises(ValueError, match="unsupported call"): + k3_embed_norm(ids, tab, w, EPS, raw) + + refused(ids8[:0], table, weight, bank[0, :0]) # N = 0 + ids65 = torch.zeros(65, dtype=torch.int32, device="cuda") + refused(ids65, table, weight, torch.empty(65, HIDDEN, dtype=torch.bfloat16, device="cuda")) + refused(ids8.to(torch.int16), table, weight, bank[0]) # id dtype + refused(ids8, table, weight, _bank(64)[0]) # raw [64, H] for 8 ids + flat = torch.empty(8 * HIDDEN + 8, dtype=torch.bfloat16, device="cuda") + refused(ids8, table, weight, flat[1 : 1 + 8 * HIDDEN].view(8, HIDDEN)) # raw 2 bytes off + refused(ids8, table, weight[: HIDDEN // 2], bank[0]) # weight [H / 2] + for hidden in (6144, 7680): # outside the norm's geometry + small = torch.zeros(16, hidden, dtype=torch.bfloat16, device="cuda") + w = torch.ones(hidden, dtype=torch.bfloat16, device="cuda") + raw = torch.empty(8, hidden, dtype=torch.bfloat16, device="cuda") + refused(ids8 % 16, small, w, raw) + assert bool((bank == SENTINEL).all()) diff --git a/tests/unittest/_torch/modeling_v2/ssm/_kda_cells.py b/tests/unittest/_torch/modeling_v2/ssm/_kda_cells.py new file mode 100644 index 000000000000..9c5d0e705b68 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/ssm/_kda_cells.py @@ -0,0 +1,330 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""What the ssm/ entry tests share: a real ``MambaHybridCacheManagerV2`` at Kimi K3's per-rank KDA shapes, its +per-layer pools, and the references of the ops' own tests (k3_kda_qkvg's rows decoded from its Lamport buffers; a +float64 plain decode; a float64 verify that keeps each request's committed history). + +Not a test file (no ``test_`` prefix): the four entry tests import it by name, as ``comm/`` imports ``_rank_job``. + +Shapes: 6 heads, K = V = 128, conv width 4 (the TP16 rank slice); conv states bf16 in the ``[q | k | v]`` layout, +SSM states fp32. The manager coalesces each slot's per-layer states, so every per-layer view is a strided slot view, +the layout the model hands these ops. +""" + +from __future__ import annotations + +import torch + +H = 6 +K = V = 128 +HK = H * K +W = 4 +CONV_DIM = 3 * HK +PROJ = 4 * HK + K + H + 2 # the fused [q | k | v | og | f_a | b | pad] rows (3208) +K_IN = 7168 +NUM_SPEC = 7 +NT = NUM_SPEC + 1 +LOWER_BOUND = -5.0 +EPS = 1e-5 +SCALE = K**-0.5 +TOL_OUT = 2e-2 # bf16 outputs +TOL_STATE = 1e-3 # fp32 states: approximate exp / rcp in the kernels and the bf16 f_b gate + + +def sm100() -> bool: + """sm_100 exactly: the architecture these entries are certified on (a missing receipt reads as unknown).""" + return torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0) + + +def load_ops(): + """Register the C++ ops and the K3 KDA CuTe DSL ops; return the k3_kda_attn op module (its constants).""" + import tensorrt_llm # noqa: F401 (registers the trtllm ops) + from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_attn import op + from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_verify import op as _verify # noqa: F401 + + return op + + +# --------------------------------------------------------------------------------------------------------------- +# The cache objects +# --------------------------------------------------------------------------------------------------------------- +def build_manager(num_layers: int, num_spec: int | None = None, max_batch_size: int = 8): + """A real MambaHybridCacheManagerV2 with ``num_layers`` KDA layers (and one attention layer). With ``num_spec``: + MTP-style speculation of ``num_spec`` drafts, the KDA replay caches and the per-token states + (``kda_token_states``) that k3_kda_verify / k3_kda_attn read and write.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + MambaHybridCacheManagerV2, + ) + 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 + + spec = MTPDecodingConfig(max_draft_len=num_spec) if num_spec else None + return MambaHybridCacheManagerV2( + mamba_d_state=K, + mamba_d_conv=W, + mamba_num_heads=H, + mamba_n_groups=H, + mamba_head_dim=V, + mamba_num_layers=num_layers, + mamba_layer_mask=[True] * num_layers + [False], + mamba_cache_dtype=torch.bfloat16, + mamba_ssm_cache_dtype=torch.float32, + kv_cache_config=KvCacheConfig(max_tokens=512, enable_block_reuse=False), + kv_cache_type=CacheTypeCpp.SELF, + num_layers=1, + num_kv_heads=4, + head_dim=64, + tokens_per_block=32, + max_seq_len=128, + max_batch_size=max_batch_size, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.HALF, + spec_config=spec, + layer_mask=[False] * num_layers + [True], + vocab_size=1024, + conv_state_layout="q_k_v", + kda_replay_num_spec=num_spec, + kda_token_states=num_spec is not None, + ) + + +def request_slots(mgr, count: int, first_id: int) -> torch.Tensor: + """State slots of ``count`` new requests (the manager's own slot assignment), int32 on the GPU.""" + ids = list(range(first_id, first_id + count)) + mgr.add_dummy_requests(ids, token_nums=[8] * count, is_gen=False) + slots = mgr.get_state_indices(ids, [False] * count) + assert len(set(slots)) == count, slots + return torch.tensor(slots, dtype=torch.int32, device="cuda") + + +def layer_pools(mgr, layer: int) -> dict: + """The manager's views of one KDA layer's pools: ``conv`` bf16 [slots, 3 HK, W - 1] (q | k | v channels) and + ``ssm`` fp32 [slots, H, V, K] (plain decode); with the replay caches also ``cs_q`` / ``cs_k`` / ``cs_v`` fp32 + [slots, HK, W - 1 + num_spec] (dim-contiguous), ``state_tok`` fp32 [slots, num_spec, H, V, K] and ``pending``, + the accepted-draft record every layer shares (``prev_num_accepted_tokens``).""" + pools = {"conv": mgr.get_conv_states(layer), "ssm": mgr.get_ssm_states(layer)} + if mgr.use_kda_replay_update: + cache = mgr.mamba_layer_cache(layer) + pools.update( + cs_q=cache.kda_conv_q, + cs_k=cache.kda_conv_k, + cs_v=cache.kda_conv_v, + state_tok=cache.kda_state_tok, + pending=cache.prev_num_accepted_tokens, + ) + return pools + + +def fill_pools(pools: dict, seed: int) -> None: + """Random finite contents in every slot of every pool (the pending record zeroed).""" + g = torch.Generator(device="cuda").manual_seed(seed) + + def rnd(t, scale): + return torch.randn(t.shape, generator=g, device="cuda") * scale + + pools["conv"].copy_(rnd(pools["conv"], 0.5).bfloat16()) + pools["ssm"].copy_(rnd(pools["ssm"], 0.05)) + for name in ("cs_q", "cs_k", "cs_v"): + if name in pools: + pools[name].copy_(rnd(pools[name], 0.5)) + if "state_tok" in pools: + pools["state_tok"].zero_() + if "pending" in pools: + pools["pending"].zero_() + + +def clone_pools(pools: dict) -> dict: + """Standalone copies with the same strides (so the same layout constraints hold), e.g. for a reference path.""" + out = {} + for name, t in pools.items(): + c = torch.empty_strided(t.shape, t.stride(), dtype=t.dtype, device=t.device) + c.copy_(t) + out[name] = c + return out + + +def snapshot(pools: dict) -> dict: + return {name: t.clone() for name, t in pools.items()} + + +def same_rows(a: dict, b: dict, rows, names) -> bool: + """Bit equality of the given slots (``rows``: list of ints) of the given pools.""" + return all(torch.equal(a[n][rows], b[n][rows]) for n in names) + + +def other_slots(num_slots: int, used) -> list: + used = set(int(s) for s in used) + return [s for s in range(num_slots) if s not in used] + + +def rel(a: torch.Tensor, b: torch.Tensor) -> float: + return ((a.float() - b.float()).abs().max() / b.float().abs().max().clamp_min(1e-6)).item() + + +# --------------------------------------------------------------------------------------------------------------- +# Weights and the projection rows +# --------------------------------------------------------------------------------------------------------------- +def make_weights(seed: int) -> dict: + """One KDA layer's decode weights at the TP16 shapes; the conv taps are bf16 values (kda_decode reads them as + bf16, the K3 kernels as fp32).""" + g = torch.Generator(device="cuda").manual_seed(seed) + + def rnd(*s, scale=1.0): + return torch.randn(*s, generator=g, device="cuda") * scale + + conv = rnd(3, HK, W, scale=0.3).bfloat16().float() + return { + "w": rnd(PROJ, K_IN, scale=0.02).bfloat16(), + "w_fb": rnd(HK, K, scale=0.05).bfloat16(), + "w_q": conv[0].contiguous(), + "w_k": conv[1].contiguous(), + "w_v": conv[2].contiguous(), + "w_t": [conv[i].t().bfloat16().contiguous() for i in range(3)], + "a_log": rnd(H, scale=0.5), + "dt_bias": rnd(HK, scale=0.5), + "onorm_w": (1 + 0.1 * rnd(V)).float(), + } + + +def decode_rows(p1: torch.Tensor, part: torch.Tensor, buf: int, n: int) -> torch.Tensor: + """The ``[q | k | v | og | f_a | b | pad]`` rows of tokens 0..n-1 from buffer ``buf`` of a k3_kda_qkvg launch: + q, k and f_a are bf16 bits; v, og and b the bf16 sum of the two fp32 K-half partials.""" + qkfa = p1.view(3, 8, 2 * HK + K)[buf, :n].view(torch.bfloat16) + parts = part.view(3, 3, 2, 8, HK)[buf].view(torch.float32) + v, og, b = ((parts[r, 0, :n] + parts[r, 1, :n]).bfloat16() for r in range(3)) + rows = torch.zeros(n, PROJ, dtype=torch.bfloat16, device=p1.device) + rows[:, : 2 * HK] = qkfa[:, : 2 * HK] + rows[:, 2 * HK : 3 * HK] = v + rows[:, 3 * HK : 4 * HK] = og + rows[:, 4 * HK : 4 * HK + K] = qkfa[:, 2 * HK :] + rows[:, 4 * HK + K : 4 * HK + K + H] = b[:, :H] + return rows + + +# --------------------------------------------------------------------------------------------------------------- +# float64 references +# --------------------------------------------------------------------------------------------------------------- +def _gate_decay(a_log, dt_bias, g): + """The lower-bound gate: decay = exp(lower_bound * sigmoid(exp(A_log) * (g + dt_bias))), per key channel.""" + xg = torch.exp(a_log.double()).unsqueeze(-1) * (g + dt_bias.double().view(H, K)) + return torch.exp(LOWER_BOUND * torch.sigmoid(xg)) + + +class F64Decode: + """The plain KDA decode in float64 torch over copies of a layer's ``conv`` / ``ssm`` pools: conv4 + SiLU, q / k + L2 norm (q scaled), beta sigmoid, the lower-bound gate, S <- S d; S <- S + beta (v - S k) k^T; o = S q, the + gated RMSNorm.""" + + def __init__(self, wt: dict, pools: dict): + self.wt = wt + self.conv = pools["conv"].double().clone() + self.state = pools["ssm"].double().clone() + + def step(self, raw_qkv, g, beta_raw, gate_raw, slots) -> torch.Tensor: + """``raw_qkv`` [n, 3 HK] (the projection's q | k | v columns), ``g`` [n, HK] (f_b's output, before dt_bias), + ``beta_raw`` [n, H], ``gate_raw`` [n, HK] (the output gate); returns [n, H, V].""" + wt = self.wt + conv_w = torch.stack([wt["w_q"], wt["w_k"], wt["w_v"]]).double() # [3, HK, W] + outs = torch.empty(raw_qkv.shape[0], H, V, dtype=torch.float64, device="cuda") + for i, s in enumerate(slots.tolist()): + new = raw_qkv[i].double().view(3, HK) + win = self.conv[s].view(3, HK, W - 1) + u = torch.cat([win, new.unsqueeze(-1)], dim=-1) # oldest first + act = (u * conv_w).sum(-1) + act = act * torch.sigmoid(act) + self.conv[s] = u[:, :, 1:].reshape(3 * HK, W - 1) + q, k, v = (act[j].view(H, K) for j in range(3)) + q = q / torch.sqrt((q * q).sum(-1, keepdim=True) + 1e-6) * SCALE + k = k / torch.sqrt((k * k).sum(-1, keepdim=True) + 1e-6) + beta = torch.sigmoid(beta_raw[i].double()) + decay = _gate_decay(wt["a_log"], wt["dt_bias"], g[i].double().view(H, K)) + st = self.state[s] * decay.unsqueeze(1) + res = (v - (st * k.unsqueeze(1)).sum(-1)) * beta.unsqueeze(-1) + st = st + res.unsqueeze(-1) * k.unsqueeze(1) + self.state[s] = st + o = (st * q.unsqueeze(1)).sum(-1) + rms = torch.rsqrt((o * o).mean(-1, keepdim=True) + EPS) + gate = torch.sigmoid(gate_raw[i].double().view(H, V)) + outs[i] = o * rms * wt["onorm_w"].double() * gate + return outs + + def step_rows(self, rows: torch.Tensor, slots) -> torch.Tensor: + """As :meth:`step` from fused projection rows (f_b in float64, rounded to bf16 as the GEMM's output is).""" + r = rows.double() + g = (r[:, 4 * HK : 4 * HK + K] @ self.wt["w_fb"].double().t()).bfloat16().double() + return self.step( + r[:, : 3 * HK], g, r[:, 4 * HK + K : 4 * HK + K + H], r[:, 3 * HK : 4 * HK], slots + ) + + +def _conv_silu(win, raw, c, w): + x = win[0][c] * w[:, 0] + win[1][c] * w[:, 1] + win[2][c] * w[:, 2] + raw[c] * w[:, 3] + return x * torch.sigmoid(x) + + +class F64Verify: + """The speculative verify in float64 over each request's committed history (raw conv inputs and the state after + the last committed token): every request's 1 + num_spec tokens from the state after its accepted drafts. The + conv caches start with pending 0, so their window columns 0..2 are the committed raw inputs.""" + + def __init__(self, wt: dict, pools: dict, slots, num_spec: int): + self.wt, self.num_spec = wt, num_spec + self.slots = slots.tolist() + self.seq = [ + [ + torch.stack([pools[c][s, :, i].double() for c in ("cs_q", "cs_k", "cs_v")]) + for i in range(3) + ] + for s in self.slots + ] + self.state = [pools["ssm"][s].double().clone() for s in self.slots] + self.last = None + + def __call__(self, rows: torch.Tensor) -> torch.Tensor: + """Outputs [N (1 + num_spec), H, V] for the fused projection rows of N requests.""" + wt, steps = self.wt, self.num_spec + 1 + r = rows.double() + g_all = (r[:, 4 * HK : 4 * HK + K] @ wt["w_fb"].double().t()).bfloat16().double() + wq, wk, wv = (wt[n].double() for n in ("w_q", "w_k", "w_v")) + out = torch.empty(rows.shape[0], H, V, dtype=torch.float64, device="cuda") + self.last = [] + for n in range(len(self.slots)): + seq, s_cur = list(self.seq[n]), self.state[n] + states, raws = [], [] + for t in range(steps): + row = n * steps + t + raw = r[row, : 3 * HK].view(3, HK) + win = seq[-3:] + q = _conv_silu(win, raw, 0, wq).view(H, K) + k = _conv_silu(win, raw, 1, wk).view(H, K) + v = _conv_silu(win, raw, 2, wv).view(H, V) + q = q * torch.rsqrt((q * q).sum(-1, keepdim=True) + 1e-6) * SCALE + k = k * torch.rsqrt((k * k).sum(-1, keepdim=True) + 1e-6) + beta = torch.sigmoid(r[row, 4 * HK + K : 4 * HK + K + H]) + decay = _gate_decay(wt["a_log"], wt["dt_bias"], g_all[row].view(H, K)) + sd = s_cur * decay[:, None, :] + vn = v - torch.einsum("hvk,hk->hv", sd, k) + s_cur = sd + beta[:, None, None] * vn[:, :, None] * k[:, None, :] + o = torch.einsum("hvk,hk->hv", s_cur, q) + rms = torch.rsqrt((o * o).mean(-1, keepdim=True) + EPS) + gate = torch.sigmoid(r[row, 3 * HK : 4 * HK].view(H, V)) + out[row] = o * rms * wt["onorm_w"].double() * gate + states.append(s_cur) + raws.append(raw) + seq.append(raw) + self.last.append((states, raws)) + return out + + def commit(self, accepted) -> None: + """The sampler accepted ``accepted[n]`` drafts: the golden token and those drafts commit.""" + for n, p in enumerate(accepted): + states, raws = self.last[n] + self.seq[n] = (self.seq[n] + raws[: p + 1])[-3:] + self.state[n] = states[p] + + +def pending_schedule(num_requests: int, num_spec: int, rnd: int) -> list: + """Accepted drafts per request after round ``rnd``: every count 0..num_spec over the rounds, different per + request.""" + return [(3 * n + 5 * rnd + 1 + (rnd * n) % 3) % (num_spec + 1) for n in range(num_requests)] diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_attn.py b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_attn.py new file mode 100644 index 000000000000..f05c0dbe78ca --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_attn.py @@ -0,0 +1,292 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the ssm/k3_kda_attn catalog entry (and its sibling k3_kda_qkvg), on a real cache object and a +caller-owned K3KdaBuffers. + +Kimi K3's fused KDA projection + speculative verify of one request's golden token and 7 drafts, at the TP16 rank +slice, on the verify pools of a real ``MambaHybridCacheManagerV2`` built with MTP-style speculation of 7 drafts, the +KDA replay caches and per-token states (``_kda_cells.build_manager``): the layer's conv caches (fp32, dim-contiguous), +its fp32 state (strided by the manager's per-slot coalescing), its per-draft states and the accepted-draft record +every layer shares (``prev_num_accepted_tokens``), at the slots the manager assigned. + +The reference, as the op's own test: the projection stream alone (``k3_kda_qkvg``, its Lamport buffers decoded into +rows) and the ``ssm/k3_kda_verify`` entry on them, on a copy of the pools: outputs and every pool bit for bit. Layer 0 +is also checked against a float64 verify over the request's committed history. Call sequences: layers x rounds on one +shared set against a set per layer and a captured round replayed, the per-CTA index across the int32 wrap; negative +controls. + +Every launch here follows a plain copy kernel (its input), as each KDA launch in the model follows other layers' +kernels: the head CTAs read the slot's pools before their grid-dependency wait, so a launch must not directly follow +another launch on the same pools. +""" + +import _kda_cells as kc +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_attn import ( + k3_kda_attn, + k3_kda_qkvg, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_buffers import ( + CTAS, + K3KdaBuffers, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_verify import k3_kda_verify + +pytestmark = pytest.mark.skipif(not kc.sm100(), reason="needs SM100 (tcgen05, TMA, clusters)") + +LAYERS = 2 +SLOT_ROUNDS = 9 # pending 0, then the schedule's 1, 6, 3, 0, 5, 2, 7, 4 +POOL_NAMES = ("cs_q", "cs_k", "cs_v", "ssm", "state_tok") + + +@pytest.fixture(scope="module") +def mgr(): + kc.load_ops() + m = kc.build_manager(LAYERS, num_spec=kc.NUM_SPEC) + try: + yield m + finally: + m.shutdown() + + +@pytest.fixture(scope="module") +def slots(mgr): + return kc.request_slots(mgr, 4, first_id=400) + + +def _layers(mgr) -> list: + return [kc.layer_pools(mgr, layer) for layer in range(LAYERS)] + + +def _fused(wt, pools, x_src, slot, bufs) -> torch.Tensor: + """One launch, after a plain copy kernel (the input).""" + x = x_src.clone() + return k3_kda_attn( + x, wt["w"], wt["w_fb"], wt["w_q"], wt["w_k"], wt["w_v"], wt["a_log"], wt["dt_bias"], wt["onorm_w"], + pools["cs_q"], pools["cs_k"], pools["cs_v"], pools["ssm"], pools["state_tok"], slot, pools["pending"], bufs, + kc.NUM_SPEC, kc.LOWER_BOUND, kc.SCALE, kc.EPS, + ) # fmt: skip + + +class _Unfused: + """k3_kda_qkvg's rows (its Lamport buffers decoded), then the ssm/k3_kda_verify entry on them.""" + + def __init__(self): + self.bufs = K3KdaBuffers.create("cuda", ctas=CTAS) + self.last_rows = None + + def rows(self, w, x) -> torch.Tensor: + buf = int(self.bufs.epoch[0].item()) + k3_kda_qkvg(x, w, self.bufs) + return kc.decode_rows(self.bufs.p1, self.bufs.part, buf, x.shape[0]) + + def __call__(self, wt, pools, x, slot) -> torch.Tensor: + self.last_rows = self.rows(wt["w"], x) + return k3_kda_verify( + self.last_rows, wt["w_fb"], wt["w_q"], wt["w_k"], wt["w_v"], wt["a_log"], wt["dt_bias"], wt["onorm_w"], + pools["cs_q"], pools["cs_k"], pools["cs_v"], pools["ssm"], pools["state_tok"], slot, pools["pending"], + kc.NUM_SPEC, kc.LOWER_BOUND, kc.SCALE, kc.EPS, None, + ) # fmt: skip + + +def _x(gen: torch.Generator) -> torch.Tensor: + return torch.randn(kc.NT, kc.K_IN, generator=gen, device="cuda").bfloat16() + + +def _accept(pools_list, slot: int, accepted: int) -> None: + """The sampler's record of the drafts it accepted: one count per slot, shared by every layer of a pool set.""" + for pools in pools_list: + pools["pending"][slot] = accepted + + +def test_attn_on_manager_pools(mgr, slots) -> None: + """Two requests, one after the other, each on its slot for SLOT_ROUNDS rounds of every layer (pending through + 0..7), all launches on one set: outputs and pools bit for bit the unfused path's on a copy; the other slots + untouched; layer 0's outputs also within fp32 tolerance of float64 over the committed history.""" + pools = _layers(mgr) + with torch.inference_mode(): + for layer, p in enumerate(pools): + kc.fill_pools(p, 1000 + layer) + copies = [kc.clone_pools(p) for p in pools] + # The record is one tensor for every layer in the manager; keep the copies' one shared too. + for c in copies[1:]: + c["pending"] = copies[0]["pending"] + wts = [kc.make_weights(500 + layer) for layer in range(LAYERS)] + bufs = K3KdaBuffers.create("cuda") + ref = _Unfused() + gen = torch.Generator(device="cuda").manual_seed(50) + for req in range(2): + s = int(slots[req]) + slot = slots[req : req + 1] + f64 = kc.F64Verify(wts[0], pools[0], slot, kc.NUM_SPEC) + for rnd in range(SLOT_ROUNDS): + for layer, (p, c) in enumerate(zip(pools, copies)): + before = kc.snapshot(p) + x = _x(gen) + got = _fused(wts[layer], p, x, slot, bufs) + want = ref(wts[layer], c, x, slot) + if layer == 0: + out_f64 = f64(ref.last_rows) + torch.cuda.synchronize() + assert torch.equal(got, want), (req, rnd, layer) + assert kc.same_rows(p, c, [s], POOL_NAMES), (req, rnd, layer) + others = kc.other_slots(p["ssm"].shape[0], [s]) + assert kc.same_rows(p, before, others, POOL_NAMES), (req, rnd, layer) + if layer == 0: + assert kc.rel(got, out_f64) <= kc.TOL_OUT, (req, rnd, kc.rel(got, out_f64)) + accepted = kc.pending_schedule(1, kc.NUM_SPEC, rnd)[0] + _accept([pools[0], copies[0]], s, accepted) + f64.commit([accepted]) + assert bool((bufs.epoch == (2 * SLOT_ROUNDS * LAYERS) % 3).all()), ( + bufs.epoch.unique().tolist() + ) + + +def test_rounds_share_one_buffer_set_and_replay(mgr, slots) -> None: + """LAYERS layers x ROUNDS rounds of one request: one set for every launch, a set per layer, and one round of every + layer captured once and replayed (the inputs rewritten in place, the record written between replays) give the + same outputs and pools bit for bit.""" + rounds = 5 + s = int(slots[2]) + slot = slots[2:3] + pools = _layers(mgr) + with torch.inference_mode(): + for layer, p in enumerate(pools): + kc.fill_pools(p, 1100 + layer) + own_pools = [kc.clone_pools(p) for p in pools] + graph_pools = [kc.clone_pools(p) for p in pools] + for copies in (own_pools, graph_pools): + for c in copies[1:]: + c["pending"] = copies[0]["pending"] + wts = [kc.make_weights(510 + layer) for layer in range(LAYERS)] + gen = torch.Generator(device="cuda").manual_seed(51) + xs = [[_x(gen) for _ in range(LAYERS)] for _ in range(rounds)] + shared = K3KdaBuffers.create("cuda") + own = [K3KdaBuffers.create("cuda") for _ in range(LAYERS)] + got, want = [], [] + for rnd in range(rounds): + for layer in range(LAYERS): + got.append(_fused(wts[layer], pools[layer], xs[rnd][layer], slot, shared).clone()) + want.append( + _fused(wts[layer], own_pools[layer], xs[rnd][layer], slot, own[layer]).clone() + ) + accepted = kc.pending_schedule(1, kc.NUM_SPEC, rnd)[0] + _accept([pools[0], own_pools[0]], s, accepted) + torch.cuda.synchronize() + # Compiles outside capture (already compiled above; kept for running this test alone). + bufs = K3KdaBuffers.create("cuda") + static = [x.clone() for x in xs[0]] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + outs = [ + _fused(wts[li], graph_pools[li], static[li], slot, bufs) for li in range(LAYERS) + ] + torch.cuda.current_stream().wait_stream(stream) + replayed = [] + for rnd in range(rounds): + for layer in range(LAYERS): + static[layer].copy_(xs[rnd][layer]) + graph.replay() + replayed.extend(o.clone() for o in outs) + _accept([graph_pools[0]], s, kc.pending_schedule(1, kc.NUM_SPEC, rnd)[0]) + torch.cuda.synchronize() + assert all(torch.equal(a, b) for a, b in zip(got, want)) + assert all(torch.equal(a, b) for a, b in zip(replayed, want)) + for p, q, r in zip(pools, own_pools, graph_pools): + assert all(torch.equal(p[n], q[n]) and torch.equal(r[n], q[n]) for n in POOL_NAMES) + + +@pytest.mark.parametrize("tokens", [1, 3, 8]) +def test_qkvg_rows_match_the_projection(tokens) -> None: + """The sibling k3_kda_qkvg: the rows decoded from its Lamport buffers are x w^T rounded to bf16 (q, k and f_a the + cluster's four K-partials summed in rank order, v, og and b the bf16 sum of two fp32 K-halves), against float64.""" + kc.load_ops() + with torch.inference_mode(): + wt = kc.make_weights(520) + gen = torch.Generator(device="cuda").manual_seed(52) + x = torch.randn(tokens, kc.K_IN, generator=gen, device="cuda").bfloat16() + bufs = K3KdaBuffers.create("cuda", ctas=CTAS) + k3_kda_qkvg(x, wt["w"], bufs) + rows = kc.decode_rows(bufs.p1, bufs.part, 0, tokens) + want = (x.double() @ wt["w"].double().t())[:, : kc.PROJ - 2] + torch.cuda.synchronize() + err = kc.rel(rows[:, : kc.PROJ - 2], want) + print(f"k3_kda_qkvg rows vs float64, T = {tokens}: rel {err:.3e}") + assert err <= 1e-2, err + assert bool((bufs.epoch == 1).all()), bufs.epoch.unique().tolist() + + +def test_epoch_wrap_on_the_object(mgr, slots) -> None: + """The set's per-CTA indices preset to 2^31 - 2, as after 2^31 launches on a device (one set serves every KDA + layer): four launches give the bits of a fresh set's and leave every index in 0..2.""" + s = int(slots[3]) + slot = slots[3:4] + p = kc.layer_pools(mgr, 0) + with torch.inference_mode(): + kc.fill_pools(p, 1200) + fresh_pools = kc.clone_pools(p) + wt = kc.make_weights(530) + gen = torch.Generator(device="cuda").manual_seed(53) + xs = [_x(gen) for _ in range(4)] + fresh = K3KdaBuffers.create("cuda") + wrapped = K3KdaBuffers.create("cuda") + wrapped.epoch.fill_(2**31 - 2) + got, want = [], [] + for rnd, x in enumerate(xs): + want.append(_fused(wt, fresh_pools, x, slot, fresh).clone()) + torch.cuda.synchronize() # the same pools again next: one launch at a time + got.append(_fused(wt, p, x, slot, wrapped).clone()) + torch.cuda.synchronize() + accepted = kc.pending_schedule(1, kc.NUM_SPEC, rnd)[0] + _accept([p, fresh_pools], s, accepted) + assert all(torch.equal(a, b) for a, b in zip(got, want)) + assert all(torch.equal(p[n], fresh_pools[n]) for n in POOL_NAMES) + assert bool(((wrapped.epoch >= 0) & (wrapped.epoch < 3)).all()), wrapped.epoch.unique().tolist() + + +def test_swapped_rounds_are_silently_wrong(mgr, slots) -> None: + """Negative control: two rounds of one request in swapped order raise nothing and give the second round's + outputs and the slot's state of a different history.""" + s = int(slots[3]) + slot = slots[3:4] + p = kc.layer_pools(mgr, 0) + with torch.inference_mode(): + kc.fill_pools(p, 1300) + swapped = kc.clone_pools(p) + wt = kc.make_weights(540) + gen = torch.Generator(device="cuda").manual_seed(54) + a, b = _x(gen), _x(gen) + bufs = K3KdaBuffers.create("cuda") + _fused(wt, p, a, slot, bufs) + torch.cuda.synchronize() + _accept([p], s, 2) + right = _fused(wt, p, b, slot, bufs).clone() + torch.cuda.synchronize() + _fused(wt, swapped, b, slot, bufs) + torch.cuda.synchronize() + _accept([swapped], s, 2) + wrong = _fused(wt, swapped, a, slot, bufs).clone() + torch.cuda.synchronize() + print(f"swapped rounds: second output rel diff {kc.rel(wrong, right):.3e}, " + f"state rel diff {kc.rel(swapped['ssm'][s], p['ssm'][s]):.3e}") # fmt: skip + assert kc.rel(swapped["ssm"][s], p["ssm"][s]) > kc.TOL_STATE + assert not torch.equal(wrong, right) + + +def test_qkvg_set_is_refused(mgr, slots) -> None: + """A set made for k3_kda_qkvg's 104 CTAs is refused by k3_kda_attn before anything is written.""" + slot = slots[3:4] + p = kc.layer_pools(mgr, 0) + with torch.inference_mode(): + kc.fill_pools(p, 1400) + before = kc.snapshot(p) + wt = kc.make_weights(550) + x = _x(torch.Generator(device="cuda").manual_seed(55)) + with pytest.raises(ValueError, match=f"epoch {CTAS}"): + _fused(wt, p, x, slot, K3KdaBuffers.create("cuda", ctas=CTAS)) + torch.cuda.synchronize() + assert all(torch.equal(p[n], before[n]) for n in POOL_NAMES) diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_decode_attn.py b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_decode_attn.py new file mode 100644 index 000000000000..319c8b4918af --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_decode_attn.py @@ -0,0 +1,358 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the ssm/k3_kda_decode_attn catalog entry, on a real cache object and a caller-owned K3KdaBuffers. + +Kimi K3's fused KDA projection + plain decode of R <= 8 requests of one token, at the TP16 rank slice, on the +plain-decode pools of a real ``MambaHybridCacheManagerV2`` (``_kda_cells.build_manager``: conv bf16 +``[slots, 2304, 3]``, fp32 state, strided by the manager's per-slot coalescing) at the slots the manager assigned. + +References, as the op's own test: the projection stream alone (``k3_kda_qkvg`` on a 104-CTA set, its Lamport buffers +decoded into rows), f_b as a bf16 ``F.linear``, then the ``ssm/kda_decode`` entry on a copy of the pools; and a +float64 decode on the same rows. Call sequences: layers x steps on one shared buffer set against a set per layer and +two sets alternating (bit for bit), one step captured and replayed, launches interleaved with ``ssm/k3_kda_attn`` on +one set, the per-CTA index across the int32 wrap; negative controls. + +Every launch here follows a plain copy kernel (its input), as each KDA launch in the model follows other layers' +kernels: the head CTAs read the slots' pools before their grid-dependency wait, so a launch must not directly follow +another launch on the same pools. +""" + +import _kda_cells as kc +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_attn import ( + k3_kda_attn, + k3_kda_qkvg, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_buffers import ( + CTAS, + FUSED_CTAS, + K3KdaBuffers, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_decode_attn import ( + k3_kda_decode_attn, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.kda_decode import kda_decode + +pytestmark = pytest.mark.skipif(not kc.sm100(), reason="needs SM100 (tcgen05, TMA, clusters)") + +LAYERS = 3 +STEPS = 3 + + +@pytest.fixture(scope="module") +def mgr(): + kc.load_ops() + m = kc.build_manager(LAYERS) + try: + yield m + finally: + m.shutdown() + + +@pytest.fixture(scope="module") +def slots(mgr): + return kc.request_slots(mgr, 8, first_id=200) + + +def _layers(mgr) -> list: + return [kc.layer_pools(mgr, layer) for layer in range(LAYERS)] + + +def _launch(wt, pools, x_src, slots, bufs) -> torch.Tensor: + """One launch, after a plain copy kernel (the input).""" + x = x_src.clone() + return k3_kda_decode_attn( + x, wt["w"], wt["w_fb"], wt["w_q"], wt["w_k"], wt["w_v"], wt["a_log"], wt["dt_bias"], wt["onorm_w"], + pools["conv"], pools["ssm"], slots, bufs, kc.LOWER_BOUND, kc.SCALE, kc.EPS, + ) # fmt: skip + + +class _Rows: + """The projection rows the fused kernel computes: the stream alone on a 104-CTA set, its buffers decoded.""" + + def __init__(self): + self.bufs = K3KdaBuffers.create("cuda", ctas=CTAS) + + def __call__(self, w, x) -> torch.Tensor: + buf = int(self.bufs.epoch[0].item()) + k3_kda_qkvg(x, w, self.bufs) + return kc.decode_rows(self.bufs.p1, self.bufs.part, buf, x.shape[0]) + + +def _native(wt, pools, rows, slots) -> torch.Tensor: + """The model's plain decode on those rows: f_b as a bf16 F.linear, then the ssm/kda_decode entry.""" + num, hk = rows.shape[0], kc.HK + + def heads(cols): + return cols.unflatten(-1, (kc.H, kc.K)).unsqueeze(0) + + g = F.linear(rows[:, 4 * hk : 4 * hk + kc.K], wt["w_fb"]) + out = torch.empty(num, 1, kc.H, kc.V, dtype=torch.bfloat16, device="cuda") + zeros = torch.zeros(hk, dtype=torch.bfloat16, device="cuda") + conv = pools["conv"] + kda_decode( + heads(rows[:, :hk]), heads(rows[:, hk : 2 * hk]), heads(rows[:, 2 * hk : 3 * hk]), + wt["w_t"][0], wt["w_t"][1], wt["w_t"][2], zeros, zeros, zeros, + conv[:, :hk], conv[:, hk : 2 * hk], conv[:, 2 * hk :], wt["a_log"], heads(g), wt["dt_bias"], + rows[:, 4 * hk + kc.K : 4 * hk + kc.K + kc.H].unsqueeze(0), heads(rows[:, 3 * hk : 4 * hk]), wt["onorm_w"], + slots, pools["ssm"], True, True, True, True, kc.LOWER_BOUND, kc.SCALE, kc.EPS, out, + ) # fmt: skip + return out.view(num, kc.H, kc.V) + + +def _x(num: int, gen: torch.Generator) -> torch.Tensor: + return torch.randn(num, kc.K_IN, generator=gen, device="cuda").bfloat16() + + +@pytest.mark.parametrize("num", range(1, 9)) +def test_decode_attn_on_manager_pools(mgr, slots, num) -> None: + """Every layer's step, all layers on one buffer set: outputs against the native path and float64, the state rows + against both (fp32 tolerance), the conv windows against the native path bit for bit; the other slots of the + layer and every slot of the other layers untouched.""" + used = slots[:num] + rows_idx = used.long() + pools = _layers(mgr) + with torch.inference_mode(): + for layer, p in enumerate(pools): + kc.fill_pools(p, 30 * num + layer) + bufs = K3KdaBuffers.create("cuda") + rows_of = _Rows() + gen = torch.Generator(device="cuda").manual_seed(40 + num) + for layer, p in enumerate(pools): + wt = kc.make_weights(100 + layer) + native = kc.clone_pools(p) + ref = kc.F64Decode(wt, p) + before = [kc.snapshot(q) for q in pools] + x = _x(num, gen) + out = _launch(wt, p, x, used, bufs) + rows = rows_of(wt["w"], x) + out_n = _native(wt, native, rows, used) + out_r = ref.step_rows(rows, used) + torch.cuda.synchronize() + assert kc.rel(out, out_n) <= kc.TOL_OUT, (layer, kc.rel(out, out_n)) + assert kc.rel(out, out_r) <= kc.TOL_OUT, (layer, kc.rel(out, out_r)) + assert kc.rel(p["ssm"][rows_idx], ref.state[rows_idx]) <= kc.TOL_STATE, layer + assert kc.rel(p["ssm"][rows_idx], native["ssm"][rows_idx]) <= kc.TOL_STATE, layer + assert torch.equal(p["conv"][rows_idx], native["conv"][rows_idx]), layer + others = kc.other_slots(p["ssm"].shape[0], used.tolist()) + assert kc.same_rows(p, before[layer], others, ("conv", "ssm")), layer + for o, q in enumerate(pools): + if o != layer: + assert all(torch.equal(q[n], before[o][n]) for n in ("conv", "ssm")), (layer, o) + assert bool((bufs.epoch == LAYERS % 3).all()), bufs.epoch.unique().tolist() + + +def _schedule(gen, num) -> list: + return [[_x(num, gen) for _ in range(LAYERS)] for _ in range(STEPS)] + + +def test_steps_share_one_buffer_set(mgr, slots) -> None: + """LAYERS layers x STEPS steps on four requests: one set for every launch (the model's layout), a set per layer, + and two sets alternating between launches give the same outputs and pools bit for bit; on the shared set every + CTA's index has moved once per launch.""" + used = slots[:4] + pools = _layers(mgr) + with torch.inference_mode(): + for layer, p in enumerate(pools): + kc.fill_pools(p, 700 + layer) + own_pools = [kc.clone_pools(p) for p in pools] + alt_pools = [kc.clone_pools(p) for p in pools] + wts = [kc.make_weights(110 + layer) for layer in range(LAYERS)] + xs = _schedule(torch.Generator(device="cuda").manual_seed(41), 4) + shared = K3KdaBuffers.create("cuda") + own = [K3KdaBuffers.create("cuda") for _ in range(LAYERS)] + alt = [K3KdaBuffers.create("cuda") for _ in range(2)] + got, want, both = [], [], [] + for s in range(STEPS): + for layer in range(LAYERS): + got.append(_launch(wts[layer], pools[layer], xs[s][layer], used, shared).clone()) + want.append( + _launch(wts[layer], own_pools[layer], xs[s][layer], used, own[layer]).clone() + ) + pick = alt[(s * LAYERS + layer) % 2] + both.append(_launch(wts[layer], alt_pools[layer], xs[s][layer], used, pick).clone()) + torch.cuda.synchronize() + assert all(torch.equal(a, b) for a, b in zip(got, want)) + assert all(torch.equal(a, b) for a, b in zip(both, want)) + for p, q, r in zip(pools, own_pools, alt_pools): + assert all(torch.equal(p[n], q[n]) and torch.equal(r[n], q[n]) for n in ("conv", "ssm")) + assert bool((shared.epoch == (STEPS * LAYERS) % 3).all()), shared.epoch.unique().tolist() + + +def test_graph_replay(mgr, slots) -> None: + """One step of every layer captured once (the inputs rewritten in place before every replay) on the shared set: + STEPS replays give the bits of the same steps run eagerly on a copy of the pools and a set of their own.""" + used = slots[:3] + pools = _layers(mgr) + with torch.inference_mode(): + for layer, p in enumerate(pools): + kc.fill_pools(p, 800 + layer) + copies = [kc.clone_pools(p) for p in pools] + wts = [kc.make_weights(120 + layer) for layer in range(LAYERS)] + xs = _schedule(torch.Generator(device="cuda").manual_seed(42), 3) + eager_bufs = K3KdaBuffers.create("cuda") + want = [[_launch(wts[li], copies[li], xs[s][li], used, eager_bufs).clone() for li in range(LAYERS)] + for s in range(STEPS)] # fmt: skip + # Compiles outside capture, on scratch pools and a scratch set. + _launch(wts[0], kc.clone_pools(pools[0]), xs[0][0], used, K3KdaBuffers.create("cuda")) + bufs = K3KdaBuffers.create("cuda") + static = [x.clone() for x in xs[0]] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + outs = [_launch(wts[li], pools[li], static[li], used, bufs) for li in range(LAYERS)] + torch.cuda.current_stream().wait_stream(stream) + got = [] + for s in range(STEPS): + for layer in range(LAYERS): + static[layer].copy_(xs[s][layer]) + graph.replay() + got.append([o.clone() for o in outs]) + torch.cuda.synchronize() + assert all(torch.equal(a, b) for s in range(STEPS) for a, b in zip(got[s], want[s])) + for p, c in zip(pools, copies): + assert all(torch.equal(p[n], c[n]) for n in ("conv", "ssm")) + + +@pytest.fixture(scope="module") +def spec_mgr(): + kc.load_ops() + m = kc.build_manager(1, num_spec=kc.NUM_SPEC) + try: + yield m + finally: + m.shutdown() + + +def test_shared_with_k3_kda_attn(mgr, slots, spec_mgr) -> None: + """Plain-decode launches and ``ssm/k3_kda_attn`` verify launches interleaved in stream order on one set (the + model's single set per device) give the bits of the same launches on sets of their own, one at a time. Both run + the same stream role and move every CTA's index once per launch, so either may follow the other.""" + vslot = kc.request_slots(spec_mgr, 1, first_id=300) + dpools = kc.layer_pools(mgr, 0) + vpools = kc.layer_pools(spec_mgr, 0) + with torch.inference_mode(): + kc.fill_pools(dpools, 900) + kc.fill_pools(vpools, 901) + own_d, own_v = kc.clone_pools(dpools), kc.clone_pools(vpools) + wt = kc.make_weights(130) + gen = torch.Generator(device="cuda").manual_seed(43) + calls = [] + for num in (1, 4, 8, 3): + calls.append(("decode", _x(num, gen), slots[:num])) + calls.append(("verify", _x(kc.NT, gen), vslot)) + + def verify(pools, x, slot, bufs): + return k3_kda_attn( + x.clone(), wt["w"], wt["w_fb"], wt["w_q"], wt["w_k"], wt["w_v"], wt["a_log"], wt["dt_bias"], + wt["onorm_w"], pools["cs_q"], pools["cs_k"], pools["cs_v"], pools["ssm"], pools["state_tok"], slot, + pools["pending"], bufs, kc.NUM_SPEC, kc.LOWER_BOUND, kc.SCALE, kc.EPS, + ) # fmt: skip + + own_bufs = {"decode": K3KdaBuffers.create("cuda"), "verify": K3KdaBuffers.create("cuda")} + want = [] + for kind, x, s in calls: + if kind == "decode": + want.append(_launch(wt, own_d, x, s, own_bufs[kind]).clone()) + else: + want.append(verify(own_v, x, s, own_bufs[kind]).clone()) + torch.cuda.synchronize() + shared = K3KdaBuffers.create("cuda") + got = [ + ( + _launch(wt, dpools, x, s, shared) + if kind == "decode" + else verify(vpools, x, s, shared) + ).clone() + for kind, x, s in calls + ] + torch.cuda.synchronize() + assert all(torch.equal(a, b) for a, b in zip(got, want)) + assert all(torch.equal(dpools[n], own_d[n]) for n in ("conv", "ssm")) + assert all( + torch.equal(vpools[n], own_v[n]) for n in ("cs_q", "cs_k", "cs_v", "ssm", "state_tok") + ) + assert bool((shared.epoch == len(calls) % 3).all()), shared.epoch.unique().tolist() + + +def test_epoch_wrap_on_the_object(mgr, slots) -> None: + """The set's per-CTA indices preset to 2^31 - 2, as after 2^31 launches on a device: four launches give the bits of + a fresh set's and leave every index in 0..2 (the kernel keeps the index, not a raw count).""" + used = slots[:2] + p = kc.layer_pools(mgr, 0) + with torch.inference_mode(): + kc.fill_pools(p, 950) + fresh_pools = kc.clone_pools(p) + wt = kc.make_weights(140) + gen = torch.Generator(device="cuda").manual_seed(44) + xs = [_x(2, gen) for _ in range(4)] + fresh = K3KdaBuffers.create("cuda") + want = [] + for x in xs: + want.append(_launch(wt, fresh_pools, x, used, fresh).clone()) + torch.cuda.synchronize() # the same pools again next: one launch at a time + wrapped = K3KdaBuffers.create("cuda") + wrapped.epoch.fill_(2**31 - 2) + got = [] + for x in xs: + got.append(_launch(wt, p, x, used, wrapped).clone()) + torch.cuda.synchronize() + assert all(torch.equal(a, b) for a, b in zip(got, want)) + assert all(torch.equal(p[n], fresh_pools[n]) for n in ("conv", "ssm")) + assert bool(((wrapped.epoch >= 0) & (wrapped.epoch < 3)).all()), wrapped.epoch.unique().tolist() + + +def test_swapped_steps_are_silently_wrong(mgr, slots) -> None: + """Negative control: two steps of one request in swapped order raise nothing and give the second step's output and + the final state of a different history; the pools carry the order, and nothing in a call can detect it.""" + used = slots[:1] + rows_idx = used.long() + p = kc.layer_pools(mgr, 0) + with torch.inference_mode(): + kc.fill_pools(p, 960) + swapped = kc.clone_pools(p) + wt = kc.make_weights(150) + gen = torch.Generator(device="cuda").manual_seed(45) + a, b = _x(1, gen), _x(1, gen) + bufs = K3KdaBuffers.create("cuda") + _launch(wt, p, a, used, bufs) + torch.cuda.synchronize() + right = _launch(wt, p, b, used, bufs).clone() + torch.cuda.synchronize() + _launch(wt, swapped, b, used, bufs) + torch.cuda.synchronize() + wrong = _launch(wt, swapped, a, used, bufs).clone() + torch.cuda.synchronize() + print(f"swapped steps: second output rel diff {kc.rel(wrong, right):.3e}, " + f"state rel diff {kc.rel(swapped['ssm'][rows_idx], p['ssm'][rows_idx]):.3e}") # fmt: skip + assert kc.rel(swapped["ssm"][rows_idx], p["ssm"][rows_idx]) > kc.TOL_STATE + assert not torch.equal(wrong, right) + + +def test_buffer_set_misuse_raises(mgr, slots) -> None: + """A set made for k3_kda_qkvg's 104 CTAs is refused before anything is written; a set cannot be made under + CUDA-graph capture (it allocates); a set for any other CTA count cannot be made.""" + used = slots[:2] + p = kc.layer_pools(mgr, 0) + with torch.inference_mode(): + kc.fill_pools(p, 970) + before = kc.snapshot(p) + wt = kc.make_weights(160) + x = _x(2, torch.Generator(device="cuda").manual_seed(46)) + with pytest.raises(ValueError, match=f"epoch {CTAS}"): + _launch(wt, p, x, used, K3KdaBuffers.create("cuda", ctas=CTAS)) + torch.cuda.synchronize() + assert all(torch.equal(p[n], before[n]) for n in ("conv", "ssm")) + with pytest.raises(ValueError, match="ctas must be"): + K3KdaBuffers.create("cuda", ctas=FUSED_CTAS - 1) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with pytest.raises(RuntimeError, match="before CUDA-graph capture"): + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + K3KdaBuffers.create("cuda") + torch.cuda.current_stream().wait_stream(stream) diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_verify.py b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_verify.py new file mode 100644 index 000000000000..247ed41d29f5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_kda_verify.py @@ -0,0 +1,202 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the ssm/k3_kda_verify catalog entry, on a real cache object. + +Kimi K3's KDA speculative verify of R requests of 1 + num_spec tokens, at the TP16 rank slice (6 heads, K = V = 128, +conv width 4), on the verify pools of a real ``MambaHybridCacheManagerV2`` built with MTP-style speculation, the KDA +replay caches and per-token states (``_kda_cells.build_manager``): every layer's conv caches (fp32, dim-contiguous), +its fp32 state (strided by the manager's per-slot coalescing) and per-draft states, and the accepted-draft record all +layers share (``prev_num_accepted_tokens``), at the slots the manager assigned. + +The reference is a float64 verify over each request's committed history (``_kda_cells.F64Verify``): every round's +outputs and the state committed after each golden token against it; the per-draft states and the conv caches are +checked through the next round, which starts from the drafts the sampler accepted. Call sequences: layers x rounds, +a captured round replayed against the same rounds run eagerly, the schedule twice; negative controls. + +Every launch here follows a plain copy kernel (its input), as each KDA launch in the model follows other layers' +kernels: the head CTAs read the slots' pools before their grid-dependency wait. +""" + +import _kda_cells as kc +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_verify import k3_kda_verify + +pytestmark = pytest.mark.skipif(not kc.sm100(), reason="needs SM100 (tcgen05, TMA, clusters)") + +LAYERS = 2 +ROUNDS = 6 +POOL_NAMES = ("cs_q", "cs_k", "cs_v", "ssm", "state_tok") +# (requests, 1 + num_spec): the model's 7 drafts, and a 3-draft manager. +CELLS = [(1, 8), (4, 8), (8, 8), (4, 4)] + + +@pytest.fixture(scope="module") +def managers(): + kc.load_ops() + built = {} + try: + for num_spec in sorted({steps - 1 for _, steps in CELLS}): + mgr = kc.build_manager(LAYERS, num_spec=num_spec) + built[num_spec] = (mgr, kc.request_slots(mgr, 8, first_id=600 + 10 * num_spec)) + yield built + finally: + for mgr, _ in built.values(): + mgr.shutdown() + + +def _layers(mgr) -> list: + return [kc.layer_pools(mgr, layer) for layer in range(LAYERS)] + + +def _verify(wt, pools, proj_src, slots, num_spec) -> torch.Tensor: + """One launch, after a plain copy kernel (its projection rows).""" + proj = proj_src.clone() + return k3_kda_verify( + proj, wt["w_fb"], wt["w_q"], wt["w_k"], wt["w_v"], wt["a_log"], wt["dt_bias"], wt["onorm_w"], + pools["cs_q"], pools["cs_k"], pools["cs_v"], pools["ssm"], pools["state_tok"], slots, pools["pending"], + num_spec, kc.LOWER_BOUND, kc.SCALE, kc.EPS, + ) # fmt: skip + + +def _proj(num: int, steps: int, gen: torch.Generator) -> torch.Tensor: + return torch.randn(num * steps, kc.PROJ, generator=gen, device="cuda").bfloat16() + + +def _accept(pools_list, slots, accepted) -> None: + for pools in pools_list: + pools["pending"][slots.long()] = torch.tensor(accepted, dtype=torch.int32, device="cuda") + + +@pytest.mark.parametrize("num,steps", CELLS, ids=[f"{r}x{t}" for r, t in CELLS]) +def test_verify_on_manager_pools(managers, num, steps) -> None: + """ROUNDS rounds of every layer with a pending count per request that changes every round: outputs and the + committed states against float64 (fp32 tolerance); the other slots of each layer untouched.""" + num_spec = steps - 1 + mgr, all_slots = managers[num_spec] + slots = all_slots[:num] + pools = _layers(mgr) + with torch.inference_mode(): + for layer, p in enumerate(pools): + kc.fill_pools(p, 2000 + 10 * num + layer) + wts = [kc.make_weights(600 + layer) for layer in range(LAYERS)] + refs = [kc.F64Verify(wts[layer], p, slots, num_spec) for layer, p in enumerate(pools)] + gen = torch.Generator(device="cuda").manual_seed(60 + num) + for rnd in range(ROUNDS): + for layer, p in enumerate(pools): + before = kc.snapshot(p) + proj = _proj(num, steps, gen) + out = _verify(wts[layer], p, proj, slots, num_spec) + want = refs[layer](proj) + torch.cuda.synchronize() + assert kc.rel(out, want) <= kc.TOL_OUT, (rnd, layer, kc.rel(out, want)) + for n, s in enumerate(slots.tolist()): + committed = refs[layer].last[n][0][0] # the state after the golden token + assert kc.rel(p["ssm"][s], committed) <= kc.TOL_STATE, (rnd, layer, n) + others = kc.other_slots(p["ssm"].shape[0], slots.tolist()) + assert kc.same_rows(p, before, others, POOL_NAMES), (rnd, layer) + accepted = kc.pending_schedule(num, num_spec, rnd) + _accept([pools[0]], slots, accepted) + for ref in refs: + ref.commit(accepted) + + +def test_layers_rounds_replay_and_repeat(managers) -> None: + """Four requests, LAYERS layers x ROUNDS rounds: one round of every layer captured once and replayed (the rows + rewritten in place, the record written between replays), and the same schedule run eagerly twice from the same + pools, give the same outputs and pools bit for bit.""" + mgr, all_slots = managers[kc.NUM_SPEC] + slots = all_slots[:4] + steps = kc.NT + pools = _layers(mgr) + with torch.inference_mode(): + for layer, p in enumerate(pools): + kc.fill_pools(p, 2100 + layer) + again = [kc.clone_pools(p) for p in pools] + graph_pools = [kc.clone_pools(p) for p in pools] + for copies in (again, graph_pools): + for c in copies[1:]: + c["pending"] = copies[0]["pending"] + wts = [kc.make_weights(610 + layer) for layer in range(LAYERS)] + gen = torch.Generator(device="cuda").manual_seed(61) + projs = [[_proj(4, steps, gen) for _ in range(LAYERS)] for _ in range(ROUNDS)] + first, second = [], [] + for rnd in range(ROUNDS): + for layer in range(LAYERS): + first.append( + _verify(wts[layer], pools[layer], projs[rnd][layer], slots, kc.NUM_SPEC).clone() + ) + second.append( + _verify(wts[layer], again[layer], projs[rnd][layer], slots, kc.NUM_SPEC).clone() + ) + accepted = kc.pending_schedule(4, kc.NUM_SPEC, rnd) + _accept([pools[0], again[0]], slots, accepted) + torch.cuda.synchronize() + static = [p.clone() for p in projs[0]] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + outs = [ + _verify(wts[li], graph_pools[li], static[li], slots, kc.NUM_SPEC) + for li in range(LAYERS) + ] + torch.cuda.current_stream().wait_stream(stream) + replayed = [] + for rnd in range(ROUNDS): + for layer in range(LAYERS): + static[layer].copy_(projs[rnd][layer]) + graph.replay() + replayed.extend(o.clone() for o in outs) + _accept([graph_pools[0]], slots, kc.pending_schedule(4, kc.NUM_SPEC, rnd)) + torch.cuda.synchronize() + assert all(torch.equal(a, b) for a, b in zip(first, second)), "the schedule twice differs" + assert all(torch.equal(a, b) for a, b in zip(replayed, first)) + for p, q, r in zip(pools, again, graph_pools): + assert all(torch.equal(p[n], q[n]) and torch.equal(r[n], q[n]) for n in POOL_NAMES) + + +def test_swapped_rounds_are_silently_wrong(managers) -> None: + """Negative control: two rounds of one request in swapped order raise nothing and give the second round's outputs + and the slot's state of a different history.""" + mgr, all_slots = managers[kc.NUM_SPEC] + slots = all_slots[:1] + s = int(slots[0]) + p = kc.layer_pools(mgr, 0) + with torch.inference_mode(): + kc.fill_pools(p, 2200) + swapped = kc.clone_pools(p) + wt = kc.make_weights(620) + gen = torch.Generator(device="cuda").manual_seed(62) + a, b = _proj(1, kc.NT, gen), _proj(1, kc.NT, gen) + _verify(wt, p, a, slots, kc.NUM_SPEC) + _accept([p], slots, [3]) + right = _verify(wt, p, b, slots, kc.NUM_SPEC).clone() + _verify(wt, swapped, b, slots, kc.NUM_SPEC) + _accept([swapped], slots, [3]) + wrong = _verify(wt, swapped, a, slots, kc.NUM_SPEC).clone() + torch.cuda.synchronize() + print(f"swapped rounds: second output rel diff {kc.rel(wrong, right):.3e}, " + f"state rel diff {kc.rel(swapped['ssm'][s], p['ssm'][s]):.3e}") # fmt: skip + assert kc.rel(swapped["ssm"][s], p["ssm"][s]) > kc.TOL_STATE + assert not torch.equal(wrong, right) + + +def test_rejects_out_of_contract(managers) -> None: + """Per-draft states for another draft count, and an int64 record, are refused before anything is written.""" + mgr, all_slots = managers[kc.NUM_SPEC] + slots = all_slots[:2] + p = kc.layer_pools(mgr, 0) + with torch.inference_mode(): + kc.fill_pools(p, 2300) + before = kc.snapshot(p) + wt = kc.make_weights(630) + proj = _proj(2, kc.NT, torch.Generator(device="cuda").manual_seed(63)) + short = dict(p, state_tok=p["state_tok"][:, : kc.NUM_SPEC - 1].contiguous()) + with pytest.raises(ValueError, match="unsupported call"): + _verify(wt, short, proj, slots, kc.NUM_SPEC) + with pytest.raises(ValueError, match="unsupported call"): + _verify(wt, dict(p, pending=p["pending"].long()), proj, slots, kc.NUM_SPEC) + torch.cuda.synchronize() + assert all(torch.equal(p[n], before[n]) for n in POOL_NAMES) diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_decode.py b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_decode.py new file mode 100644 index 000000000000..035eb7ee3869 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_kda_decode.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the ssm/kda_decode catalog entry, on a real cache object. + +Kimi K3's plain decode shape (6 heads, K = V = 128, conv width 4, the lower-bound gate, the output norm, beta sigmoid +in the kernel) on the pools of a real ``MambaHybridCacheManagerV2`` (``_kda_cells.build_manager``): its per-layer +conv views (``[q | k | v]`` sections of one bf16 ``[slots, 2304, 3]`` pool) and fp32 state views, strided by the +manager's per-slot coalescing, at the slots the manager assigned. The reference is a float64 torch decode +(``_kda_cells.F64Decode``). B runs 1..8: on sm_100 B <= 5 selects the four-CTA cluster kernel and 6 <= B <= 24 the +legacy compact-heads kernel (the one #19830 fixes). +""" + +import _kda_cells as kc +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.kda_decode import kda_decode + +assert torch.cuda.is_available(), "kda_decode requires a CUDA device" +pytestmark = pytest.mark.skipif(not kc.sm100(), reason="certified on SM100 only") + +LAYERS = 3 +STEPS = 3 + + +@pytest.fixture(scope="module") +def mgr(): + kc.load_ops() + m = kc.build_manager(LAYERS) + try: + yield m + finally: + m.shutdown() + + +@pytest.fixture(scope="module") +def slots(mgr): + return kc.request_slots(mgr, 8, first_id=100) + + +def _heads(cols: torch.Tensor) -> torch.Tensor: + return cols.unflatten(-1, (kc.H, kc.K)).unsqueeze(0) + + +def _inputs(num: int, gen: torch.Generator) -> dict: + def rnd(*s): + return torch.randn(*s, generator=gen, device="cuda").bfloat16() + + return { + "raw": rnd(num, 3 * kc.HK), + "g": rnd(num, kc.HK), + "beta": rnd(num, kc.H), + "gate": rnd(num, kc.HK), + } + + +def _call(wt, pools, inp, slots, out=None, use_lower_bound=True) -> torch.Tensor: + """The entry on one layer's pools, as the model calls it: packed conv sections updated in place, indexed state.""" + num = slots.numel() + raw, conv = inp["raw"], pools["conv"] + zeros = torch.zeros(kc.HK, dtype=torch.bfloat16, device="cuda") + if out is None: + out = torch.empty(num, 1, kc.H, kc.V, dtype=torch.bfloat16, device="cuda") + kda_decode( + _heads(raw[:, : kc.HK]), _heads(raw[:, kc.HK : 2 * kc.HK]), _heads(raw[:, 2 * kc.HK :]), + wt["w_t"][0], wt["w_t"][1], wt["w_t"][2], zeros, zeros, zeros, + conv[:, : kc.HK], conv[:, kc.HK : 2 * kc.HK], conv[:, 2 * kc.HK :], + wt["a_log"], _heads(inp["g"]), wt["dt_bias"], inp["beta"].unsqueeze(0), _heads(inp["gate"]), + wt["onorm_w"], slots, pools["ssm"], True, True, use_lower_bound, True, kc.LOWER_BOUND, kc.SCALE, kc.EPS, out, + ) # fmt: skip + return out.view(num, kc.H, kc.V) + + +def _layers(mgr) -> list: + return [kc.layer_pools(mgr, layer) for layer in range(LAYERS)] + + +@pytest.mark.parametrize("num", range(1, 9)) +def test_decode_on_manager_pools(mgr, slots, num) -> None: + """Every layer's step against float64: outputs and the state rows at fp32 tolerance, the conv windows bit for + bit (raw bf16 inputs), the other slots of the layer and every slot of the other layers untouched.""" + used = slots[:num] + pools = _layers(mgr) + with torch.inference_mode(): + for layer, p in enumerate(pools): + kc.fill_pools(p, 10 * num + layer) + gen = torch.Generator(device="cuda").manual_seed(num) + for layer, p in enumerate(pools): + wt = kc.make_weights(100 + layer) + ref = kc.F64Decode(wt, p) + before = [kc.snapshot(q) for q in pools] + inp = _inputs(num, gen) + out = _call(wt, p, inp, used) + want = ref.step(inp["raw"], inp["g"], inp["beta"], inp["gate"], used) + torch.cuda.synchronize() + rows = used.long() + assert kc.rel(out, want) <= kc.TOL_OUT, (layer, kc.rel(out, want)) + assert kc.rel(p["ssm"][rows], ref.state[rows]) <= kc.TOL_STATE, layer + assert torch.equal(p["conv"][rows], ref.conv[rows].bfloat16()), layer + others = kc.other_slots(p["ssm"].shape[0], used.tolist()) + assert kc.same_rows(p, before[layer], others, ("conv", "ssm")), layer + for o, q in enumerate(pools): + if o != layer: + assert all(torch.equal(q[n], before[o][n]) for n in ("conv", "ssm")), (layer, o) + + +def test_steps_on_layers_and_graph_replay(mgr, slots) -> None: + """LAYERS layers x STEPS decode steps on four requests: the steps captured once in a CUDA graph (inputs and slots + rewritten in place before each replay) give the bits of the same steps run eagerly on a copy of the pools.""" + used = slots[:4] + pools = _layers(mgr) + with torch.inference_mode(): + for layer, p in enumerate(pools): + kc.fill_pools(p, 500 + layer) + copies = [kc.clone_pools(p) for p in pools] + wts = [kc.make_weights(200 + layer) for layer in range(LAYERS)] + gen = torch.Generator(device="cuda").manual_seed(7) + steps = [[_inputs(4, gen) for _ in range(LAYERS)] for _ in range(STEPS)] + want = [ + [_call(wts[li], copies[li], steps[s][li], used).clone() for li in range(LAYERS)] + for s in range(STEPS) + ] + static = [{n: t.clone() for n, t in steps[0][li].items()} for li in range(LAYERS)] + s_in = used.clone() + outs = [ + torch.empty(4, 1, kc.H, kc.V, dtype=torch.bfloat16, device="cuda") + for _ in range(LAYERS) + ] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + for li in range(LAYERS): + _call(wts[li], pools[li], static[li], s_in, outs[li]) + torch.cuda.current_stream().wait_stream(stream) + got = [] + for s in range(STEPS): + for li in range(LAYERS): + for n, t in static[li].items(): + t.copy_(steps[s][li][n]) + graph.replay() + got.append([o.view(4, kc.H, kc.V).clone() for o in outs]) + torch.cuda.synchronize() + assert all(torch.equal(a, b) for s in range(STEPS) for a, b in zip(got[s], want[s])) + for p, c in zip(pools, copies): + assert torch.equal(p["ssm"], c["ssm"]) and torch.equal(p["conv"], c["conv"]) + + +def test_swapped_steps_are_silently_wrong(mgr, slots) -> None: + """Negative control: two decode steps of one request applied in swapped order raise nothing and give the second + step's output and the final state of a different history (the pools carry the order; nothing in the call can + detect it).""" + used = slots[:1] + p = kc.layer_pools(mgr, 0) + with torch.inference_mode(): + kc.fill_pools(p, 900) + start = kc.clone_pools(p) + wt = kc.make_weights(300) + gen = torch.Generator(device="cuda").manual_seed(11) + a, b = _inputs(1, gen), _inputs(1, gen) + _call(wt, p, a, used) + right = _call(wt, p, b, used).clone() + right_state = p["ssm"][used.long()].clone() + swapped = kc.clone_pools(start) + _call(wt, swapped, b, used) + wrong = _call(wt, swapped, a, used).clone() + torch.cuda.synchronize() + rows = used.long() + print(f"swapped steps: second output rel diff {kc.rel(wrong, right):.3e}, " + f"state rel diff {kc.rel(swapped['ssm'][rows], right_state):.3e}") # fmt: skip + assert kc.rel(swapped["ssm"][rows], right_state) > kc.TOL_STATE + assert not torch.equal(wrong, right) + + +def test_rejects_out_of_contract(mgr, slots) -> None: + """A state view that is not 16-byte aligned, int64 slot indices, and the gate without its lower bound (the op + supports only apply_onorm, use_lower_bound and apply_beta_sigmoid all on): each raises before the pools are + written.""" + used = slots[:2] + p = kc.layer_pools(mgr, 0) + with torch.inference_mode(): + kc.fill_pools(p, 950) + before = kc.snapshot(p) + wt = kc.make_weights(400) + inp = _inputs(2, torch.Generator(device="cuda").manual_seed(13)) + ssm = p["ssm"] + shifted = dict(p, ssm=ssm.as_strided(ssm.shape, ssm.stride(), ssm.storage_offset() + 1)) + with pytest.raises(RuntimeError, match="16B-aligned"): + _call(wt, shifted, inp, used) + with pytest.raises(RuntimeError, match="int32"): + _call(wt, p, inp, used.long()) + with pytest.raises( + RuntimeError, match="only supports apply_onorm=true, use_lower_bound=true" + ): + _call(wt, p, inp, used, use_lower_bound=False) + torch.cuda.synchronize() + assert all(torch.equal(p[n], before[n]) for n in ("conv", "ssm")) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py index 1f7e0435d5b3..b62a1dad85fe 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py @@ -17,6 +17,7 @@ from __future__ import annotations +import ast import re from pathlib import Path @@ -215,3 +216,73 @@ def test_targets_do_not_share_files(): f"directory for {tail!r}; targets may import the " f"catalog and nothing else" ) + + +#: tensorrt_llm imports that compute nothing, so a target's generic path need not declare them: types and configs +#: it reads, enums it passes, the metadata it is handed and the views it takes of it, the graph-capture context it +#: asks about, its registration and its logger. +_NON_COMPUTE_IMPORTS = frozenset( + { + "AllReduceStrategy", + "AttentionMetadata", + "AuxStreamType", + "MambaHybridCacheManagerV2", + "Mapping", + "ModelConfig", + "MoESchedulerKind", + "QuantAlgo", + "QuantConfig", + "SiTuActivation", + "is_in_breakable_cuda_graph", + "k3_mla_decode_view", + "logger", + "register_auto_model", + } +) + + +def _declared_tuple(tree: ast.Module, name: str): + """The string tuple assigned to ``name`` at module level, or None when the module does not assign it.""" + for node in tree.body: + if isinstance(node, ast.Assign) and any( + isinstance(target, ast.Name) and target.id == name for target in node.targets + ): + return tuple(ast.literal_eval(node.value)) + return None + + +def test_uncertified_generic_calls_name_every_stock_import(): + """A target whose generic path runs stock code (``UNCERTIFIED_GENERIC_CALLS``) declares every tensorrt_llm + name it imports outside the catalog, module level or function local, unless the name computes nothing; and it + declares nothing it does not import, so the list cannot go stale as entries replace stock calls.""" + checked = 0 + for arch in _ARCHS: + routing = routing_module(arch) + for name, dotted in routing.TARGET_MODULES.items(): + tree = ast.parse(_module_path(dotted).read_text()) + declared = _declared_tuple(tree, "UNCERTIFIED_GENERIC_CALLS") + if declared is None: + continue + checked += 1 + imported = set() + for node in ast.walk(tree): + if ( + isinstance(node, ast.ImportFrom) + and node.level == 0 + and node.module is not None + and node.module.split(".")[0] == "tensorrt_llm" + and ".modeling_v2.catalog" not in node.module + ): + imported.update( + f"{node.module}.{alias.name}" + for alias in node.names + if alias.name not in _NON_COMPUTE_IMPORTS + ) + assert imported <= set(declared), ( + f"{name}: imported but not in UNCERTIFIED_GENERIC_CALLS: {sorted(imported - set(declared))}" + ) + assert set(declared) <= imported, ( + f"{name}: UNCERTIFIED_GENERIC_CALLS names what it does not import: " + f"{sorted(set(declared) - imported)}" + ) + assert checked, "no target declares UNCERTIFIED_GENERIC_CALLS" diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_gemv.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_gemv.py new file mode 100644 index 000000000000..f87c8c6a6832 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_gemv.py @@ -0,0 +1,315 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Kimi K3 target's decode GEMVs, LM head and embedding on the catalog's single-GPU entries (``decode_gemv.py`` +of ``kimi_k3_mxfp4__sm_100__tp16_moetp4ep4``), on one GPU. + +* Each GEMV site at every row count it takes (1..8, and 9..64 where k3_ctm_gemv_wide takes it): the bits of its + catalog entry's call, within 8e-3 of ``max |ref|`` of a float64 product (the sigmoid columns within 1e-2 of the + sigmoid of it); declined above its rows, at another shape or dtype, and under capture before an eager call. +* The LM head through ``K3LogitsProcessor`` on a real ``LMHead`` (one rank): fp32 logits from ``k3_head_gemv`` + within the same bound, the stock processor's rows selected, the stock path above 8 rows and before the state is + built. +* The embedding: bit-identical to ``nn.Embedding`` followed by the stock RMSNorm, the rows in the bank's slot 0. +* Layer 0's dense MLP at 1..8 rows: the bits of its three entries chained, within 3e-2 of ``max |ref|`` of the + float64 SituAndMul MLP; declined above 8 rows, at linear_beta 0 and under capture before an eager call. +* CUDA-graph replays give the eager bits. +""" + +import types + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.k3_situ_mul import k3_situ_mul +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_long import ( + k3_ctm_gemv_long, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_wide import ( + k3_ctm_gemv_wide, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_decode_gemv import k3_decode_gemv +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4 import ( # noqa: E501 + decode_gemv, +) +from tensorrt_llm._torch.modules.embedding import LMHead +from tensorrt_llm._torch.modules.linear import TensorParallelMode +from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm.mapping import Mapping + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), + reason="the K3 decode kernels run on sm_100 only", +) + +TOL = 8e-3 # max |y - ref| / max |ref|, ref a float64 product +VOCAB_SHARD, HIDDEN = 10240, 7168 +ROWS = range(1, 9) + + +def _weight(n, k, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(n, k, generator=g, device="cuda") * 0.02).to(torch.bfloat16) + + +def _rows(m, k, seed): + g = torch.Generator(device="cuda").manual_seed(1000 + seed) + return torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + + +def _check_product(y, x, w, sig_col0=-1): + """Within TOL of the float64 product (columns from ``sig_col0`` on: within 1e-2 of its sigmoid).""" + ref = x.double() @ w.double().t() + sig = ref.shape[1] if sig_col0 < 0 else sig_col0 + lin = ((y[:, :sig].double() - ref[:, :sig]).abs().max() / ref[:, :sig].abs().max()).item() + assert lin <= TOL, lin + if sig < ref.shape[1]: + err = (y[:, sig:].double() - ref[:, sig:].sigmoid()).abs().max().item() + assert err <= 1e-2, err + + +def _rel_err(y, x, w): + ref = x.double() @ w.double().t() + return ((y.double() - ref).abs().max() / ref.abs().max()).item() + + +def _bits(t): + return t.contiguous().view(torch.int16) + + +def _entry(spec, rows, x, w): + """The site's catalog entry called directly: what ``project`` must reproduce bit for bit.""" + if rows > decode_gemv.MAX_ROWS or spec.small == "wide": + return k3_ctm_gemv_wide(x, w, sig_col0=spec.sig_col0) + if spec.small == "decode": + return k3_decode_gemv(x, w) + return k3_ctm_gemv_long( + x, w, sig_col0=spec.sig_col0, split=spec.split, ring=spec.ring, push=spec.push + ) + + +def _site_rows(spec): + return list(ROWS) + ([9, 16, 24, 40, 64] if spec.wide else []) + + +@pytest.mark.parametrize("site", list(decode_gemv.SITES)) +def test_site(site): + """Every row count the site takes runs its catalog entry: its bits, within TOL of the float64 product.""" + spec = decode_gemv.SITES[site] + w = _weight(spec.n, spec.k, seed=len(site)) + gemvs = decode_gemv.K3DecodeGemvs.create(None, sites=[site]) + for m in _site_rows(spec): + x = _rows(m, spec.k, seed=m) + y = gemvs.project(site, x, w) + assert y is not None and y.shape == (m, spec.n), (site, m) + assert torch.equal(_bits(y), _bits(_entry(spec, m, x, w))), (site, m) + _check_product(y, x, w, spec.sig_col0) + + +@pytest.mark.parametrize("site", ["kv_a", "mla_ag"]) +def test_site_declines(site): + """More rows than the site takes, another weight shape, a non-bf16 input, and under capture before any eager + call: None, nothing launched.""" + spec = decode_gemv.SITES[site] + w = _weight(spec.n, spec.k, seed=1) + gemvs = decode_gemv.K3DecodeGemvs.create(None, sites=[site]) + too_many = decode_gemv.WIDE_ROWS + 1 if spec.wide else decode_gemv.MAX_ROWS + 1 + assert gemvs.project(site, _rows(too_many, spec.k, seed=9), w) is None + assert ( + gemvs.project(site, _rows(4, spec.k, seed=4), _weight(spec.n + 128, spec.k, seed=2)) is None + ) + assert gemvs.project(site, _rows(4, spec.k, seed=4).float(), w) is None + fresh = decode_gemv.K3DecodeGemvs() + x = _rows(4, spec.k, seed=4) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + y = fresh.project(site, x, w) + assert y is None + + +@pytest.mark.parametrize("rows", [4, 16]) +def test_site_capture_replays_the_eager_bits(rows): + """A captured call (the long kernel at 4 rows, the wide one at 16, the gate columns through the sigmoid), its + input rewritten in place before each replay, gives the eager call's bits.""" + spec = decode_gemv.SITES["mla_ag"] + w = _weight(spec.n, spec.k, seed=3) + gemvs = decode_gemv.K3DecodeGemvs.create(None, sites=["mla_ag"]) + x = _rows(rows, spec.k, seed=4) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + y = gemvs.project("mla_ag", x, w) + assert y is not None + for seed in (5, 6): + x.copy_(_rows(rows, spec.k, seed=seed)) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(_bits(y), _bits(gemvs.project("mla_ag", x, w))) + + +def _lm_head(): + head = LMHead( + VOCAB_SHARD, + HIDDEN, + dtype=torch.bfloat16, + mapping=Mapping(), + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=True, + ).cuda() + with torch.no_grad(): + head.weight.copy_(_weight(VOCAB_SHARD, HIDDEN, seed=7)) + return head + + +def test_lm_head(): + """At most 8 rows: fp32 logits from k3_head_gemv within TOL; 9 rows and an unbuilt processor: the stock bits.""" + head = _lm_head() + stock = LogitsProcessor() + processor = decode_gemv.K3LogitsProcessor(stock) + for m in (1, 9): + x = _rows(m, HIDDEN, seed=m) + assert torch.equal(processor(x, head, None, True), stock(x, head, None, True)) + processor.gemvs = decode_gemv.K3DecodeGemvs.create(head, sites=()) + assert processor.gemvs.head_workspace is not None + for m in ROWS: + x = _rows(m, HIDDEN, seed=m) + logits = processor(x, head, None, True) + assert logits.dtype == torch.float32 and logits.shape == (m, VOCAB_SHARD) + assert _rel_err(logits, x, head.weight) <= TOL, m + x = _rows(9, HIDDEN, seed=9) + assert torch.equal(processor(x, head, None, True), stock(x, head, None, True)) + + +def test_lm_head_selects_the_stock_rows(): + """Without context logits, the stock processor's last-token selection feeds the head.""" + head = _lm_head() + processor = decode_gemv.K3LogitsProcessor(LogitsProcessor()) + processor.gemvs = decode_gemv.K3DecodeGemvs.create(head, sites=()) + x = _rows(7, HIDDEN, seed=3) + metadata = types.SimpleNamespace( + seq_lens_cuda=torch.tensor([3, 1, 3], dtype=torch.int32, device="cuda") + ) + logits = processor(x, head, metadata, False) + want = processor(x[[2, 3, 6]].contiguous(), head, None, True) + assert torch.equal(logits, want) + + +def test_lm_head_capture_replays_the_eager_bits(): + head = _lm_head() + gemvs = decode_gemv.K3DecodeGemvs.create(head, sites=()) + x = _rows(8, HIDDEN, seed=8) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + y = gemvs.lm_head_logits(x, head) + assert y is not None + x.copy_(_rows(8, HIDDEN, seed=11)) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(_bits(y), _bits(gemvs.lm_head_logits(x, head))) + + +@pytest.mark.parametrize("n", [1, 8, 64]) +def test_embed_norm(n): + """Layer 0's normed input is bit-identical to nn.Embedding followed by the stock RMSNorm; the bank's slot 0 holds + the embedding rows.""" + vocab = 4096 + table = _weight(vocab, HIDDEN, seed=21) + norm = RMSNorm(hidden_size=HIDDEN, eps=1e-5, dtype=torch.bfloat16).cuda() + with torch.no_grad(): + norm.weight.copy_(1 + _weight(1, HIDDEN, seed=22)[0]) + embedding = nn.Embedding(vocab, HIDDEN, dtype=torch.bfloat16).cuda() + with torch.no_grad(): + embedding.weight.copy_(table) + g = torch.Generator(device="cuda").manual_seed(n) + ids = torch.randint(0, vocab, (n,), generator=g, device="cuda", dtype=torch.int32) + gemvs = decode_gemv.K3DecodeGemvs() + bank = table.new_empty(3, n, HIDDEN) + normed = gemvs.embed_norm(ids, table, norm, bank) + assert normed is not None + raw = embedding(ids) + assert torch.equal(_bits(bank[0]), _bits(raw)) + assert torch.equal(_bits(normed), _bits(norm(raw))) + + +def test_embed_norm_capture(): + """Under capture a token count that ran eagerly replays the eager bits; one that never ran is declined.""" + vocab = 4096 + table = _weight(vocab, HIDDEN, seed=31) + norm = RMSNorm(hidden_size=HIDDEN, eps=1e-5, dtype=torch.bfloat16).cuda() + gemvs = decode_gemv.K3DecodeGemvs() + ids = torch.arange(4, dtype=torch.int32, device="cuda") + bank = table.new_empty(2, 4, HIDDEN) + eager = gemvs.embed_norm(ids, table, norm, bank).clone() + other_ids = torch.arange(5, dtype=torch.int32, device="cuda") + other_bank = table.new_empty(2, 5, HIDDEN) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = gemvs.embed_norm(ids, table, norm, bank) + declined = gemvs.embed_norm(other_ids, table, norm, other_bank) + assert captured is not None and declined is None + graph.replay() + torch.cuda.synchronize() + assert torch.equal(_bits(captured), _bits(eager)) + + +SITU = (4.0, 25.0) # the Kimi K3 checkpoint's (beta, linear_beta) + + +def _dense_weights(): + gate_up, down = decode_gemv.SITES["dense_gate_up"], decode_gemv.SITES["dense_down"] + return _weight(gate_up.n, gate_up.k, seed=41), _weight(down.n, down.k, seed=42) + + +def _dense_reference(x, w_gu, w_down, beta, linear_beta): + gu = x.double() @ w_gu.double().t() + g, u = gu.chunk(2, dim=1) + a = beta * torch.tanh(g / beta) * torch.sigmoid(g) + v = u if linear_beta is None else linear_beta * torch.tanh(u / linear_beta) + return (a * v) @ w_down.double().t() + + +@pytest.mark.parametrize("situ", [SITU, (1.0, None)], ids=["k3", "default"]) +def test_dense_mlp(situ): + """At 1..8 rows: the bits of the gate_up GEMV, k3_situ_mul and the down GEMV called directly, within 3e-2 of the + float64 MLP.""" + w_gu, w_down = _dense_weights() + gemvs = decode_gemv.K3DecodeGemvs.create(None, sites=["dense_gate_up", "dense_down"]) + gate_up, down = decode_gemv.SITES["dense_gate_up"], decode_gemv.SITES["dense_down"] + for m in ROWS: + x = _rows(m, HIDDEN, seed=m) + y = gemvs.dense_mlp(x, w_gu, w_down, *situ) + assert y is not None and y.shape == (m, down.n), m + gu = k3_ctm_gemv_long(x, w_gu, split=gate_up.split, ring=gate_up.ring) + want = k3_ctm_gemv_long(k3_situ_mul(gu, *situ), w_down, split=down.split, ring=down.ring) + assert torch.equal(_bits(y), _bits(want)), m + ref = _dense_reference(x, w_gu, w_down, *situ) + err = ((y.double() - ref).abs().max() / ref.abs().max()).item() + assert err <= 3e-2, (m, err) + + +def test_dense_mlp_declines(): + """9 rows, linear_beta 0 (the activation would run it as 1), and under capture before any eager call: None.""" + w_gu, w_down = _dense_weights() + gemvs = decode_gemv.K3DecodeGemvs.create(None, sites=["dense_gate_up", "dense_down"]) + assert gemvs.dense_mlp(_rows(9, HIDDEN, seed=9), w_gu, w_down, *SITU) is None + assert gemvs.dense_mlp(_rows(4, HIDDEN, seed=4), w_gu, w_down, 4.0, 0.0) is None + fresh = decode_gemv.K3DecodeGemvs() + x = _rows(4, HIDDEN, seed=4) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + y = fresh.dense_mlp(x, w_gu, w_down, *SITU) + assert y is None + + +def test_dense_mlp_capture_replays_the_eager_bits(): + w_gu, w_down = _dense_weights() + gemvs = decode_gemv.K3DecodeGemvs.create(None, sites=["dense_gate_up", "dense_down"]) + x = _rows(8, HIDDEN, seed=8) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + y = gemvs.dense_mlp(x, w_gu, w_down, *SITU) + assert y is not None + x.copy_(_rows(8, HIDDEN, seed=12)) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(_bits(y), _bits(gemvs.dense_mlp(x, w_gu, w_down, *SITU))) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_step.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_step.py new file mode 100644 index 000000000000..345f803fac8f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_step.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""How the Kimi K3 target classifies a step for its decode kernels (host-side, no GPU): ``decode_step`` of +``kimi_k3_mxfp4__sm_100__tp16_moetp4ep4``.""" + +import types + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4.modeling import ( # noqa: E501 + DecodeStep, + decode_step, +) + +pytestmark = pytest.mark.cpu_only + + +def _metadata(seq_lens, num_contexts=0): + return types.SimpleNamespace( + num_contexts=num_contexts, + num_generations=len(seq_lens) - num_contexts, + seq_lens=torch.tensor(seq_lens, dtype=torch.int32), + ) + + +@pytest.mark.parametrize( + "requests,tokens", + [(1, 1), (1, 8), (8, 1), (2, 4), (4, 2), (3, 1), (5, 1), (2, 8), (4, 8), (8, 8), (8, 2)], +) +def test_pure_decode_steps(requests, tokens): + """R <= 8 generation requests of the same T <= 8 tokens: a decode step, small up to one token tile, wide above.""" + step = decode_step(_metadata([tokens] * requests), requests * tokens) + assert step == DecodeStep(requests * tokens, requests, tokens) + assert step.decode + assert step.small == (requests * tokens <= 8) + assert step.wide == (requests * tokens > 8) + + +@pytest.mark.parametrize( + "seq_lens,num_contexts,rows", + [ + ([5], 1, 5), # a short prefill + ([3, 1], 1, 4), # a short mixed step + ([4, 4], 0, 4), # rows that are not the step's tokens + ([3, 2], 0, 5), # a ragged decode step + ], +) +def test_short_steps_are_small_only(seq_lens, num_contexts, rows): + """Other steps of at most 8 tokens take the token-count kernels only.""" + step = decode_step(_metadata(seq_lens, num_contexts), rows) + assert step == DecodeStep(rows) + assert step.small and not step.decode and not step.wide + + +@pytest.mark.parametrize( + "seq_lens,num_contexts,rows", + [ + ([1024], 1, 1024), # prefill + ([100, 8], 1, 108), # mixed step + ([1] * 9, 0, 9), # more requests than the kernels take + ([16], 0, 16), # more tokens per request than the kernels take + ([8, 4], 0, 12), # ragged decode step + ([4, 4], 0, 16), # rows that are not the step's tokens + ([], 0, 0), # empty step + ], +) +def test_other_steps_take_the_generic_path(seq_lens, num_contexts, rows): + assert decode_step(_metadata(seq_lens, num_contexts), rows) is None diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py index f40cb6428713..25fe8e786b66 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -28,6 +28,7 @@ ModelingV2Mode, modeling_v2_resolve, ) +from tensorrt_llm._torch.configs.kimi_k3 import KimiK3Config from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_auto import AutoModelForCausalLM from tensorrt_llm._torch.models.modeling_utils import ( @@ -35,6 +36,8 @@ get_registered_model_class, ) from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo _SM103 = (10, 3) @@ -235,3 +238,160 @@ def test_an_unknown_mode_raises_rather_than_falling_back(monkeypatch): monkeypatch.setenv(MODELING_V2_ENV, "yes") with pytest.raises(ValueError, match="not a modeling_v2 mode"): ModelingV2Mode.from_env() + + +_SM100 = (10, 0) +_TP16_MOETP4EP4 = dict(world_size=16, tp_size=16, moe_tp_size=4, moe_ep_size=4) +_K3_TARGET = "ModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4" + + +def _k3_config(text_as_dict=False, **text_overrides): + """The shape Kimi K3's own config.json declares: a vision-language wrapper + whose language model is ``text_config``.""" + text = dict( + model_type="kimi_linear", + num_hidden_layers=93, + hidden_size=7168, + num_experts=896, + routed_expert_hidden_size=3584, + ) + text.update(text_overrides) + return PretrainedConfig( + architectures=["KimiK3ForConditionalGeneration"], + model_type="kimi_k3", + text_config=text if text_as_dict else PretrainedConfig(**text), + ) + + +# The quantization the MXFP4 checkpoint declares in its config.json, inside ``text_config``. +_K3_CHECKPOINT_QUANTIZATION = { + "quant_method": "compressed-tensors", + "format": "mxfp4-pack-quantized", + "config_groups": { + "group_0": { + "format": "mxfp4-pack-quantized", + "input_activations": None, + "output_activations": None, + "targets": ["Linear"], + "weights": { + "group_size": 32, + "num_bits": 4, + "strategy": "group", + "symmetric": True, + "type": "float", + }, + } + }, + "ignore": [ + "re:.*self_attn.*", + "re:.*shared_experts.*", + r"re:.*mlp\.(gate|up|gate_up|down)_proj.*", + "re:.*lm_head.*", + "re:.*vision_tower.*", + "re:.*mm_projector.*", + ], + "kv_cache_scheme": None, +} + + +def _k3_checkpoint_quant_config(): + """What the engine reads for the MXFP4 checkpoint: ``KimiK3Config`` surfaces the text config's + declaration, and ``ModelConfig`` parses it.""" + config = KimiK3Config( + text_config=dict(quantization_config=_K3_CHECKPOINT_QUANTIZATION), + architectures=["KimiK3ForConditionalGeneration"], + ) + quant_config, layer_quant_config = ModelConfig.load_hf_quant_config( + config.quantization_config, "TRTLLM" + ) + assert layer_quant_config is None + return quant_config + + +def _k3_model_config(pretrained_config, quant_config=None, **mapping_kwargs): + """A model config of the MXFP4 checkpoint's quantization (or ``quant_config``).""" + return ModelConfig( + pretrained_config=pretrained_config, + mapping=Mapping(**mapping_kwargs), + quant_config=quant_config if quant_config is not None else _k3_checkpoint_quant_config(), + ) + + +@pytest.fixture +def _on_sm100(monkeypatch): + """Route as if this were a GB200.""" + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: _SM100) + + +def test_kimi_k3_checkpoint_reads_as_w4a16_mxfp4(): + quant_config = _k3_checkpoint_quant_config() + assert quant_config.quant_algo == QuantAlgo.W4A16_MXFP4 + assert quant_config.group_size == 32 + + +@pytest.mark.usefixtures("_on_sm100") +@pytest.mark.parametrize("mode", ["auto", "require"]) +@pytest.mark.parametrize("text_as_dict", [False, True], ids=["text-config", "text-dict"]) +def test_kimi_k3_tp16_moetp4ep4_matches(monkeypatch, mode, text_as_dict): + _set_mode(monkeypatch, mode) + config = _k3_model_config(_k3_config(text_as_dict=text_as_dict), **_TP16_MOETP4EP4) + assert modeling_v2_resolve(config) == _K3_TARGET + + +@pytest.mark.usefixtures("_on_sm100") +def test_kimi_k3_target_registers_and_counts_as_external(): + config = _k3_model_config(_k3_config(), **_TP16_MOETP4EP4) + cls = get_registered_model_class(modeling_v2_resolve(config)) + assert cls is not None and cls.__name__ == _K3_TARGET + assert not _is_builtin_model_class(cls) + + +@pytest.mark.usefixtures("_on_sm100") +@pytest.mark.parametrize( + "quant_algo", + [QuantAlgo.MIXED_PRECISION, None], + ids=["nvfp4-requant", "unquantized"], +) +def test_kimi_k3_other_quantizations_do_not_match(monkeypatch, quant_algo): + """The NVFP4 requant (MIXED_PRECISION) and an unquantized checkpoint have the MXFP4 checkpoint's shape; the + quantization is what keeps them out of a target whose loader reads packed MXFP4 experts.""" + config = _k3_model_config(_k3_config(), QuantConfig(quant_algo=quant_algo), **_TP16_MOETP4EP4) + assert modeling_v2_resolve(config) is None + + _set_mode(monkeypatch, "require") + with pytest.raises(ValueError, match="quant"): + modeling_v2_resolve(config) + + +@pytest.mark.usefixtures("_on_sm100") +@pytest.mark.parametrize( + "text_overrides, mapping_kwargs, missed", + [ + # a Kimi-family checkpoint of another depth + (dict(num_hidden_layers=61), _TP16_MOETP4EP4, "shape"), + # route B's expert split: experts tensor-parallel 16 ways, no target yet + (dict(), dict(world_size=16, tp_size=16, moe_tp_size=16, moe_ep_size=1), "parallel"), + # attention data parallelism splits the requests, not the heads + (dict(), dict(_TP16_MOETP4EP4, enable_attention_dp=True), "parallel"), + # one tray instead of four + (dict(), dict(world_size=4, tp_size=4, moe_tp_size=1, moe_ep_size=4), "parallel"), + ], +) +def test_kimi_k3_near_misses_do_not_match(monkeypatch, text_overrides, mapping_kwargs, missed): + config = _k3_model_config(_k3_config(**text_overrides), **mapping_kwargs) + assert modeling_v2_resolve(config) is None + + _set_mode(monkeypatch, "require") + with pytest.raises(ValueError, match=missed): + modeling_v2_resolve(config) + + +def test_kimi_k3_on_another_sm_does_not_match(monkeypatch): + """The autouse fixture routes as a GB300 (sm 10.3); the K3 target is sm 10.0 only.""" + config = _k3_model_config(_k3_config(), **_TP16_MOETP4EP4) + assert modeling_v2_resolve(config) is None + + _set_mode(monkeypatch, "require") + with pytest.raises(ValueError) as excinfo: + modeling_v2_resolve(config) + assert "sm" in str(excinfo.value) and "(10, 3)" in str(excinfo.value) diff --git a/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_proxy_fences.py b/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_proxy_fences.py new file mode 100644 index 000000000000..b1a83d52af7f --- /dev/null +++ b/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_proxy_fences.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Source check of the cross-proxy fences in the Kimi K3 attention-residual kernels (PTX ISA, memory consistency +model, proxies). + +The consumers read the V (and delta) slots with generic-proxy shared loads and release each slot with an mbarrier +arrive on ``bar_consumed``. The producer then refills the slot with ``cp.async.bulk``, an async-proxy write. Accesses +to one location through two proxies need a cross-proxy fence: every reading thread issues +``fence.proxy.async.shared::cta`` after its last read and before the release. Without it the order rests on how the +compiler schedules the instructions. No test of values can see that, so the kernels' sources are read. + + pytest test_attn_res_proxy_fences.py +""" + +import pathlib +import re + +import pytest + +pytestmark = pytest.mark.cpu_only + +_SRC = pathlib.Path(__file__).resolve().parents[5] / "cpp/tensorrt_llm/kernels/kimiK3AttnRes" + +# The kernels whose consumers release cp.async.bulk-filled slots, with the number of releases each file has. +RELEASES = {"attnResFwd.cu": 3, "attnResFwdPersistentFused.cu": 1} + +RELEASE = re.compile(r"\b(cute::arrive_barrier|mbarrier_arrive)\(\s*plan\.bar_consumed\b") +FENCE = 'asm volatile("fence.proxy.async.shared::cta;" ::: "memory");' +# What may sit between the fence and the release: the warp sync that gathers every lane's reads before lane 0 releases, +# and the lane-0 guard. Any other statement could hold a read after the fence. +BETWEEN = re.compile(r"^(__syncwarp\(\);|if \(lane == 0\)|\{)$") + + +def _code(line): + return line.split("//", 1)[0].strip() + + +def violations(lines): + """(line number, release) of every ``bar_consumed`` release whose closest preceding statement, past the warp sync + and the lane-0 guard, is not the cross-proxy fence.""" + out = [] + for i, line in enumerate(lines): + if not RELEASE.search(_code(line)): + continue + for j in range(i - 1, -1, -1): + code = _code(lines[j]) + if not code or BETWEEN.match(code): + continue + if code != FENCE: + out.append((i + 1, line.strip())) + break + else: + out.append((i + 1, line.strip())) + return out + + +@pytest.mark.parametrize("name", sorted(RELEASES)) +def test_slot_release_follows_cross_proxy_fence(name): + lines = (_SRC / name).read_text().splitlines() + releases = [i for i, line in enumerate(lines) if RELEASE.search(_code(line))] + assert len(releases) == RELEASES[name], ( + f"{name}: {len(releases)} bar_consumed releases, expected {RELEASES[name]} " + "(update RELEASES if the kernels changed)" + ) + assert any("cp.async.bulk.shared::cta.global" in line for line in lines), ( + f"{name}: no cp.async.bulk refill; the fence may no longer be needed" + ) + assert violations(lines) == [], ( + f"{name}: slot releases without fence.proxy.async.shared::cta right before them: {violations(lines)}" + ) + + +def test_check_flags_a_release_without_the_fence(): + """The check itself: the release pattern without the fence, as before the fix, is reported.""" + lines = [ + "// Every lane's reads of the chunk's slots before lane 0 releases them.", + "__syncwarp();", + "if (lane == 0)", + "{", + " cute::arrive_barrier(plan.bar_consumed[chunk_slot]);", + "}", + ] + assert violations(lines) == [(5, "cute::arrive_barrier(plan.bar_consumed[chunk_slot]);")] + assert violations([FENCE] + lines) == [] + assert violations([FENCE, "float x = buf[0];"] + lines) == [ + (7, "cute::arrive_barrier(plan.bar_consumed[chunk_slot]);") + ]