diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/kernels/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/kernels/__init__.py similarity index 68% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/kernels/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/kernels/__init__.py index 3811a6f831f0..a5adae93fbad 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/kernels/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/kernels/__init__.py @@ -1,3 +1,3 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Kernels used only by this target, called through its ops.""" +"""Kernels used only by this family, called through its ops.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/kernels/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/kernels/moe/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/kernels/moe/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/kernels/moe/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/kernels/moe/blockscaled_contiguous_grouped_gemm_finalize_fusion.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/kernels/moe/blockscaled_contiguous_grouped_gemm_finalize_fusion.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/kernels/moe/blockscaled_contiguous_grouped_gemm_finalize_fusion.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/kernels/moe/blockscaled_contiguous_grouped_gemm_finalize_fusion.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/__init__.py similarity index 66% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/__init__.py index 1059445198db..a8aa08cbff1d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/__init__.py @@ -1,3 +1,3 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Ops used only by this target. Shared ops stay in the catalog.""" +"""Operations of the qwen4_exp family's target.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/fused_sigmoid_mul_inplace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/fused_sigmoid_mul_inplace.py new file mode 100644 index 000000000000..98189e91514c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/fused_sigmoid_mul_inplace.py @@ -0,0 +1,227 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""In-place sigmoid output gating: ``attention_output *= sigmoid(gate)``. + +The gated attention's output gate, between the attention core and o_proj, in +one Triton kernel (`torch.ops.trtllm.fused_sigmoid_mul_inplace`). Called by the +qwen4_exp target (`qwen3_8_flash_next_nvfp4__sm_103__default`) at every layout -- its +trunk, gate off the fused input-projection row, and its MTP layer, gate as the strided +half of the q projection. +""" + +from __future__ import annotations + +import torch + +# Registered by the module that defines it: after `import +# tensorrt_llm._torch.custom_ops` alone, torch.ops.trtllm has no +# `fused_sigmoid_mul_inplace` (measured). +import tensorrt_llm._torch.attention.kernels.fused_qk_norm_rope_gate # noqa: F401 +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The checkpoint's gated attention: 24 q heads of 256, q projection [T, 24, 512]. +_CKPT = dict(heads=24, head_dim=256, dtype=torch.bfloat16, gate="interleaved") + +#: The default target at tp4 / tep4: a rank's fused input projection row is +#: [q 6 x 256 | k 256 | v 256 | gate 6 x 256 | index scratch 768] = 4352 columns, +#: so the gate is the 1536-wide window at column qkv_width = 1536 + 2 * 256 = 2048. +_TP4_ROW = dict( + heads=6, + head_dim=256, + dtype=torch.bfloat16, + gate="fused_row", + row_stride=4352, + gate_col=2048, +) + +#: The default target at dep4: tp1's widths, 14080 columns, the gate at 7168. +_DEP4_ROW = dict( + heads=24, + head_dim=256, + dtype=torch.bfloat16, + gate="fused_row", + row_stride=14080, + gate_col=7168, +) + +#: The default target's MTP layer head-sharded at tep4: q_proj 12288 -> 3072 columns a +#: rank, read as [T, 6, 512], the gate the second 256 of each head. +_MTP_TEP4 = dict(heads=6, head_dim=256, dtype=torch.bfloat16, gate="interleaved") + + +class FusedSigmoidMulInplace(OpWrapper): + """`attention_output[t, i] *= sigmoid(gate[t, i // D, i % D])` in place, `D = gate.shape[-1]`. + + Returns `attention_output`, the tensor it wrote. Both operands are widened to + fp32 and the product is rounded once to `attention_output.dtype`. `gate` is + read through its own token and head strides, so the gate half of an + interleaved `[T, heads, 2 * D]` q projection is passed as a strided view with + no repack (bitwise equal to its compacted copy, measured). A 2-D `gate` of + `attention_output`'s shape is one head spanning the row. + + Fusion boundary: sigmoid and multiply only. The caller owns producing both + operands and everything downstream (the o projection). `gate` is read-only. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + # The target at tp1: hidden 6144, the trunk's gate off the 14080-wide fused + # row (the test's default row stride and gate column), the MTP layer's + # off its [T, 24, 512] q_proj. + Cell( + why="decode 128 tokens: 2-D gate off the fused projection row (stride 14080)", + spec=dict(tokens=128, heads=24, head_dim=256, dtype=torch.bfloat16, gate="fused_row"), + ), + Cell( + why="prefill 8192 tokens: 2-D gate off the fused projection row, 2048-wide blocks", + spec=dict(tokens=8192, heads=24, head_dim=256, dtype=torch.bfloat16, gate="fused_row"), + ), + Cell( + why="MTP decode: 1 token, 24x256 gate as the strided half of q_proj", + spec=dict(tokens=1, **_CKPT), + ), + # The default target's trunk. Hidden 1536 is not a multiple of the kernel's + # block, so every tp4 / tep4 row ends in 512 masked lanes: two 1024-wide + # blocks below 1024 rows, one 2048-wide block from 1024 rows. + Cell( + why="default tp4/tep4 decode, one token: 6 heads off the 4352-wide row", + spec=dict(tokens=1, **_TP4_ROW), + ), + Cell( + why="default tp4/tep4 largest captured decode batch: 128 rows, 6 heads", + spec=dict(tokens=128, **_TP4_ROW), + ), + Cell( + why="default tep4 MTP-7 verify 128 x 8 = 1024 rows: the first 2048-wide block count", + spec=dict(tokens=1024, **_TP4_ROW), + ), + Cell( + why="default tp4/tep4 prefill chunk at max_num_tokens: 8192 rows, 6 heads", + spec=dict(tokens=8192, **_TP4_ROW), + ), + Cell( + why="default dep4 rank's largest captured decode batch: 32 rows, 24 heads", + spec=dict(tokens=32, **_DEP4_ROW), + ), + # The default target's MTP layer at tep4, on the same 1536-wide masked rows. + Cell( + why="default MTP drafter, 1 row: 6-head gate half of the [T, 6, 512] q_proj", + spec=dict(tokens=1, **_MTP_TEP4), + ), + Cell( + why="default MTP drafter at 128 requests: 128 rows, 6 heads, 1024-wide blocks", + spec=dict(tokens=128, **_MTP_TEP4), + ), + Cell( + why="default MTP-3 verify 128 x 4 = 512 rows, 6 heads", + spec=dict(tokens=512, **_MTP_TEP4), + ), + Cell( + why="default MTP-7 verify 128 x 8 = 1024 rows: the first 2048-wide block count", + spec=dict(tokens=1024, **_MTP_TEP4), + ), + Cell( + why="default MTP prefill chunk at max_num_tokens: 8192 rows, 6 heads", + spec=dict(tokens=8192, **_MTP_TEP4), + ), + ) + + note = """ + Metadata consumed: none. Source: `torch.ops.trtllm.fused_sigmoid_mul_inplace`, + a Python `torch.library.custom_op` registered by + `tensorrt_llm/_torch/attention/kernels/fused_qk_norm_rope_gate.py` around the + Triton kernel `_fused_sigmoid_mul_kernel` there. One kernel, JIT-compiled for + the running device at the first call of each (hidden, head_dim, block) + specialization, no autotuner: grid `(T, cdiv(hidden, block))`, + `block = min(next_pow2(hidden), 1024 if T < 1024 else 2048)`, 4 warps. The + cells ran on sm_103. + + Call sites, each the attention output gate between the attention core and + o_proj: + - The target at tp1 (`qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default`): the + trunk's gate as the 2-D window of the fused input projection row + (14080 wide, the gate at column 7168, hidden 6144) at decode and + prefill row counts; its MTP layer's gate as the strided second half of + each head of a `[T, 24, 512]` q_proj. + - The default target's trunk + (`qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default`): + `DecodeTarget._attention_block` (decode, 1..128 rows at tp4 / tep4, + 1..32 at a dep4 rank), its MTP verify (up to 128 x 8 = 1024 rows, tep4, + draft 7) and `PrefillTarget._attention_block` (one chunk of up to 8192 + rows, `max_num_tokens`). `gate` is the 2-D window of the fused input + projection row at column qkv_width: at tp4 / tep4 the row is 4352 wide + (q 1536 | k 256 | v 256 | gate 1536 | index scratch 768), the gate at + column 2048 and hidden 1536; at dep4 tp1's 14080-wide row, gate at + 7168, hidden 6144. + - The default target's MTP layer, head-sharded at tep4 (and tp4): q_proj + 12288 -> 3072 columns a rank, read as `[T, 6, 512]`, `gate` the second + 256 of each head (token stride 3072, head stride 512), hidden 1536. + Rows: the drafter 1..128, the depth-3 verify 128 x 4 = 512, the depth-7 + verify 128 x 8 = 1024, a prefill chunk of up to 8192 + (`max_num_tokens`). + Every row of the default target's 6-head forms ends in masked lanes (see + the cells); the test places `attention_output` at the head of a larger + buffer and holds the elements past its last row unchanged. + + Preconditions, every one enforced loudly by the op's own bare `assert` + (an AssertionError with no message): `attention_output` 2-D with + `stride(-1) == 1` (row stride free); `gate` 2-D of `attention_output`'s + shape, or 3-D with `shape[0] == T` and `shape[1] * shape[2] == hidden`, + `stride(-1) == 1`. A CPU tensor raises ValueError from Triton (measured). + Hence no `is_valid`. Cross-device operands were not probed. + + The pre-call contents of `attention_output` are destroyed; the gate is + left bitwise unchanged. + + Numerics, measured on GB300 at every cell: bf16 bitwise equal to the fp32 + reference except at most ~2 in 10^6 elements one ulp off (`tl.sigmoid`'s + fast exp). Against a float64 product every cell sits at bf16 output + rounding, the op and the reference alike (relative L2 1.64e-3 to 1.69e-3). + Torch's default bf16 tolerance holds, nothing loosened. A wrong read lands + far outside it (1.6 to 3.8 max abs): the q half instead of the gate half, + the fused row's q window instead of its gate window, or at the 6-head 3-D + gate the q half of each head or the flat row with the head stride ignored; + the test holds the default band refusing the last two. So the column offset + and the per-head indexing are under test. + + Certified: bf16 operands, a compact `attention_output`; the 2-D window of + the fused projection row at hidden 6144 (row stride 14080, gate column + 7168) at 32, 128 and 8192 rows, and at hidden 1536 (row stride 4352, gate + column 2048) at 1, 128, 1024 and 8192 rows; the strided q_proj half (3-D + gate) at 24 heads (`[T, 24, 512]`, hidden 6144) at 1 token, and at 6 heads + (a compact `[T, 6, 512]`, hidden 1536) at 1, 128, 512, 1024 and 8192 rows; + CUDA-graph capture and replay of the 1536-wide fused-row form at 8 rows and + of the 6-head 3-D gate at 1, 128 and 1024 rows. + + Hard limits + T * row stride token offset computed in int64 boundary 2^63 not probed + in-row offset int32; < hidden (6144) at the 2-D form, + head * head stride + dim < 2 * hidden at the 3-D form + (largest run 8192 x 14080 = 1.2e8 elements) + + Not certified: 0 tokens; rows past 8192; the 3-D gate at a token stride + other than its own q_proj's (heads x 2 x head_dim, e.g. a gate inside a + fused q|k|v row); the 24-head strided q_proj half above 1 token (a dep4 + rank's MTP layer: drafter 1..32 rows and its verify); a strided + `attention_output`; a gate dtype other than `attention_output.dtype`; + fp16/fp32 operands; hidden sizes other than 1536 and 6144; CUDA-graph + capture of the 6144-wide forms; torch.compile. + """ + + def __call__(self, attention_output: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + torch.ops.trtllm.fused_sigmoid_mul_inplace(attention_output, gate) + return attention_output + + def reference(self, attention_output: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + """What `attention_output` holds after the call; neither input is mutated. + + `gate.reshape(T, hidden)` is the per-head `(t, i // D, i % D)` read, for + the 2-D and 3-D forms alike. + """ + num_tokens, hidden = attention_output.shape + g = gate.reshape(num_tokens, hidden).float() + return (attention_output.float() * torch.sigmoid(g)).to(attention_output.dtype) + + +fused_sigmoid_mul_inplace = FusedSigmoidMulInplace() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/gelu_erf.py similarity index 66% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/gelu_erf.py index 07150279c931..a31ca88bb4a3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/gelu_erf.py @@ -3,12 +3,16 @@ """Non-gated exact (erf) GELU -- what `torch.nn.GELU()` is -- the vision patch merger's activation, which no config key names. -Tier-3 composed: one `torch.nn.functional.gelu(x)` call, `approximate="none"`. +Native torch (rung 4 of the sourcing ladder): one `torch.nn.functional.gelu(x)` +call, `approximate="none"` -- one elementwise CUDA kernel, the same call the +stock tower makes (`modeling_qwen3vl.py`'s merger `nn.GELU()`). `torch.ops.trtllm.moe_gelu` computes this function bitwise but is not a substitute: its `tile_idx_to_mn_limit` / `num_non_exiting_tiles` operands must be int32 *device* tensors (on the host it raises nothing and returns garbage, -nan, or an illegal memory access), so wrapping it means launching two fills -before the call; it also rejects fp32 and non-2-D inputs. +nan, or an illegal memory access), and it rejects fp32 and non-2-D inputs. +FlashInfer's `activation` module has only gated forms (`gelu_and_mul`, +`gelu_tanh_and_mul`, `silu_and_mul`), as has +`torch.ops.trtllm.flashinfer_gelu_tanh_and_mul`. """ from __future__ import annotations @@ -54,6 +58,10 @@ class GeluErf(OpWrapper): why="merger activand, 1 merged row x 4608 bf16", spec=dict(shape=(1, MERGER_WIDTH), dtype=torch.bfloat16, lo=-16.0, hi=16.0), ), + Cell( + why="the merger-built activand at the tower call ceiling: 8192 merged rows x 4608", + spec=dict(shape=(8192, MERGER_WIDTH), dtype=torch.bfloat16, dist="merger"), + ), ) note = """ @@ -61,14 +69,21 @@ class GeluErf(OpWrapper): Certified on contiguous bf16 `[rows, 4608]` activands (what the merger's fc1 returns) at 1, 512 and 8192 rows (8192, from 32768 patch rows, the most one - tower call carries). - - Identity is pinned by fp32 only (the identity test). - Max |erf GELU - tanh GELU| is 4.73e-4 at x = -2.699; in fp32 the tanh form - lands ~1400x the correct path's deviation away, 31x outside the default - gate. In bf16 one ulp at the output scale exceeds that gap (tanh variant - 8.24e-3 vs correct 7.77e-3 max abs), so a bf16 match does not say which GELU - was computed. + tower call carries); 512 and 8192 rows also on the merger-built + distribution. The default target's call site is the vision tower's patch + merger, eager and rank-invariant: every rank runs the whole tower at tp4, + tep4 and dep4, so the shapes are tp1's. + + Measured on GB300 against `reference` (fp64, rounded once to bf16): the + merger-built activand is bitwise equal or nearly so; the uniform [-16, 16] + spread differs in ~12% of elements. Every difference is at most 2.4e-7 + absolute, in the negative tail where the output is near zero. One + compiled-in torch elementwise kernel per call: no JIT path. + + Identity is pinned by fp32 only (the identity test): max |erf GELU - tanh + GELU| is 4.73e-4 (x = -2.699), 31x outside the default fp32 gate. In bf16 + one ulp at the output scale exceeds that gap, so a bf16 match does not say + which GELU was computed. Measured, not certified: `-inf -> nan` (the closed form does the same); `+inf -> +inf`; large negatives give `-0.0`; `finfo.max` comes back diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/gelu_tanh.py similarity index 55% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/gelu_tanh.py index e9cd2ebec23a..e28a9c0b4d6b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/gelu_tanh.py @@ -3,12 +3,14 @@ """Non-gated tanh GELU -- HuggingFace's `gelu_pytorch_tanh` -- the vision tower's block-MLP activation. -Tier-3 composed: one `torch.nn.functional.gelu(x, approximate="tanh")` call. -No registered op computes a plain non-gated GELU over an unquantized tensor: -`flashinfer_gelu_tanh_and_mul` is the gated form (and silently returns a -half-width `gelu(a) * b` when handed this tower's non-gated activand), +Tier 4 (native torch), a macro entry of one call: `torch.nn.functional.gelu(x, +approximate="tanh")`, the aten elementwise kernel. No registered op computes a +plain non-gated GELU over an unquantized tensor: `flashinfer_gelu_tanh_and_mul` (and FlashInfer's +`gelu_tanh_and_mul` / `gelu_and_mul`) is the gated form, and silently returns a +half-width `gelu(a) * b` when handed this tower's non-gated activand. `moe_gelu` is the erf GELU behind MoE tile metadata, and the -`cute_dsl_nvfp4_dense_gemm_gelu*` ops are NVFP4 GEMM epilogues. +`cute_dsl_nvfp4_dense_gemm_gelu*` ops are NVFP4 GEMM epilogues (this tower's +fc1 is bf16). """ from __future__ import annotations @@ -46,6 +48,10 @@ class GeluTanh(OpWrapper): why="activand built through fc1 at the checkpoint's measured scale", spec=dict(shape=(512, INTERMEDIATE_SIZE), dtype=torch.bfloat16, dist="fc1"), ), + Cell( + why="fc1-scaled activand at the tower call ceiling, 32768 patch rows", + spec=dict(shape=(32768, INTERMEDIATE_SIZE), dtype=torch.bfloat16, dist="fc1"), + ), Cell( why="block MLP activand at the tower call ceiling, 32768 patch rows x 4304 bf16", spec=dict(shape=(32768, INTERMEDIATE_SIZE), dtype=torch.bfloat16, lo=-16.0, hi=16.0), @@ -57,27 +63,44 @@ class GeluTanh(OpWrapper): ) note = """ - Metadata consumed: none. Out-of-place; `x` is not mutated. + Metadata consumed: none. Out-of-place; `x` is not mutated. Call site: the + vision tower's block MLP, `gelu_tanh(cublas_mm(layer_norm(h), fc1, fc1_b))`, + eager only, at most 32768 patch rows a call (`_VISION_CALL_MAX_PATCHES`). + The tower is rank-invariant, so every member (tp4, tep4, dep4) calls it at + the unsplit width 4304. - Certified on contiguous bf16 `[rows, 4304]` activands (what the block MLP's - fc1 returns) at 1, 512 and 32768 rows (32768, the tower call ceiling). + Source: the aten CUDA kernel compiled into torch (no JIT path), one launch a + call, computing in fp32 and rounding once to bf16. - Identity is pinned by fp32 only (the identity test). - Max |tanh GELU - erf GELU| is 4.73e-4 (|x| ~ 2.7); in fp32 the erf form - lands ~1600x the correct path's deviation away, 31x outside the default - gate. In bf16 one ulp at the output scale exceeds that gap (erf variant - 8.19e-3 vs correct 7.75e-3 max abs), so a bf16 match does not say which GELU - was computed. + Certified on contiguous bf16 `[rows, 4304]` activands (what the block MLP's + fc1 returns) at 1, 512 and 32768 rows (32768, the tower call ceiling): + fc1-scaled activands (x std ~1.8, |x| <= 8.8) at 512 and 32768 rows, + uniform [-16, 16] at 1 and 32768. Gate: torch's bf16 default against the + fp64 reference. + + Error, measured on GB300 at every cell: within one bf16 ulp of the reference + everywhere except x in [-7.19, -4.56], where fp32 cancellation in `1 + tanh` + leaves the output up to 254 ulp of its own size off (abs error <= 7.6e-8 at + |ref| <= 3.6e-6). HF's `gelu_pytorch_tanh` is this same call. x >= 0 is + bitwise the reference. Against the unrounded fp64 form the fc1 cells sit at + bf16 output rounding (relative L2 1.4e-3, max abs 7.755e-3: half an ulp at + |y| in [2, 4)). + + Identity is pinned by fp32 only (the identity test): max |tanh GELU - erf + GELU| is 4.73e-4 (|x| ~ 2.7), 31x outside the default fp32 gate. In bf16 one + ulp at the output scale exceeds that gap, so a bf16 match does not say which + GELU was computed. Measured, not certified: `-inf -> nan` (the closed form does the same: -inf times a factor saturated to 0); `+inf -> +inf`; large negatives give `-0.0`; `finfo.max` comes back unchanged in bf16/fp16/fp32 (the cube is not - taken in the input dtype); CUDA-graph capturable; the output is a fresh + taken in the input dtype); CUDA-graph capturable, replay bitwise eager at 1, + 512 and 32768 rows (the call site is eager); the output is a fresh allocation. Not certified: fp32 beyond the identity test, non-contiguous inputs, fp16, - other widths, float8 operands. No index-typed argument; numel past 2^31 - not probed. + other widths, float8 operands, row counts between the cells. No index-typed + argument; numel past 2^31 (498951 rows of 4304) not probed. """ def __call__(self, x: torch.Tensor) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/silu_and_mul_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/silu_and_mul_rows.py similarity index 64% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/silu_and_mul_rows.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/silu_and_mul_rows.py index cae68c64e1a9..74088ea716fd 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/silu_and_mul_rows.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/activation/silu_and_mul_rows.py @@ -3,10 +3,9 @@ """SiLU-and-mul over the leading `2 * inter` columns of each row, read through the row stride: the shared expert's activation taken out of its fused gate|up|gate-logit GEMM output. -Tier 5 (Triton written for this target). `flashinfer_silu_and_mul` addresses row `t` at -`t * 2 * inter`, so its input must be exactly the gate|up columns, contiguous; here they are -the head of a wider row (the same GEMM also forms the per-token gate logit) and are read in -place. +Tier 5: a Triton kernel. `flashinfer_silu_and_mul` addresses row `t` at `t * 2 * inter`, so +its input must be exactly the gate|up columns, contiguous; here they are the head of a wider +row (the same GEMM also forms the per-token gate logit) and are read in place. """ from __future__ import annotations @@ -50,24 +49,38 @@ class SiluAndMulRows(OpWrapper): Cell(why="prefill chunk 16384 rows of the fused 1288-wide output", spec=dict(rows=16384)), Cell(why="short prefill step, 128 rows of the fused output", spec=dict(rows=128)), Cell(why="zero rows: no launch", spec=dict(rows=0)), + Cell(why="decode, one request: the shared expert's single row", spec=dict(rows=1)), + Cell( + why="decode, 32 rows: a dep4 rank's local share of a 128-request batch", + spec=dict(rows=32), + ), + Cell(why="MTP-7 verify: 128 requests x 8 rows", spec=dict(rows=1024)), + Cell(why="prefill chunk at max_num_tokens 8192", spec=dict(rows=8192)), ) note = """ - Metadata consumed: none; every row independent. Call site: the prefill target's MoE - block, on the shared expert's fused `[gate | up | gate logit | pad]` GEMM output. + Metadata consumed: none; every row independent. Call sites: the decode and prefill + targets' MoE block, on the replicated shared expert's fused + `[gate | up | gate logit | pad]` GEMM output, so inter 640 and the 1288-wide row at + every member (tp4, tep4, dep4). Decode rows: 1..128 per rank at tp4/tep4, 1..32 local + at dep4, 1024 at an MTP-7 verify of 128 requests; prefill chunks of up to + max_num_tokens. Preconditions: 2-D `x` with unit column stride, `x.shape[1] >= 2 * inter`; bf16 or fp16. - Launch: `(rows, cdiv(inter, 1024))` programs of 4 warps. GB300, 16384 rows of the - 1288-wide output: 16.4 us (`flashinfer_silu_and_mul` on a contiguous copy of the gate|up - columns: 20.5 us for the kernel, 62.1 with the copy). + Source: Triton, JIT-compiled at first call for the running device; the cells ran that + path on sm_103 (GB300). Launch: `(rows, cdiv(inter, 1024))` programs of 4 warps, the + same kernel at every row count (no path switch). Row offsets are int64, so no index + product overflows inside the envelope (2^31 elements of the 1288-wide input is 1.67M + rows). Numerics: `g / (1 + exp(-g)) * u` in fp32 (Triton's `exp` and division), rounded once; - bitwise `flashinfer_silu_and_mul` on the same values (16384 x 640 random bf16 values, - the test's cells). + bitwise `flashinfer_silu_and_mul` on the same values at every non-empty cell, including + 0, -0, 88 and -100. - Certified: bf16 at inter 640 on the 1288-wide fused output, 0, 128 and 16384 rows; - CUDA-graph replay at 128 rows. Not certified: fp16, other widths or row counts. + Certified: bf16 at inter 640 on the 1288-wide fused output, 0, 1, 32, 128, 1024, 8192 + and 16384 rows; CUDA-graph replay at 1, 32 and 128 rows. Not certified: fp16, other + widths; row counts between the cells run the same kernel but were not driven. """ def __call__(self, x: torch.Tensor, inter: int) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_canonical_selection.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_canonical_selection.py new file mode 100644 index 000000000000..7ec7256b8eff --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_canonical_selection.py @@ -0,0 +1,608 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA block selection, made canonical: one answer per row, whichever tied columns the +selection engine happened to keep. + +Tier 5 (Triton). No registered op rewrites a selection, and none +of the top-k engines (`indexer_topk_prefill`, the CuTe DSL kernels, FlashInfer's top-k) +takes a tie-break rule or fixes the order of its output. A row's k-th value is read off +the engine's picks; the window is counted against it once and compacted into the row in +a second, early-exiting pass. Only the window's columns are read, and nothing +`[rows, cols]`-sized is allocated. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: This checkpoint's block budget: `indexer_budget` 2048 tokens over +#: `indexer_compress_ratio` 4. +_BLOCK_TOPK = 512 +#: Score-row width the target hands over: a 260-page block table at 32 tokens +#: per block, compressed 4:1. +_SCORE_COLUMNS = 260 * 32 // 4 +#: The score-row width at the 256k-token context: 65536 compressed blocks. +_LONG_SCORE_COLUMNS = 262144 // 4 +#: From this many rows on, one program per row keeps the GPU busy and rereads its window +#: from L2 (`_canonical_rows_kernel`). Below it a row's window is cut into segments of +#: `_SEGMENT` columns, each its own program, so a decode-sized launch is not one program +#: streaming a 256 KiB row. +_SPLIT_ROWS = 1024 +_SEGMENT = 4096 +#: Columns per step of the counting pass, of the writing pass (whose in-block prefix sum +#: costs more per column), and warps per program, tuned on GB300 for both launch shapes. +_ROWS_COUNT_BLOCK = 2048 +_SPLIT_COUNT_BLOCK = 4096 +_WRITE_BLOCK = 1024 +_NUM_WARPS = 4 + + +def _cell(why: str, rows, windows, logits: str, tail: str = "nan", **extra) -> Cell: + return Cell( + why=why, + spec=dict( + rows=rows, + cols=extra.pop("cols", _SCORE_COLUMNS), + index_topk=_BLOCK_TOPK, + windows=windows, + logits=logits, + tail=tail, + selection=extra.pop("selection", "random"), + **extra, + ), + ) + + +@triton.jit +def _kth_value( + row_logits, + row_indices, + start, + num_cols, + stride_col, + stride_index, + K: tl.constexpr, + K_P2: tl.constexpr, +): + """The smallest score the engine kept, read at the clamped column `start + pick`: an + empty slot (`pick < 0`) reads +inf, and a NaN pick makes it NaN, which no score + compares to, so the row comes out empty.""" + slots = tl.arange(0, K_P2) + in_k = slots < K + picks = tl.load(row_indices + slots * stride_index, mask=in_k, other=-1) + columns = tl.minimum(tl.maximum(picks + start, 0), num_cols - 1) + picked = tl.load( + row_logits + columns.to(tl.int64) * stride_col, + mask=in_k & (picks >= 0), + other=float("inf"), + ).to(tl.float32) + kth = tl.min(picked, axis=0) + picked_nan = tl.max((picked != picked).to(tl.int32), axis=0) + return tl.where(picked_nan > 0, float("nan"), kth) + + +@triton.jit +def _load_window(row_logits, cols, lo, hi, stride_col): + # NaN past the window: it compares false, so it is never counted or kept. + return tl.load( + row_logits + cols.to(tl.int64) * stride_col, + mask=(cols >= lo) & (cols < hi), + other=float("nan"), + ).to(tl.float32) + + +@triton.jit +def _count(x, kth): + """Columns of `x` above the k-th value and tied with it (float equality: -0.0 == +0.0, + NaN equals nothing).""" + return tl.sum((x > kth).to(tl.int32), axis=0), tl.sum((x == kth).to(tl.int32), axis=0) + + +@triton.jit +def _write_block( + row_indices, + stride_index, + cols, + start, + x, + kth, + drop, + kept, + ties_seen, + K: tl.constexpr, + WRITE_BLOCK: tl.constexpr, +): + """Compact one block of the window into the row; returns the updated running counts. + + Kept: every column above the k-th value, and every tie whose 0-based rank among the + window's ties is at least `drop` -- the lowest-index `drop` ties are the surplus. The + j-th kept column, ascending, goes to slot j as a window-local offset; slots from k on + are not written. One prefix sum carries both counts, above in the high half-word.""" + tl.static_assert(WRITE_BLOCK < 65536) + is_above = (x > kth).to(tl.int32) + is_tied = (x == kth).to(tl.int32) + packed = is_above * 65536 + is_tied + before = tl.cumsum(packed, axis=0) - packed + ties_before = ties_seen + (before & 65535) + kept_ties = tl.maximum(ties_seen - drop, 0) + slot = kept + (before >> 16) + tl.maximum(ties_before - drop, 0) - kept_ties + keep = (is_above != 0) | ((is_tied != 0) & (ties_before >= drop)) + tl.store(row_indices + slot * stride_index, cols - start, mask=keep & (slot < K)) + block = tl.sum(packed, axis=0) + ties_seen += block & 65535 + kept += (block >> 16) + tl.maximum(ties_seen - drop, 0) - kept_ties + return kept, ties_seen + + +@triton.jit +def _canonical_rows_kernel( + logits_ptr, + starts_ptr, + ends_ptr, + indices_ptr, + num_cols, + stride_starts, + stride_ends, + stride_logits_row, + stride_logits_col, + stride_indices_row, + stride_indices_col, + K: tl.constexpr, + K_P2: tl.constexpr, + COUNT_BLOCK: tl.constexpr, + WRITE_BLOCK: tl.constexpr, +): + # One program per row: both passes, the second reading the window back from L2. Each + # loop loads its next block before reducing the current one. Every pick is read + # before the k-th value's reduction completes and every store depends on it, so the + # in-place rewrite never overwrites a pick another thread has yet to read. + row = tl.program_id(0).to(tl.int64) + row_logits = logits_ptr + row * stride_logits_row + row_indices = indices_ptr + row * stride_indices_row + start = tl.load(starts_ptr + row * stride_starts).to(tl.int32) + end = tl.load(ends_ptr + row * stride_ends).to(tl.int32) + kth = _kth_value( + row_logits, row_indices, start, num_cols, stride_logits_col, stride_indices_col, K, K_P2 + ) + # The window's columns inside the matrix, in blocks aligned to the block size so the + # loads vectorize. + lo = tl.maximum(start, 0) + hi = tl.minimum(end, num_cols) + + above = tl.zeros((), dtype=tl.int32) + tied = tl.zeros((), dtype=tl.int32) + first = (lo // COUNT_BLOCK) * COUNT_BLOCK + x = _load_window(row_logits, first + tl.arange(0, COUNT_BLOCK), lo, hi, stride_logits_col) + for c0 in range(first, hi, COUNT_BLOCK): + c0 = tl.multiple_of(c0, COUNT_BLOCK) + cols = c0 + COUNT_BLOCK + tl.arange(0, COUNT_BLOCK) + x_next = _load_window(row_logits, cols, lo, hi, stride_logits_col) + block_above, block_tied = _count(x, kth) + above += block_above + tied += block_tied + x = x_next + # A window of at most k has no surplus. + drop = tl.maximum(above + tied - K, 0) + + kept = tl.zeros((), dtype=tl.int32) + ties_seen = tl.zeros((), dtype=tl.int32) + c0 = (lo // WRITE_BLOCK) * WRITE_BLOCK + x = _load_window(row_logits, c0 + tl.arange(0, WRITE_BLOCK), lo, hi, stride_logits_col) + while (c0 < hi) & (kept < K): + c0 = tl.multiple_of(c0, WRITE_BLOCK) + cols = c0 + tl.arange(0, WRITE_BLOCK) + x_next = _load_window(row_logits, cols + WRITE_BLOCK, lo, hi, stride_logits_col) + kept, ties_seen = _write_block( + row_indices, + stride_indices_col, + cols, + start, + x, + kth, + drop, + kept, + ties_seen, + K, + WRITE_BLOCK, + ) + x = x_next + c0 += WRITE_BLOCK + slots = tl.arange(0, K_P2) + tl.store(row_indices + slots * stride_indices_col, -1, mask=(slots < K) & (slots >= kept)) + + +@triton.jit +def _split_kth_kernel( + logits_ptr, + starts_ptr, + indices_ptr, + kth_ptr, + num_cols, + stride_starts, + stride_logits_row, + stride_logits_col, + stride_indices_row, + stride_indices_col, + K: tl.constexpr, + K_P2: tl.constexpr, +): + # Its own launch: the segments' programs that write the row come after every pick of + # it has been read. + row = tl.program_id(0).to(tl.int64) + start = tl.load(starts_ptr + row * stride_starts).to(tl.int32) + kth = _kth_value( + logits_ptr + row * stride_logits_row, + indices_ptr + row * stride_indices_row, + start, + num_cols, + stride_logits_col, + stride_indices_col, + K, + K_P2, + ) + tl.store(kth_ptr + row, kth) + + +@triton.jit +def _split_count_kernel( + logits_ptr, + starts_ptr, + ends_ptr, + kth_ptr, + counts_ptr, + num_rows, + num_cols, + num_segments, + stride_starts, + stride_ends, + stride_logits_row, + stride_logits_col, + SEGMENT: tl.constexpr, + COUNT_BLOCK: tl.constexpr, +): + # One program per (row, segment): the segment's window columns above and tied with + # the k-th value, into `counts[0 | 1, row, segment]`. + row = tl.program_id(0).to(tl.int64) + segment = tl.program_id(1) + row_logits = logits_ptr + row * stride_logits_row + seg_lo = segment * SEGMENT + start = tl.load(starts_ptr + row * stride_starts).to(tl.int32) + end = tl.load(ends_ptr + row * stride_ends).to(tl.int32) + lo = tl.maximum(tl.maximum(start, 0), seg_lo) + hi = tl.minimum(tl.minimum(end, num_cols), seg_lo + SEGMENT) + kth = tl.load(kth_ptr + row) + above = tl.zeros((), dtype=tl.int32) + tied = tl.zeros((), dtype=tl.int32) + for c0 in range(seg_lo, hi, COUNT_BLOCK): + c0 = tl.multiple_of(c0, COUNT_BLOCK) + x = _load_window(row_logits, c0 + tl.arange(0, COUNT_BLOCK), lo, hi, stride_logits_col) + block_above, block_tied = _count(x, kth) + above += block_above + tied += block_tied + tl.store(counts_ptr + row * num_segments + segment, above) + tl.store(counts_ptr + (num_rows + row) * num_segments + segment, tied) + + +@triton.jit +def _split_write_kernel( + logits_ptr, + starts_ptr, + ends_ptr, + kth_ptr, + counts_ptr, + indices_ptr, + num_rows, + num_cols, + num_segments, + stride_starts, + stride_ends, + stride_logits_row, + stride_logits_col, + stride_indices_row, + stride_indices_col, + K: tl.constexpr, + K_P2: tl.constexpr, + SEGMENT: tl.constexpr, + WRITE_BLOCK: tl.constexpr, + COUNTS_BLOCK: tl.constexpr, +): + # One program per (row, segment): the row's totals and the counts of the segments + # before this one place the segment's kept columns; segment 0 also writes the `-1`s. + row = tl.program_id(0).to(tl.int64) + segment = tl.program_id(1) + row_logits = logits_ptr + row * stride_logits_row + row_indices = indices_ptr + row * stride_indices_row + start = tl.load(starts_ptr + row * stride_starts).to(tl.int32) + end = tl.load(ends_ptr + row * stride_ends).to(tl.int32) + seg_lo = segment * SEGMENT + lo = tl.maximum(tl.maximum(start, 0), seg_lo) + hi = tl.minimum(tl.minimum(end, num_cols), seg_lo + SEGMENT) + + above = tl.zeros((), dtype=tl.int32) + tied = tl.zeros((), dtype=tl.int32) + above_before = tl.zeros((), dtype=tl.int32) + ties_before = tl.zeros((), dtype=tl.int32) + for j0 in range(0, num_segments, COUNTS_BLOCK): + j = j0 + tl.arange(0, COUNTS_BLOCK) + seg_above = tl.load(counts_ptr + row * num_segments + j, mask=j < num_segments, other=0) + seg_tied = tl.load( + counts_ptr + (num_rows + row) * num_segments + j, mask=j < num_segments, other=0 + ) + above += tl.sum(seg_above, axis=0) + tied += tl.sum(seg_tied, axis=0) + above_before += tl.sum(tl.where(j < segment, seg_above, 0), axis=0) + ties_before += tl.sum(tl.where(j < segment, seg_tied, 0), axis=0) + drop = tl.maximum(above + tied - K, 0) + if segment == 0: + total = above + tl.maximum(tied - drop, 0) + slots = tl.arange(0, K_P2) + tl.store(row_indices + slots * stride_indices_col, -1, mask=(slots < K) & (slots >= total)) + + kth = tl.load(kth_ptr + row) + kept = above_before + tl.maximum(ties_before - drop, 0) + ties_seen = ties_before + c0 = seg_lo + while (c0 < hi) & (kept < K): + c0 = tl.multiple_of(c0, WRITE_BLOCK) + cols = c0 + tl.arange(0, WRITE_BLOCK) + x = _load_window(row_logits, cols, lo, hi, stride_logits_col) + kept, ties_seen = _write_block( + row_indices, + stride_indices_col, + cols, + start, + x, + kth, + drop, + kept, + ties_seen, + K, + WRITE_BLOCK, + ) + c0 += WRITE_BLOCK + + +class QsaCanonicalSelection(OpWrapper): + """Rewrite an exact top-k selection, in place, into the window's canonical selection. + + For row `r` with window `[s, e) = [row_starts[r], row_ends[r])` of `logits` and + `k = indices.shape[1]`, `indices[r]` holds on entry what a top-k engine wrote: + window-local offsets of an exact top-k of the window, in any order, `-1` past + `min(k, e - s)`. Let `v` be the smallest selected score. On return the row is + + * every window column scoring strictly above `v`, then the highest-index window + columns scoring exactly `v` until there are `k` -- the k largest + `(score, column)` pairs; or the whole window when `e - s <= k`; + * as window-local offsets in ascending order, `-1` in every remaining slot. + + So the row is a function of `(logits, row_starts, row_ends, k)` alone: two engines, + or two launches of one engine, that broke a tie at the k-th value differently produce + the same row. On a row with no tie at the k-th value this is the engine's own set, + sorted. Ties are float equality: `-0.0` and `+0.0` are one tie class. Returns + `indices`. + + Fusion boundary: the rewrite only. Selecting is the engine's call before it; the + engine must be exact -- `v` is read off its picks, and a pick below the true k-th + value would admit more than `k` columns above it. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The target calls this right after `qsa_topk_select`, at `k = 512` over the + # `[rows, cols]` score matrix whose columns past each `row_end` are allocator + # garbage (the cells poison them). `selection` names what produced the input row: + # `random` is an exact top-k that keeps a random subset of the tie class at the k-th + # value, in a random order (the adversary); `radix` is the engine the target calls. + CELLS: tuple[Cell, ...] = ( + _cell( + "decode: 1, 128 and 256 rows (MTP3 and MTP7 at batch 32)", + (1, 128, 256), + ("positions_random", 2051, 8191), + "relu", + ), + _cell( + "prefill: 8192 rows from position 0", + 8192, + ("chunk", 0), + "relu", + ), + _cell( + "input from the radix engine (qsa_topk_select), zero ties", + 64, + ("positions_random", 2051, 8191), + "relu_sparse", + selection="radix", + ), + _cell( + "decode at up to 256k context: 65536-wide table, relu-zero ties", + (1, 16, 64), + ("positions_random", 2051, 262143), + "relu_sparse", + cols=_LONG_SCORE_COLUMNS, + ), + _cell( + "prefill: a 2048-row chunk ending at the 256k context", + 2048, + ("chunk", 262144 - 2048), + "relu", + cols=_LONG_SCORE_COLUMNS, + ), + ) + + note = """ + Metadata consumed: none. Stateless; every input is an argument, and `indices` is + both the engine's selection read and the canonical row written. + + Preconditions (none checked): + - `indices` holds an EXACT top-k of each window: the smallest selected score is + the window's k-th largest. `qsa_topk_select` is exact; an approximate selection + would be rewritten into a row with more than `k` columns at or above its k-th + pick, truncated to the lowest `k` column indices -- wrong, silently. + - `0 <= row_starts[r] <= row_ends[r] <= logits.shape[1]`. Anything outside the + window is never read for a decision, so garbage, NaN and +inf there are fine + (the cells poison it). + - Scores inside the window are not NaN: a NaN column is never kept, and a NaN pick + empties the row. (The engine returns the identity for a window of at most k + without reading it; this rewrite reads it.) + + Cost: no host read, no data-dependent shape, so it can be captured into a CUDA graph + (the entry's test replays one). A row reads its k picks, its window once to count + and again, up to its k-th kept column, to write; columns outside the window are not + read, so narrowing `logits` to a bound on the windows (`logits[:, :bound]`) gives + the same rows and saves nothing. From 1024 rows: one launch, one program per row. + Below: three launches -- the k-th values, then counts and writes over 4096-column + segments, each its own program -- through `[rows]` fp32 and `[2, rows, cols / 4096]` + int32 scratch. + """ + + def __call__( + self, + logits: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + indices: torch.Tensor, + ) -> torch.Tensor: + num_rows, num_cols = logits.shape + index_topk = indices.shape[1] + if num_rows == 0 or index_topk == 0: + return indices + # A one-element window broadcasts over every row -- the hazard `is_valid` + # refuses -- as a stride-0 view rather than a read past the tensor's end. + row_starts = row_starts.expand(num_rows) + row_ends = row_ends.expand(num_rows) + k_p2 = triton.next_power_of_2(index_topk) + if num_rows >= _SPLIT_ROWS: + _canonical_rows_kernel[(num_rows,)]( + logits, + row_starts, + row_ends, + indices, + num_cols, + row_starts.stride(0), + row_ends.stride(0), + logits.stride(0), + logits.stride(1), + indices.stride(0), + indices.stride(1), + K=index_topk, + K_P2=k_p2, + COUNT_BLOCK=_ROWS_COUNT_BLOCK, + WRITE_BLOCK=_WRITE_BLOCK, + num_warps=_NUM_WARPS, + ) + return indices + num_segments = triton.cdiv(num_cols, _SEGMENT) + kth = torch.empty(num_rows, dtype=torch.float32, device=logits.device) + counts = torch.empty(2, num_rows, num_segments, dtype=torch.int32, device=logits.device) + _split_kth_kernel[(num_rows,)]( + logits, + row_starts, + indices, + kth, + num_cols, + row_starts.stride(0), + logits.stride(0), + logits.stride(1), + indices.stride(0), + indices.stride(1), + K=index_topk, + K_P2=k_p2, + num_warps=_NUM_WARPS, + ) + _split_count_kernel[(num_rows, num_segments)]( + logits, + row_starts, + row_ends, + kth, + counts, + num_rows, + num_cols, + num_segments, + row_starts.stride(0), + row_ends.stride(0), + logits.stride(0), + logits.stride(1), + SEGMENT=_SEGMENT, + COUNT_BLOCK=_SPLIT_COUNT_BLOCK, + num_warps=_NUM_WARPS, + ) + _split_write_kernel[(num_rows, num_segments)]( + logits, + row_starts, + row_ends, + kth, + counts, + indices, + num_rows, + num_cols, + num_segments, + row_starts.stride(0), + row_ends.stride(0), + logits.stride(0), + logits.stride(1), + indices.stride(0), + indices.stride(1), + K=index_topk, + K_P2=k_p2, + SEGMENT=_SEGMENT, + WRITE_BLOCK=_WRITE_BLOCK, + COUNTS_BLOCK=64, + num_warps=_NUM_WARPS, + ) + return indices + + def reference( + self, + logits: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + indices: torch.Tensor, + ) -> torch.Tensor: + """The k largest `(score, column)` pairs of each window, ascending by column, + `-1`-padded: a top-k over one int64 key per column whose high word orders the + scores as fp32 does and whose low word is the column. `indices` supplies only the + width `k`; its contents are ignored.""" + num_rows, num_cols = logits.shape + index_topk = indices.shape[1] + columns = torch.arange(num_cols, device=logits.device)[None, :] + starts, ends = row_starts.long()[:, None], row_ends.long()[:, None] + in_window = (columns >= starts) & (columns < ends) + # `+ 0.0` folds -0.0 into +0.0, so the two are one tie class. + bits = (logits.float() + 0.0).view(torch.int32).long() + order = torch.where(bits < 0, bits ^ 0x7FFFFFFF, bits) + key = torch.where(in_window, (order << 32) | columns, torch.iinfo(torch.int64).min) + width = min(index_topk, num_cols) + top = torch.topk(key, width, dim=1).indices + taken = torch.gather(in_window, 1, top) + pad = torch.iinfo(torch.int64).max + ordered = torch.sort(torch.where(taken, top, pad), dim=1).values + out = torch.where(ordered == pad, -1, ordered - starts).to(torch.int32) + if width < index_topk: + out = torch.nn.functional.pad(out, (0, index_topk - width), value=-1) + return out + + def is_valid( + self, + logits: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + indices: torch.Tensor, + ) -> None: + # A one-element `row_starts` / `row_ends`, or a one-row `indices`, broadcasts over + # every row of `logits` without an error: a [1] `row_starts` against 4 rows + # returns every row framed by row 0's start. + num_rows = logits.shape[0] + assert row_starts.shape == (num_rows,) and row_ends.shape == (num_rows,), ( + f"row_starts and row_ends need one entry per logits row ({num_rows}); got " + f"{tuple(row_starts.shape)} and {tuple(row_ends.shape)}" + ) + assert indices.dim() == 2 and indices.shape[0] == num_rows, ( + f"indices must be [{num_rows}, k]; got {tuple(indices.shape)}" + ) + + +qsa_canonical_selection = QsaCanonicalSelection() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_expand_block_indices.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_expand_block_indices.py similarity index 74% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_expand_block_indices.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_expand_block_indices.py index ba1cf5acd32e..2133bf12c615 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_expand_block_indices.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_expand_block_indices.py @@ -20,16 +20,28 @@ #: First positions of 32 MTP verify blocks in the sparse branch (KV past 2048). _VERIFY_STARTS = tuple(2100 + (499 * g) % 6000 for g in range(32)) +#: The longest sequence the default target serves: `max_seq_len` 262144, with +#: `max_num_tokens` 8192 and chunked prefill. +_ENVELOPE_SEQ_LEN = 262144 + +#: First positions of 128 MTP7 verify blocks, spread over the sparse branch up to +#: the envelope (the last block ends at 262107). +_VERIFY_STARTS_128 = tuple(2100 + (2039 * g) % 260000 for g in range(128)) + def _spec(batch, **overrides) -> dict: """A cell spec at this checkpoint's geometry and the engine's metadata dtypes. `batch` is one of: ("sweep", n) n rows at 4096 + 13 i, length = position + 1 + ("decode", n, top) n one-row requests spread over [1, top), + every tail width, length = position + 1 ("serving", requests, context_lens) one packed step: context chunks, then two MTP rows per remaining request ("verify", starts, p) one MTP verify block of p rows per start, every row's length the block's end + ("chunk", start, n) one context chunk of n rows from `start`, + every row's length the chunk's end """ base = dict( batch=batch, @@ -65,6 +77,9 @@ class QsaExpandBlockIndices(OpWrapper): Fusion boundary: index arithmetic and the tail only. Scoring, Top-K, the per-row causal group bound and the geometry are the caller's. + + Called by the qwen4_exp target at every layout; its row counts and sequence + lengths are among the cells. """ ARCHS = frozenset({Arch.SM_103}) @@ -84,6 +99,20 @@ class QsaExpandBlockIndices(OpWrapper): why="warm-up capture: 2 x 4096 rows, the [8192, 512] buffer's height", spec=_spec(("serving", 2, (4096, 4096))), ), + # ---- the default target's decode, verify and prefill extents (tp4 / tep4 / dep4 + # alike: one selection per query row, no head or rank axis). + Cell( + why="decode batch 128: the first whole-row launch, KV up to the 262144 envelope", + spec=_spec(("decode", 128, _ENVELOPE_SEQ_LEN)), + ), + Cell( + why="verify MTP7 at 128 requests: 1024 rows, KV up to the 262144 envelope", + spec=_spec(("verify", _VERIFY_STARTS_128, 8)), + ), + Cell( + why="context chunk of max_num_tokens 8192 rows ending a 262144-token prompt", + spec=_spec(("chunk", _ENVELOPE_SEQ_LEN - 8192, 8192)), + ), ) note = """ @@ -111,16 +140,30 @@ class QsaExpandBlockIndices(OpWrapper): the test pins the width-vs-geometry rejection. A one-element strided metadata view reports contiguous and is accepted -- benign at one row. - Launch: one program per (row, 256 columns) below 128 rows, one per row - from 128 up. No autotune. CUDA-graph capturable (captured at 256 rows); a - replay recomputes from the live buffers, and the captured row count fixes - the launch shape. - - Hard limits (read from the kernel, not probed): the row offset is - `program_id * stride(0)` in int32, so `rows * stride(0) >= 2^31` (4 Mi - rows at width 512) would wrap. Certified extent: 8192 rows x 2051 columns. + Source: a Triton kernel JIT-compiled on first call for the device it runs + on; the cells were driven through that JIT path on sm_103 (GB300). - Not certified: multiple devices, `torch.compile`, rows above 8192. + Launch: one program per (row, 256 columns) below 128 rows, one per row + from 128 up. No autotune. CUDA-graph capturable (captured at 1, 128, 256 + and 1024 rows: both launch shapes); a replay recomputes from the live + buffers, and the captured row count fixes the launch shape. + + Rank-invariant: one selection per query row, no head or rank axis, so tp4, + tep4 and dep4 call it with the same geometry; only the row count and the + positions move. + + Hard limits (read from the kernel, not probed): both row strides are + `tl.constexpr` and the row offset is `program_id * stride` in int32. The + output's stride is `final_topk` (2051), so the store wraps first, at + `rows > 2^31 / 2051` (about 1.05 M rows); the input read wraps at + `rows * stride(0) >= 2^31` (4 Mi rows at width 512). Token ids are int32, + exact far past the envelope. The envelope reaches neither: `max_num_tokens` + 8192 rows, positions below `max_seq_len` 262144. + Certified extent: 8192 rows x 2051 columns, positions up to 262143. + + Not certified: multiple devices, `torch.compile`, rows above 8192, and the + column-split launch at 2..127 rows (only its one-row shape is a cell; + decode batches 2..120 are captured there). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_mtp_reuse_topk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_mtp_reuse_topk.py similarity index 71% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_mtp_reuse_topk.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_mtp_reuse_topk.py index 582e514682ea..3ed07e74c1c1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_mtp_reuse_topk.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_mtp_reuse_topk.py @@ -17,7 +17,20 @@ _BLOCK_TOPK = 512 #: `(slots, this layer's slot)` of the shared buffers: with the MTP layer in the #: target's cache manager, 12 trunk full-attention layers plus the MTP layer. +#: This is dep4's layout (attention DP keeps the draft layer in the trunk's manager). _COMBINED_SLOTS = (13, 12) +#: tp4 / tep4's layout: attention TP puts the draft layer in a separate draft KV +#: cache, so the metadata sizes the buffers by the larger manager (the trunk's 12 +#: sparse layers) and the draft loop reads the draft manager's only slot, 0. +_SEPARATE_DRAFT_SLOTS = (12, 0) + +#: The engines' `max_batch_size` at every member and variant, and so the requests +#: one draft step carries (one row each) and the shared buffers' height +#: (`max_num_sequences` = `max_batch_size * max_beam_width`, beam width 1). +_MAX_BATCH = 128 +#: The longest sequence the default target serves: `max_seq_len` 262144 (16384 +#: on the MTP member). +_ENVELOPE_SEQ_LEN = 262144 def _spec(contexts, **overrides) -> dict: @@ -25,12 +38,14 @@ def _spec(contexts, **overrides) -> dict: `contexts` says what each request's captured row and current row are: ("mixed", B) B requests; the first 16 at fixed boundary positions - (empty, short, 511, 512 and full rows), the rest seeded + (empty, short, 511, 512 and full rows), plus one at the + last position `max_seq_len` allows, the rest seeded ("full", B) captured positions >= 2047: all 512 slots valid ("explicit", P) one request per captured position in P A request captured at position `p` (its last accepted row) holds `(p + 1) // 4` complete groups; its draft row `d` positions later sees - `(p + d + 1) // 4`. `drafted` cycles over the requests. + `(p + d + 1) // 4`. `drafted` cycles over the requests. Positions stay below + `max_seq_len - d`. """ base = dict( contexts=contexts, @@ -41,10 +56,24 @@ def _spec(contexts, **overrides) -> dict: block_topk=_BLOCK_TOPK, metadata_dtype=torch.int32, graph_batches=(), + max_seq_len=8192, ) return {**base, **overrides} +def _default(drafted, shared_rows) -> dict: + """The default target's draft step: 128 requests, tp4 / tep4's separate-draft + slots, positions over the 262144 envelope, captured at 1 and 128 requests.""" + return _spec( + ("mixed", _MAX_BATCH), + drafted=drafted, + shared_rows=shared_rows, + slots=_SEPARATE_DRAFT_SLOTS, + graph_batches=(1, _MAX_BATCH), + max_seq_len=_ENVELOPE_SEQ_LEN, + ) + + class QsaMtpReuseTopk(OpWrapper): """Rebuild each query row's group selection from its request's captured row, with the groups that completed since the capture written into it. @@ -78,6 +107,8 @@ class QsaMtpReuseTopk(OpWrapper): into the shared buffers, the per-row request map and visible counts, and the expansion into token indices plus the incomplete causal tail (`qsa_expand_block_indices`) are the caller's. + + Called by the qwen4_exp target at every layout. """ ARCHS = frozenset({Arch.SM_103}) @@ -103,6 +134,24 @@ class QsaMtpReuseTopk(OpWrapper): why="CUDA graph, drafted 1..6: captured at 1, 8 and 32 requests, new-data replay", spec=_spec(("mixed", 32), drafted=(1, 2, 3, 4, 5, 6), graph_batches=(1, 8, 32)), ), + # ---- the default target (tp4 / tep4 / dep4 alike: one selection row per + # request, no head or rank axis): its 128-request draft steps. + Cell( + why="128 requests 3 positions past the capture, 512-row shared buffer (tep4 slots)", + spec=_default(drafted=(3,), shared_rows=512), + ), + Cell( + why="128 requests 1 position past the capture, 256-row shared buffer (tep4 slots)", + spec=_default(drafted=(1,), shared_rows=256), + ), + Cell( + why="128 requests 7 positions past the capture, 1024-row shared buffer (tep4 slots)", + spec=_default(drafted=(7,), shared_rows=1024), + ), + Cell( + why="served mtp7 step: 128 requests filling the engine's 128-row buffer, drafted 1..6", + spec=_default(drafted=(1, 2, 3, 4, 5, 6), shared_rows=_MAX_BATCH), + ), ) note = """ @@ -147,11 +196,38 @@ class QsaMtpReuseTopk(OpWrapper): a 2-D `shared_visible_blocks` (the whole `[slots, seqs]` buffer is read as slot 0's counts). The test pins the last. - Certified: int32 ids and metadata, `block_topk` 512, up to 32 requests, the - identity request map over the target cache manager's 13-slot buffers, CUDA - graph capture at 1, 8 and 32 requests with a new-data replay. - Not certified: a permuted request map, a separate draft manager's one-slot - buffers. + Source: tier 2, an in-tree Triton kernel in the QSA attention backend, + JIT-compiled on first call for the device it runs on; the cells were driven + through that JIT path on sm_103 (GB300). One program per output row, no + autotune, no PDL. + + The default target (tp4 / tep4 / dep4) calls it with the same geometry on every + rank: one selection row per request, no head or rank axis (the indexer is + replicated). The slots are `_SEPARATE_DRAFT_SLOTS` under attention TP (tp4, + tep4) and `_COMBINED_SLOTS` under attention DP. The buffers are `max_num_sequences` rows + tall, `max_batch_size * max_beam_width` (`build_attention_metadata`): 128 at + every member. The 256-, 512- and 1024-row cells are taller than any engine + allocates and are certified as stated. A draft step carries at most 128 + requests, one row each. + + The `drafted` axis is the offset past the capture (above), not + `max_draft_len`: mtp7's offsets 1..6 are the "served mtp7 step" cell, and the + offset-7 cell is beyond any served draft step. Every cell holds a full row + gaining the widest fresh count its offsets allow. + + Hard limits (read from the kernel, not probed): both row offsets are int64 + (`request.to(tl.int64) * BLOCK_TOPK`, `row.to(tl.int64) * BLOCK_TOPK`), so no + buffer size the envelope reaches wraps them. Ids and visible counts are int32; + positions below `max_seq_len` 262144 give at most 65536 groups. `block_topk` + must be a power of two (`tl.arange`). + + Certified: int32 ids and metadata, `block_topk` 512, the identity request map, + up to 128 requests in 8- to 1024-row shared buffers (128 requests filling a + 128-row buffer included), slots (13, 12) and (12, 0), captured positions up to + the 262144 envelope, CUDA graph capture at 1, 8, 32 and 128 requests with a + new-data replay. + Not certified: a permuted request map, more than one row per request, int64 + metadata, 128 requests at dep4's (13, 12) slots (MTP is served at tep4). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_nvfp4_pages_to_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_nvfp4_pages_to_bf16.py similarity index 99% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_nvfp4_pages_to_bf16.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_nvfp4_pages_to_bf16.py index 03f261bf15ca..32f1b145a305 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_nvfp4_pages_to_bf16.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_nvfp4_pages_to_bf16.py @@ -16,8 +16,9 @@ import triton.language as tl from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -from .qsa_paged_sparse_gqa_nvfp4 import _e2m1_to_f32 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4 import ( + _e2m1_to_f32, +) #: Values per e4m3 block scale, e2m1 values per byte, tokens per V scale group. _SF_VEC = 16 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_nvfp4_verify_ptx.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_nvfp4_verify_ptx.py similarity index 98% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_nvfp4_verify_ptx.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_nvfp4_verify_ptx.py index d6b416ac80b1..d184882c004c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_nvfp4_verify_ptx.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_nvfp4_verify_ptx.py @@ -22,7 +22,10 @@ from cutlass.cute.nvgpu import tcgen05 from cutlass.cutlass_dsl import T, dsl_user_op -from .qsa_paged_sparse_gqa_nvfp4_decode import _asm, _ir +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4_decode import ( # noqa: E501 + _asm, + _ir, +) # --------------------------------------------------------------------------------------- # scalar and memory helpers diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_index_scores.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_index_scores.py new file mode 100644 index 000000000000..6c74ae3a15c7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_index_scores.py @@ -0,0 +1,1158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA compressed-block index scoring, read straight out of the paged index-K side cache.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# Approved exception: QSA's selection path is Triton called as Python, with no +# torch.ops boundary. The engine's own sparse hook (`select_qsa_paged_tokens` +# in `_torch/attention/backends/sparse/qsa/module.py`) calls this same callable. +from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import ( + triton_qsa_paged_index_scores as _triton_qsa_paged_index_scores, +) + +_INTEGER_DTYPES = (torch.int32, torch.int64) + +#: At or below this many rows a call is decode-sized: the engine's own threshold +#: between its one-tile-a-program and its deeper-tile launches. +_DECODE_ROW_LIMIT = 256 +#: Dot columns of one visible-split program: whole query rows of `heads` columns each +#: (`tl.dot` needs at least 16), so the rows of one request share each key gather. +_SPLIT_DOT_WIDTH = 16 +#: Score columns a tile gathers keys for. +_SPLIT_BLOCK_N = 64 +#: Programs a launch holds per SM, empty request slots included. One row tile (one +#: request's verify rows: a chain of dependent loads) wants a tile a program, up to the +#: 1024 tiles of a 262143-position window; wider batches stream several requests at once. +_SPLIT_PROGRAMS_PER_SM = 16 +_SPLIT_PROGRAMS_PER_SM_ONE_TILE = 32 +_SPLIT_NUM_WARPS = 2 +_SM_COUNT: dict = {} + + +# Metadata views can begin at any row of a mixed context/decode step. Their +# alignment does not change these scalar loads or the score computation. +@triton.jit( + do_not_specialize=["num_rows", "split"], + do_not_specialize_on_alignment=["query_positions", "request_indices"], +) +def _qsa_index_scores_visible_split_kernel( + q, + index_cache, + block_table, + query_positions, + request_indices, + output, + score_scale, + num_rows, + split, + q_stride_row: tl.constexpr, + q_stride_head: tl.constexpr, + q_stride_dim: tl.constexpr, + cache_stride_page: tl.constexpr, + cache_stride_token: tl.constexpr, + cache_stride_dim: tl.constexpr, + block_table_stride_request: tl.constexpr, + block_table_stride_page: tl.constexpr, + output_stride_row: tl.constexpr, + output_stride_block: tl.constexpr, + NUM_INDEX_HEADS: tl.constexpr, + INDEX_HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + TOKENS_PER_BLOCK: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + MAX_COMPRESSED_BLOCKS: tl.constexpr, + ROWS_PER_TILE: tl.constexpr, + QUERY_TILE: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_N: tl.constexpr, + NUM_STAGES: tl.constexpr, + USE_PDL: tl.constexpr, +): + """Visible-only scores of one request's rows among ROWS_PER_TILE consecutive rows (in + QUERY_TILE dot lanes). Program (t, g, s) takes the g-th distinct request of row tile t, + in row order, and walks its column tiles s, s + split, ... up to its visible bound: the + grid is sized by the device, not by the table, a tile's keys are gathered once for every + row of their request in the row tile, and different requests run side by side. + """ + if USE_PDL: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + row_base = tl.program_id(0) * ROWS_PER_TILE + group_slot = tl.program_id(1) + first_tile = tl.program_id(2) + + # One dot column per (query row, index head) pair; q is read before the metadata + # decides whether this program has work, so the two loads overlap. + dot_columns = tl.arange(0, BLOCK_H) + column_row = dot_columns // NUM_INDEX_HEADS + column_head = dot_columns % NUM_INDEX_HEADS + rows = row_base + column_row + present = (column_row < ROWS_PER_TILE) & (rows < num_rows) + dim_offsets = tl.arange(0, BLOCK_D) + query_values = tl.load( + q + + rows[None, :] * q_stride_row + + dim_offsets[:, None] * q_stride_dim + + column_head[None, :] * q_stride_head, + mask=(dim_offsets[:, None] < INDEX_HEAD_DIM) & present[None, :], + other=0.0, + ) + + # Each tile row (lane): its request and visible bound. + lanes = tl.arange(0, QUERY_TILE) + lane_rows = row_base + lanes + lane_present = (lanes < ROWS_PER_TILE) & (lane_rows < num_rows) + lane_requests = tl.load(request_indices + lane_rows, mask=lane_present, other=-1).to(tl.int64) + lane_positions = tl.load(query_positions + lane_rows, mask=lane_present, other=0) + lane_visible = tl.where( + lane_present, tl.minimum((lane_positions + 1) // COMPRESS_RATIO, MAX_COMPRESSED_BLOCKS), 0 + ) + # A lane leads its request's group when no earlier lane holds the same request; groups + # are numbered by their leader's lane. + earlier_same = ( + (lane_requests[:, None] == lane_requests[None, :]) + & (lanes[:, None] < lanes[None, :]) + & lane_present[:, None] + ) + leads = lane_present & (tl.max(earlier_same.to(tl.int32), axis=0) == 0) + rank = tl.cumsum(leads.to(tl.int32), axis=0) - 1 + request = tl.max(tl.where(leads & (rank == group_slot), lane_requests, -1)) + lane_group = lane_present & (lane_requests == request) + group_visible = tl.max(tl.where(lane_group, lane_visible, 0)) + # No such group (request -1) leaves group_visible at 0. + if first_tile * BLOCK_N >= group_visible: + return + + column_offsets = tl.arange(0, BLOCK_N) + group_tiles = tl.cdiv(group_visible, BLOCK_N) + for tile in tl.range(first_tile, group_tiles, split, num_stages=NUM_STAGES): + block_columns = tile * BLOCK_N + column_offsets + valid_blocks = block_columns < group_visible + anchor_positions = block_columns * COMPRESS_RATIO + COMPRESS_RATIO - 1 + logical_pages = anchor_positions // TOKENS_PER_BLOCK + token_in_page = anchor_positions % TOKENS_PER_BLOCK + physical_pages = tl.load( + block_table + + request * block_table_stride_request + + logical_pages * block_table_stride_page, + mask=valid_blocks, + other=0, + ).to(tl.int64) + keys = tl.load( + index_cache + + physical_pages[:, None] * cache_stride_page + + token_in_page[:, None] * cache_stride_token + + dim_offsets[None, :] * cache_stride_dim, + mask=valid_blocks[:, None] & (dim_offsets[None, :] < INDEX_HEAD_DIM), + other=0.0, + eviction_policy="evict_first", + ) + positive_scores = tl.maximum(tl.dot(keys, query_values, out_dtype=tl.float32), 0.0) + # Each row's heads are adjacent dot columns: sum them per row. + scores = tl.sum(tl.reshape(positive_scores, [BLOCK_N, QUERY_TILE, NUM_INDEX_HEADS]), axis=2) + tl.store( + output + + lane_rows[None, :] * output_stride_row + + block_columns[:, None] * output_stride_block, + scores * score_scale, + mask=lane_group[None, :] & (block_columns[:, None] < lane_visible[None, :]), + ) + + +def _sm_count(device: torch.device) -> int: + index = device.index if device.index is not None else torch.cuda.current_device() + if index not in _SM_COUNT: + _SM_COUNT[index] = torch.cuda.get_device_properties(index).multi_processor_count + return _SM_COUNT[index] + + +def _rows_per_tile(rows: int, query_tile: int) -> int: + """Rows a program's dot lanes hold: the whole batch when it fits one dot tile, else a + full tile. + """ + return min(rows, query_tile) + + +def _takes_visible_split(q: torch.Tensor, only_visible_blocks: bool, prefetch: bool) -> bool: + """The visible-split kernel serves decode-sized, visible-only calls whose head count + divides its dot width; `prefetch` keeps the engine's PDL prefetch kernel. + """ + heads = q.shape[1] if q.ndim == 3 else 0 + return ( + only_visible_blocks + and not prefetch + and q.ndim == 3 + and q.shape[0] <= _DECODE_ROW_LIMIT + and 0 < heads <= _SPLIT_DOT_WIDTH + and _SPLIT_DOT_WIDTH % heads == 0 + ) + + +def _qsa_index_scores_visible_split( + q: torch.Tensor, + index_cache: torch.Tensor, + block_table: torch.Tensor, + query_positions: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + compress_ratio: int, + use_pdl: bool, + *, + dot_width: int = _SPLIT_DOT_WIDTH, + rows_per_tile: int | None = None, + block_n: int = _SPLIT_BLOCK_N, + programs_per_sm: int | None = None, + num_warps: int = _SPLIT_NUM_WARPS, + num_stages: int = 2, +) -> torch.Tensor: + """One launch of the visible-split kernel after the callable's own checks.""" + if index_cache.ndim != 4 or index_cache.shape[2] != 1: + raise ValueError("QSA index cache must have shape [pages, tokens, 1, dim]") + if block_table.ndim != 2 or block_table.dtype != torch.int32: + raise ValueError("QSA index scoring requires a two-dimensional int32 block table") + rows, num_index_heads, index_head_dim = q.shape + if index_head_dim <= 0 or index_cache.shape[3] != index_head_dim: + raise ValueError("QSA index query and cache dimensions do not match") + if q.dtype != index_cache.dtype: + raise ValueError("QSA index query and cache dtypes must match") + if query_positions.shape != (rows,) or request_indices.shape != (rows,): + raise ValueError("QSA index-scoring metadata must have one entry per row") + if query_positions.dtype not in _INTEGER_DTYPES or request_indices.dtype not in _INTEGER_DTYPES: + raise ValueError("QSA index-scoring metadata must use integer storage") + if tokens_per_block <= 0 or compress_ratio < 2: + raise ValueError("QSA index scoring requires positive pages and compression ratio >= 2") + for t in (index_cache, block_table, query_positions, request_indices): + if t.device != q.device: + raise ValueError("QSA paged index scoring tensors must share one CUDA device") + max_compressed_blocks = block_table.shape[1] * tokens_per_block // compress_ratio + if max_compressed_blocks <= 0: + raise ValueError("QSA index scoring requires a non-empty block table") + output = torch.empty((rows, max_compressed_blocks), dtype=torch.float32, device=q.device) + if rows == 0: + return output + query_tile = dot_width // num_index_heads + if rows_per_tile is None: + rows_per_tile = _rows_per_tile(rows, query_tile) + row_tiles = triton.cdiv(rows, rows_per_tile) + if programs_per_sm is None: + programs_per_sm = ( + _SPLIT_PROGRAMS_PER_SM_ONE_TILE if row_tiles == 1 else _SPLIT_PROGRAMS_PER_SM + ) + # Every (row tile, group slot) pair gets `split` programs; slots past a tile's request + # count exit at once. + split = min( + triton.cdiv(max_compressed_blocks, block_n), + max(1, triton.cdiv(programs_per_sm * _sm_count(q.device), row_tiles * rows_per_tile)), + ) + _qsa_index_scores_visible_split_kernel[(row_tiles, rows_per_tile, split)]( + q, + index_cache, + block_table, + query_positions, + request_indices, + output, + index_head_dim**-0.5, + rows, + split, + q.stride(0), + q.stride(1), + q.stride(2), + index_cache.stride(0), + index_cache.stride(1), + index_cache.stride(3), + block_table.stride(0), + block_table.stride(1), + output.stride(0), + output.stride(1), + NUM_INDEX_HEADS=num_index_heads, + INDEX_HEAD_DIM=index_head_dim, + BLOCK_D=max(16, triton.next_power_of_2(index_head_dim)), + TOKENS_PER_BLOCK=tokens_per_block, + COMPRESS_RATIO=compress_ratio, + MAX_COMPRESSED_BLOCKS=max_compressed_blocks, + ROWS_PER_TILE=rows_per_tile, + QUERY_TILE=query_tile, + BLOCK_H=dot_width, + BLOCK_N=block_n, + NUM_STAGES=num_stages, + USE_PDL=use_pdl, + num_warps=num_warps, + launch_pdl=use_pdl, + ) + return output + + +@triton.jit +def _qsa_index_scores_group_kernel( + q, + index_cache, + block_table, + query_positions, + request_indices, + output, + score_scale, + rows, + q_stride_row: tl.constexpr, + q_stride_head: tl.constexpr, + q_stride_dim: tl.constexpr, + cache_stride_page: tl.constexpr, + cache_stride_token: tl.constexpr, + cache_stride_dim: tl.constexpr, + block_table_stride_request: tl.constexpr, + block_table_stride_page: tl.constexpr, + output_stride_row: tl.constexpr, + output_stride_block: tl.constexpr, + NUM_INDEX_HEADS: tl.constexpr, + INDEX_HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + TOKENS_PER_BLOCK: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + MAX_COMPRESSED_BLOCKS: tl.constexpr, + GROUP_ROWS: tl.constexpr, + BLOCK_G: tl.constexpr, + BLOCK_N: tl.constexpr, + TILES_PER_PROGRAM: tl.constexpr, + USE_PDL: tl.constexpr, +): + """The engine's row-wise kernel (visible blocks only) over a run of ``GROUP_ROWS`` + consecutive rows a program: when the run's rows name one request (a verify step's + rows), one key gather feeds every row's heads as the dot's columns; otherwise each + row gathers its own request's keys. Each score is the same dot, ReLU and head sum as + the row-wise kernel's, bit for bit.""" + BLOCK_H: tl.constexpr = BLOCK_G * NUM_INDEX_HEADS + if USE_PDL: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + group = tl.program_id(0) + slots = tl.arange(0, BLOCK_G) + group_rows = group * GROUP_ROWS + slots + row_ok = (slots < GROUP_ROWS) & (group_rows < rows) + lead = tl.load(request_indices + group * GROUP_ROWS).to(tl.int64) + requests = tl.load(request_indices + group_rows, mask=row_ok, other=0).to(tl.int64) + requests = tl.where(row_ok, requests, lead) + positions = tl.load(query_positions + group_rows, mask=row_ok, other=0) + visible = tl.where(row_ok, (positions + 1) // COMPRESS_RATIO, 0) + max_visible = tl.max(visible, axis=0) + first_tile = tl.program_id(1) * TILES_PER_PROGRAM + if first_tile * BLOCK_N >= max_visible: + return + last_tile = tl.minimum(first_tile + TILES_PER_PROGRAM, tl.cdiv(max_visible, BLOCK_N)) + shared = tl.max((requests != lead).to(tl.int32), axis=0) == 0 + + dim_offsets = tl.arange(0, BLOCK_D) + # Dot column c is head c % heads of the run's row c // heads. + columns = tl.arange(0, BLOCK_H) + column_slot = columns // NUM_INDEX_HEADS + column_rows = group * GROUP_ROWS + column_slot + column_ok = (column_slot < GROUP_ROWS) & (column_rows < rows) + query_values = tl.load( + q + + column_rows[None, :] * q_stride_row + + dim_offsets[:, None] * q_stride_dim + + (columns % NUM_INDEX_HEADS)[None, :] * q_stride_head, + mask=(dim_offsets[:, None] < INDEX_HEAD_DIM) & column_ok[None, :], + other=0.0, + ) + dim_ok = dim_offsets[None, :] < INDEX_HEAD_DIM + block_offsets = tl.arange(0, BLOCK_N) + for tile in tl.range(first_tile, last_tile): + block_columns = tile * BLOCK_N + block_offsets + valid_blocks = (block_columns < MAX_COMPRESSED_BLOCKS) & (block_columns < max_visible) + anchor_positions = block_columns * COMPRESS_RATIO + COMPRESS_RATIO - 1 + logical_pages = anchor_positions // TOKENS_PER_BLOCK + token_in_page = anchor_positions % TOKENS_PER_BLOCK + if shared: + physical_pages = tl.load( + block_table + + lead * block_table_stride_request + + logical_pages * block_table_stride_page, + mask=valid_blocks, + other=0, + ).to(tl.int64) + keys = tl.load( + index_cache + + physical_pages[:, None] * cache_stride_page + + token_in_page[:, None] * cache_stride_token + + dim_offsets[None, :] * cache_stride_dim, + mask=valid_blocks[:, None] & dim_ok, + other=0.0, + eviction_policy="evict_first", + ) + per_head_scores = tl.dot(keys, query_values, out_dtype=tl.float32) + else: + per_head_scores = tl.zeros((BLOCK_N, BLOCK_H), dtype=tl.float32) + for slot in tl.static_range(GROUP_ROWS): + request = tl.sum(tl.where(slots == slot, requests, 0), axis=0) + physical_pages = tl.load( + block_table + + request * block_table_stride_request + + logical_pages * block_table_stride_page, + mask=valid_blocks, + other=0, + ).to(tl.int64) + keys = tl.load( + index_cache + + physical_pages[:, None] * cache_stride_page + + token_in_page[:, None] * cache_stride_token + + dim_offsets[None, :] * cache_stride_dim, + mask=valid_blocks[:, None] & dim_ok, + other=0.0, + ) + per_head_scores = tl.where( + (column_slot == slot)[None, :], + tl.dot(keys, query_values, out_dtype=tl.float32), + per_head_scores, + ) + # The row-wise kernel's head sum is (h0 + h1) + (h2 + h3): pairs, then the pairs. + pairs = tl.sum( + tl.reshape( + tl.maximum(per_head_scores, 0.0), (BLOCK_N, BLOCK_G, NUM_INDEX_HEADS // 2, 2) + ), + axis=3, + ) + scores = tl.sum(pairs, axis=2) + tl.store( + output + + group_rows[None, :] * output_stride_row + + block_columns[:, None] * output_stride_block, + scores * score_scale, + mask=(block_columns[:, None] < visible[None, :]) & row_ok[None, :], + ) + + +#: Index-head counts the request-group kernel serves: its head sum is the row-wise kernel's +#: at up to 8 heads, two a thread then across the lanes; at 16 the row-wise kernel holds +#: four a thread. +_GROUP_HEADS = (2, 4, 8) +_GROUP_BLOCK_N = 128 + + +def _takes_request_groups( + q: torch.Tensor, + only_visible_blocks: bool, + context_rows: bool, + prefetch: bool, + rows_per_request: int, +) -> bool: + """The request-group kernel serves verify-shaped, visible-only calls: runs of + `rows_per_request` rows of one request at 2, 4 or 8 index heads. A prefill-sized + `context_rows` call keeps the engine's query-tiled kernel.""" + return ( + rows_per_request > 1 + and only_visible_blocks + and not prefetch + and q.ndim == 3 + and q.shape[1] in _GROUP_HEADS + and not (context_rows and q.shape[0] > _DECODE_ROW_LIMIT) + ) + + +def _qsa_index_scores_request_groups( + q: torch.Tensor, + index_cache: torch.Tensor, + block_table: torch.Tensor, + query_positions: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + compress_ratio: int, + use_pdl: bool, + rows_per_request: int, +) -> torch.Tensor: + """One launch of the request-group kernel after the callable's own checks.""" + if index_cache.ndim != 4 or index_cache.shape[2] != 1: + raise ValueError("QSA index cache must have shape [pages, tokens, 1, dim]") + if block_table.ndim != 2 or block_table.dtype != torch.int32: + raise ValueError("QSA index scoring requires a two-dimensional int32 block table") + rows, num_index_heads, index_head_dim = q.shape + if index_head_dim <= 0 or index_cache.shape[3] != index_head_dim: + raise ValueError("QSA index query and cache dimensions do not match") + if q.dtype != index_cache.dtype: + raise ValueError("QSA index query and cache dtypes must match") + if query_positions.shape != (rows,) or request_indices.shape != (rows,): + raise ValueError("QSA index-scoring metadata must have one entry per row") + if query_positions.dtype not in _INTEGER_DTYPES or request_indices.dtype not in _INTEGER_DTYPES: + raise ValueError("QSA index-scoring metadata must use integer storage") + if tokens_per_block <= 0 or compress_ratio < 2: + raise ValueError("QSA index scoring requires positive pages and compression ratio >= 2") + for t in (index_cache, block_table, query_positions, request_indices): + if t.device != q.device: + raise ValueError("QSA paged index scoring tensors must share one CUDA device") + max_compressed_blocks = block_table.shape[1] * tokens_per_block // compress_ratio + if max_compressed_blocks <= 0: + raise ValueError("QSA index scoring requires a non-empty block table") + output = torch.empty((rows, max_compressed_blocks), dtype=torch.float32, device=q.device) + if rows == 0: + return output + # A verify step's rows of one request read the same keys: one gather a run of rows, + # wider tiles, and a deeper column range once there are runs enough to fill the GPU. + groups = triton.cdiv(rows, rows_per_request) + tiles = 1 if groups <= 1 else 2 if groups <= 4 else 4 + grid = (groups, triton.cdiv(max_compressed_blocks, _GROUP_BLOCK_N * tiles)) + _qsa_index_scores_group_kernel[grid]( + q, + index_cache, + block_table, + query_positions, + request_indices, + output, + index_head_dim**-0.5, + rows, + q.stride(0), + q.stride(1), + q.stride(2), + index_cache.stride(0), + index_cache.stride(1), + index_cache.stride(3), + block_table.stride(0), + block_table.stride(1), + output.stride(0), + output.stride(1), + NUM_INDEX_HEADS=num_index_heads, + INDEX_HEAD_DIM=index_head_dim, + BLOCK_D=max(16, triton.next_power_of_2(index_head_dim)), + TOKENS_PER_BLOCK=tokens_per_block, + COMPRESS_RATIO=compress_ratio, + MAX_COMPRESSED_BLOCKS=max_compressed_blocks, + GROUP_ROWS=rows_per_request, + # The dot's columns: every row's heads, at least the 16 the row-wise kernel pads to. + BLOCK_G=max(triton.next_power_of_2(rows_per_request), 16 // num_index_heads, 1), + BLOCK_N=_GROUP_BLOCK_N, + TILES_PER_PROGRAM=tiles, + USE_PDL=use_pdl, + num_warps=4, + launch_pdl=use_pdl, + ) + return output + + +_CHECKPOINT = dict( + heads=4, + head_dim=128, + tokens_per_block=32, + compress_ratio=4, + dtype=torch.bfloat16, + # `qsa_logical_positions` is int64, `qsa_req_idx_per_token` int32. + positions_dtype=torch.int64, + request_dtype=torch.int32, + # The target's q is a head slice of its packed [T, 4 + 1 + 1, 128] index + # scratch buffer, so its row stride is 768 rather than 512. + q_layout="scratch_slice", +) + + +def _spec(**overrides) -> dict: + """A cell spec: this checkpoint's geometry unless the cell says otherwise.""" + return {**_CHECKPOINT, **overrides} + + +#: The serving tables: [8, 260] on the warm-up capture path and [32, 260] in mixed +#: steps. +_WIDTH = 260 + + +#: The same plane under an NVFP4 KV pool: the QSA position page joins each physical +#: page after the index buffers (`page_pad` elements), so the trunk's view is +#: `[pages, 32, 1, 128]` at page stride 12 x 4096 + 192 = 49344. +_TRUNK_POOL_NVFP4 = dict(pool_buffers=12, page_pad=192) + +#: Core's plane (bf16 K/V) at every four-rank layout, as the stock manager allocates it: +#: the 12 attention layers' `[32, 1, 128]` INDEX_KEY buffers coalesce into one pool, layer k at +#: byte 8192 k of a 98304-byte page, so one layer's view is `[pages, 32, 1, 128]` at page stride +#: 12 x 4096 = 49152. The cells take the last layer's plane. +_TRUNK_POOL_BF16 = dict(pool_buffers=12, page_pad=0) + +#: Pages in one rank's pool: the stock engine's `maxNumBlocks` at core's configuration (bf16 K/V +#: and recurrent state, KV fraction 0.6, max_batch_size 128, max_seq_len 9216, the n-gram table +#: host-resident). The largest per member; the index plane shares the attention pages. +_TEP4_POOL_PAGES = 258983 +_DEP4_POOL_PAGES = 120530 + +#: `QSAAttentionMetadata.qsa_block_table` is `[max_num_requests, max_blocks_per_seq]`: 128 rows +#: on a tep4 or dep4 rank (the engine's max_num_requests), and ceil((max_seq_len + generation +#: headroom) / 32) padded to a multiple of 4 -- 292 at core's 9216, 516 at MTP's 16384. +_TABLE_SLOTS = 128 +_CORE_TABLE_WIDTH = 292 +_MTP_TABLE_WIDTH = 516 + +#: The longest histories each envelope serves: the last position below max_seq_len. +_CORE_LAST_POSITION = 9216 - 1 +_MTP_LAST_POSITION = 16384 - 1 + +#: The target's AgentX server at tp1 (max_seq_len 262144): an 8200-page table (65600 score +#: columns, 1025 tiles of 64), histories near 100k tokens and to the end. +_AGENTX_TABLE_WIDTH = 8200 +_AGENTX_CONTEXT_LAST = 100000 - 1 +_AGENTX_LAST_POSITION = 262144 - 1 + + +#: The tune's table: max_seq_len 262144 is 8192 pages, and MTP's headroom makes it 8196, +#: 65568 columns. +_TUNE_TABLE_WIDTH = 8196 +_TUNE_LAST_POSITION = 262144 - 1 +#: Windows of 0 (positions 0..2), 1, 8, 32 and 33 blocks, ~20k-94k decode contexts, and the +#: envelope's last position. +_TUNE_DEP4_POSITIONS = ( + 94000, 0, 2, 3, 31, 127, 131, 20000, 34800, 49600, 64400, 65535, 79200, 88888, 131071, + _TUNE_LAST_POSITION, +) # fmt: skip + + +def _decode_positions(requests: int, last: int, rows_per_request: int = 1) -> tuple[int, ...]: + """Each request's first row position; the first request's last row sits at `last`.""" + first = last - rows_per_request + 1 + return (first,) + tuple(1 + (7919 * g) % first for g in range(1, requests)) + + +class QsaPagedIndexScores(OpWrapper): + """Score every causally visible compressed block of each query row. + + One Triton launch after allocating the output: the entry's own visible-split + kernel (`_qsa_index_scores_visible_split_kernel`) for ungrouped decode-sized + visible-only calls, its request-group kernel (`_qsa_index_scores_group_kernel`) for + verify-shaped ones, otherwise `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels. + triton_qsa_paged_index_scores`. For row `r` and compressed block `b` of request + `request_indices[r]`:: + + anchor = b * compress_ratio + compress_ratio - 1 + key_b = index_cache[block_table[req, anchor // tpb], anchor % tpb, 0] + out[r, b] = sum_h relu(q[r, h] . key_b) * index_head_dim ** -0.5 + + The block's key is whatever the compression step left at the group's LAST + position; nothing is pooled here. Row `r` sees `(query_positions[r] + 1) // + compress_ratio` blocks (complete groups only). Columns at or past that are + `-inf` when `only_visible_blocks` is False and **not written** -- the fresh + allocation's garbage -- when it is True. The column count is the block + table's width, `block_table.shape[1] * tpb // compress_ratio`. + + Kernel choice. With `rows_per_request=1`, `only_visible_blocks`, no `prefetch`, + at most 256 rows and an index-head count dividing 16, the visible-split kernel + runs: a `(row tiles, request slots, split)` grid sized by the device (16 programs + an SM, 32 when the batch is one tile) whose programs stride over each request's + visible column tiles, one key gather serving every row of that request among + `16 // heads` consecutive rows -- so the grid does not scale with the table's + columns. Otherwise the engine chooses a kernel: `context_rows` is a layout + declaration ("consecutive rows mostly share a request"), not a phase one. With + more than 256 rows and an index-head count that divides 128 into more than one + row (1, 2, 4, ..., 64), it selects the query-tiled kernel sharing one key gather + across `128 // heads` rows; otherwise the row-wise kernel runs. Before the engine's + choice, `rows_per_request` above 1 declares runs of that many rows of one request + from row 0 (a verify step's `predicted_tokens_per_seq`): with `only_visible_blocks`, + no `prefetch` and 2, 4 or 8 index heads, the entry's request-group kernel shares one + key gather across each run (a run whose rows name different requests is gathered + row by row). Each cell names the kernel it was + measured on. The request-group kernel is bitwise the row-wise one; the + visible-split and query-tiled kernels are not. + + Fusion boundary: the table walk, the gather, the dots, relu, head sum, + scale and causal bound are inside. Producing `q` (norm, RoPE), writing and + compressing the index cache, and everything after the scores are not. + Returns a fresh fp32 `[rows, columns]`; nothing is written in place. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="NVFP4 trunk plane, page stride 49344: decode MTP 4 requests x 2 rows", + spec=_spec( + layout="decode", + gen_positions=(7000, 5000, 4200, 2100), + rows_per_request=2, + slots=32, + pages=_WIDTH, + only_visible_blocks=True, + context_rows=False, + call_rows_per_request=2, + branch="request_group", + **_TRUNK_POOL_NVFP4, + ), + ), + Cell( + why="decode MTP7: 32 slots x 8 rows = 256 rows, the last decode-sized count", + spec=_spec( + layout="decode", + gen_positions=tuple(100 + (499 * g) % 8000 for g in range(32)), + rows_per_request=8, + slots=32, + pages=_WIDTH, + only_visible_blocks=True, + context_rows=False, + call_rows_per_request=8, + branch="request_group", + ), + ), + # The target's AgentX verify steps at tp1 (MTP3: 4 rows a request) over the NVFP4 trunk + # plane and its 8200-page table, NaN on every page no request maps. + Cell( + why="AgentX MTP3 verify: 16 requests x 4 rows to 99999, 65600 columns", + spec=_spec( + layout="decode", + gen_positions=_decode_positions(16, _AGENTX_CONTEXT_LAST, rows_per_request=4), + rows_per_request=4, + call_rows_per_request=4, + slots=32, + pages=_AGENTX_TABLE_WIDTH, + pool_pages=60000, + only_visible_blocks=True, + context_rows=False, + use_pdl=True, + branch="request_group", + **_TRUNK_POOL_NVFP4, + ), + ), + Cell( + why="AgentX MTP3 verify: 16 requests x 4 rows to 262143, 65600 columns", + spec=_spec( + layout="decode", + gen_positions=_decode_positions(16, _AGENTX_LAST_POSITION, rows_per_request=4), + rows_per_request=4, + call_rows_per_request=4, + slots=32, + pages=_AGENTX_TABLE_WIDTH, + pool_pages=140000, + only_visible_blocks=True, + context_rows=False, + use_pdl=True, + branch="request_group", + **_TRUNK_POOL_NVFP4, + ), + ), + Cell( + why="AgentX MTP3 verify: 1 request x 4 rows to 262143, 65600 columns", + spec=_spec( + layout="decode", + gen_positions=_decode_positions(1, _AGENTX_LAST_POSITION, rows_per_request=4), + rows_per_request=4, + call_rows_per_request=4, + slots=32, + pages=_AGENTX_TABLE_WIDTH, + pool_pages=60000, + only_visible_blocks=True, + context_rows=False, + use_pdl=True, + branch="request_group", + **_TRUNK_POOL_NVFP4, + ), + ), + Cell( + why="capture: 8191 rows over 8 ragged requests, visible-only", + spec=_spec( + layout="capture", + lengths=(1024, 512, 2048, 256, 1500, 700, 1151, 1000), + slots=8, + pages=_WIDTH, + only_visible_blocks=True, + context_rows=True, + branch="query_tile", + ), + ), + # The four-rank layouts' sites over one rank's measured pool: the call-site flags each + # site passes (the decode block's PDL launch, with prefetch at up to 4 rows), the + # engine's [128, width] table with requests on seeded rows, and pages drawn over the + # whole pool, both ends included. + Cell( + why="tep4 decode: one request at position 9215, prefetch, 258983-page pool", + spec=_spec( + layout="decode", + gen_positions=_decode_positions(1, _CORE_LAST_POSITION), + rows_per_request=1, + slots=_TABLE_SLOTS, + pages=_CORE_TABLE_WIDTH, + pool_pages=_TEP4_POOL_PAGES, + only_visible_blocks=True, + context_rows=False, + use_pdl=True, + prefetch=True, + branch="row_wise", + **_TRUNK_POOL_BF16, + ), + ), + Cell( + why="tep4 decode: 128 requests x 1 row, the largest captured batch", + spec=_spec( + layout="decode", + gen_positions=_decode_positions(128, _CORE_LAST_POSITION), + rows_per_request=1, + slots=_TABLE_SLOTS, + pages=_CORE_TABLE_WIDTH, + pool_pages=_TEP4_POOL_PAGES, + only_visible_blocks=True, + context_rows=False, + use_pdl=True, + branch="visible_split", + **_TRUNK_POOL_BF16, + ), + ), + Cell( + why="tep4 MTP7 verify: 128 requests x 8 rows = 1024, table 516, to 16383", + spec=_spec( + layout="decode", + gen_positions=_decode_positions(128, _MTP_LAST_POSITION, rows_per_request=8), + rows_per_request=8, + slots=_TABLE_SLOTS, + pages=_MTP_TABLE_WIDTH, + pool_pages=_TEP4_POOL_PAGES, + only_visible_blocks=True, + context_rows=False, + use_pdl=True, + branch="row_wise", + **_TRUNK_POOL_BF16, + ), + ), + Cell( + why="tep4 context: one fresh 8192-token prompt (max_num_tokens)", + spec=_spec( + layout="capture", + lengths=(8192,), + slots=_TABLE_SLOTS, + pages=_CORE_TABLE_WIDTH, + pool_pages=_TEP4_POOL_PAGES, + only_visible_blocks=True, + context_rows=True, + branch="query_tile", + **_TRUNK_POOL_BF16, + ), + ), + Cell( + why="dep4 decode: 32 requests x 1 row over a 120530-page pool", + spec=_spec( + layout="decode", + gen_positions=_decode_positions(32, _CORE_LAST_POSITION), + rows_per_request=1, + slots=_TABLE_SLOTS, + pages=_CORE_TABLE_WIDTH, + pool_pages=_DEP4_POOL_PAGES, + only_visible_blocks=True, + context_rows=False, + use_pdl=True, + branch="visible_split", + **_TRUNK_POOL_BF16, + ), + ), + Cell( + why="envelope: an 8192-row chunk at positions 1024..9215, table column 287", + spec=_spec( + layout="decode", + gen_positions=(_CORE_LAST_POSITION + 1 - 8192,), + rows_per_request=8192, + slots=_TABLE_SLOTS, + pages=_CORE_TABLE_WIDTH, + pool_pages=_TEP4_POOL_PAGES, + only_visible_blocks=True, + context_rows=True, + branch="query_tile", + **_TRUNK_POOL_BF16, + ), + ), + # The tune's envelope (max_seq_len 262144): an 8196-page table, 65568 columns. + Cell( + why="tune ll MTP3 verify: 1 request x 4 rows at 84000..84003, 8196-page table", + spec=_spec( + layout="decode", + gen_positions=(84000,), + rows_per_request=4, + slots=_TABLE_SLOTS, + pages=_TUNE_TABLE_WIDTH, + pool_pages=_TEP4_POOL_PAGES, + only_visible_blocks=True, + context_rows=False, + branch="visible_split", + **_TRUNK_POOL_BF16, + ), + ), + Cell( + why="tune draft: 1 row at 262143, the envelope's last position", + spec=_spec( + layout="decode", + gen_positions=(_TUNE_LAST_POSITION,), + rows_per_request=1, + slots=_TABLE_SLOTS, + pages=_TUNE_TABLE_WIDTH, + pool_pages=_TEP4_POOL_PAGES, + only_visible_blocks=True, + context_rows=False, + branch="visible_split", + **_TRUNK_POOL_BF16, + ), + ), + Cell( + why="tune MTP7 verify: 4 requests x 8 rows, the first to 262143", + spec=_spec( + layout="decode", + gen_positions=_decode_positions(4, _TUNE_LAST_POSITION, rows_per_request=8), + rows_per_request=8, + slots=_TABLE_SLOTS, + pages=_TUNE_TABLE_WIDTH, + pool_pages=_TEP4_POOL_PAGES, + only_visible_blocks=True, + context_rows=False, + use_pdl=True, + branch="visible_split", + **_TRUNK_POOL_BF16, + ), + ), + Cell( + why="tune dep4 decode: 16 requests x 1 row, empty to 262143-position windows", + spec=_spec( + layout="decode", + gen_positions=_TUNE_DEP4_POSITIONS, + rows_per_request=1, + slots=_TABLE_SLOTS, + pages=_TUNE_TABLE_WIDTH, + pool_pages=_DEP4_POOL_PAGES, + only_visible_blocks=True, + context_rows=False, + branch="visible_split", + **_TRUNK_POOL_BF16, + ), + ), + Cell( + why="tune context step: 2 rows of one request at 84000 (the snapshot split)", + spec=_spec( + layout="decode", + gen_positions=(84000,), + rows_per_request=2, + slots=_TABLE_SLOTS, + pages=_TUNE_TABLE_WIDTH, + pool_pages=_TEP4_POOL_PAGES, + only_visible_blocks=True, + context_rows=True, + branch="visible_split", + **_TRUNK_POOL_BF16, + ), + ), + Cell( + why="small context: fresh 200- and 56-token prompts, 256 rows from position 0", + spec=_spec( + layout="capture", + lengths=(200, 56), + slots=_TABLE_SLOTS, + pages=_TUNE_TABLE_WIDTH, + pool_pages=_DEP4_POOL_PAGES, + only_visible_blocks=True, + context_rows=True, + branch="visible_split", + **_TRUNK_POOL_BF16, + ), + ), + ) + + note = """ + Source: the entry's own Triton kernel (`_qsa_index_scores_visible_split_kernel`, + this file) and the engine's beside the wrapped callable + (`_qsa_paged_index_scores_kernel`, `_qsa_paged_index_scores_qtile_kernel`, + `_qsa_paged_index_scores_group_kernel`), + JIT-compiled for the running device on first call; the cells ran that JIT + path on sm_103. + + Visible-split kernel (decode-sized visible-only calls). Its grid is + `(row tiles, request slots, split)`: a row tile is `min(rows, 16 // heads)` + consecutive rows, program `(t, g, s)` takes the g-th distinct request of + tile t (a slot past the tile's request count exits at once) and its column + tiles s, s + split, ... of 64 blocks up to that request's bound, `split` + fixed by the SM count (16 programs an SM, 32 for a one-tile batch) and the + table, never by the positions -- so a captured graph follows any window + (pinned by `test_captured_visible_split_follows_the_positions`). The dot is + `[64, 128] x [128, rows x heads]` with fp32 accumulation, relu, a per-row + head sum and the scale, as in the engine's kernels; the columns written are + exactly each row's visible ones. `use_pdl` waits on its predecessor before the first read and + releases dependents at once; `prefetch` keeps the engine's kernel (the + visible-split one does not read before the wait). + + Metadata consumed: none -- every input is an argument. What the target + passes: `block_table` is `QSAAttentionMetadata.qsa_block_table` (raw + pool-page ids, int32, -1 past each sequence); `query_positions` / + `request_indices` are `qsa_logical_positions` / `qsa_req_idx_per_token`; + `index_cache` is the layer's bf16 index-K plane (`get_index_k_buffer`). + + Preconditions (not checked): every block-table entry the causal bound + reaches is a valid page, and `request_indices[r]` is a row of the table -- + the loads are masked by the bound, not by the pool size, so a violation is + an out-of-bounds read (not probed). The compression step must already have + run: an uncompressed anchor scores as a raw index key. `rows == 0` returns + an empty `[0, columns]` tensor without launching. + + The op validates its own metadata loudly (`ValueError`s; the test pins the + one-entry-per-row metadata refusal), so `is_valid` adds nothing. + + `only_visible_blocks=True` really leaves garbage in the whole unscored tail + (measured, not pinned by a test). A consumer that scans `[row_start, + row_end)` -- `qsa_topk_select` -- is unaffected. + + fp32 operands are accepted but `tl.dot` rounds them to a tensor-core + format: 8.1e-3 max abs against an fp32 reference on scores to 7.8 (measured), + outside the test's band. No cell covers fp32; the side cache is bf16. + + CUDA graphs: the test captures the decode launch (`use_pdl`, `prefetch`) at + 4 rows and replays it after `q` and the anchor keys were rewritten in + place; every replay matches an eager call bitwise. The same holds captured + at each measured-pool decode and verify cell (1, 128 and 1024 rows at tep4, + 32 at dep4, the AgentX MTP3 verify cells' 4 and 64 rows, with their sites' flags) + after q and every mapped key changed. + + The request-group kernel (`rows_per_request` > 1): each program takes one run of + rows and 128 blocks a tile (1, 2 or 4 tiles a program at 1, up to 4, and more runs); + the dot's columns are the run's rows' heads, padded to 16 as the row-wise kernel + pads its heads, and the head sum is the row-wise kernel's pairwise order, so every + score is bitwise the row-wise kernel's -- the test compares the bit patterns at 2, + 3, 4 and 8 rows a run, 2, 4 and 8 heads, partial last runs and runs that mix + requests; 16 heads stay row-wise. + The run reads up to its longest row's visible blocks and stores each row's own. + + At the measured-pool cells: at most 2.9e-6 against the reference (scores to + 11.3), relative L2 below 6.4e-8. + + `prefetch` (with `use_pdl` and `only_visible_blocks`, at generation-sized + batches) reads the pages and compressed keys before the PDL wait and q after + it. Sound only behind the decode step's own prologue + (`qwen4_exp_attn_decode_prologue`), which writes no compressed key but the last + visible block's, at a row whose position ends a group; that key is read again + after the wait. + + `prefetch` takes effect only where the call is one tile a program -- 256 rows or + fewer -- and is otherwise dropped silently: a 1024-row verify step passing it + runs the plain PDL kernel. + + Every cell runs `only_visible_blocks=True` (the target never passes False), + 4 index heads x 128, compress ratio 4, 32-token pages, bf16 operands, int64 + positions and int32 request indices. The NVFP4-pool page stride (49344) is + certified on the row-wise kernel at the tp1 cell's small pool and on the + request-group kernel at the AgentX cells' 60000- and 140000-page pools. + + The four-rank layouts' cells run over one rank's measured pool, core's coalesced + plane at page stride 49152 (`page_stride` is a compile-time constant, so this + is its own kernel variant): 258983 pages at tep4, 120530 at dep4, the plane + poisoned with NaN outside the mapped pages, which include page 0 and the + pool's last page. The last page's elements sit near 1.27e10 elements from the + plane's base: both kernels widen the page and request index to int64 before + the stride product, and the cells prove it. + + Index arithmetic left in int32: `row * q_stride_row` and the output's + `row * output_stride_row` (program ids times compile-time strides). The output + term passes 2^31 at rows x columns >= 2^31 -- 919k rows at 2336 columns; the + envelope's 8192 rows x 65568 columns (max_seq_len 262144) is 5.4e8. Positions + are read at their own width (int64 here), so the fp32 2^24 bound does not apply. + + The visible-split kernel's index arithmetic: pages and request rows are widened + to int64 before the stride products (the measured-pool cells read the pool's last + page); `lane_rows * output_stride_row` stays int32, at most 256 x 65568 here. + + Certified: up to 8192 rows (one request's 8192-row context, fresh or ending at + position 9215, on the query-tile kernel; 1024 verify rows on the row-wise + kernel with 8 tiles a program), tables of 128 x 292 and 128 x 516 pages, + positions to 16383; on the visible-split kernel also the tune's 128 x 8196 + table (65568 columns) with positions to 262143 at 1, 2, 4, 16 and 32 rows + (empty windows, 1 and 2 rows a request, 8-row requests across two tiles, small + context steps from position 0); on the request-group kernel 2, 4 and 8 rows a + request, up to 256 rows, a 32 x 8200 table and positions to 262143. Not certified: + more than 8192 rows, tables wider than 516 pages on the row-wise and query-tile + kernels (the tune's context-sized calls), any other head geometry, compress ratio, + page size or operand dtype, a context step's mixed request rows on the pool + layout, the fp8 side-plane stride at a measured pool, and the `-inf` padded mode. + """ + + def __call__( + self, + q: torch.Tensor, + index_cache: torch.Tensor, + block_table: torch.Tensor, + query_positions: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + compress_ratio: int, + only_visible_blocks: bool, + context_rows: bool, + use_pdl: bool = False, + prefetch: bool = False, + rows_per_request: int = 1, + ) -> torch.Tensor: + if rows_per_request < 1: + raise ValueError("QSA index scoring requires rows_per_request >= 1") + if rows_per_request == 1 and _takes_visible_split(q, only_visible_blocks, prefetch): + return _qsa_index_scores_visible_split( + q, + index_cache, + block_table, + query_positions, + request_indices, + tokens_per_block, + compress_ratio, + use_pdl, + ) + if _takes_request_groups(q, only_visible_blocks, context_rows, prefetch, rows_per_request): + return _qsa_index_scores_request_groups( + q, + index_cache, + block_table, + query_positions, + request_indices, + tokens_per_block, + compress_ratio, + use_pdl, + rows_per_request, + ) + return _triton_qsa_paged_index_scores( + q=q, + index_cache=index_cache, + block_table=block_table, + query_positions=query_positions, + request_indices=request_indices, + tokens_per_block=tokens_per_block, + compress_ratio=compress_ratio, + only_visible_blocks=only_visible_blocks, + context_rows=context_rows, + use_pdl=use_pdl, + prefetch=prefetch, + ) + + def reference( + self, + q: torch.Tensor, + index_cache: torch.Tensor, + block_table: torch.Tensor, + query_positions: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + compress_ratio: int, + only_visible_blocks: bool, + context_rows: bool, + use_pdl: bool = False, + prefetch: bool = False, + rows_per_request: int = 1, + ) -> torch.Tensor: + """fp32 scores, `-inf` past each row's bound -- or NaN there, meaning + "unspecified", when `only_visible_blocks` is True. `context_rows` and + `rows_per_request` are ignored: they pick a kernel, not a result. + + Walked in row chunks only to bound the gathered keys, which reach + 8.7 GiB in fp32 at 8191 rows x 2080 columns. + """ + rows, _, head_dim = q.shape + num_columns = block_table.shape[1] * tokens_per_block // compress_ratio + blocks = torch.arange(num_columns, device=q.device) + anchors = blocks * compress_ratio + compress_ratio - 1 + logical_pages, in_page = anchors // tokens_per_block, anchors % tokens_per_block + table = block_table.long() + requests = request_indices.long() + scores = torch.empty(rows, num_columns, dtype=torch.float32, device=q.device) + chunk = max(1, 2**26 // max(num_columns * head_dim, 1)) + for lo in range(0, rows, chunk): + hi = min(lo + chunk, rows) + # Pages past a sequence are -1; their columns are masked below. + physical = table[requests[lo:hi]][:, logical_pages].clamp_min(0) + keys = index_cache[physical, in_page.expand(hi - lo, num_columns)][:, :, 0, :] + per_head = torch.einsum("rhd,rbd->rbh", q[lo:hi].float(), keys.float()) + scores[lo:hi] = torch.relu(per_head).sum(-1) * (head_dim**-0.5) + visible = blocks[None, :] < ((query_positions.long() + 1) // compress_ratio)[:, None] + tail = float("nan") if only_visible_blocks else float("-inf") + return scores.masked_fill(~visible, tail) + + +qsa_paged_index_scores = QsaPagedIndexScores() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_index_store.py similarity index 76% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_index_store.py index 0100f49724d9..a014e6a74314 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_index_store.py @@ -108,8 +108,7 @@ def _qsa_paged_index_store_kernel( def _launch_rows(rows: int) -> Tuple[int, int]: """Rows per program and warps: one row per program while the grid is small, - a tile of rows once there are enough rows to fill the machine. Measured on - GB300 against 1/4/8/16/32 rows per program.""" + a tile of rows once there are enough rows to fill the machine.""" if rows <= 256: return 1, 1 return 16, 4 @@ -151,6 +150,40 @@ def _spec(segments, pages, **overrides) -> dict: return {**_CKPT, "segments": segments, "pages": pages, **overrides} +#: The default target (tp4 / tep4 / dep4) at one rank. Its fused input projection row +#: `[q | k | v | gate | index q, k | pad of one index head]` is 4352 wide under attention TP (6 q +#: and 1 KV heads of 256: 34 heads of 128, the raw index key head 32) and 14080 under attention DP +#: (110 heads, key head 108). Its bf16 index pool is the stock manager's at every member: the 12 +#: sparse layers' INDEX_KEY buffers coalesce into one pool, so a layer's plane is `[pages, 32, 1, +#: 128]` at page stride 12 x 4096 = 49152 elements, layer `j` at `j` x 8192 bytes into the slot, +#: and the position plane is its own contiguous `[pages, 32, 3]` int32 pool (`cache_layout` +#: "coalesced", the last layer's plane). `pages` is the tep4 pool: the stock tep4 engine's +#: (bf16 KV) `kv_max_tokens` 7866720 a rank, 245835 pages of 32. The table +#: is `qsa_block_table`, `[max_batch_size, max_blocks_per_seq]`: 128 x 292 (`max_seq_len` 9216) or +#: 128 x 516 under MTP (16384). `straddle` puts the batch's last row on the pool's last page, past +#: 2^31 elements of the index plane. +_DEFAULT = dict( + _CKPT, + key_layout="fused_row", + cache_layout="coalesced", + layers=12, + layer=11, + page_stride=49152, + slots=128, + width=292, + pages=245835, + straddle=True, +) + +#: 128 decode requests, one row each, at positions on both sides of page edges; the last at +#: 9215 (`max_seq_len` - 1, table column 287). +_DECODE_128 = tuple((31 + 72 * i, 1) for i in range(127)) + ((9215, 1),) +#: 32 decode requests (one dep4 rank's share of 128), the last at 9215. +_DECODE_32 = tuple((31 + 287 * i, 1) for i in range(31)) + ((9215, 1),) +#: 128 MTP7 verify blocks of 8 rows, some across a page edge, up to position 16165. +_MTP7_128 = tuple((29 + 127 * i, 8) for i in range(128)) + + class QsaPagedIndexStore(OpWrapper): """Write each row's raw index key and position coordinates into its paged slot. @@ -168,13 +201,11 @@ class QsaPagedIndexStore(OpWrapper): rotation itself. Returns `(index_cache, position_cache)`, both written in place. - Sourcing tier 5: one Triton launch written for this target - (`_qsa_paged_index_store_kernel`). Nothing in the engine performs this - store for a packed, ragged batch: the registered index-cache ops write fp8 - planes or norm and rotate on the way in, `triton_qsa_decode_pre_indexer` - is one row per request, and `QSAIndexer.update_cache_and_compress` bundles - the compression. Native torch needs a block-table gather, a floor-divide, - a remainder, two indexed writes and their index casts (7 launches). + Sourcing: one Triton launch (`_qsa_paged_index_store_kernel`). Nothing in + the engine performs this store for a packed, ragged batch: the registered + index-cache ops write fp8 planes or norm and rotate on the way in, + `triton_qsa_decode_pre_indexer` is one row per request, and + `QSAIndexer.update_cache_and_compress` bundles the compression. Fusion boundary: page lookup, slot arithmetic, both scatters. The index projection, the query's norm and rotation, group compression, scoring, @@ -208,6 +239,31 @@ class QsaPagedIndexStore(OpWrapper): hole="unallocated_page", ), ), + Cell( + why="default tep4 decode: one row at position 9215 on the pool's last page, token_k " + "off the 4352-wide fused row, the coalesced bf16 index plane", + spec=dict(_DEFAULT, member="tep4", segments=((9215, 1),), row_width=4352), + ), + Cell( + why="default tep4 decode: 128 requests one row each (the largest captured batch), " + "4352-wide row", + spec=dict(_DEFAULT, member="tep4", segments=_DECODE_128, row_width=4352), + ), + Cell( + why="default tep4 MTP7 verify: 128 requests x 8 rows (1024), 16 rows per program, " + "a 516-column table", + spec=dict(_DEFAULT, member="tep4", segments=_MTP7_128, row_width=4352, width=516), + ), + Cell( + why="default tp4 step of 8192 rows (max_num_tokens): a context from 0 and a decode " + "row at 9215, 4352-wide row", + spec=dict(_DEFAULT, member="tp4", segments=((0, 8191), (9215, 1)), row_width=4352), + ), + Cell( + why="default dep4 decode: 32 requests one row each on a 128-row table, token_k off " + "the 14080-wide row", + spec=dict(_DEFAULT, member="dep4", segments=_DECODE_32, row_width=14080), + ), ) note = """ @@ -228,6 +284,20 @@ class QsaPagedIndexStore(OpWrapper): `[pages, 32, 3]` at (24672, 3, 1); the cell checks that the other index-K roles of each page survive the call. + The default target (`_DEFAULT` cells) under bf16 KV, read off the stock + manager at tp4, tep4 and dep4 (the test re-reads it): each sparse layer's + `get_index_k_buffer(i)` is `[pages, 32, 1, 128]` at strides (49152, 128, + 128, 1), the 12 layers' planes 8192 bytes apart in one coalesced pool, and + `get_qsa_position_buffer()` a contiguous `[pages, 32, 3]` int32 pool. Under + fp8 and NVFP4 KV both are the shared-page layout above (49344 / 24672) at + every member. The cells drive the last layer's plane of a 245835-page pool + (the stock tep4 engine's), its other 11 planes checked untouched, with the + batch's last row on the pool's last page: index-plane offsets past 2^31 + elements (1.2e10). `token_k` is the index-K head off the fused input + projection row, 4352 wide under attention TP (head 32) and 14080 under + attention DP (head 108); the table is `[128, 292]` (`max_seq_len` 9216) or + `[128, 516]` (MTP's `max_seq_len` 16384). + Preconditions: - destination slots are unique across the batch: two rows naming one (page, slot) both store, unordered. The engine's mapping cannot @@ -264,10 +334,13 @@ class QsaPagedIndexStore(OpWrapper): them) share that mask and have no cell. Not certified: an fp8 index plane, keys other than bf16, any index geometry - other than 1 head x 128 at 32-token pages, decode batches other than 32 MTP7 - verify blocks, strided metadata or tables, - stride-0 coordinates, multiple devices, concurrent streams onto one pool, - duplicate destinations. + other than 1 head x 128 at 32-token pages, decode batches other than 1, 32 + and 128 requests and verify batches other than 32 and 128 MTP7 blocks, + strided metadata or tables, stride-0 coordinates, multiple devices, + concurrent streams onto one pool, duplicate destinations. The tp4 and dep4 + pools are driven at tep4's page count, not their own: both hold at most + that many pages a rank (the same or larger weights a rank, the same or + larger KV page). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store.py similarity index 58% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store.py index be0791c3ccd6..5d5af38707ad 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store.py @@ -43,12 +43,64 @@ #: page edge and some inside one. _MTP7_32 = tuple((29 + 61 * i, 8) for i in range(32)) +#: One rank's QSA K/V pool as the stock `Qwen4ExpHybridCacheManager` lays it out at the default target's +#: members: a page holds the K and V role planes of all 12 attention layers (24 role pages of +#: `[hkv, 32, 256]`), so a layer's `get_buffers(i, kv_layout="HND")` view steps 24 role pages a +#: page. `pool="engine"` cells drive layer 47's view (`layer_slot` 11, the page's last planes) at +#: an engine's own page count. `slots` x `width` is the engine's `qsa_block_table`: +#: `max_num_sequences` 128 by `max_blocks_per_seq` 292 at `max_seq_len` 9216 (516 at 16384). +_ENGINE = dict(_CKPT, pool="engine", role_pages=24, layer_slot=11, slots=128, width=292) + +#: A tp4 / tep4 rank: one KV head, k and v at heads 6 and 7 of the fused projection row +#: `[q 6 | k 1 | v 1 | gate 6 | index 3]` (17 heads of 256, row stride 4352). +_TP4 = dict(_ENGINE, hkv=1, heads_q=6, row_pad_heads=9, row_stride=4352) + +#: A dep4 rank (attention DP, every head): the tp1 row `[q 24 | k 2 | v 2 | gate 24 | index 3]` +#: (55 heads, row stride 14080) over the kv_fp8 variant's e4m3 pool. +_DEP4_FP8 = dict(_ENGINE, row_pad_heads=27, row_stride=14080, cache_dtype=torch.float8_e4m3fn) + +#: Attention pool pages of an engine at KV fraction 0.6, bf16 recurrent state, `max_batch_size` +#: 128 and `max_seq_len` 9216: its `maxNumBlocks` less the recurrent pool's slots, which that +#: total also counts. tep4, bf16 KV: 258983 - 1874. dep4, kv_fp8: 172579 - 1217. +PAGES_TEP4 = 257109 +PAGES_DEP4_FP8 = 171362 + +#: 128 decode rows, one a request: page-aligned, mid-page, page-crossing and the last position +#: `max_seq_len` 9216 admits, then a stride-577 walk over the rest. +_DECODE_128 = ((0, 1), (31, 1), (32, 1), (9215, 1)) + tuple( + ((i * 577 + 45) % 9216, 1) for i in range(124) +) + +#: 32 decode rows, one a request, the last at position 9215. +_DECODE_32 = tuple(((i * 293 + 5) % 9216, 1) for i in range(31)) + ((9215, 1),) + +#: 128 requests of one MTP7 verify block each (8 rows, 1024), the last ending at position 16383 +#: (`max_seq_len` 16384 under MTP); starts at every offset in a page, so blocks straddle pages. +_MTP7_128 = tuple(((29 + 127 * i) % 16376, 8) for i in range(127)) + ((16376, 8),) + +#: An 8192-row (`max_num_tokens`) context step: a chunk continuing from mid-page to position +#: 9215, and a request's first 1023 tokens. +_PREFILL_8192 = ((2047, 7169), (0, 1023)) + def _spec(segments, pages, **overrides) -> dict: """`segments` is `(first position, rows)` per request, in row order.""" return {**_CKPT, "segments": segments, "pages": pages, **overrides} +def _engine(base: dict, member: str, segments, pages: int, **overrides) -> dict: + """An engine-pool cell: `rows` is the step's row count, `member` whose pool `pages` is.""" + rows = sum(n for _, n in segments) + return { + **base, + "member": member, + "segments": segments, + "pages": pages, + "rows": rows, + **overrides, + } + + def _bits(tensor: torch.Tensor) -> torch.Tensor: """The same storage as integers: fp8 has no comparison ops, and -0.0 != 0.0.""" as_int = {1: torch.uint8, 2: torch.int16, 4: torch.int32}[tensor.element_size()] @@ -100,9 +152,37 @@ class QsaPagedKvStore(OpWrapper): why="fp8 decode MTP7: 32 requests x 8 rows (256), blocks straddling pages", spec=_spec(_MTP7_32, pages=1200, slots=40, cache_dtype=_F8), ), + # The default target's members over the engine's own pool: one rank's page count, the + # 24-role-page stride, layer 47's view, a 128-row table. + Cell( + why="tep4 decode, 1 row: hkv 1 off the 4352-wide row, position 9215, " + "257109-page bf16 pool", + spec=_engine(_TP4, "tep4", ((9215, 1),), PAGES_TEP4), + ), + Cell( + why="tep4 decode, 128 rows (max_batch_size), aligned, mid and page-crossing", + spec=_engine(_TP4, "tep4", _DECODE_128, PAGES_TEP4), + ), + Cell( + why="tep4 MTP7 verify, 128 requests x 8 rows (1024) to position 16383, table width 516", + spec=_engine(_TP4, "tep4", _MTP7_128, PAGES_TEP4, width=516), + ), + Cell( + why="tp4 context, 8192 rows (max_num_tokens): a mid-page chunk to 9215 and a " + "request's first 1023", + spec=_engine(_TP4, "tp4", _PREFILL_8192, PAGES_TEP4), + ), + Cell( + why="dep4 kv_fp8 decode, 32 rows: hkv 2 off the 14080-wide row into a " + "171362-page e4m3 pool", + spec=_engine(_DEP4_FP8, "dep4", _DECODE_32, PAGES_DEP4_FP8), + ), ) note = """ + Source: the in-tree Triton kernel `_qsa_paged_kv_store_kernel`, JIT-compiled + on its first call; `ARCHS` is the architecture it compiled for and ran on. + Metadata consumed: none -- every input is an argument. The target passes `qsa_req_idx_per_token` (int32, the row's batch index) and `qsa_logical_positions` (int64) sliced to the step's rows, @@ -126,22 +206,38 @@ class QsaPagedKvStore(OpWrapper): table (not pinned by a test). `rows == 0` returns before launching. Strides and the pool/table extents are Triton constexprs: each distinct - layout or pool size compiles its own kernel. CUDA-graph capturable - (captured at 256 rows, 32 MTP7 verify blocks); a replay follows positions - rewritten in place. + layout or pool size compiles its own kernel (once per engine, at its + pool's page count). CUDA-graph capturable (captured at 256 rows, 32 MTP7 + verify blocks, and at 128 tep4 decode rows over the 257109-page pool); a + replay follows positions rewritten in place. + + The engine's pool, as the stock manager lays it out at tp4 / tep4 / dep4: + one page holds the K and V planes of all 12 attention layers, so a + layer's K view steps 24 role pages a page -- 196608 elements at one KV + head, 393216 at two, bf16 and fp8 alike -- and its page count is the + engine's (`PAGES_TEP4`, `PAGES_DEP4_FP8`). The `pool="engine"` cells drive + layer 47's view with the pool's last page among the touched ones, past + 2^32 elements of the view's base (byte 101099339776 at tep4, 67382247424 + at dep4 kv_fp8). The kernel loads the page id as int64 before its stride + product, and every such row lands bitwise. Hard limits: both metadata values are cast to int32 inside the kernel - (`.to(tl.int32)`); an int64 position at `2^32 + 5` was probed once (not - pinned by a test) and lands on slot 5 of page column 0 -- written, not - skipped. The engine's positions stay below 2^14. + (`.to(tl.int32)`); an int64 position of `2^32 + 5` lands on slot 5 of page + column 0 -- written, not skipped (not pinned by a test). The engine's + positions stay below 2^14 at `max_seq_len` 16384 with MTP. Not certified: duplicate destinations, multiple devices, a scaled fp8 pool (there is no scale argument), fp8 conversion of values past 448, the - packed-QKV-row (MTP) layout over a bf16 pool, fp16 rows, head geometry other than - this checkpoint's (D 256, Hkv 2, 32-token pages), int32 positions or - int64 request indices, head-major rows, sliced or transposed tables, - `rows == 0`, and of the rows that resolve to no slot any but an - unallocated (`-1`) table entry (the reference skips each). + packed-QKV-row (MTP draft layer) layout over a bf16 pool or at one KV head, + fp16 rows, head geometry other than D 256 with Hkv 1 or 2 and 32-token + pages, int32 positions or int64 request indices, head-major rows, sliced + or transposed tables, `rows == 0`, and of the rows that resolve to no slot + any but an unallocated (`-1`) table entry (the reference skips each). + Pools at a page count not driven here: tp4's own engine (tep4's count + stands in: a tp4 rank's pool has the same per-page shape), the MTP + engine's at tep4 (its draft layer has a separate cache, so its trunk + count differs; the non-MTP count stands in), and dep4's bf16 pool (119680 + pages). Positions past 16383 and tables wider than 516 columns. """ def __call__( @@ -207,9 +303,8 @@ def is_valid( block_table: torch.Tensor, tokens_per_block: int, ) -> None: - # The kernel reads these two as flat unit-stride buffers. Measured: a - # column of a 2-wide buffer is read through its base pointer and the - # rows land at other rows' slots, with nothing raised. + # The kernel reads these two as flat unit-stride buffers; a strided view + # lands rows at other rows' slots, with nothing raised. assert request_indices.is_contiguous(), "request_indices must be a contiguous 1-D tensor" assert logical_positions.is_contiguous(), ( "logical_positions must be a contiguous 1-D tensor" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store_nvfp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store_nvfp4.py similarity index 89% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store_nvfp4.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store_nvfp4.py index 73e6ec1bebf7..719c548b5fd2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store_nvfp4.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_kv_store_nvfp4.py @@ -60,6 +60,19 @@ _DECODE = ((0, 1), (31, 1), (32, 1), (33, 1), (100, 1), (255, 1), (7, 1), (64, 1)) +#: One dep4 rank's kv_nvfp4 pool as the stock engine sizes it (NVFP4 K/V, chunked prefill): +#: 213295 slots, its `maxNumBlocks`; the trunk layout of 24 role pages a slot, layer 11's +#: planes (the last attention layer's); and the `qsa_block_table` of +#: `max_num_requests` 128 rows by `max_blocks_per_seq` 292 columns (max_seq_len 9216). +#: `last_page` puts that request's last page on the pool's last slot, so the stores land past +#: 2^31 bytes of both the data and the scale plane. The pool alone is 47 GB, so `large`. +_DEP4_PAGES = 213295 +_DEP4 = dict(slots=128, width=292, slot_pages=24, layer=11, last_page=0, large=True) +#: 32 decode rows of 32 requests: max_seq_len - 1 first, page edges, then a stride over the table. +_DEP4_DECODE = ((9215, 1), (0, 1), (31, 1), (32, 1), (33, 1)) + tuple( + (293 * i + 17, 1) for i in range(1, 28) +) + def _spec(segments, pages, **overrides) -> dict: """`segments` is `(first position, rows)` per request, in row order.""" @@ -391,6 +404,21 @@ class QsaPagedKvStoreNvfp4(OpWrapper): "P=1 and P=3) at unit second-level scales: data and both scale layouts bitwise equal", spec=_spec(((0, 115), (0, 101), (0, 36)), pages=96, thop_append=True), ), + Cell( + why="dep4 kv_nvfp4 decode: one row at position 9215 (max_seq_len - 1) on the last " + "of the rank's 213295 pool slots, 128 x 292 table", + spec=_spec(((9215, 1),), pages=_DEP4_PAGES, **_DEP4), + ), + Cell( + why="dep4 kv_nvfp4 decode: 32 requests of the 128 table rows at positions 0..9215 " + "across page edges, on the rank's 213295-slot pool", + spec=_spec(_DEP4_DECODE, pages=_DEP4_PAGES, **_DEP4), + ), + Cell( + why="dep4 kv_nvfp4 context chunk: 8192 rows (max_num_tokens), one prompt's chunk at " + "2048..9215 and a new prompt's first 1024, on the rank's 213295-slot pool", + spec=_spec(((2048, 7168), (0, 1024)), pages=_DEP4_PAGES, **_DEP4), + ), ) note = """ @@ -431,9 +459,20 @@ class QsaPagedKvStoreNvfp4(OpWrapper): tokens_per_block. A served call puts the width and each stride in one fixed class (the k/v head slices of the packed QKV row, the planes' slot, head and token strides, the contiguous table's) and varies the page count and the - table height. The cells run (other, other) and (multiple of 16, multiple of - 16); the other combinations, a table height of 1 among them, are not - certified. + table height. The cells run (other, other), (multiple of 16, multiple of 16) + and, at dep4, (other, multiple of 16); (multiple of 16, other) and a table + height of 1 are not certified. + + The default target at dep4 with NVFP4 K/V: a rank holds both KV heads in + the trunk layout, data `[slots, 24, 2, 32, 128]` and scales + `[slots, 24, 2, 32, 16]`, attention layer j's K and V at role pages 2j and + 2j + 1 (the stock manager's `get_buffers` / `get_block_scale_buffers` + views). The stock engine at that configuration (chunked prefill on, + max_num_tokens 8192) has `maxNumBlocks` 213295 a rank, and the dep4 cells + drive that many slots under its 128 x 292 `qsa_block_table`, one + request's last page on the pool's last slot: stores land 4.19e10 bytes + into the data plane and 5.24e9 into the scale plane. A cell's pool is + 47 GB, so the dep4 cells are `large`. Not finite values: a NaN or +-inf in a block is not certified (the e2m1 and e4m3 conversions saturate, the block maximum is the kernel's `max`). diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa.py new file mode 100644 index 000000000000..b65b55eac0cf --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa.py @@ -0,0 +1,484 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA sparse GQA: attention over an explicitly selected set of paged K/V tokens.""" + +from __future__ import annotations + +import math +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# Approved exception: QSA's sparse path never reaches the thop attention op's +# sparse argument group, so there is no torch.ops surface for this computation. +from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import ( + can_fuse_qsa_splitk_output_gate as _can_fuse_qsa_splitk_output_gate, +) +from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import ( + triton_qsa_paged_sparse_gqa as _triton_qsa_paged_sparse_gqa, +) + +_CHECKPOINT = dict( + q_heads=24, + kv_heads=2, + head_dim=256, + tokens_per_block=32, + compress_ratio=4, + token_topk=2048, + cache_dtype=torch.bfloat16, + q_dtype=torch.bfloat16, + # q is the first 24 heads of the packed [T, 24 + 2 + 2, 256] QKV buffer; + # "fused_row" is the target's fused input projection row, 55 heads of 256. + q_layout="qkv_slice", + request_dtype=torch.int32, + # The engine's table: 260 pages per slot, -1 past each sequence. + pages=260, +) + +_SPLIT_K = ("splitk", "merge") +_PREFILL = ("prefill",) +_FP8 = dict(cache_dtype=torch.float8_e4m3fn) + + +def _spec(**overrides) -> dict: + """A cell spec: this checkpoint's geometry unless the cell says otherwise. + + `segments` is `(first position, rows)` per request, in row order; + `prefill_pair` says whether `query_positions` / `compress_ratio` are + passed, which the target does exactly when the step has a context request. + """ + return {**_CHECKPOINT, **overrides} + + +def _decode(segments, launches=_SPLIT_K, **overrides) -> dict: + base = dict(segments=segments, slots=32, prefill_pair=False, launches=launches) + return _spec(**{**base, **overrides}) + + +#: The engine's trunk layout (`Qwen4ExpHybridCacheManager.get_buffers`): a lifecycle slot +#: holds the 24 role pages of the 12 attention layers, K then V each, and a layer's +#: `k_cache` / `v_cache` are its two planes, so one pool page advances 24 role pages. +#: Layer 11's planes are the slot's last. +_TRUNK = dict(slot_pages=24, layer=11) + +#: The element offset the default target's cells straddle: the requests' pages lie on both +#: sides of it in the layer's K and V planes. An engine's pool reaches far past it (tep4's +#: bf16 engine: 245835 pages, 4.8e10 elements of trunk). +_OFFSET_EDGE = 2**31 + +#: One tp4 / tep4 rank (attention TP 4): 6 q heads and one of the 2 KV heads, q the first 6 +#: heads of the packed [T, 6 + 1 + 1, 256] QKV buffer; a `max_batch_size` 128-row table, +#: `ceil((max_seq_len + 1) / 32)` padded to a multiple of 4 columns: 292 at max_seq_len 9216, +#: 516 at MTP's 16384. +_TP4 = dict(q_heads=6, kv_heads=1, slots=128, pages=292, **_TRUNK) +_TP4_MTP = dict(_TP4, pages=516) +#: One dep4 rank (attention DP): tp1's 24 / 2 heads and a 32-row table (128 requests over 4 +#: ranks); a bf16 pool, or the kv_fp8 variant's e4m3 pool. +_DEP4 = dict(q_heads=24, kv_heads=2, slots=32, pages=292, **_TRUNK) +_DEP4_FP8 = dict(_DEP4, **_FP8) + +#: q and the output gate off one rank's fused input projection row, `[q | k | v | gate | +#: index q, k | pad]` in heads of 256 (tp1's `weights._attn_in_w`): 6 + 1 + 1 + 6 + 3 = 17 +#: heads (row stride 4352) at a tp4 / tep4 rank, 24 + 2 + 2 + 24 + 3 = 55 (14080) at dep4. +_TP4_ROW = dict(q_layout="fused_row", gate="fused_row", row_heads=17) +_DEP4_ROW = dict(q_layout="fused_row", gate="fused_row", row_heads=55) + + +def _straddling(segments, kv_heads: int) -> int: + """The pool page count that puts a cell's requests -- the pool's last pages, as the + test lays them out -- on both sides of `_OFFSET_EDGE` elements of the K plane (half + of them below it, or the whole pool when they need more than twice the pages below).""" + tpb, d, slot_pages, layer = 32, 256, _TRUNK["slot_pages"], _TRUNK["layer"] + needed = sum(math.ceil((s + n) / tpb) for s, n in segments) + 7 + role_page = kv_heads * tpb * d + edge_page = math.ceil((_OFFSET_EDGE - 2 * layer * role_page) / (slot_pages * role_page)) + return max(edge_page + needed // 2, needed) + + +def _default(segments, member: dict, launches, **overrides) -> dict: + """A cell of the default target's call: `member`'s heads, table and trunk layout, its pool + straddling `_OFFSET_EDGE`.""" + return _decode( + segments, + launches=launches, + rows=sum(n for _, n in segments), + pool_pages=_straddling(segments, member["kv_heads"]), + offset_edge=_OFFSET_EDGE, + **{**member, **overrides}, + ) + + +def _spread(count: int, rows: int, last_start: int) -> tuple: + """`count` requests of `rows` rows: one starting at `last_start` (its last row at the + envelope's last position), the others spread from position 0 up.""" + span = last_start + 1 + return tuple(((5 + 1129 * g) % span, rows) for g in range(count - 1)) + ((last_start, rows),) + + +class QsaPagedSparseGqa(OpWrapper): + """Sparse GQA over the selected tokens of an HND paged K/V pool. + + Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels. + triton_qsa_paged_sparse_gqa`. For row `r`, query head `h`, KV head + `h // (Hq // Hkv)` (contiguous groups, not interleaved), over the columns + `j` of `selected_tokens[r]` that name a resident token:: + + t = selected_tokens[r, j] # logical position + page = block_table[request_indices[r], t // tokens_per_block] + s_j = (q[r, h] . k_cache[page, kv, t % tpb]) * softmax_scale + out[r, h] = sum_j softmax(s)_j * v_cache[page, kv, t % tpb] + + A negative index, a page past the table's width, a negative or + out-of-pool table entry, and a request index outside the table's rows all + contribute nothing; a row left with nothing returns zeros. Indices are not + deduplicated: a token named twice is attended twice. Returns a fresh, + contiguous `[rows, Hq, D]` tensor of `q.dtype`, even for a strided `q`. + + Two kernels, chosen inside the call: with `query_positions` and + `compress_ratio` given and `rows * Hkv > 512`, a causal-prefix prefill + kernel; otherwise split-K, plus a merge kernel when it used more than one + split. Each cell names the launches it was measured on. + + **The prefill pair is a bound, not a mask.** The prefill kernel reads only + `min(visible // cr, block_topk) * cr + visible % cr` leading columns of a + row (`visible = query_positions[r] + 1`, rounded up to a 64-wide tile) and + nothing after them. That is exactly what QSA's own top-k + expansion emit, + packed at the front; a selection with valid entries past the bound is + silently truncated. The split-K kernel ignores the pair. + + **The output gate rides the merge when `output_gate` is passed.** Each + merged head row is rounded to `q.dtype`, multiplied in fp32 by + `sigmoid(output_gate)` and rounded again: the arithmetic of the ungated + call followed by `fused_sigmoid_mul_inplace`. `output_gate` is `[rows, + Hq * D]` or `[rows, Hq, D]` in `q.dtype`, unit stride in its last dim, any + row stride (the gate columns of the fused projection row). Only a split-K + launch with more than one split has a merge to carry it; `fuses_output_gate` + says whether a call does, and the op raises `ValueError` if handed a gate + it cannot carry. + + Fusion boundary: the block-table gather, scale, online softmax, P.V, the + split-K merge and, when passed, the output gate are inside. Deciding what + to attend to, appending K/V to the pool, and q's projection, norm and RoPE + are not. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + # ---- generation: no prefill pair, split-K --------------------------- + Cell( + why="decode MTP7: 32 slots x 8 rows = 256 rows, 512 programs, 4-way split-K, " + "q and the gate off the fused row, the gate fused into the merge", + spec=_decode( + tuple((100 + (499 * g) % 8000, 8) for g in range(32)), + q_layout="fused_row", + gate="fused_row", + ), + ), + # ---- context: the prefill pair is passed ---------------------------- + Cell( + why="fp8 context 272 rows from 2032: causal-prefix prefill kernel", + spec=_decode(((2032, 272),), launches=_PREFILL, prefill_pair=True, **_FP8), + ), + Cell( + why="q off the fused projection row (55 x 256), 8191-row prefill", + spec=_decode(((0, 8191),), launches=_PREFILL, prefill_pair=True, q_layout="fused_row"), + ), + # ---- the default target: one rank's heads, table and trunk-layout pool ------ + Cell( + why="tp4/tep4 decode, 1 row at position 9215 (max_seq_len, table column 287): " + "6 q / 1 kv heads, 64-way split-K + merge with PDL, [128, 292] table", + spec=_default(_spread(1, 1, 9215), _TP4, _SPLIT_K, pdl=True), + ), + Cell( + why="tp4/tep4 decode, 128 requests x 1 row (the largest captured batch): " + "6 q / 1 kv heads, 8-way split-K + merge, [128, 292] table", + spec=_default(_spread(128, 1, 9215), _TP4, _SPLIT_K), + ), + Cell( + why="tep4 MTP7 verify, 128 requests x 8 rows = 1024: 6 q / 1 kv heads, one split " + "and no merge, [128, 516] table (max_seq_len 16384), rows to position 16383", + spec=_default(_spread(128, 8, 16376), _TP4_MTP, ("splitk",), rows_per_request=8), + ), + Cell( + why="tp4/tep4 8192-row step (max_num_tokens) with the prefill pair: a context " + "request 0..8064 and 127 generation rows to position 9215, all on the prefill " + "kernel, 6 q / 1 kv heads", + spec=_default(((0, 8065),) + _spread(127, 1, 9215), _TP4, _PREFILL, prefill_pair=True), + ), + Cell( + why="dep4 kv_fp8 decode, 32 requests x 1 row (a rank's batch): 24 q / 2 kv heads " + "over an e4m3 pool, 8-way split-K + merge, [32, 292] table", + spec=_default(_spread(32, 1, 9215), _DEP4_FP8, _SPLIT_K), + ), + # ---- the gated decode / verify site: q and the gate off the fused row ------ + Cell( + why="tep4 gated decode 1 row in a CUDA graph: q and the gate off the 17-head " + "(4352-wide) fused row, 64-way split-K + gated merge with PDL, captured, refilled " + "and replayed", + spec=_default( + _spread(1, 1, 9215), _TP4, _SPLIT_K, splits=64, pdl=True, capture=True, **_TP4_ROW + ), + ), + Cell( + why="tep4 gated decode 3 rows (PDL's 2..4-row branch): 64-way split-K + gated merge " + "with PDL, 17-head fused row", + spec=_default(_spread(3, 1, 9215), _TP4, _SPLIT_K, splits=64, pdl=True, **_TP4_ROW), + ), + Cell( + why="tep4 gated decode 16 rows (the 5..31-row branch): 32-way split-K + gated merge, " + "17-head fused row", + spec=_default(_spread(16, 1, 9215), _TP4, _SPLIT_K, splits=32, **_TP4_ROW), + ), + Cell( + why="tep4 gated decode 128 x 1 row in a CUDA graph (the largest captured batch): " + "8-way split-K + gated merge, 17-head fused row, captured, refilled and replayed", + spec=_default( + _spread(128, 1, 9215), _TP4, _SPLIT_K, splits=8, capture=True, **_TP4_ROW + ), + ), + Cell( + why="tep4 MTP3 gated verify 128 x 4 rows = 512 (the 257..512-row branch): 4-way " + "split-K + gated merge, [128, 516] table, rows to position 16383, 17-head fused row", + spec=_default( + _spread(128, 4, 16380), + _TP4_MTP, + _SPLIT_K, + splits=4, + rows_per_request=4, + **_TP4_ROW, + ), + ), + Cell( + why="dep4 core gated decode 32 x 1 row: 24 q / 2 kv heads over a bf16 pool, q and " + "the gate off the 55-head (14080-wide) fused row, 8-way split-K + gated merge", + spec=_default(_spread(32, 1, 9215), _DEP4, _SPLIT_K, splits=8, **_DEP4_ROW), + ), + Cell( + why="dep4 kv_fp8 context 8192 rows from position 0 with the prefill pair: 24 q / 2 " + "kv heads over an e4m3 pool on the prefill kernel, [32, 292] table", + spec=_default(((0, 8192),), _DEP4_FP8, _PREFILL, prefill_pair=True), + ), + Cell( + why="tep4 8192-row context step (max_num_tokens) with the prefill pair, q off the " + "17-head fused row: a context request 0..8064 and 127 generation rows to position " + "9215, all on the prefill kernel", + spec=_default( + ((0, 8065),) + _spread(127, 1, 9215), + _TP4, + _PREFILL, + prefill_pair=True, + q_layout="fused_row", + row_heads=17, + ), + ), + ) + + note = """ + Metadata consumed: none -- every input is an argument. What the target + passes: `block_table` is `qsa_block_table`, physical page ids into dim 0 + of the caches (the *unscaled* table, not logical == physical); + `selected_tokens` is `qsa_expand_block_indices`' output, logical + positions, `-1`-padded; `k_cache` / `v_cache` are the layer pool's role + planes. + + Preconditions: + - Whenever the prefill kernel can run, each selection row is packed at + the front with exactly the bound's count of valid entries (see the + class docstring). QSA's producer -- `qsa_topk_select` with zero + `row_starts` and `row_ends = (pos + 1) // cr`, then the expansion -- + satisfies it (measured, including with NaN/inf in the scores). + Checking it needs the contents, so it is not an `is_valid` guard. + - The prefill pair is a property of the launch: passed because some + rows are context, it puts the generation rows on the prefill kernel + too, and the bound applies to them (the default 8192-row cell drives + 127 generation rows through it). + - There is no KV scale argument: an fp8 pool is read at face value (each code + widens exactly to bf16 under a bf16 q), so it must hold K/V written at unit scale. + - `D` must be a power of two (192 is refused); `Hq % Hkv == 0`. + + The op validates its own inputs loudly (`ValueError`s; the test pins the + refusal of positions passed without a compress ratio; a selection width + the compress ratio cannot produce is refused on the prefill kernel only), + so `is_valid` adds nothing. `rows == 0` returns an empty `[0, Hq, D]` + tensor (not certified). + + PDL: `TRTLLM_ENABLE_PDL` (default on) enables programmatic dependent + launch for multi-split launches of 4 rows or fewer; the default 1-row + cell runs one (split-K and merge both launched with PDL, asserted). At tp1 + the target sends gated decode batches that small to + `qsa_paged_sparse_gqa_decode`. Calls are bitwise deterministic (pinned + on a 32-split launch and its merge). + + The decode site carries up to 8 rows per request (max_draft_len 7): the + decode cell runs 32 requests of 8 rows, 512 `rows * Hkv` programs and a + 4-way split-K launch with the output gate in its merge. + + The default target (tp4 / tep4 / dep4). A tp4 or tep4 rank holds 6 q heads + and one KV head (`Hkv` 1, a group of 6 padded to 16 MMA rows), so its + launch follows `rows`: 64-way split-K up to 4 rows, 32-way below 32, + 8-way to 256, 4-way to 512, then one split and no merge (the 1024-row + MTP7 verify step, where the output gate cannot ride a merge and + `fuses_output_gate` is False); with the prefill pair, the prefill kernel + past 512 rows. A dep4 rank holds tp1's 24 / 2 heads, 32 requests, and + under kv_fp8 an e4m3 pool, read on the split-K kernel. Its cells drive + the engine's trunk layout (a page advances 24 role pages; layer 11's + planes) over `[128, 292]`, `[128, 516]` (MTP's max_seq_len 16384) and + `[32, 292]` tables, with pools whose mapped pages straddle 2^31 elements of + the K and V planes: both kernels widen the page id to int64 before its + stride product. Against `reference` over the five ungated cells: at most + 2.00 bf16 ulp of the row scale, relative RMS at most 0.572 ulp. + + The gated decode and verify site. q and the output gate are views of one + rank's fused input projection row (`row_heads`): 17 heads of 256 at tp4 / + tep4 (q heads 0..5, the gate 8..13, row stride 4352), 55 at dep4 (q 0..23, + the gate 28..51, 14080). `q_stride_row` is a Triton constexpr, so each row + width compiles its own kernels. Cells assert their split count, so every + 6/1 launch branch with a merge carries the gate: 64-way split with PDL at + 1 and 3 rows, 32-way at 16, 8-way at 128 (decode), 4-way at 512 (MTP3 + verify, 128 x 4 rows to position 16383). dep4's decode runs it at 24/2 + over a bf16 pool (8-way, 32 rows). + The 1-row and 128-row cells are captured in a CUDA graph (the split-K + partials come from the graph's pool); a replay over q, the gate, the + selection and K/V refilled in place takes the eager call's launches (PDL + included) and matches an eager call on the new inputs bitwise. + A gated output is rounded twice (the attention, then the gated + product), each at most one bf16 ulp of the value rounded, so its element + error is bounded near the band's 4 ulp rather than 2: at most 3.32 ulp of + the row scale and 0.667 ulp relative RMS over the seven gated cells + (tp1's included); the same calls ungated at most 2.00. + Context rows: dep4's kv_fp8 context step runs + 8192 rows from position 0 on the prefill kernel over an e4m3 pool, and + tep4's 8192-row step (max_num_tokens) with q off the 17-head row. Both at + most 2.00 ulp and 0.573 ulp RMS. + + Pool page count. `NUM_CACHE_PAGES` is a constexpr that only bounds the + page id (`physical_pages < NUM_CACHE_PAGES`). The cells' pools straddle + 2^31 elements at sizes a normal cell holds (5592 to 33167 pages). The + engines' K/V views hold more: 257109 pages (tep4 bf16), 119680 (dep4 bf16) + and 171362 (dep4 kv_fp8), 94.2, 87.7 and 62.8 GiB as trunk pools. Probed + at those exact counts, each matched its cell's numbers; no cell carries + them. + + Not certified: head geometries other than 24/2 and 6/1 x 256; + `tokens_per_block` other than 32; more than 8192 rows; capture of the + one-split launch (the 1024-row MTP7 verify, which allocates only its + output), of the 512-row MTP3 verify and of the prefill kernel; dep4's + gated decode below 16 rows (its 64- and 32-way splits at 24/2); an fp8 + pool at 6/1 (no tp4 / tep4 member stores one) and on the prefill kernel + with q off the 55-head row (the fp8 context cell reads the packed + `[q | k | v]` slice); positions past 16383; rows whose selection + resolves no token (all padding, a foreign + request, unmapped pages); an fp16 pool or q; standalone contiguous K/V + buffers; int64 request indices. Hard limits: not probed; the largest + certified extent is 8192 rows x 2051 selected tokens, a [128, 516] table + and a 33167-page trunk pool. + """ + + def __call__( + self, + *, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + selected_tokens: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + softmax_scale: float, + query_positions: Optional[torch.Tensor] = None, + compress_ratio: Optional[int] = None, + output_gate: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return _triton_qsa_paged_sparse_gqa( + q=q, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + selected_tokens=selected_tokens, + request_indices=request_indices, + tokens_per_block=tokens_per_block, + softmax_scale=softmax_scale, + query_positions=query_positions, + compress_ratio=compress_ratio, + output_gate=output_gate, + ) + + @staticmethod + def fuses_output_gate( + q: torch.Tensor, + k_cache: torch.Tensor, + selected_tokens: torch.Tensor, + output_gate: torch.Tensor, + ) -> bool: + """Whether a call with these operands and no prefill pair carries `output_gate`.""" + return _can_fuse_qsa_splitk_output_gate(q, k_cache, selected_tokens, output_gate) + + def reference( + self, + *, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + selected_tokens: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + softmax_scale: float, + query_positions: Optional[torch.Tensor] = None, + compress_ratio: Optional[int] = None, + output_gate: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """fp32 sparse GQA over every resident selected token, rounded once; + with `output_gate`, then gated in fp32 and rounded again. + + `query_positions` and `compress_ratio` are ignored: on a prefix-packed + selection the prefill kernel's bound reads exactly the valid entries. + Walked 32 rows at a time to bound the gathered K/V. + """ + rows, num_q_heads, head_dim = q.shape + num_pages, num_kv_heads = k_cache.shape[0], k_cache.shape[1] + groups = num_q_heads // num_kv_heads + num_requests, table_width = block_table.shape + requests = request_indices.long() + out = torch.empty(rows, num_q_heads, head_dim, dtype=torch.float32, device=q.device) + for start in range(0, rows, 32): + stop = min(start + 32, rows) + chunk = selected_tokens[start:stop].long() + logical_page = chunk.clamp_min(0) // tokens_per_block + in_page = chunk.clamp_min(0) % tokens_per_block + request = requests[start:stop] + owned = (request >= 0) & (request < num_requests) + physical = block_table[ + request.clamp(0, num_requests - 1)[:, None].expand_as(logical_page), + logical_page.clamp(max=table_width - 1), + ].long() + resident = ( + (chunk >= 0) + & owned[:, None] + & (logical_page < table_width) + & (physical >= 0) + & (physical < num_pages) + ) + physical = physical.clamp(0, num_pages - 1) + # What a non-resident column gathers (a clamped page) is not part of the answer: + # zeroed, so a pool holding NaN outside the mapped pages cannot reach the sum. + absent = ~resident[:, :, None, None] + keys = k_cache[physical, :, in_page, :].float().masked_fill(absent, 0.0) + values = v_cache[physical, :, in_page, :].float().masked_fill(absent, 0.0) + grouped = q[start:stop].reshape(-1, num_kv_heads, groups, head_dim).float() + scores = torch.einsum("bhgd,bkhd->bhgk", grouped, keys) * softmax_scale + scores = scores.masked_fill(~resident[:, None, None, :], float("-inf")) + # A row with nothing resident softmaxes over all -inf: NaN, i.e. zeros. + weights = torch.nan_to_num(torch.softmax(scores, dim=-1), nan=0.0) + pooled = torch.einsum("bhgk,bkhd->bhgd", weights, values) + out[start:stop] = pooled.reshape(stop - start, num_q_heads, head_dim) + if output_gate is None: + return out.to(q.dtype) + gate = output_gate.reshape(rows, num_q_heads, head_dim).float() + return (out.to(q.dtype).float() * torch.sigmoid(gate)).to(q.dtype) + + +qsa_paged_sparse_gqa = QsaPagedSparseGqa() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_decode.py similarity index 99% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_decode.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_decode.py index b6180bcb7566..7047ce0ec03b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_decode.py @@ -51,13 +51,14 @@ from cutlass.cute.runtime import from_dlpack, make_ptr from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa import ( + qsa_paged_sparse_gqa as _composition, +) # Approved exception (as for `qsa_paged_sparse_gqa`): QSA's sparse path has no torch.ops # surface; its row-count PDL policy is reached as Python. from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import _qsa_pdl_enabled -from .qsa_paged_sparse_gqa import qsa_paged_sparse_gqa as _composition - #: Output elements a merge CTA covers (four adjacent ones a lane of each warp). CHUNK = 128 #: Splits the merge reads: two LSEs a lane. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4.py similarity index 91% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4.py index 75e9c55e17a8..734370a39e0b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4.py @@ -48,7 +48,7 @@ layer=0, ) -#: The prefill kernel's (BLOCK_N, warps): on GB300 8192 context rows 6.5 ms, 4096 rows 2.9 ms. +#: The prefill kernel's (BLOCK_N, warps). _PREFILL_TILE = (32, 2) _SPLIT_K = ("splitk", "merge") @@ -71,6 +71,30 @@ def _decode(segments, launches=_SPLIT_K, **overrides) -> dict: return _spec(**{**base, **overrides}) +#: One dep4 rank's engine under `kv_nvfp4` (attention DP, so tp1's 24/2 x 256 geometry): +#: `max_batch_size` 32 a rank and `max_seq_len` 9216 give a `[32, 292]` table +#: (`ceil((9216 + 1) / 32)` rounded up to a multiple of 4); the trunk layout's 24 role pages +#: a slot, layer 11's planes the last of them. +_DEP4 = dict(pages=292, slots=32, slot_pages=24, layer=11) + +#: The byte offset the dep4 cells' pools straddle: one page of layer 11's scale plane +#: advances 24 x 2 x 32 x 16 bytes, so from page 87382 on it lies past 2^31 bytes (the data +#: plane, 8x wider, from page 10923). A dep4 pool holds far more pages than either. +_OFFSET_EDGE = 2**31 + + +def _dep4_context(segments, pool_pages) -> dict: + """A context step of dep4's kv_nvfp4 engine: the prefill pair, the prefill kernel.""" + return _decode( + segments, + launches=_PREFILL, + prefill_pair=True, + pool_pages=pool_pages, + offset_edge=_OFFSET_EDGE, + **_DEP4, + ) + + @triton.jit def _e2m1_to_f32(code): """e2m1 code (0..15, bit 3 the sign) -> its value, exactly.""" @@ -604,8 +628,7 @@ def _nvfp4_merge_splitk_kernel( def _splitk_launch_config( rows: int, num_kv_heads: int, selected_width: int ) -> tuple[int, int, int]: - """(BLOCK_N, splits, warps) by programs = rows x Hkv, the fastest of a sweep on GB300 at - 2051 selected tokens a row (1 row 12.6 us, 128 rows 160 us, 512 rows 560 us).""" + """(BLOCK_N, splits, warps) by programs = rows x Hkv.""" base_programs = rows * num_kv_heads if base_programs <= 4: block_n, target_splits, num_warps = 32, 64, 4 @@ -935,6 +958,18 @@ class QsaPagedSparseGqaNvfp4(OpWrapper): why="captured decode: G 1, 8 and 32 requests x 1 and 8 rows, replay == eager", spec=_decode((), slots=32, capture=dict(groups=(1, 8, 32), rows_per_request=(1, 8))), ), + Cell( + why="dep4 8192-row context chunk (max_num_tokens): a request at 1152..9215 " + "(max_seq_len, table column 287) and one from 0, [32, 292] table, trunk layout, " + "pages straddling 2^31 bytes of the scale plane; 87408 pages", + spec=_dep4_context(((1152, 8064), (0, 128)), pool_pages=87408), + ), + Cell( + why="dep4 512-row context step (the chunked variant's max_num_tokens): four " + "requests' chunks, one ending at max_seq_len, one crossing the 2048-token " + "selection budget; 87401 pages (not a multiple of 16)", + spec=_dep4_context(((8960, 256), (0, 160), (2000, 64), (5000, 32)), pool_pages=87401), + ), ) note = """ @@ -948,7 +983,11 @@ class QsaPagedSparseGqaNvfp4(OpWrapper): Call sites (the target's `modeling.py`): the context/mixed forward passes the prefill pair on every step, the decode forward never does; both pass `q` as - the 24-head slice of the packed QKV row (row stride 7168). + the 24-head slice of the packed QKV row (row stride 7168). In the default + target only dep4's `kv_nvfp4` engine binds an NVFP4 pool, with tp1's 24/2 x 256 + geometry (attention DP): its context steps, up to `max_num_tokens` 8192 rows (512 + in the chunked variant), over a `[32, 292]` table (`max_batch_size` 32 a rank, + `max_seq_len` 9216) and the trunk layout's 24 role pages a slot. Launch shapes at Hkv 2 (programs = rows x 2): split-K at 32-token tiles 64-way up to 2 rows, 16-way up to 8, 8-way up to 16 and 4-way up to 32 @@ -961,8 +1000,7 @@ class QsaPagedSparseGqaNvfp4(OpWrapper): Each byte is decoded once, into the even and the odd halves of the head dimension, which the kernels keep apart: q's halves against K's on two MMAs, P against V's halves into two accumulators, the output stored back - interleaved. On GB300 at 2051 selected tokens a row: 1 row 12.6 us, 128 rows - 160 us, against 10.4 / 108 us for QSA's bf16 kernel over the same values. + interleaved. Equivalence: the test also runs QSA's bf16 sparse GQA over a bf16 pool holding the dequantized values; on every cell the two outputs agree within @@ -973,12 +1011,25 @@ class QsaPagedSparseGqaNvfp4(OpWrapper): [32, 260] table; a replay after new operands are written in place is bitwise the eager call on them. - Compiled variants: the kernels take the pool's page count (`num_pages`), the + Compiled variants: the kernels are Triton, JIT-compiled on first call for the + device that runs them; every cell ran that path on sm_103 (GB300). They take + the pool's page count (`num_pages`), the table's row count (`num_requests`) and the split-K kernel the step's row count as plain ints, and Triton compiles each apart at 1, at multiples of 16 and otherwise. Every cell runs a 32-row table (a multiple of 16); the other - table-height classes are not certified. Every other plain int -- the strides, the table width 260 and - the selection's row stride 2051 -- is fixed at the served value the cells use. + table-height classes are not certified. The page count is driven at both the + multiple-of-16 class and the other (the dep4 cells: 87408 and 87401). Every other + plain int -- the strides, the table widths 260 (tp1) and 292 (dep4) and the + selection's row stride 2051 -- is fixed at the served values the cells use. + + Large pools: every page index is widened to 64 bits before its product with a + plane's page stride. In the trunk layout one page of a layer's scale plane + advances 24 x 2 x 32 x 16 bytes and of its data plane 8x that, so a pool past + 87381 pages addresses both planes past 2^31 bytes; the dep4 cells place their + requests' pages on both sides of that edge in the scale plane (the test asserts + it), with every scale outside them NaN. Against `reference`: at most 2.00 bf16 + ulp of a row's magnitude and 0.25 ulp relative RMS at 8192 rows, 1.98 / 0.24 + ulp at 512 rows, inside the test's band of 4 per element and 2 RMS. Preconditions: - Whenever the prefill kernel can run, each selection row is packed at the @@ -999,11 +1050,13 @@ class QsaPagedSparseGqaNvfp4(OpWrapper): Every row, request, page and position index is widened to 64 bits before any product. Not certified: head geometries other than 24/2 x 256; `tokens_per_block` other than 32; more than 8192 rows; capture of the - context/mixed site; a pool of one slot; second-level scales other than 1.0; - the prefill pair on the split-K kernel (256 context rows or fewer); one split - without a merge; the trunk layout on the prefill kernel; rows whose selection - resolves no token (all padding, a foreign request, unmapped pages); an fp16 - or unsliced q; int64 request indices. + context/mixed site; a pool of one slot; second-level scales other than 1.0 + (the target binds ones); the prefill pair on the split-K kernel (256 context + rows or fewer); one split without a merge; rows whose selection resolves no + token (all padding, a foreign request, unmapped pages); an fp16 or unsliced q; + int64 request indices; dep4's own pool page count (the cells' pools + are sized to the 2^31-byte edge rather than to an engine's allocation, which + is larger); positions past `max_seq_len` 9216. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4_decode.py similarity index 86% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4_decode.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4_decode.py index b25c1c1b63a6..7dee4f99f9c0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4_decode.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """QSA sparse GQA over an NVFP4 paged K/V pool, decode rows: CuTe DSL split-K + merge. -Tier 5 (CuTe DSL written for this target), both launches. One CTA (eight warps) serves one +Both launches are CuTe DSL. One CTA (eight warps) serves one (row, KV head, split) and walks the split's tiles of 64 selected tokens with an online softmax; the split count is chosen by row count so the grid stays about two waves deep (33 splits of one tile at 1 row, one split of 33 tiles from 128 rows up). @@ -45,11 +45,12 @@ from cutlass.cutlass_dsl import T, dsl_user_op from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4 import ( + qsa_paged_sparse_gqa_nvfp4 as _composition, +) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr -from .qsa_paged_sparse_gqa_nvfp4 import qsa_paged_sparse_gqa_nvfp4 as _composition - TILE = 64 WARPS = 8 THREADS = WARPS * 32 @@ -265,22 +266,6 @@ def _dq_word(word, scale2, *, loc=None, ip=None): return _unpack(res, 4, cutlass.Int32, loc, ip) -@dsl_user_op -def _e2m1_half(half, *, loc=None, ip=None): - """Two e2m1 bytes (the low 16 bits: 4 values, low nibble the even one) -> 2 f16x2 - (values 0-1, values 2-3), unscaled.""" - res = _asm( - llvm.StructType.get_literal([T.i32()] * 2), - [_ir(half, cutlass.Int32, loc, ip)], - "{ .reg .b8 b0, b1, b2, b3; mov.b32 {b0, b1, b2, b3}, $2; " - "cvt.rn.f16x2.e2m1x2 $0, b0; cvt.rn.f16x2.e2m1x2 $1, b1; }", - "=r,=r,r", - loc, - ip, - ) - return _unpack(res, 2, cutlass.Int32, loc, ip) - - @dsl_user_op def _e4m3x2(half, *, loc=None, ip=None): """Two e4m3 bytes (the low 16 bits) -> f16x2, the low byte in the low half.""" @@ -296,38 +281,6 @@ def _e4m3x2(half, *, loc=None, ip=None): ) -@dsl_user_op -def _hmul2(a, b, *, loc=None, ip=None): - return cutlass.Int32( - _asm( - T.i32(), - [_ir(a, cutlass.Int32, loc, ip), _ir(b, cutlass.Int32, loc, ip)], - "mul.rn.f16x2 $0, $1, $2;", - "=r,r,r", - loc, - ip, - ) - ) - - -@dsl_user_op -def _prmt(a, b, sel, *, loc=None, ip=None): - return cutlass.Int32( - _asm( - T.i32(), - [ - _ir(a, cutlass.Int32, loc, ip), - _ir(b, cutlass.Int32, loc, ip), - _ir(sel, cutlass.Int32, loc, ip), - ], # fmt: skip - "prmt.b32 $0, $1, $2, $3;", - "=r,r,r,r", - loc, - ip, - ) - ) - - @dsl_user_op def _lds_u16(saddr, *, loc=None, ip=None): return cutlass.Int32( @@ -357,22 +310,6 @@ def _stg4(gaddr, a, b, c, d, *, loc=None, ip=None): ) -@dsl_user_op -def _scale2(sbyte, *, loc=None, ip=None): - """An e4m3 byte (in the low 8 bits) -> f16x2 {s, s}.""" - return cutlass.Int32( - _asm( - T.i32(), - [_ir(sbyte, cutlass.Int32, loc, ip)], - "{ .reg .b32 t; .reg .b16 h; and.b32 t, $1, 255; mul.lo.u32 t, t, 257; " - "cvt.u16.u32 h, t; cvt.rn.f16x2.e4m3x2 $0, h; }", - "=r,r", - loc, - ip, - ) - ) - - @dsl_user_op def _scale2s(word, sel, *, loc=None, ip=None): """The e4m3 byte `word.byte[sel & 3]` (sel = k * 0x11) -> f16x2 {s, s}: one prmt + one cvt.""" @@ -745,16 +682,21 @@ def kernel( cute.arch.sync_threads() # ---- q of this KV head: rows g and g + 8 of the m16 tile, as permuted A fragments -- - # thread (g, c) holds dims 32m + 8c .. 32m + 8c + 7 of heads g and g + 8, m = 0..7 + # thread (g, c) holds dims 32m + 8c .. 32m + 8c + 7 of heads g and g + 8, m = 0..7; a + # row past the group (g + 8 always below 16 heads, g itself below 8) holds zeros and + # stores nothing qa = cute.make_rmem_tensor(cute.make_layout((8, 4), stride=(4, 1)), cutlass.Int32) qb = cute.make_rmem_tensor(cute.make_layout((8, 4), stride=(4, 1)), cutlass.Int32) qs0, qs1 = self.qs q_row = q_base + (cutlass.Int64(row) * qs0 + cutlass.Int64(kvh * group) * qs1) * 2 for m in cutlass.range_constexpr(8): off = cutlass.Int64((32 * m + 8 * c) * 2) - w4 = _ldg4(q_row + cutlass.Int64(g) * (qs1 * 2) + off) + hac = cutlass.min(g, group - 1) + w4 = _ldg4(q_row + cutlass.Int64(hac) * (qs1 * 2) + off) for e in cutlass.range_constexpr(4): - qa[m, e] = _bf2_to_f2(w4[e]) + qa[m, e] = cutlass.Int32( + cutlass.select_(g < group, _bf2_to_f2(w4[e]), cutlass.Int32(0)) + ) hb = g + 8 hbc = cutlass.min(hb, group - 1) w4b = _ldg4(q_row + cutlass.Int64(hbc) * (qs1 * 2) + off) @@ -972,6 +914,7 @@ def kernel( ih = gv / lh hg = kvh * group + g hh = kvh * group + g + 8 + okg = g < group okh = (g + 8) < group for nt in cutlass.range_constexpr(4): d = 32 * warp + 8 * nt + 2 * c @@ -982,13 +925,14 @@ def kernel( if const_expr(splits == 1): ag = out_base + (cutlass.Int64(row) * heads + hg) * (HEAD_DIM * 2) + d * 2 ah = out_base + (cutlass.Int64(row) * heads + hh) * (HEAD_DIM * 2) + d * 2 - if const_expr(self.gated): - gw = _ldg1( - gate_base + (cutlass.Int64(row) * gate_stride + hg * HEAD_DIM + d) * 2 - ) - og0 = _round_bf16(og0) * _sigmoid(_bf2_lo(gw)) - og1 = _round_bf16(og1) * _sigmoid(_bf2_hi(gw)) - _stg1(ag, _bf2_pack(og0, og1)) + if okg: + if const_expr(self.gated): + gw = _ldg1( + gate_base + (cutlass.Int64(row) * gate_stride + hg * HEAD_DIM + d) * 2 + ) + og0 = _round_bf16(og0) * _sigmoid(_bf2_lo(gw)) + og1 = _round_bf16(og1) * _sigmoid(_bf2_hi(gw)) + _stg1(ag, _bf2_pack(og0, og1)) if okh: if const_expr(self.gated): gw2 = _ldg1( @@ -1008,7 +952,8 @@ def kernel( + ((cutlass.Int64(split) * rows + row) * heads + hh) * (HEAD_DIM * 4) + d * 4 ) - _stg2f(pg_, og0, og1) + if okg: + _stg2f(pg_, og0, og1) if okh: _stg2f(ph_, oh0, oh1) if const_expr(splits > 1): @@ -1017,10 +962,11 @@ def kernel( lse_g = neg_inf if lg > 0.0: lse_g = st[0] + cute.math.log2(lg) - _stg1( - lse_base + ((cutlass.Int64(split) * rows + row) * heads + hg) * 4, - _f32_bits(lse_g), - ) + if okg: + _stg1( + lse_base + ((cutlass.Int64(split) * rows + row) * heads + hg) * 4, + _f32_bits(lse_g), + ) if okh: lse_h = neg_inf if lh > 0.0: @@ -1201,8 +1147,8 @@ def kernel( _COMPILED: dict = {} -#: Splits by (rows x KV heads), measured on GB300 at 2051 selected tokens (sweep_splits.py): -#: about 250-290 CTAs where the merge is cheap, a single split from 128 rows up. +#: Splits by (rows x KV heads): about 250-290 CTAs where the merge is cheap, a single split +#: from 128 rows up. _SPLIT_TABLE = ((4, 33), (8, 11), (16, 17), (32, 9), (64, 4), (128, 2)) @@ -1390,6 +1336,49 @@ def _decode(segments, **overrides) -> dict: return dict(segments=segments, **overrides) +#: One tp4 / tep4 rank's heads: 6 q heads over 1 KV head, fewer q heads a KV head than the +#: m16 tile's eight first rows. +_TP4_HEADS = dict(q_heads=6, kv_heads=1) + + +#: One dep4 rank's pool under `kv_nvfp4` (attention DP, so tp1's 24/2 x 256 geometry), as the +#: stock engine allocates it (`max_batch_size` 128, `max_seq_len` 9216, +#: `free_gpu_memory_fraction` 0.6, chunked prefill): a K/V view of 211797 pages of 32 tokens a +#: rank (the manager's 213295 blocks less the 1498 recurrent-state slots its second pool group +#: holds), the engine's `qsa_block_table` of `max_num_sequences` 128 rows by 292 columns +#: (`ceil((9216 + 1) / 32)` rounded up to a multiple of 4), the trunk layout's 24 role pages a +#: slot with layer 11's planes the last of them. The requests' pages are scattered over the +#: whole pool, its last page among them. +_DEP4_NVFP4 = dict(slots=128, pages=292, slot_pages=24, layer=11, pool_pages=211797, spread=True) + +#: The byte offsets the dep4 cells' pages must straddle in both planes: one page of layer 11's +#: scale plane advances 24 x 2 x 32 x 16 bytes and of its data plane 8x that, so a 211797-page +#: pool addresses the scale plane to 5.2 GB and the data plane to 41.6 GB. +_OFFSET_EDGES = (2**31, 2**32) + + +def _dep4_decode(positions) -> dict: + """A dep4 decode step: one row a request at each position, with the output gate, over + the pool above; captured, as the engine runs it. About 44 GiB of pool (`large`).""" + return _decode( + tuple((p, 1) for p in positions), + gated=True, + capture=True, + large=True, + offset_edges=_OFFSET_EDGES, + **_DEP4_NVFP4, + ) + + +#: Decode positions of a dep4 step's requests: the last token `max_seq_len` 9216 allows (table +#: column 287), a first token, the last position whose history fits inside the 2048-token +#: budget, then a spread across the rest. +def _dep4_positions(requests: int) -> tuple[int, ...]: + edge = (9215, 0, 2047) + spread = tuple((100 + 499 * r) % 9216 for r in range(requests - len(edge))) + return (edge + spread)[:requests] + + class QsaPagedSparseGqaNvfp4Decode(OpWrapper): """`qsa_paged_sparse_gqa_nvfp4` for decode rows (no prefill pair): sparse GQA over the explicitly selected tokens of an NVFP4 paged pool, `[rows, Hq, D]` bf16. With @@ -1420,6 +1409,42 @@ class QsaPagedSparseGqaNvfp4Decode(OpWrapper): pool_pages=2048, ), ), + Cell( + why="tp4 rank heads (6 q over 1 KV): 1 row, 33-way split + merge", + spec=_decode(((4000, 1),), **_TP4_HEADS), + ), + Cell( + why="tp4 rank heads (6 q over 1 KV): 128 rows one a request, 2-way split + gated " + "merge, captured", + spec=_decode( + tuple((100 + (499 * r) % 8000, 1) for r in range(128)), + gated=True, + capture=True, + **_TP4_HEADS, + ), + ), + Cell( + why="tp4 rank heads (6 q over 1 KV): 160 rows (20 x 8), one split, gate fused into " + "the split kernel", + spec=_decode( + tuple((100 + (499 * g) % 8000, 8) for g in range(20)), gated=True, **_TP4_HEADS + ), + ), + Cell( + why="dep4 kv_nvfp4 decode: 1 row at position 9215, 33-way split, gate fused into " + "the merge, [128, 292] table, 211797-page pool, captured", + spec=_dep4_decode(_dep4_positions(1)), + ), + Cell( + why="dep4 kv_nvfp4 decode: 32 requests (a fourth of the slots), 4-way split + " + "gated merge, 211797-page pool, captured", + spec=_dep4_decode(_dep4_positions(32)), + ), + Cell( + why="dep4 kv_nvfp4 decode: 128 requests (max_batch_size, every slot), one split " + "with the gate fused, 211797-page pool, captured", + spec=_dep4_decode(_dep4_positions(128)), + ), ) note = """ @@ -1430,19 +1455,53 @@ class QsaPagedSparseGqaNvfp4Decode(OpWrapper): Launch shapes: splits by (rows x KV heads) from `_SPLIT_TABLE` (33 splits of one 64-token tile at one row, a single split of 33 tiles from 128 rows up), each split followed by the - merge. GB300, 2051 selected tokens a row, one layer: 1 row 6.7 us, 4 rows 8.5, 32 rows 22.3, - 128 rows 67.8, 256 rows 133, 512 rows 231, 1024 rows 449 -- against 9.0 / 18.8 / 49.5 / - 158 / 292 / 557 / 1036 us for the Triton op. A one-split launch (65 rows and up) runs the - kernel without the f16 V tile: 72 rows 65.1 -> 61.7 us, 128 rows 71.6 -> 66.6, 192 rows - 111.4 -> 101.6, 256 rows 137.5 -> 128.3 (12 chained calls, gated), bitwise the standard - kernel. + merge. A one-split launch (65 rows and up) runs the kernel without the f16 V tile, bitwise + the standard kernel. Numerics: dequantized K/V values are exact in f16; scores on f16 MMAs with fp32 accumulation; P rounded to f16 (as the Triton op); the second-level scales act on the scores and the normalized output. Within the Triton op's band (4 ulp element, 2 ulp RMS) - of the float64 reference; not bitwise the Triton op. Deterministic; graph-capturable + of the fp32 reference; not bitwise the Triton op. Deterministic; graph-capturable (row count is a runtime value; one compile per split variant). + The CuTe kernel puts a KV head's q heads in the m16 tile's rows g and g + 8; a group of + fewer than eight q heads (tp4 and tep4: 6 over 1 KV head a rank) leaves rows past it + zero and unstored. The tp4-heads cells drive 1, 128 and 160 rows at that geometry (33 + splits, 2 splits with the gated merge, captured, and one split with the gate fused). + + dep4's `kv_nvfp4` engine binds tp1's 24/2 x 256 geometry (attention DP): its decode + steps, one row a request with the output gate, up to `max_batch_size` 128 rows a rank + (the captured batches 1..128), over the engine's `qsa_block_table` ([128, 292] at + `max_seq_len` 9216) and one rank's pool as the stock engine allocates it: a K/V view of + 211797 pages of 32 tokens (the manager's 213295 blocks less its 1498 recurrent-state + slots), the trunk layout's 24 role pages a slot. The dep4 cells drive 1, 32 and 128 rows + there (33 splits, 4 splits, one split with the gate fused into the split kernel), each + also captured and replayed at its own compile variant. Against `reference`, in bf16 ulp + (2^-7) of a (row, head)'s magnitude, max / relative RMS: 1 row 0.55 / 0.12, 32 rows 1.15 + / 0.03, 128 rows 1.21 / 0.05; QSA's bf16 kernel over the same values lies within 1.00 + per element of each. Each dep4 cell holds about 44 GiB of pool (41.6 GB of data planes, + 5.2 GB of scales), past a normal cell's 40 GiB: `large`. + + Compiled variants: CuTe DSL, JIT-compiled on the first call of each key for the device + that runs it; every cell ran that path on sm_103 (GB300). The key holds the pool's page + count, the table's shape, the planes' strides, the split plan and whether the gate is + fused, so the dep4 cells compile the variants a dep4 engine at those settings runs; any + other page count or table shape compiles a variant of its own. + + Large pools: a page index is widened to 64 bits before its product with a plane's page + stride. A 211797-page pool addresses layer 11's scale planes to 5.2 GB and its data + planes to 41.6 GB; the dep4 cells scatter their requests' pages over the whole pool, its + last page among them, so both planes are read on both sides of 2^31 and 2^32 bytes (the + test asserts it), with every scale outside them NaN. Hard limit, read from the source and + not probed: the token table's entry `page * tokens_per_block + slot` is 32-bit, so a pool + of 2^26 pages or more at 32 tokens a page wraps it (1.6 TB of scale planes at this + geometry, beyond any device). + + Not certified: the target engine's own dep4 pool (the cells take the stock engine's page + count); dep4 with MTP (the draft layer's two more role pages a slot, and the verify + entry); positions past `max_seq_len` 9216 (a wider table is another compiled variant); + second-level scales other than 1.0 (the target binds ones); int64 request indices. + Preconditions: bf16 q with head_dim 256; K and V planes (and their scale planes) share strides; power-of-two tokens per block; int32 row-contiguous block table; contiguous int32 selections. @@ -1493,7 +1552,7 @@ def reference( softmax_scale: float, output_gate=None, ) -> torch.Tensor: - """`qsa_paged_sparse_gqa_nvfp4`'s float64 reference, then the gate in fp32 over its + """`qsa_paged_sparse_gqa_nvfp4`'s fp32 reference, then the gate in fp32 over its bf16 rounding, rounded once.""" out = _composition.reference( q, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4_verify.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4_verify.py similarity index 99% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4_verify.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4_verify.py index a71036db0466..4b4aec3eff8f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4_verify.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_paged_sparse_gqa_nvfp4_verify.py @@ -64,23 +64,7 @@ from cutlass.cute.typing import Float32, Int32, Int64 from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl - -from .qsa_nvfp4_verify_ptx import ( - ctz, - gemm_ts, - groups8, - ldg2, - popc, - red_or_shared, - spread4, - stg2, - stg4, - stg4f, - tcgen05_fence_after, - tcgen05_fence_before, -) -from .qsa_paged_sparse_gqa_nvfp4_decode import ( +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4_decode import ( # noqa: E501 MAX_SPLITS, _bf2_hi, _bf2_lo, @@ -103,6 +87,22 @@ nvfp4_sparse_gqa_decode, qsa_paged_sparse_gqa_nvfp4_decode, ) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +from .qsa_nvfp4_verify_ptx import ( + ctz, + gemm_ts, + groups8, + ldg2, + popc, + red_or_shared, + spread4, + stg2, + stg4, + stg4f, + tcgen05_fence_after, + tcgen05_fence_before, +) try: from cutlass.memory import SmemAllocator diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_prefill_compress.py similarity index 61% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_prefill_compress.py index 415956a05579..0a051c7f90c0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_prefill_compress.py @@ -1,6 +1,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""QSA block compression of the paged index-key side cache: completed groups to anchors.""" +"""QSA block compression of the paged index-key side cache: completed groups to anchors. + +Called by the qwen4_exp target (`qwen3_8_flash_next_nvfp4__sm_103__default`) at every +layout. +""" from __future__ import annotations @@ -14,9 +18,8 @@ # kernel plane has no torch.ops surface. from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels -#: The mRoPE sections the checkpoint declares; the target passes `None` today -#: (`_QSA_MROPE_SECTION`) and these are certified, on the NVFP4 trunk-plane cell, -#: for when it flips. +#: The mRoPE sections the checkpoint declares. The target passes `None` +#: (`_QSA_MROPE_SECTION`); these are certified on the NVFP4 trunk-plane cell. MROPE_SECTION = (11, 11, 10) #: This checkpoint's index geometry and the engine's QSA metadata: index head @@ -40,6 +43,26 @@ _MTP7_32 = tuple((29 + 61 * i, 8) for i in range(32)) +#: The default target's index pool at bf16 KV, as the stock `Qwen4ExpHybridCacheManager` lays it out on +#: every member and rank: KV cache manager V2 coalesces the 12 sparse layers' same-size index +#: pages into one pool, so a layer's `get_index_k_buffer` view is `[pages, 32, 1, 128]` with page +#: stride 12 x 4096 elements, starting `ordinal x 4096` into each slot. The position plane is its +#: own contiguous `[pages, 32, 3]` int32 pool. +INDEX_ROLES = 12 + +#: A rank's attention pool, in pages: the stock engine's `maxNumBlocks` at core config (bf16 KV, +#: fraction 0.6, host-resident n-gram table), measured per member. The manager's index view holds +#: slightly fewer pages than its primary block count, so these bound it from above. +TEP4_PAGES = 258983 +DEP4_PAGES = 120530 + +#: `qsa_block_table` is `[max_num_sequences, max_blocks_per_seq]`: 128 seats (`max_batch_size`, per +#: rank under attention DP), 292 pages a sequence at max_seq_len 9216 and 516 at the MTP member's 16384. +SEATS = 128 +CORE_WIDTH = 292 +MTP_WIDTH = 516 + + def _spec(segments, pages, **overrides) -> dict: """`segments` is `(first position, rows)` per request, in row order.""" return {**_CKPT, "segments": segments, "pages": pages, **overrides} @@ -50,6 +73,39 @@ def _mrope(segments, pages, **overrides) -> dict: return _spec(segments, pages, mrope_section=MROPE_SECTION, **overrides) +def _pool(segments, pages, width, ordinal, **overrides) -> dict: + """The default target's call over a rank's coalesced index pool: layer `ordinal`'s plane, 128 seats, + and request 0's last page the pool's last (`top_page`).""" + return _spec( + segments, + pages, + slots=SEATS, + width=width, + layout="coalesced", + roles=INDEX_ROLES, + ordinal=ordinal, + top_page=True, + **overrides, + ) + + +def _decode(requests: int, last: int) -> tuple: + """`requests` decode rows spread over positions 3..`last`; every other one an anchor + (`4k + 3`), the others at `4k + 0..2`, the last request an anchor at `last`.""" + span = (last - 3) // 4 + segments = [] + for r in range(requests): + k = r * span // (requests - 1) + offset = 3 if r % 2 == 0 or r == requests - 1 else r % 3 + segments.append((4 * k + offset, 1)) + return tuple(segments) + + +#: 128 MTP7 verify blocks (8 rows, two anchors each) at starts spread to the MTP member's last +#: position, 16383: some blocks inside one page, some across a page edge. +_MTP7_128 = tuple((r * (16384 - 8) // 127, 8) for r in range(128)) + + class QsaPrefillCompress(OpWrapper): """Replace each completed group's anchor slot with its pooled, normed, rotated key. @@ -73,8 +129,7 @@ class QsaPrefillCompress(OpWrapper): else, including the non-anchor members' raw keys, is untouched. Not idempotent: the anchor slot is the group's last input and its only - output, so a second call over the same rows re-pools a compressed anchor - (measured: all 48 of 48 anchors moved, up to 2.94). + output, so a second call over the same rows re-pools a compressed anchor. Fusion boundary: pool, Gemma norm, position gather, rotation, store. The raw-key and position stores (`qsa_paged_index_store`, which must run @@ -99,6 +154,24 @@ class QsaPrefillCompress(OpWrapper): why="NVFP4 trunk planes: 12 index roles and the position page share each page", spec=_mrope(((0, 64),) * 2, pages=4, layout="shared_page", roles=12), ), + # The default target's calls over a rank's measured pool, every page reachable: the last + # page's plane sits past 2^31 elements of the pool's base. + Cell( + why="tep4 decode: 128 requests x 1 row, the 258983-page coalesced pool, ordinal 11", + spec=_pool(_decode(128, 9215), TEP4_PAGES, CORE_WIDTH, 11), + ), + Cell( + why="tep4 MTP7 verify: 128 requests x 8 rows = 1024, positions to 16383", + spec=_pool(_MTP7_128, TEP4_PAGES, MTP_WIDTH, 5), + ), + Cell( + why="tep4 prefill at max_num_tokens: 8192 rows (6144 + 2047 + 1, all from 0)", + spec=_pool(((0, 6144), (0, 2047), (0, 1)), TEP4_PAGES, CORE_WIDTH, 0, rounding_flips=5), + ), + Cell( + why="dep4 decode: 32 requests x 1 row over a rank's 120530-page pool, 128 seats", + spec=_pool(_decode(32, 9215), DEP4_PAGES, CORE_WIDTH, 11), + ), ) note = """ @@ -114,13 +187,13 @@ class QsaPrefillCompress(OpWrapper): - every member's raw key and the first member's coordinates are in the caches before the call; this call stores neither. - every logical position is non-negative. `L = -1` satisfies the anchor - test and addresses positions -4..-1: measured, an illegal memory - access that poisons the process. A value, so not guarded; pad with 0 - (measured an exact no-op). + test and addresses positions -4..-1, an illegal memory access that + poisons the process. A value, so not guarded; pad with 0 (an exact + no-op). - `logical_positions`, `request_indices`, `k_norm_weight` and the last dim of `cos_sin` are read with an assumed unit stride while the - callable checks only shape: each was measured accepted and wrong - (max abs 4.72, 6.20, 0.80, 7.16). `is_valid` refuses them. + callable checks only shape, so a strided one is accepted and wrong. + `is_valid` refuses them. The callable rejects loudly (not pinned by a test): metadata shape/dtype mismatches, a non-int32 or non-2-D table, more than one index head, a position plane with fewer than 3 axes or not int32, a non-bf16 cache, page sizes that disagree with @@ -129,34 +202,60 @@ class QsaPrefillCompress(OpWrapper): `mrope_section` that is not 3 entries summing to `R`. `rows == 0` returns before launching. - `cos_sin`'s dtype is the rotation's precision: bf16 is accepted and was - measured 1.4% of elements off by up to 1.56e-2 -- not certified. Text - rows reduce mRoPE to plain RoPE bit for bit (measured, not pinned by a - test). Plain `@triton.jit`, no - autotune; CUDA-graph capturable at a fixed row count (captured at 256 - rows, 32 MTP7 verify blocks). + `cos_sin`'s dtype is the rotation's precision: bf16 is accepted, puts some + elements off by up to 1.56e-2, and is not certified. Text rows reduce + mRoPE to plain RoPE bit for bit (measured, not pinned by a test). Plain + `@triton.jit`, no autotune, JIT-compiled at first call; the cells drove + that JIT path on sm_103 (GB300). Every stride is a + `tl.constexpr`, so each pool layout (contiguous, the NVFP4 shared page, + the default target's coalesced bf16 pool) is its own compiled specialization, and + each is driven by a cell. CUDA-graph capturable at a fixed row count + (captured at 256 rows, 32 MTP7 verify blocks, and at 128 decode rows over + the default target's tep4 pool). Under an NVFP4 KV pool both caches are strided views of one per-slot allocation (each physical page holds the index roles, then the position page); the trunk-plane cell checks the other index roles and the position page untouched. - The kernel's `tl.rsqrt` is the hardware approximation, an fp32 ulp from the - reference's: every cell is bitwise the reference but the 32-block MTP7 one, where - 1 of 8192 anchor elements rounds one bf16 step away. + The `coalesced` cells hold the default target's bf16-KV index pool (layout + at `INDEX_ROLES`, read off the stock manager at tep4 and dep4 rank 0) at a + rank's page count -- tep4 258983, dep4 120530, an upper bound on the index + view -- with request 0's last page the pool's last, written at 258982 x + 49152 elements from the base, past 2^31. The table is the engine's + `[128, 292]` (core) or `[128, 516]` (MTP, max_seq_len 16384), 128 seats at + 128 and at 32 requests. The other 11 layers' planes are checked untouched. - Hard limits (read from the kernel, not probed): positions, requests and - page ids are widened to int64 before any address arithmetic, so no int32 - boundary inside the call. Certified extent: positions to 1928 through a - [40, 260] table over 1100 pages. + The kernel's `tl.rsqrt` is the hardware approximation, an fp32 ulp from the + reference's: every cell is bitwise the reference but two. In the 32-block + MTP7 one, 1 of 8192 anchor elements rounds one bf16 step away. In the 8192-row + tep4 prefill, 5 of 262016 do, every one an exact bf16 tie in the reference's + fp32 value (low 16 bits 0x8000). The nearest wrong variants (sequential bf16 + pooling, coordinates from the group's last member, the norm weight without + its `+1`) fall far outside that budget. + + Hard limits (read from the kernel; the page offset also driven): positions, + requests and page ids are widened to int64 before any address arithmetic, + so no int32 boundary inside the call -- the coalesced cells address pages + past 2^31 elements of the pool's base. Certified extent: positions to 16383 + through a [128, 516] table over 258983 pages; 8192 rows in one launch. + + Each rank's call is local to its device (the indexer is replicated, no + collective), so a one-device cell stands for any rank's call at its pool. Not certified: `compress_ratio` other than 4, `D` other than 128, a bf16 - or fp16 `cos_sin`, positions past 1928 or the table's rows, mRoPE outside - the NVFP4 trunk-plane cell, the NVFP4 draft planes (one index role a page), multiple devices, + or fp16 `cos_sin`, positions past 16383 or the table's rows (max_seq_len + 262144 included), mRoPE outside the NVFP4 + trunk-plane cell, the NVFP4 draft planes (one index role a page), `tokens_per_block` other than 32 (no group straddles a page), eps other than 1e-6, int32 positions or int64 request indices, a strided view of the index cache, position plane or table other than the NVFP4 pool's shared - page, two rows naming one anchor in a launch, and `rows == 0`. + page and the default target's coalesced bf16 pool, the kv_fp8 / kv_nvfp4 dep4 + pools at their measured page count (their packed `QSA_SIDE_PLANES` layout + is the shared-page cell's, at 4 pages), the MTP member's own pool count and + its separate draft cache (the 1024-row cell runs on the core tep4 pool), a + tp4 pool (never booted; the indexer is replicated, so tep4's stands in), two + rows naming one anchor in a launch, and `rows == 0`. """ def __call__( @@ -263,7 +362,7 @@ def is_valid( mrope_section: Optional[Tuple[int, int, int]], ) -> None: # Read by the kernel with an assumed unit stride; the callable checks - # only their shapes. Each measured accepted and silently wrong. + # only their shapes. Each is accepted and silently wrong. assert logical_positions.is_contiguous(), "logical_positions must be contiguous" assert request_indices.is_contiguous(), "request_indices must be contiguous" assert k_norm_weight.is_contiguous(), "k_norm_weight must be contiguous" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rollback_planes.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rollback_planes.py new file mode 100644 index 000000000000..c7e0ff27e19d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rollback_planes.py @@ -0,0 +1,241 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The speculative rollback of every attention layer's QSA index-K plane in one launch each way: +gather the verify rows' slots of all planes before the trunk writes them, and after the +draft loop write the gathered values back into the slots of the rows the sampler rejected. + +Tier 5 (Triton, called by the qwen4_exp target, `qwen3_8_flash_next_nvfp4__sm_103__default`, +at every layout). The planes are separate buffers of the cache manager (one per +attention layer, each `[pages, tokens_per_block, heads, dim]` with a free page stride and a +contiguous tail); the kernels reach them through a table of their base addresses and page +strides, one program per (plane, row). +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: This checkpoint's index plane row: one index KV head of 128 bf16. +ROW = 128 + +#: The index planes of core's bf16 K/V pool at every four-rank layout, as the stock +#: `Qwen4ExpHybridCacheManager` (KV cache manager V2) allocates them: the 12 attention layers' +#: `[32, 1, 128]` INDEX_KEY buffers coalesce into one pool, layer k's buffer at element 4096 k of +#: a 49152-element page, so every plane is a `[pages, 32, 1, 128]` view at page stride 49152. +#: 257109 pages is a tep4 rank's attention pool at core's configuration, which the index pool +#: shares: the stock engine's `maxNumBlocks` 258983 less the 1874 recurrent-state slots that +#: total also counts. The largest measured per member. +_TEP4_POOL = dict(planes=12, layout="coalesced", pool_pages=257109) + + +def _verify(requests: int, p: int) -> dict: + """A verify step's rows: `requests` generation requests of `p` rows (`max_draft_len + 1`).""" + return dict(_TEP4_POOL, rows=requests * p, requests=requests, p=p, rejected="bool") + + +@triton.jit +def _gather_kernel(ptrs, page_strides, pages, within, out, rows, ROW: tl.constexpr): + plane = tl.program_id(0) + r = tl.program_id(1) + base = tl.load(ptrs + plane).to(tl.pointer_type(tl.bfloat16)) + stride = tl.load(page_strides + plane) + page = tl.load(pages + r) + tok = tl.load(within + r) + offs = tl.arange(0, ROW) + vals = tl.load(base + page * stride + tok * ROW + offs) + tl.store(out + (plane.to(tl.int64) * rows + r) * ROW + offs, vals) + + +@triton.jit +def _restore_kernel(ptrs, page_strides, pages, within, rejected, old, rows, ROW: tl.constexpr): + plane = tl.program_id(0) + r = tl.program_id(1) + if tl.load(rejected + r) != 0: + base = tl.load(ptrs + plane).to(tl.pointer_type(tl.bfloat16)) + stride = tl.load(page_strides + plane) + page = tl.load(pages + r) + tok = tl.load(within + r) + offs = tl.arange(0, ROW) + vals = tl.load(old + (plane.to(tl.int64) * rows + r) * ROW + offs) + tl.store(base + page * stride + tok * ROW + offs, vals) + + +def plane_table(planes: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + """`(base addresses, page strides)` of the planes, int64 on their device (built once per + cache-manager bind: the planes do not move).""" + dev = planes[0].device + for t in planes: + assert t.dtype == torch.bfloat16 and t.dim() == 4, "index planes [pages, tpb, H, D] bf16" + assert t.shape[2] * t.shape[3] == ROW, "one 128-wide row a token" + assert t.stride(1) == ROW and t.stride(2) == t.shape[3] and t.stride(3) == 1, ( + "contiguous [tpb, H, D] tail" + ) + ptrs = torch.tensor([t.data_ptr() for t in planes], dtype=torch.int64, device=dev) + strides = torch.tensor([t.stride(0) for t in planes], dtype=torch.int64, device=dev) + return ptrs, strides + + +class QsaRollbackGather(OpWrapper): + """`torch.stack([plane[pages, within] for plane in planes])` as `[L, R, 128]` bf16, with + the planes given as `plane_table(planes)`: every row's slot of every plane, in plane + order. Nothing is written but the output. + + Fusion boundary: the gather of all planes; the slots (`pages`, `within`) are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + # Both halves are certified at the same cells: the target's at tp1, and its MTP verify + # rows at `max_batch_size` 128 over a tep4 rank's pool. + CELLS: tuple[Cell, ...] = ( + Cell(why="12 planes, 4 verify rows (MTP3, one request)", spec=dict(planes=12, rows=4)), + Cell(why="12 planes, 512 verify rows (MTP3, 128 requests)", spec=dict(planes=12, rows=512)), + Cell( + why="tep4 MTP1 verify: 128 requests x 2 rows (256) over the 257109-page coalesced pool", + spec=_verify(128, 2), + ), + Cell( + why="tep4 MTP7 verify: 128 requests x 8 rows (1024, the envelope's largest), same pool", + spec=_verify(128, 8), + ), + ) + note = """ + Metadata consumed: none. Call site: `SpeculativeRollback.begin`, before the verify step's + trunk writes any plane; `pages` / `within` are the generation rows' QSA page and offset + (`block_table[req, pos // 32]`, `pos % 32`), int64 `[R]`, R = requests x (max_draft_len + 1). + One program per (plane, row), one 256-byte row each. Preconditions: the planes as + `plane_table` accepts them (bf16 NHD views, any page stride); every (page, within) inside + the pool. + + Triton, JIT-compiled at first call for the running device (sm_103 here). + + Addressing is 64-bit: `pages` and the page strides are int64, so a slot past 2^31 elements + from a plane's base is reached. The tep4 cells gather from (and restore into) the pool's + last page, 257108 x 49152 = 1.26e10 elements past every plane's base, and its first. + + Layouts certified: separate contiguous planes and the packed side-plane page of an FP8 / + NVFP4 pool (page stride 12 x 4096 + 192) at the tp1 cells' 4 and 512 rows over 600 pages; + core's coalesced bf16 pool (page stride 49152) at 256 and 1024 rows over 257109 pages. + `test_cuda_graph_replay` captures the 1024-row call and replays it at new slots. + + Hard limits (not probed): rows is the launch grid's y extent, at most 65535; out-of-pool or + duplicated slots are the caller's. Not certified: a pool page count measured under the MTP + unit's own engine (`max_seq_len` 16384; the cells use core's tep4 count), and the + side-plane layout at 256 / 1024 rows. + """ + + def __call__(self, ptrs, page_strides, pages, within, num_planes: int) -> torch.Tensor: + rows = pages.shape[0] + out = torch.empty([num_planes, rows, ROW], dtype=torch.bfloat16, device=pages.device) + if rows: + _gather_kernel[(num_planes, rows)]( + ptrs, page_strides, pages, within, out, rows, ROW=ROW, num_warps=1 + ) + return out + + def reference(self, planes, pages, within, num_planes: int) -> torch.Tensor: + """Every row's slot of the first `num_planes` planes, by torch's advanced indexing. + + Takes the planes themselves where `__call__` takes `plane_table(planes)`: torch + cannot dereference the raw base addresses the kernel is handed, so the reference reads + the tensors the table was built from. A copy, exact. + """ + return torch.stack([t[pages, within].reshape(-1, ROW) for t in planes[:num_planes]]) + + def is_valid(self, ptrs, page_strides, pages, within, num_planes: int) -> None: + assert ptrs.dtype == torch.int64 and ptrs.numel() == num_planes, ( + "ptrs must be plane_table's int64 table of exactly num_planes addresses; a " + "num_planes past it reads addresses beyond the table" + ) + assert page_strides.dtype == torch.int64 and page_strides.numel() == num_planes, ( + "page_strides must be plane_table's int64 table of num_planes strides" + ) + assert pages.dtype == torch.int64 and within.dtype == torch.int64, ( + "pages and within must be int64" + ) + assert pages.shape == within.shape and pages.dim() == 1, ( + "pages and within must be one [R] pair: rows are counted off pages, so a shorter " + "within is read past its end" + ) + + +class QsaRollbackRestore(OpWrapper): + """For every row `r` with `rejected[r]` and every plane `l`: `plane_l[pages[r], within[r]] + = old[l, r]` -- the inverse of `qsa_rollback_gather` over the rejected rows. Rows not + rejected are not touched. + + Fusion boundary: the restore of all planes; `rejected` is the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + # Both halves are certified at the same cells: the target's at tp1, and its MTP verify + # rows at `max_batch_size` 128 over a tep4 rank's pool. + CELLS: tuple[Cell, ...] = ( + Cell(why="12 planes, 4 verify rows (MTP3, one request)", spec=dict(planes=12, rows=4)), + Cell(why="12 planes, 512 verify rows (MTP3, 128 requests)", spec=dict(planes=12, rows=512)), + Cell( + why="tep4 MTP1 verify: 128 requests x 2 rows (256) over the 257109-page coalesced pool", + spec=_verify(128, 2), + ), + Cell( + why="tep4 MTP7 verify: 128 requests x 8 rows (1024, the envelope's largest), same pool", + spec=_verify(128, 8), + ), + ) + note = """ + Metadata consumed: none. Call site: `SpeculativeRollback.commit_speculative_states`, + after the draft loop: `old` is the begin step's `qsa_rollback_gather` output, `rejected` + the bool `[R]` mask `ordinal >= accepted` (int32 / int8 / bool, nonzero = restore; the tp1 + cells drive int32, the tep4 cells the call site's bool). + + Returns nothing: it writes the planes, which it holds only as `plane_table` addresses. + Triton, JIT-compiled at first call for the running device (sm_103 here). + + Every cell checks that the rejected rows' slots hold `old`, the others the step's values, + and that nothing else in the pool was written. Addressing, layouts and the CUDA-graph + replay as `qsa_rollback_gather`'s note states; the same limits. + + Silent, refused by `is_valid`: a non-contiguous `old` (read as dense `[L, R, 128]`; + measured wrong in `test_guard_refuses_a_strided_old`). + """ + + def __call__(self, ptrs, page_strides, pages, within, rejected, old) -> None: + num_planes, rows, _ = old.shape + if rows: + _restore_kernel[(num_planes, rows)]( + ptrs, page_strides, pages, within, rejected, old, rows, ROW=ROW, num_warps=1 + ) + + def reference(self, planes, pages, within, rejected, old) -> torch.Tensor: + """What each row's slot of each plane holds after the call, `[L, R, 128]`: `old` where + `rejected`, else the slot's value now. Read before the call, as the planes stand. + + Takes the planes where `__call__` takes `plane_table(planes)`, for the reason + `QsaRollbackGather.reference` states. Every location outside these slots is unchanged, + which the test checks over the whole pool. + """ + now = torch.stack([t[pages, within].reshape(-1, ROW) for t in planes[: old.shape[0]]]) + return torch.where(rejected.bool()[None, :, None], old, now) + + def is_valid(self, ptrs, page_strides, pages, within, rejected, old) -> None: + assert old.dtype == torch.bfloat16 and old.dim() == 3 and old.shape[2] == ROW, ( + "old must be qsa_rollback_gather's [L, R, 128] bf16 output" + ) + assert old.is_contiguous(), ( + "old must be contiguous; the kernel reads it as dense [L, R, 128] and a strided " + "view restores the wrong values" + ) + assert ptrs.numel() == old.shape[0] and page_strides.numel() == old.shape[0], ( + "plane_table must hold one address and one stride per plane of old" + ) + assert pages.shape == (old.shape[1],) and within.shape == (old.shape[1],), ( + "pages and within must hold one slot per row of old" + ) + assert rejected.shape == (old.shape[1],), "rejected must hold one flag per row of old" + + +qsa_rollback_gather = QsaRollbackGather() +qsa_rollback_restore = QsaRollbackRestore() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rowtile_attention_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rowtile_attention_cute.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rowtile_attention_cute.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_rowtile_attention_cute.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_sparse_gqa_rowtile.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_sparse_gqa_rowtile.py new file mode 100644 index 000000000000..a7b5ad5e6a8e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_sparse_gqa_rowtile.py @@ -0,0 +1,938 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA sparse GQA for context batches: several query rows per program, K/V read once per tile. + +The first launch turns each row's selection (the columns the prefill bound +lets it read) into a bitmap over logical token positions. Three plan launches +form, per group of consecutive rows of one request, the union of its rows' +resident selected tokens as a packed list of physical cache slots, and each +row's bitmap re-indexed onto that list. The row-tile launch is persistent: one +program per SM takes (group, KV head) items longest first and walks each packed +list in 64-slot tiles, each tile's K/V loaded once for every row and head of the +group, each row masked to its own bits. A group whose rows span requests +(generation rows of a mixed batch) walks each row's own selected columns +instead, as the per-row kernel does. +""" + +from __future__ import annotations + +import math + +import cutlass +import cutlass.cute as cute +import torch +import triton +import triton.language as tl +from cutlass._mlir.dialects import llvm +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa import ( + qsa_paged_sparse_gqa as _per_row_entry, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr + +from .qsa_rowtile_attention_cute import BLOCK_M, BLOCK_N +from .qsa_rowtile_attention_cute import eligible as _cute_eligible +from .qsa_rowtile_attention_cute import num_ctas as _rowtile_ctas +from .qsa_rowtile_attention_cute import rowtile_attention_cute as _rowtile_attention_cute + +_INTEGER_DTYPES = (torch.int32, torch.int64) +# The per-row kernel's scan tile: the prefill bound is rounded up to it, so the +# bitmap reads exactly the columns that kernel reads. +_BOUND_TILE = 64 +# Programs of the lone-row pass per KV head; each walks the lone list with this stride. +_LONE_PROGRAMS = 256 # capacity: fixed: a launch width; the lone list may be any length +# Bitmap words one program of the plan's union and slot passes covers (a grid dimension). +_CHUNK_WORDS = 32 # capacity: fixed: a launch tile; the grid covers every word +_PLAN_WARPS = 4 +_LONE_WARPS = 2 +_LONE_STAGES = 2 + + +# -------------------------------------------------------------------------------------- +# The bitmap pass (CuTe DSL, Tier 6). Row `r`: one bit per readable selected token (a +# column below the per-row kernel's prefill bound naming a token >= 0 whose page is inside +# the table) set in `bits[r]`, and the highest such token and the number of columns read +# in `row_info[r]`. A CTA owns one row: the row's selection is staged in shared memory with +# coalesced loads, each thread walks `_BITS_RUN` consecutive columns, ORs the bits of one +# 64-token word in a register and flushes the word into a shared-memory bitmap with one +# `red.shared.or.b64` when the run leaves it; then the CTA writes every word of the row and +# `row_info`. Any column order gives the same words (OR is idempotent and commutative); the +# order only decides how few flushes a run makes, and the selection arrives in token order. +# No memset of the bitmap is needed. +# -------------------------------------------------------------------------------------- + +_BITS_THREADS = 128 +#: Consecutive selected columns a thread walks. +_BITS_RUN = 16 + + +@dsl_user_op +def _red_or_shared(addr, val, *, loc=None, ip=None): + llvm.inline_asm( + T.i32(), + [cutlass.Int32(addr).ir_value(loc=loc, ip=ip), cutlass.Int64(val).ir_value(loc=loc, ip=ip)], + "{ red.shared.or.b64 [$1], $2; mov.u32 $0, 0; }", + "=r,r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +class _RowBitsKernel: + """The bitmap pass's `bits` and `row_info`, one CTA a row.""" + + def __init__(self, topk, block_topk, compress_ratio, tokens_per_block, table_width, num_words): + self.topk, self.block_topk, self.cr = topk, block_topk, compress_ratio + self.tpb, self.width, self.num_words = tokens_per_block, table_width, num_words + self.staged = topk + topk // _BITS_RUN + 1 + + @cute.jit + def __call__( + self, + sel_t: cute.Tensor, + qp_ptr: cute.Pointer, + bits_ptr: cute.Pointer, + info_ptr: cute.Pointer, + rows: cutlass.Int32, + stream, + ): + qp = cute.make_tensor(qp_ptr, cute.make_layout((rows,), stride=(1,))) + bits = cute.make_tensor( + bits_ptr, cute.make_layout((rows, self.num_words), stride=(self.num_words, 1)) + ) + info = cute.make_tensor(info_ptr, cute.make_layout((rows, 2), stride=(2, 1))) + # `sel_t` arrives as a tensor: the TVM-FFI entry takes its stream from a GPU tensor. + self.kernel(sel_t, qp, bits, info).launch( + grid=[rows, 1, 1], + block=[_BITS_THREADS, 1, 1], + smem=self.num_words * 8 + self.staged * 4 + 128, + stream=stream, + ) + + @cute.kernel + def kernel(self, mSel: cute.Tensor, mQp: cute.Tensor, mBits: cute.Tensor, mInfo: cute.Tensor): + tidx, _, _ = cute.arch.thread_idx() + row, _, _ = cute.arch.block_idx() + smem = cutlass.utils.SmemAllocator() + words = smem.allocate_tensor( + cutlass.Int64, cute.make_layout((self.num_words,), stride=(1,)), byte_alignment=16 + ) + red = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((4,), stride=(1,)), byte_alignment=16 + ) + # The row's selection, one padding slot every `_BITS_RUN` columns (conflict-free runs). + sel_s = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((self.staged,), stride=(1,)), byte_alignment=16 + ) + words_base = cutlass.Int32(words.iterator.toint()) + for i in cutlass.range_constexpr((self.num_words + _BITS_THREADS - 1) // _BITS_THREADS): + w = tidx + i * _BITS_THREADS + if w < self.num_words: + words[w] = cutlass.Int64(0) + + # The columns the per-row kernel reads: the prefill bound rounded up to its tile. + visible = cutlass.Int32(mQp[row]) + 1 + selected_blocks = cutlass.min(visible // self.cr, cutlass.Int32(self.block_topk)) + valid = selected_blocks * self.cr + visible % self.cr + bounded = cutlass.min(valid, cutlass.Int32(self.topk)) + columns_read = cutlass.min( + (bounded + _BOUND_TILE - 1) // _BOUND_TILE * _BOUND_TILE, cutlass.Int32(self.topk) + ) + columns_read = cutlass.max(columns_read, cutlass.Int32(0)) + for i in cutlass.range_constexpr((self.topk + _BITS_THREADS - 1) // _BITS_THREADS): + c = tidx + i * _BITS_THREADS + if c < columns_read: + sel_s[c + c // _BITS_RUN] = cutlass.Int32(mSel[row, c]) + cute.arch.sync_threads() + + highest = cutlass.Int32(-1) + runs = (self.topk + _BITS_THREADS * _BITS_RUN - 1) // (_BITS_THREADS * _BITS_RUN) + for run in cutlass.range_constexpr(runs): + c0 = (run * _BITS_THREADS + tidx) * _BITS_RUN + if c0 < columns_read: + cur_w = cutlass.Int32(-1) + cur_m = cutlass.Int64(0) + for k in cutlass.range_constexpr(_BITS_RUN): + c = c0 + k + t = cutlass.Int32(-1) + if c < columns_read: + t = sel_s[c + c // _BITS_RUN] + if t >= 0: + if t // self.tpb < self.width: + w = t >> 6 + if w != cur_w: + if cur_w >= 0: + _red_or_shared(words_base + cur_w * 8, cur_m) + cur_w = w + cur_m = cutlass.Int64(0) + cur_m = cur_m | (cutlass.Int64(1) << cutlass.Int64(t & 63)) + highest = cutlass.max(highest, t) + if cur_w >= 0: + _red_or_shared(words_base + cur_w * 8, cur_m) + for off in cutlass.range_constexpr(5): + highest = cutlass.max(highest, cute.arch.shuffle_sync_bfly(highest, offset=1 << off)) + if tidx % 32 == 0: + red[tidx // 32] = highest + cute.arch.sync_threads() + if tidx == 0: + mInfo[row, 0] = cutlass.max(cutlass.max(red[0], red[1]), cutlass.max(red[2], red[3])) + mInfo[row, 1] = columns_read + for i in cutlass.range_constexpr((self.num_words + _BITS_THREADS - 1) // _BITS_THREADS): + w = tidx + i * _BITS_THREADS + if w < self.num_words: + mBits[row, w] = words[w] + + +_ROW_BITS: dict = {} + + +def _row_bits( + selected_tokens, query_positions, compress_ratio, tokens_per_block, table_width, num_words +): + """`(bits [rows, num_words] int64, row_info [rows, 2] int32)` of the bitmap pass over a + selection with unit column stride.""" + rows, final_topk = selected_tokens.shape + block_topk = (final_topk - compress_ratio + 1) // compress_ratio + device = selected_tokens.device + row_info = torch.empty((rows, 2), dtype=torch.int32, device=device) + query_positions = query_positions.contiguous() + bits = torch.empty((rows, num_words), dtype=torch.int64, device=device) + key = ( + final_topk, compress_ratio, tokens_per_block, table_width, + query_positions.dtype, selected_tokens.dtype, + ) # fmt: skip + args = (selected_tokens, query_positions.data_ptr(), bits.data_ptr(), row_info.data_ptr()) + fn = _ROW_BITS.get(key) + if fn is None: + gm = cute.AddressSpace.gmem + qty = cutlass.Int64 if query_positions.dtype == torch.int64 else cutlass.Int32 + fn = cute.compile( + _RowBitsKernel( + final_topk, block_topk, compress_ratio, tokens_per_block, table_width, num_words + ), + from_dlpack(selected_tokens, assumed_align=4).mark_layout_dynamic(leading_dim=1), + make_ptr(qty, args[1], gm, assumed_align=4), + make_ptr(cutlass.Int64, args[2], gm, assumed_align=16), + make_ptr(cutlass.Int32, args[3], gm, assumed_align=8), + cutlass.Int32(1), + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--opt-level 2 --enable-tvm-ffi", + ) + _ROW_BITS[key] = fn + fn(*args, rows) + return bits, row_info + + +@triton.jit +def _or_combine(a, b): + return a | b + + +@triton.jit +def _compress_bits(x, m): + """Per word, the bits of `x` at the set positions of `m`, packed to the low + end in order (a parallel bit extract; Hacker's Delight `compress`, 64-bit). + `x` and `m` are uint64 and broadcast together.""" + x = x & m + mk = ~m << 1 + for i in tl.static_range(6): + mp = mk ^ (mk << 1) + mp = mp ^ (mp << 2) + mp = mp ^ (mp << 4) + mp = mp ^ (mp << 8) + mp = mp ^ (mp << 16) + mp = mp ^ (mp << 32) + mv = mp & m + m = (m ^ mv) | (mv >> (1 << i)) + t = x & mv + x = (x ^ t) | (t >> (1 << i)) + mk = mk & ~mp + return x + + +@triton.jit +def _popc64(x): + """Set bits of each uint64 of `x`, as int32.""" + return tl.inline_asm_elementwise( + "popc.b64 $0, $1;", "=r,l", [x], dtype=tl.int32, is_pure=True, pack=1 + ) + + +@triton.jit +def _bits_below(n): + """uint64 with the low `n` bits set, `0 <= n <= 64`.""" + ones = (n * 0 - 1).to(tl.int64).to(tl.uint64, bitcast=True) + return tl.where(n >= 64, ones, ~(ones << tl.minimum(n, 63).to(tl.uint64))) + + +@triton.jit +def _group_rows(group, request_indices, row_info, rows, NUM_REQUESTS, + ROWS_PER_TILE: tl.constexpr, ROWS_POW2: tl.constexpr): # fmt: skip + """A group's members, which of them exist, whether they are of one request (and which), + and the bitmap words its union spans (0 for a group of several requests).""" + offsets = tl.arange(0, ROWS_POW2) + members = group * ROWS_PER_TILE + offsets + live = (offsets < ROWS_PER_TILE) & (members < rows) + requests = tl.load(request_indices + members, mask=live, other=-1).to(tl.int64) + lowest = tl.min(tl.where(live, requests, NUM_REQUESTS), axis=0) + highest_request = tl.max(tl.where(live, requests, -1), axis=0) + highest_token = tl.max(tl.load(row_info + members * 2, mask=live, other=-1), axis=0) + one_request = (lowest == highest_request) & (lowest >= 0) & (lowest < NUM_REQUESTS) + words_needed = tl.where(one_request, (highest_token + 64) // 64, 0) + return offsets, members, live, one_request, lowest, words_needed + + +@triton.jit +def _rowtile_union_kernel( + request_indices, row_info, bits, block_table, union_words, chunk_counts, rows, + bits_stride_row, block_table_stride_request, block_table_stride_page, + union_stride_group, counts_stride_group, + NUM_REQUESTS, NUM_CACHE_PAGES, + ROWS_PER_TILE: tl.constexpr, ROWS_POW2: tl.constexpr, + PAGE_TABLE_WIDTH: tl.constexpr, TOKENS_PER_BLOCK: tl.constexpr, CHUNK_WORDS: tl.constexpr, + WORD_PAGES: tl.constexpr, +): # fmt: skip + """Per (group, chunk of `CHUNK_WORDS` bitmap words): the union of the group's rows' bits + restricted to resident tokens (a mapped page inside the pool), into `union_words`, and + its token count into `chunk_counts`. A chunk past the group's highest token, or a group + of several requests, counts zero and writes no word.""" + group = tl.program_id(0).to(tl.int64) + chunk = tl.program_id(1) + _, members, live, _, lowest, words_needed = _group_rows( + group, request_indices, row_info, rows, NUM_REQUESTS, ROWS_PER_TILE, ROWS_POW2 + ) + start = chunk * CHUNK_WORDS + count = tl.full([], 0, tl.int32) + if start < words_needed: + words = start + tl.arange(0, CHUNK_WORDS) + in_range = words < words_needed + row_words = tl.load( + bits + members[:, None] * bits_stride_row + words[None, :], + mask=live[:, None] & in_range[None, :], + other=0, + ).to(tl.uint64, bitcast=True) # fmt: skip + selected = tl.reduce(row_words, 0, _or_combine) + # The pages a word's 64 tokens span (at most WORD_PAGES), each as the word's bits it + # covers; a page is looked up only where the union selects one of its tokens. + first_token = words.to(tl.int32) * 64 + logical_pages = first_token[:, None] // TOKENS_PER_BLOCK + tl.arange(0, WORD_PAGES)[None, :] + low = tl.maximum(logical_pages * TOKENS_PER_BLOCK - first_token[:, None], 0) + high = tl.minimum((logical_pages + 1) * TOKENS_PER_BLOCK - first_token[:, None], 64) + page_bits = tl.where(low < high, _bits_below(high) & ~_bits_below(low), 0) + page_bits = page_bits & selected[:, None] + page_base = block_table + lowest * block_table_stride_request + pages = tl.load( + page_base + tl.minimum(logical_pages, PAGE_TABLE_WIDTH - 1) * block_table_stride_page, + mask=(page_bits != 0) & (logical_pages < PAGE_TABLE_WIDTH), + other=-1, + ) # fmt: skip + kept = (pages >= 0) & (pages < NUM_CACHE_PAGES) & (logical_pages < PAGE_TABLE_WIDTH) + resident = tl.reduce(tl.where(kept, page_bits, 0), 1, _or_combine) + tl.store(union_words + group * union_stride_group + words, + resident.to(tl.int64, bitcast=True), mask=in_range) # fmt: skip + count = tl.sum(_popc64(resident), axis=0) + tl.store(chunk_counts + group * counts_stride_group + chunk, count) + + +@triton.jit +def _rowtile_group_kernel( + request_indices, row_info, chunk_counts, chunk_bases, group_plan, lone_rows, lone_count, + rows, counts_stride_group, num_chunks, NUM_REQUESTS, + ROWS_PER_TILE: tl.constexpr, ROWS_POW2: tl.constexpr, + BLOCK_N: tl.constexpr, CHUNKS_POW2: tl.constexpr, +): # fmt: skip + """Per group: each chunk's first list position (`chunk_bases`), and `group_plan[group]` = + (tiles of the list for the row-tile pass, or -1; list length). A group of several + requests appends its rows to the lone list instead.""" + group = tl.program_id(0).to(tl.int64) + offsets, members, live, one_request, _, _ = _group_rows( + group, request_indices, row_info, rows, NUM_REQUESTS, ROWS_PER_TILE, ROWS_POW2 + ) + chunks = tl.arange(0, CHUNKS_POW2) + counts = tl.load(chunk_counts + group * counts_stride_group + chunks, + mask=chunks < num_chunks, other=0) # fmt: skip + tl.store(chunk_bases + group * counts_stride_group + chunks, + tl.cumsum(counts, axis=0) - counts, mask=chunks < num_chunks) # fmt: skip + listed = tl.sum(counts, axis=0) + tl.store(group_plan + group * 2, tl.where(one_request, tl.cdiv(listed, BLOCK_N), -1)) + tl.store(group_plan + group * 2 + 1, listed) + if not one_request: + base = tl.atomic_add(lone_count, tl.sum(live.to(tl.int32), axis=0)) + tl.store(lone_rows + base + offsets, members.to(tl.int32), mask=live) + + +@triton.jit +def _rowtile_slots_kernel( + request_indices, row_info, bits, block_table, union_words, chunk_counts, chunk_bases, + group_plan, slots, packed_bits, rows, + bits_stride_row, packed_stride_row, block_table_stride_request, block_table_stride_page, + union_stride_group, counts_stride_group, slots_stride_group, + NUM_REQUESTS, + ROWS_PER_TILE: tl.constexpr, ROWS_POW2: tl.constexpr, + TOKENS_PER_BLOCK: tl.constexpr, CHUNK_WORDS: tl.constexpr, WORD_PAGES: tl.constexpr, +): # fmt: skip + """Per (group, chunk) of a row-tile group: the chunk's listed tokens as physical cache + slots (`page * tokens_per_block + offset`) at their list positions in `slots[group]`, + and each row's bits re-indexed onto the list ORed into `packed_bits[row]` (bit `j` set + when the row selects the list's `j`-th token): a per-word parallel bit extract.""" + group = tl.program_id(0).to(tl.int64) + chunk = tl.program_id(1) + count = tl.load(chunk_counts + group * counts_stride_group + chunk) + if (tl.load(group_plan + group * 2) >= 0) & (count > 0): + _, members, live, _, lowest, words_needed = _group_rows( + group, request_indices, row_info, rows, NUM_REQUESTS, ROWS_PER_TILE, ROWS_POW2 + ) + words = chunk * CHUNK_WORDS + tl.arange(0, CHUNK_WORDS) + in_range = words < words_needed + resident = tl.load(union_words + group * union_stride_group + words, mask=in_range, + other=0).to(tl.uint64, bitcast=True) # fmt: skip + word_count = _popc64(resident) + first = (tl.load(chunk_bases + group * counts_stride_group + chunk) + + tl.cumsum(word_count, axis=0) - word_count) # fmt: skip + # The pages the chunk's words span, looked up once a word (as in the union pass). + first_token = words.to(tl.int32) * 64 + first_page = first_token // TOKENS_PER_BLOCK + word_pages = first_page[:, None] + tl.arange(0, WORD_PAGES)[None, :] + low = tl.maximum(word_pages * TOKENS_PER_BLOCK - first_token[:, None], 0) + high = tl.minimum((word_pages + 1) * TOKENS_PER_BLOCK - first_token[:, None], 64) + page_bits = tl.where(low < high, _bits_below(high) & ~_bits_below(low), 0) + pages_of_word = tl.load( + block_table + lowest * block_table_stride_request + + word_pages * block_table_stride_page, + mask=(page_bits & resident[:, None]) != 0, + other=0, + ) # fmt: skip + lanes = tl.arange(0, 64) + lanes64 = lanes.to(tl.uint64) + kept = ((resident[:, None] >> lanes64[None, :]) & 1) != 0 + rank = first[:, None] + _popc64(resident[:, None] & _bits_below(lanes)[None, :]) + tokens = first_token[:, None] + lanes[None, :] + which = tokens // TOKENS_PER_BLOCK - first_page[:, None] + pages = tl.zeros([CHUNK_WORDS, 64], tl.int32) + for k in tl.static_range(WORD_PAGES): + page_k = tl.sum(tl.where(tl.arange(0, WORD_PAGES)[None, :] == k, pages_of_word, 0), 1) + pages = tl.where(which == k, page_k[:, None], pages) + tl.store(slots + group * slots_stride_group + rank, + pages * TOKENS_PER_BLOCK + tokens % TOKENS_PER_BLOCK, mask=kept) # fmt: skip + row_words = tl.load( + bits + members[:, None] * bits_stride_row + words[None, :], + mask=live[:, None] & in_range[None, :], + other=0, + ).to(tl.uint64, bitcast=True) # fmt: skip + packed = _compress_bits(row_words, resident[None, :]) + shift = (first % 64).to(tl.uint64) + destination = packed_bits + members[:, None] * packed_stride_row + (first // 64)[None, :] + writes = live[:, None] & in_range[None, :] & (packed != 0) + tl.atomic_or(destination, (packed << shift[None, :]).to(tl.int64, bitcast=True), + mask=writes, sem="relaxed") # fmt: skip + # Bits past the destination word spill into the next one (never for shift 0). + spill = (packed >> 1) >> (63 - shift[None, :]) + tl.atomic_or(destination + 1, spill.to(tl.int64, bitcast=True), + mask=writes & ((first % 64) + word_count > 64)[None, :], sem="relaxed") # fmt: skip + + +@triton.jit +def _rowtile_lone_kernel( + q, k_cache, v_cache, block_table, selected_tokens, request_indices, row_info, + lone_rows, lone_count, output, softmax_scale, rows, + q_stride_row: tl.constexpr, q_stride_head: tl.constexpr, q_stride_dim: tl.constexpr, + k_stride_page: tl.constexpr, k_stride_head: tl.constexpr, + k_stride_token: tl.constexpr, k_stride_dim: tl.constexpr, + v_stride_page: tl.constexpr, v_stride_head: tl.constexpr, + v_stride_token: tl.constexpr, v_stride_dim: tl.constexpr, + block_table_stride_request: tl.constexpr, block_table_stride_page: tl.constexpr, + selected_stride_row: tl.constexpr, selected_stride_token: tl.constexpr, + output_stride_row: tl.constexpr, output_stride_head: tl.constexpr, + output_stride_dim: tl.constexpr, + NUM_CACHE_PAGES, NUM_REQUESTS, LONE_PROGRAMS, + PAGE_TABLE_WIDTH: tl.constexpr, + GROUP_SIZE: tl.constexpr, TOKENS_PER_BLOCK: tl.constexpr, ROWS_PER_TILE: tl.constexpr, + ROWS_POW2: tl.constexpr, BLOCK_H: tl.constexpr, BLOCK_N: tl.constexpr, + HEAD_DIM: tl.constexpr, USE_PDL: tl.constexpr, + LONE_STAGES: tl.constexpr, +): # fmt: skip + """Rows of groups the row-tile pass does not serve: each its own columns, gathered.""" + if USE_PDL: + # Independent of the row-tile pass (disjoint rows): let it start now. + tl.extra.cuda.gdc_launch_dependents() + kv_head = tl.program_id(1) + for slot in range(tl.program_id(0), tl.load(lone_count), LONE_PROGRAMS): + _rowtile_lone_row( + q, k_cache, v_cache, block_table, selected_tokens, request_indices, row_info, output, + softmax_scale, tl.load(lone_rows + slot).to(tl.int64), kv_head, + q_stride_row, q_stride_head, q_stride_dim, + k_stride_page, k_stride_head, k_stride_token, k_stride_dim, + v_stride_page, v_stride_head, v_stride_token, v_stride_dim, + block_table_stride_request, block_table_stride_page, + selected_stride_row, selected_stride_token, + output_stride_row, output_stride_head, output_stride_dim, + NUM_CACHE_PAGES, NUM_REQUESTS, PAGE_TABLE_WIDTH, GROUP_SIZE, TOKENS_PER_BLOCK, + BLOCK_H, BLOCK_N, HEAD_DIM, LONE_STAGES, + ) # fmt: skip + + +@triton.jit +def _rowtile_lone_row( + q, k_cache, v_cache, block_table, selected_tokens, request_indices, row_info, output, + softmax_scale, row, kv_head, + q_stride_row, q_stride_head, q_stride_dim, + k_stride_page, k_stride_head, k_stride_token, k_stride_dim, + v_stride_page, v_stride_head, v_stride_token, v_stride_dim, + block_table_stride_request, block_table_stride_page, + selected_stride_row, selected_stride_token, + output_stride_row, output_stride_head, output_stride_dim, + NUM_CACHE_PAGES, NUM_REQUESTS, PAGE_TABLE_WIDTH: tl.constexpr, + GROUP_SIZE: tl.constexpr, TOKENS_PER_BLOCK: tl.constexpr, BLOCK_H: tl.constexpr, + BLOCK_N: tl.constexpr, HEAD_DIM: tl.constexpr, LONE_STAGES: tl.constexpr, +): # fmt: skip + request = tl.load(request_indices + row).to(tl.int64) + valid_request = (request >= 0) & (request < NUM_REQUESTS) + safe_request = tl.minimum(tl.maximum(request, 0), NUM_REQUESTS - 1) + page_base = block_table + safe_request * block_table_stride_request + columns_read = tl.load(row_info + row * 2 + 1) + head_offsets = tl.arange(0, BLOCK_H) + dim_offsets = tl.arange(0, HEAD_DIM) + query_heads = kv_head * GROUP_SIZE + head_offsets + head_live = head_offsets < GROUP_SIZE + q_tile = tl.load( + q + row * q_stride_row + query_heads[:, None] * q_stride_head + + dim_offsets[None, :] * q_stride_dim, + mask=head_live[:, None], + other=0.0, + ) # fmt: skip + accumulator = tl.zeros([BLOCK_H, HEAD_DIM], tl.float32) + running_max = tl.full([BLOCK_H], -1.0e30, tl.float32) + running_sum = tl.zeros([BLOCK_H], tl.float32) + token_offsets = tl.arange(0, BLOCK_N) + for tile in tl.range(0, tl.cdiv(columns_read, BLOCK_N), num_stages=LONE_STAGES): + columns = tile * BLOCK_N + token_offsets + tokens = tl.load( + selected_tokens + row * selected_stride_row + columns * selected_stride_token, + mask=columns < columns_read, + other=-1, + ).to(tl.int32) + safe_tokens = tl.maximum(tokens, 0) + logical_pages = safe_tokens // TOKENS_PER_BLOCK + token_in_page = safe_tokens % TOKENS_PER_BLOCK + usable = valid_request & (tokens >= 0) & (logical_pages < PAGE_TABLE_WIDTH) + physical_pages = tl.load( + page_base + tl.minimum(logical_pages, PAGE_TABLE_WIDTH - 1) * block_table_stride_page, + mask=usable, + other=-1, + ).to(tl.int64) + usable = usable & (physical_pages >= 0) & (physical_pages < NUM_CACHE_PAGES) + safe_pages = tl.maximum(physical_pages, 0) + keys = tl.load( + k_cache + safe_pages[None, :] * k_stride_page + kv_head * k_stride_head + + token_in_page[None, :] * k_stride_token + dim_offsets[:, None] * k_stride_dim, + mask=usable[None, :], + other=0.0, + ).to(q_tile.dtype) # fmt: skip + values = tl.load( + v_cache + safe_pages[:, None] * v_stride_page + kv_head * v_stride_head + + token_in_page[:, None] * v_stride_token + dim_offsets[None, :] * v_stride_dim, + mask=usable[:, None], + other=0.0, + ).to(q_tile.dtype) # fmt: skip + scores = tl.dot(q_tile, keys) * (softmax_scale * 1.4426950408889634) + scores = tl.where(usable[None, :], scores, -float("inf")) + next_max = tl.maximum(running_max, tl.max(scores, axis=1)) + correction = tl.math.exp2(running_max - next_max) + probabilities = tl.where(usable[None, :], tl.math.exp2(scores - next_max[:, None]), 0.0) + accumulator = tl.dot( + probabilities.to(q_tile.dtype), values, accumulator * correction[:, None] + ) + running_sum = running_sum * correction + tl.sum(probabilities, axis=1) + running_max = next_max + normalized = tl.where( + running_sum[:, None] > 0, accumulator / tl.maximum(running_sum[:, None], 1.0e-20), 0.0 + ) + tl.store( + output + row * output_stride_row + query_heads[:, None] * output_stride_head + + dim_offsets[None, :] * output_stride_dim, + normalized.to(output.dtype.element_ty), + mask=head_live[:, None], + ) # fmt: skip + + +def _check_inputs(q, k_cache, v_cache, block_table, selected_tokens, request_indices, + tokens_per_block, softmax_scale, query_positions, compress_ratio): # fmt: skip + if q.ndim != 3: + raise ValueError("QSA queries must have shape [rows, query_heads, head_dim]") + if k_cache.ndim != 4 or v_cache.shape != k_cache.shape: + raise ValueError("QSA K/V caches must have matching HND paged shapes") + if block_table.ndim != 2 or block_table.dtype != torch.int32: + raise ValueError("QSA sparse GQA requires a two-dimensional int32 block table") + if selected_tokens.ndim != 2 or selected_tokens.dtype != torch.int32: + raise ValueError("QSA selected tokens must be a two-dimensional int32 tensor") + rows, num_q_heads, head_dim = q.shape + if selected_tokens.shape[0] != rows or selected_tokens.shape[1] <= 0: + raise ValueError("QSA selected-token rows must match non-empty query rows") + if request_indices.shape != (rows,) or request_indices.dtype not in _INTEGER_DTYPES: + raise ValueError("QSA request indices must have one integer entry per row") + if not request_indices.is_contiguous(): + raise ValueError("QSA request indices must be contiguous") + if tokens_per_block <= 0 or k_cache.shape[2] != tokens_per_block: + raise ValueError("QSA cache token dimension must match positive tokens_per_block") + if k_cache.shape[3] != head_dim or k_cache.dtype != v_cache.dtype: + raise ValueError("QSA query and K/V cache head geometry is incompatible") + if q.dtype != torch.bfloat16: + raise ValueError("QSA queries must use BF16") + if k_cache.dtype != torch.bfloat16: + raise ValueError("QSA K/V cache must use BF16") + if selected_tokens.stride(1) != 1: + raise ValueError("QSA selected tokens must have unit column stride") + num_kv_heads = k_cache.shape[1] + if num_kv_heads <= 0 or num_q_heads <= 0 or num_q_heads % num_kv_heads: + raise ValueError("QSA query heads must be divisible by non-zero KV heads") + if k_cache.shape[0] <= 0 or block_table.shape[0] <= 0 or block_table.shape[1] <= 0: + raise ValueError("QSA sparse GQA requires non-empty cache and block-table geometry") + if k_cache.shape[0] * tokens_per_block >= 2**31: + raise ValueError("QSA row-tile cache slots must fit int32") + if not math.isfinite(softmax_scale) or softmax_scale <= 0: + raise ValueError("QSA sparse GQA softmax scale must be finite and positive") + if query_positions is None or compress_ratio is None: + raise ValueError("the row-tile kernel needs the prefill pair (positions, compress ratio)") + if query_positions.shape != (rows,) or query_positions.dtype not in _INTEGER_DTYPES: + raise ValueError("QSA query positions must be one integer entry per row") + if compress_ratio <= 0: + raise ValueError("QSA compression ratio must be positive") + token_topk = selected_tokens.shape[1] - compress_ratio + 1 + if token_topk <= 0 or token_topk % compress_ratio: + raise ValueError("QSA selected-token width is incompatible with its compression ratio") + devices = {t.device for t in (q, k_cache, v_cache, block_table, selected_tokens, + request_indices, query_positions)} # fmt: skip + if len(devices) != 1 or next(iter(devices)).type != "cuda": + raise ValueError("QSA row-tile sparse GQA needs every tensor on one CUDA device") + if num_q_heads // num_kv_heads > BLOCK_M: + raise ValueError("a row's query heads per KV head must fit the 128-row tile") + if not _cute_eligible(q, k_cache, v_cache): + raise ValueError( + "the row-tile pass needs SM103, head_dim 256, unit-stride heads (q's heads packed) " + "and 16-byte-aligned rows" + ) + + +def qsa_rowtile_sparse_gqa( + *, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + selected_tokens: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + softmax_scale: float, + query_positions: torch.Tensor, + compress_ratio: int, +) -> torch.Tensor: + """Bitmap pass, plan pass, lone-row pass, row-tile pass (`qsa_rowtile_attention_cute`); + returns `[rows, Hq, D]` contiguous.""" + _check_inputs(q, k_cache, v_cache, block_table, selected_tokens, request_indices, + tokens_per_block, softmax_scale, query_positions, compress_ratio) # fmt: skip + rows, num_q_heads, head_dim = q.shape + num_kv_heads = k_cache.shape[1] + group_size = num_q_heads // num_kv_heads + output = torch.empty((rows, num_q_heads, head_dim), dtype=q.dtype, device=q.device) + if rows == 0: + return output + table_width = block_table.shape[1] + num_words = triton.cdiv(table_width * tokens_per_block, 64) + bits, row_info = _row_bits( + selected_tokens, query_positions, compress_ratio, tokens_per_block, table_width, num_words + ) + rows_per_tile = BLOCK_M // group_size + rows_pow2 = triton.next_power_of_2(rows_per_tile) + kv_strides = ( + k_cache.stride(0), k_cache.stride(1), k_cache.stride(2), k_cache.stride(3), + v_cache.stride(0), v_cache.stride(1), v_cache.stride(2), v_cache.stride(3), + ) # fmt: skip + q_strides = (q.stride(0), q.stride(1), q.stride(2)) + common = (*q_strides, *kv_strides, block_table.stride(0), block_table.stride(1)) + geometry = dict( + NUM_CACHE_PAGES=k_cache.shape[0], NUM_REQUESTS=block_table.shape[0], + PAGE_TABLE_WIDTH=table_width, GROUP_SIZE=group_size, + TOKENS_PER_BLOCK=tokens_per_block, ROWS_PER_TILE=rows_per_tile, ROWS_POW2=rows_pow2, + HEAD_DIM=head_dim, + ) # fmt: skip + out_strides = (output.stride(0), output.stride(1), output.stride(2)) + num_groups = triton.cdiv(rows, rows_per_tile) + # A group's list holds at most its rows' columns: the slot lists and the re-indexed + # bitmaps are that wide (in whole 64-slot tiles), not the table's width. + list_words = triton.cdiv(min(num_words * 64, rows_per_tile * selected_tokens.shape[1]), 64) + num_chunks = triton.cdiv(num_words, _CHUNK_WORDS) + slots = torch.empty((num_groups, list_words * 64), dtype=torch.int32, device=q.device) + packed_bits = torch.zeros((rows, list_words), dtype=torch.int64, device=q.device) + union_words = torch.empty((num_groups, num_words), dtype=torch.int64, device=q.device) + # [0]: per chunk, its resident tokens; [1]: its first list position. + chunk_info = torch.empty((2, num_groups, num_chunks), dtype=torch.int32, device=q.device) + group_plan = torch.empty((num_groups, 2), dtype=torch.int32, device=q.device) + lone_rows = torch.empty((rows,), dtype=torch.int32, device=q.device) + # [0]: the lone list's length; [1]: the row-tile pass's work-ticket counter. + counters = torch.zeros((2,), dtype=torch.int32, device=q.device) + lone_count = counters[:1] + rows_meta = dict(ROWS_PER_TILE=rows_per_tile, ROWS_POW2=rows_pow2) + # The pages an aligned 64-token word spans. + if 64 % tokens_per_block == 0: + word_pages = 64 // tokens_per_block + elif tokens_per_block % 64 == 0: + word_pages = 1 + else: + word_pages = triton.cdiv(64, tokens_per_block) + 1 + word_meta = dict(TOKENS_PER_BLOCK=tokens_per_block, CHUNK_WORDS=_CHUNK_WORDS, + WORD_PAGES=triton.next_power_of_2(word_pages)) # fmt: skip + _rowtile_union_kernel[(num_groups, num_chunks)]( + request_indices, row_info, bits, block_table, union_words, chunk_info[0], rows, + bits.stride(0), block_table.stride(0), block_table.stride(1), + union_words.stride(0), chunk_info.stride(1), + block_table.shape[0], k_cache.shape[0], + PAGE_TABLE_WIDTH=table_width, num_warps=_PLAN_WARPS, **rows_meta, **word_meta, + ) # fmt: skip + _rowtile_group_kernel[(num_groups,)]( + request_indices, row_info, chunk_info[0], chunk_info[1], group_plan, lone_rows, + lone_count, rows, chunk_info.stride(1), num_chunks, block_table.shape[0], + BLOCK_N=BLOCK_N, CHUNKS_POW2=triton.next_power_of_2(num_chunks), + num_warps=_PLAN_WARPS, **rows_meta, + ) # fmt: skip + _rowtile_slots_kernel[(num_groups, num_chunks)]( + request_indices, row_info, bits, block_table, union_words, chunk_info[0], chunk_info[1], + group_plan, slots, packed_bits, rows, + bits.stride(0), packed_bits.stride(0), block_table.stride(0), block_table.stride(1), + union_words.stride(0), chunk_info.stride(1), slots.stride(0), + block_table.shape[0], + num_warps=_PLAN_WARPS, **rows_meta, **word_meta, + ) # fmt: skip + use_pdl = get_env_enable_pdl() + _rowtile_lone_kernel[(min(rows, _LONE_PROGRAMS), num_kv_heads)]( + q, k_cache, v_cache, block_table, selected_tokens, request_indices, row_info, + lone_rows, lone_count, output, softmax_scale, rows, + *common, selected_tokens.stride(0), selected_tokens.stride(1), *out_strides, + BLOCK_H=max(16, triton.next_power_of_2(group_size)), BLOCK_N=64, + LONE_PROGRAMS=min(rows, _LONE_PROGRAMS), USE_PDL=use_pdl, LONE_STAGES=_LONE_STAGES, + num_warps=_LONE_WARPS, **geometry, + ) # fmt: skip + _rowtile_attention_cute[(_rowtile_ctas(num_groups * num_kv_heads, q.device),)]( + q, k_cache, v_cache, slots, packed_bits, group_plan, output, counters[1:], softmax_scale, + rows, + group_size=group_size, rows_per_tile=rows_per_tile, tokens_per_block=tokens_per_block, + use_pdl=use_pdl, + ) # fmt: skip + return output + + +_CHECKPOINT = dict( + q_heads=24, + kv_heads=2, + head_dim=256, + tokens_per_block=32, + compress_ratio=4, + token_topk=2048, + cache_dtype=torch.bfloat16, + q_dtype=torch.bfloat16, + # K and V are the two role planes of one [pages, 2, kv, tpb, D] pool. + cache_layout="pool_planes", + # q is the first 24 heads of the packed [T, 24 + 2 + 2, 256] QKV buffer. + q_layout="qkv_slice", + request_dtype=torch.int32, + # The engine's table at max_seq_len 10240: 320 pages per slot, -1 past each sequence. + pages=320, +) + + +def _spec(segments, slots=32, **overrides) -> dict: + """`segments` is `(first position, rows)` per request, in row order.""" + return {**_CHECKPOINT, "segments": tuple(segments), "slots": slots, **overrides} + + +def _decodes(count, rows_each=1, first=8200): + return [(first + (37 * g) % 1000, rows_each) for g in range(count)] + + +class QsaSparseGqaRowtile(OpWrapper): + """Sparse GQA over the selected tokens of an HND paged K/V pool, for a + launch that carries the prefill pair (`query_positions`, `compress_ratio`). + + For row `r`, query head `h`, KV head `h // (Hq // Hkv)`, over the columns + `j < bound(r)` of `selected_tokens[r]` that name a resident token:: + + t = selected_tokens[r, j] # logical position + page = block_table[request_indices[r], t // tokens_per_block] + s_j = (q[r, h] . k_cache[page, kv, t % tpb]) * softmax_scale + out[r, h] = sum_j softmax(s)_j * v_cache[page, kv, t % tpb] + + `bound(r)` is the per-row kernel's prefill bound -- `min(visible // cr, + block_topk) * cr + visible % cr` columns with `visible = query_positions[r] + + 1`, rounded up to 64 -- so a row reads exactly the columns + `qsa_paged_sparse_gqa` reads on a context launch. Negative indices, pages + past the table, unmapped or out-of-pool pages and foreign requests + contribute nothing; a row left with nothing returns zeros. Returns a + fresh, contiguous `[rows, Hq, D]` tensor of `q.dtype`. + + Six launches. A bitmap pass: one bit per selected logical token per row. + A plan per group of 10 consecutive rows of one request, in three passes + over (group, 32 bitmap words) programs: the union of the rows' selected + tokens whose page is mapped (a page lookup per word's page, not per + token) and its count per chunk; a per-group scan of those counts (list + length and tiles); the union listed in logical order as physical cache + slots, and each row's bitmap re-indexed onto that list (a per-word + parallel bit extract ORed into place). Every group of one request goes to + the row-tile pass, however long its list: its list is never longer than + its rows' own columns combined, and a listed tile costs no more than a + per-row tile. A group whose rows span requests joins the lone + list. A lone-row pass (256 programs + per KV head striding the list) serves listed rows over their own columns, + as the per-row kernel does. The row-tile pass is persistent (one program + per SM drawing (group, KV head) items from a ticket counter, the plan's + last groups first); per item it loads each listed tile's K/V once for the + 120 (row, head) pairs and masks each row to its own re-indexed bits, and + one item's epilogue overlaps the next item's first MMAs. It is launched + programmatically dependent on the lone pass, whose rows are disjoint, so + the two overlap. That pass is `qsa_rowtile_attention_cute` (tcgen05 MMAs, + S / P / O in tensor memory, a cp.async gather of the K/V rows, a TMA store + of O), so the entry serves SM103, bf16 q and K/V and D = 256. + + Fusion boundary: as `qsa_paged_sparse_gqa` -- the block-table gather, + scale, online softmax and P.V are inside; deciding what to attend to, + appending K/V and q's projection, norm, RoPE and gate are not. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="con 1 prefill: 8192 rows of one request from 0, across the budget", + spec=_spec([(0, 8192)]), + ), + Cell( + why="sparse union: 4096 rows from 4096 ranking one score vector", + spec=_spec([(4096, 4096)], shared_scores=True), + ), + Cell( + why="con 128 mixed: 3 chunks (16240 rows) + 52 decodes on the lone pass", + spec=_spec([(0, 8192), (2192, 6000), (6144, 2048)] + _decodes(52), slots=128), + ), + ) + + note = """ + Metadata consumed: none -- every input is an argument, the same arguments + `qsa_paged_sparse_gqa` takes, with the prefill pair required. What the + target passes: `qsa_block_table` (physical page ids), the expansion's + `-1`-padded logical selections, the layer pool's K/V role planes, and the + logical positions with the checkpoint's compress ratio. + + Preconditions (neither is checkable without reading the contents): + - Within the bound, a row names each token at most once. The row-tile + pass attends to the *set* of tokens a row names; the per-row kernel + attends to every column, so a repeated index would weigh twice there + and once here. QSA's top-k + expansion never repeats one (distinct + blocks, and a tail outside every complete block): measured zero + repeats on four layers of the target's own 8192-token prefill. + Rows served by the lone-row pass keep per-column semantics. + - As for the per-row kernel: each row is prefix-packed with exactly the + bound's count of valid entries; entries past the 64-rounded bound + are ignored. + + Work: a row-tile group walks its packed list, so its MMA work grows + with the union of its rows' selections rather than with the group's + highest position. On the target's own selections (8192-token prefill, + layers 0/3/6/11) that union is 13-32% shorter than every 64-token tile up + to the highest token, and the walk loads ~5x fewer K/V tiles than the + per-row kernel. Measured on GB300 (eager, whole call) at 8192 context + rows of real selections: 1544-1565 us on layers 0-6 and 1274 us on layer + 11, against 1711-1715 us for a walk of every tile up to the highest token + and 2953 us for `qsa_paged_sparse_gqa`. Uniformly random selections + leave almost no gap in the union, and the call costs what the full walk + does (16292-row mixed step: 3859 vs 3836 us). Below a few hundred rows + it is slower than the per-row kernel (generation-only launches 1.4-2x); + the target's served context launches are thousands of rows. The tensor-core + row-tile pass alone, on six dumped 16384-row context launches of the + target (GB300, layers 0-5 of one step): 581-775 us, against 694-886 us for + the one-program-per-item pass it replaced (the same output, bit for bit). + + Compiled once per (group size, rows per tile, tokens per block, PDL): + nothing in the row-tile pass's compile key depends on the row count or + the plan. + + Scratch per call: the bitmap, `rows x ceil(table_width * tpb / 64)` int64 + words (21 MB at 16384 rows and 320 pages), and per group its union over + as many words (2 MB); the re-indexed bitmaps, `rows x ceil(10 x + selection width / 64)` int64 words (42 MB at 16384 rows), and the slot + lists, `ceil(rows / 10) x` 64 times that many int32 slots (27 MB); the + per-chunk counts, the plan and the lone list. Calls allocate, + so capture in a CUDA graph is not certified; this entry serves the + prefill path, which runs eager. Results are deterministic: the bitmap + OR, the re-indexed bitmap OR and the lone list order do not change any + row's arithmetic. + + Not certified: int64 request indices or positions; rows whose selection + resolves no token (all padding, a foreign request, unmapped pages); + `rows == 0`. + """ + + def __call__( + self, + *, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + selected_tokens: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + softmax_scale: float, + query_positions: torch.Tensor, + compress_ratio: int, + ) -> torch.Tensor: + return qsa_rowtile_sparse_gqa( + q=q, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + selected_tokens=selected_tokens, + request_indices=request_indices, + tokens_per_block=tokens_per_block, + softmax_scale=softmax_scale, + query_positions=query_positions, + compress_ratio=compress_ratio, + ) + + def reference( + self, + *, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + selected_tokens: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + softmax_scale: float, + query_positions: torch.Tensor, + compress_ratio: int, + ) -> torch.Tensor: + """The reference `qsa_paged_sparse_gqa` certifies against: fp32 sparse + GQA over every resident selected token, rounded once. The prefill pair + is ignored, as there: on a prefix-packed selection the bound reads + exactly the valid entries.""" + return _per_row_entry.reference( + q=q, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + selected_tokens=selected_tokens, + request_indices=request_indices, + tokens_per_block=tokens_per_block, + softmax_scale=softmax_scale, + query_positions=query_positions, + compress_ratio=compress_ratio, + ) + + +qsa_sparse_gqa_rowtile = QsaSparseGqaRowtile() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_topk_expand_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_topk_expand_decode.py new file mode 100644 index 000000000000..cf74ab9cefec --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_topk_expand_decode.py @@ -0,0 +1,1556 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA's block top-k and its expansion to token indices, one launch per batch of rows. + +A CuTe DSL kernel, an in-tree callable. One CTA per row holds four consecutive scores a +thread and narrows to the k-th largest by radix passes over the scores' bit patterns +(shared-memory histograms of 11 bits), starting below the bits every visible score +shares. It stops as soon as the bin holding the threshold has at most `CAP` keys and +ranks those exactly by pairwise comparison (ties to the higher block, the canonical +selection's rule). Every selected block is then placed in block order by one block-wide +scan, and expanded to token indices with the causal tail, as the expansion entry does. + +Scores are a sum of ReLUs times a positive scale, so non-negative: their fp32 bit +patterns, read as int32, order like their values (a negative zero is taken as zero, a +NaN ranks above everything). + +A table wider than one CTA's four keys a thread (`MAX_WIDTH`, e.g. the 65536 columns of a +262144-token `max_seq_len`) takes `_TopkExpandWideKernel`, still one CTA a row and one +launch: the row's window is staged in shared memory, a bound every top-k key clears is +read off the lanes' maxima, the keys above it are compacted in block order, and the same +selection runs over that short list (or over the window, slab by slab, when the list +cannot hold them). +""" + +from __future__ import annotations + +import cutlass +import cutlass.cute as cute +import torch +from cutlass import const_expr +from cutlass.cute.nvgpu import cpasync +from cutlass.cute.runtime import from_dlpack + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_expand_block_indices import ( + qsa_expand_block_indices as _expand, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr + +#: Consecutive scores a thread holds (one 16-byte load). +KPT = 4 +#: Score columns a row may have: `KPT` a thread, at most 1024 threads a CTA. +MAX_WIDTH = KPT * 1024 +#: CTAs an SM the register allocation is held to. At the 2592-wide score table a row is 672 +#: threads, which the per-subpartition register file holds once an SM at the compiler's own +#: 48 registers a thread; at three (32 a thread) rows overlap each other's barriers and +#: shared-memory round trips. +MIN_BLOCKS = 3 +#: Radix passes at most (11 bits each; 31 magnitude bits). +_PASSES = 3 +_HIST = 2048 +_I32_MAX = 0x7FFFFFFF +#: At most this many keys in the threshold's bin are ranked pairwise instead of +#: narrowed by another pass. +CAP = 128 +#: Threads that own the 2048 bins of a scan, four each. +_SCAN_THREADS = _HIST // 4 +#: Consecutive threads that rank one candidate of the threshold's bin, each over every +#: `_RANK_LANES`-th candidate (`_SCAN_THREADS` threads cover `CAP` candidates). +_RANK_LANES = _SCAN_THREADS // CAP + +#: Tables wider than `MAX_WIDTH` take the wide kernel: one CTA a row (1024 threads by +#: default) walks its visible window in slabs of `KPT` columns a thread. +_WIDE_THREADS = 1024 +#: The widest table the wide kernel serves (the row count times the row stride stays an +#: Int32 product for 8192 rows). +MAX_WIDE_WIDTH = 1 << 17 +#: Shared memory of an SM on sm_100/sm_103 (228 KiB), the most one CTA may hold (227 KiB), +#: and what the driver reserves a CTA (1 KiB). +_SMEM_SM = 233472 +_SMEM_LIMIT = 232448 +_SMEM_RESERVED = 1024 +#: The wide kernel's fixed shared state in 4-byte words: the histogram, the candidates, +#: three sets of warp totals, the scan's double-buffered warp totals and the state words; +#: and the alignment padding of its allocations. +_WIDE_FIXED_WORDS = _HIST + 2 * CAP + 3 * 32 + 64 + 16 +_WIDE_PAD = 10 * 128 + + +def _threads(width: int) -> int: + return max(_SCAN_THREADS, -(-width // (KPT * 32)) * 32) + + +class _TopkExpandKernel: + def __init__(self, *, width, key_stride, block_topk, compress, token_topk, pdl): + self.width, self.key_stride, self.k, self.cr, self.token_topk, self.pdl = ( + width, + key_stride, + block_topk, + compress, + token_topk, + pdl, + ) + self.final = token_topk + compress - 1 + self.threads = _threads(width) + self.warps = self.threads // 32 + self.flag_words = -(-width // 32) + + @cute.jit + def __call__( + self, + keys_t: cute.Tensor, + end_ptr: cute.Pointer, + pos_ptr: cute.Pointer, + len_ptr: cute.Pointer, + topk_ptr: cute.Pointer, + sel_ptr: cute.Pointer, + rows: cutlass.Int32, + stream, + ): + # The keys arrive as a tensor, as the TVM-FFI entry takes its stream from a GPU + # tensor; their layout is rebuilt with the static width and row stride the 16-byte + # loads need. The row count is the call's. + mKeys = cute.make_tensor( + keys_t.iterator, cute.make_layout((rows, self.width), stride=(self.key_stride, 1)) + ) + per_row = cute.make_layout((rows,), stride=(1,)) + mEnd = cute.make_tensor(end_ptr, per_row) + mPos = cute.make_tensor(pos_ptr, per_row) + mLen = cute.make_tensor(len_ptr, per_row) + mTopk = cute.make_tensor(topk_ptr, cute.make_layout((rows, self.k), stride=(self.k, 1))) + mSel = cute.make_tensor( + sel_ptr, cute.make_layout((rows, self.final), stride=(self.final, 1)) + ) + ld_keys = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), cutlass.Int32, num_bits_per_copy=32 * KPT + ) + self.kernel(mKeys, mEnd, mPos, mLen, mTopk, mSel, ld_keys).launch( + grid=[rows, 1, 1], + block=[self.threads, 1, 1], + smem=(_HIST + self.flag_words + 2 * CAP + 64 + 16) * 4 + 128, + stream=stream, + use_pdl=self.pdl, + min_blocks_per_mp=MIN_BLOCKS, + ) + + @cute.kernel + def kernel(self, mKeys, mEnd, mPos, mLen, mTopk, mSel, ld_keys: cute.CopyAtom): + tidx, _, _ = cute.arch.thread_idx() + row, _, _ = cute.arch.block_idx() + warp = cute.arch.warp_idx() + lane = cute.arch.lane_idx() + nt: cutlass.Constexpr = self.threads + nw: cutlass.Constexpr = self.warps + k: cutlass.Constexpr = self.k + cr: cutlass.Constexpr = self.cr + smem = cutlass.utils.SmemAllocator() + hist = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((_HIST,), stride=(1,)), byte_alignment=16 + ) + flags = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((self.flag_words,), stride=(1,)), byte_alignment=16 + ) + cand_key = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((CAP,), stride=(1,)), byte_alignment=16 + ) + cand_idx = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((CAP,), stride=(1,)), byte_alignment=16 + ) + wa = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((32,), stride=(1,)), byte_alignment=16 + ) + wb = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((32,), stride=(1,)), byte_alignment=16 + ) + state = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((16,), stride=(1,)), byte_alignment=16 + ) + + # 0. Shared state, before the wait (nothing global is read here). + for b in cutlass.range_constexpr(-(-_HIST // nt)): + if tidx + b * nt < _HIST: + hist[tidx + b * nt] = cutlass.Int32(0) + for b in cutlass.range_constexpr(-(-self.flag_words // nt)): + if tidx + b * nt < self.flag_words: + flags[tidx + b * nt] = cutlass.Int32(0) + if tidx == 0: + state[5] = cutlass.Int32(0) # candidates compacted + # The row's metadata, host-filled: read before the wait too. + n = mEnd[row] + seqlen = mLen[row] + visible = cutlass.Int32(mPos[row]) + 1 + if const_expr(self.pdl): + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + + # 1. Four consecutive keys a thread (a column past the window is no candidate: + # -1); the largest and smallest key in the window. + e0 = tidx * KPT + keys = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32) + keys.fill(-1) + if e0 < self.width: + cute.copy(ld_keys, cute.logical_divide(mKeys, (None, KPT))[row, (None, tidx)], keys) + kmax = cutlass.Int32(-1) + kmin = cutlass.Int32(_I32_MAX) + for c in cutlass.range_constexpr(KPT): + b = keys[c] + if b < 0: # -0.0 (or a negative): rank as zero + b = cutlass.Int32(0) + if e0 + c >= n: + b = cutlass.Int32(-1) + keys[c] = b + kmax = cutlass.max(kmax, b) + if b >= 0: + kmin = cutlass.min(kmin, b) + kmax = cute.arch.warp_reduction_max(kmax) + kmin = cute.arch.warp_reduction(kmin, lambda x, y: cutlass.min(x, y)) + if lane == 0: + wa[warp] = kmax + wb[warp] = kmin + cute.arch.sync_threads() + gmax = cutlass.Int32(-1) + gmin = cutlass.Int32(_I32_MAX) + if lane < nw: + gmax = wa[lane] + gmin = wb[lane] + gmax = cute.arch.warp_reduction_max(gmax) + gmin = cute.arch.warp_reduction(gmin, lambda x, y: cutlass.min(x, y)) + + # 2. Every key in the window shares the bits above the highest bit where the + # largest and the smallest differ (`hb`); radix passes walk down from there, 11 bits + # at a time, until the threshold's bin holds at most CAP keys or no bits are left. + diff = gmax ^ gmin + hb = cutlass.Int32(-1) + if diff != 0: + hb = 31 - cutlass.Int32(cute.arch.clz(diff)) + low = (cutlass.Int64(1) << cutlass.Int64(hb + 1)) - 1 + mask0 = cutlass.Int32(cutlass.Int64(_I32_MAX) & (cutlass.Int64(-1) ^ low)) + if tidx == 0: + state[0] = gmax & mask0 # threshold prefix + state[1] = mask0 # its mask + state[2] = cutlass.Int32(k) # still needed at the threshold + state[3] = n # keys in the threshold's bin + state[4] = cutlass.Int32(1) # narrowing done + state[6] = cutlass.max(hb - 10, 0) # the next pass's shift + if n > k: + if hb >= 0: + state[4] = cutlass.Int32(0) + cute.arch.sync_threads() + for _p in cutlass.range_constexpr(_PASSES): + if state[4] == 0: + prefix = state[0] + mask = state[1] + need = state[2] + shift = state[6] + for c in cutlass.range_constexpr(KPT): + if keys[c] >= 0: + if (keys[c] & mask) == prefix: + digit = (keys[c] >> shift) & (_HIST - 1) + cute.arch.atomic_add( + hist.iterator + digit, cutlass.Int32(1), scope="cta" + ) + cute.arch.sync_threads() + # Thread t < 512 owns bins 2047-4t .. 2044-4t, counted from the top; each is + # cleared for the next pass once read. + h = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), cutlass.Int32) + h.fill(0) + own = cutlass.Int32(0) + if tidx < _SCAN_THREADS: + for j in cutlass.range_constexpr(4): + h[j] = hist[_HIST - 1 - 4 * tidx - j] + hist[_HIST - 1 - 4 * tidx - j] = cutlass.Int32(0) + own = own + h[j] + incl = own + for off in (1, 2, 4, 8, 16): + up = cute.arch.shuffle_sync_up(incl, offset=off) + if lane >= off: + incl = incl + up + if lane == 31: + wa[warp] = incl + cute.arch.sync_threads() + above = cutlass.Int32(0) + if lane < warp: + above = wa[lane] + before = cute.arch.warp_reduction_sum(above) + incl - own # keys in the bins above + if tidx < _SCAN_THREADS: + if before < need: + if need <= before + own: + run = before + for j in cutlass.range_constexpr(4): + if run < need: + if need <= run + h[j]: + state[0] = prefix | ((_HIST - 1 - 4 * tidx - j) << shift) + state[1] = mask | ((_HIST - 1) << shift) + state[2] = need - run + state[3] = h[j] + state[4] = cutlass.Int32( + cutlass.select_((h[j] <= CAP) | (shift == 0), 1, 0) + ) + state[6] = cutlass.max(shift - 11, 0) + run = run + h[j] + cute.arch.sync_threads() + + prefix = state[0] + mask = state[1] + need = state[2] + at_thr = state[3] + # 3. At most CAP keys in the threshold's bin: rank them exactly, ties to the + # higher block, and flag the `need` best. + if n > k: + if at_thr <= CAP: + for c in cutlass.range_constexpr(KPT): + if keys[c] >= 0: + if (keys[c] & mask) == prefix: + slot = cute.arch.atomic_add( + state.iterator + 5, cutlass.Int32(1), scope="cta" + ) + cand_key[slot] = keys[c] + cand_idx[slot] = e0 + c + cute.arch.sync_threads() + # Candidate `ci` is ranked by `_RANK_LANES` consecutive lanes, lane `sub` + # counting the better ones among candidates sub, sub + _RANK_LANES, ...; + # butterfly shuffles add the counts up. + ci = tidx // _RANK_LANES + sub = tidx % _RANK_LANES + rank = cutlass.Int32(0) + it = cutlass.Int32(0) + if ci < at_thr: + kt = cand_key[ci] + it = cand_idx[ci] + for j in cutlass.range(sub, at_thr, _RANK_LANES, unroll=4): + kj = cand_key[j] + ij = cand_idx[j] + better = (kj > kt) | ((kj == kt) & (ij > it)) + rank = rank + cutlass.Int32(cutlass.select_(better, 1, 0)) + for off in cutlass.range_constexpr(_RANK_LANES.bit_length() - 1): + rank = rank + cute.arch.shuffle_sync_bfly(rank, offset=1 << off) + if ci < at_thr: + if sub == 0: + if rank < need: + cute.arch.atomic_or( + flags.iterator + (it >> 5), + cutlass.Int32(1) << (it & 31), + scope="cta", + ) + cute.arch.sync_threads() + + # 4. Block order: one scan of (selected, tied at the threshold) counts. A tied key + # is left only when the passes ran out of bits with more than CAP keys tied; the + # last `need` of those in block order are taken. + take = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32) + tied = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32) + own = cutlass.Int32(0) + for c in cutlass.range_constexpr(KPT): + t_c = cutlass.Int32(0) + q_c = cutlass.Int32(0) + if keys[c] >= 0: + if n <= k: + t_c = cutlass.Int32(1) + else: + mk = keys[c] & mask + if mk > prefix: + t_c = cutlass.Int32(1) + if mk == prefix: + if at_thr > CAP: + q_c = cutlass.Int32(1) + else: + e = e0 + c + t_c = (flags[e >> 5] >> (e & 31)) & 1 + take[c] = t_c + tied[c] = q_c + own = own + t_c + (q_c << 16) + incl = own + for off in (1, 2, 4, 8, 16): + up = cute.arch.shuffle_sync_up(incl, offset=off) + if lane >= off: + incl = incl + up + if lane == 31: + wb[warp] = incl + cute.arch.sync_threads() + above = cutlass.Int32(0) + if lane < warp: + above = wb[lane] + pre = cute.arch.warp_reduction_sum(above) + incl - own + taken_before = pre & 0xFFFF + tied_before = pre >> 16 + # The tied keys left out, the lowest-index ones (none unless the threshold's bin + # holds more keys than are still needed). + skip = cutlass.max(at_thr - need, 0) + for c in cutlass.range_constexpr(KPT): + e = e0 + c + sel = cutlass.Int32(0) + pos = cutlass.Int32(0) + if tied[c] == 1: + if tied_before >= skip: + sel = cutlass.Int32(1) + pos = taken_before + tied_before - skip + tied_before = tied_before + 1 + if take[c] == 1: + sel = cutlass.Int32(1) + pos = taken_before + cutlass.max(tied_before - skip, 0) + taken_before = taken_before + 1 + if sel == 1: + mTopk[row, pos] = e + for j in cutlass.range_constexpr(cr): + tok = e * cr + j + val = cutlass.Int32(-1) + if tok < seqlen: + val = tok + mSel[row, pos * cr + j] = val + + # 5. Padding of the block list; the causal tail after the selected tokens, then -1. + total = n + if n > k: + total = cutlass.Int32(k) + for c in cutlass.range_constexpr(-(-k // nt)): + s = tidx + c * nt + if s < k: + if s >= total: + mTopk[row, s] = cutlass.Int32(-1) + tail_start = (visible // cr) * cr + tail_count = visible - tail_start + body = total * cr + if body > self.token_topk: + body = cutlass.Int32(self.token_topk) + for c in cutlass.range_constexpr(-(-self.final // nt)): + col = tidx + c * nt + if col < self.final: + if col >= body: + off = col - body + val = cutlass.Int32(-1) + if off < cr - 1: + if off < tail_count: + if tail_start + off < seqlen: + val = tail_start + off + mSel[row, col] = val + + +def _cache_slabs(width: int, threads: int, ctas: int) -> int: + """Slabs of the wide kernel's window cache: as many of the table's as the shared memory + left after the fixed state and the candidate list hold with `ctas` CTAs an SM (11 of + 4096 columns at one CTA of 1024 threads: a window of up to 45056 blocks, 180224 tokens, + is read from global once).""" + slab = KPT * threads + per_cta = min(_SMEM_LIMIT, _SMEM_SM // ctas - _SMEM_RESERVED) + room = per_cta - (_WIDE_FIXED_WORDS + 2 * slab) * 4 - _WIDE_PAD + return max(0, min(-(-width // slab), room // (slab * 4))) + + +class _TopkExpandWideKernel: + """`_TopkExpandKernel`'s selection over a table wider than one CTA's registers. The row's + visible window (`row_ends`, not the table width) is walked in slabs, `KPT` consecutive + keys a thread a slab; its first `cache_slabs` slabs are read from global once (cp.async + into shared memory) and every walk reads them from there, a later slab from global (L2). + + Prefilter: every lane's largest key is a distinct key of the window, so with `per_warp` + the `k / warps` rounded up, the smallest over the warps of the `per_warp`-th largest lane + maximum of each warp (`lo`) has at least `k` keys at or above it. The `k`-th largest key, + every key tied with it and everything above are among those, so the keys below `lo` are + no candidates and the selection is unchanged. + + Two walks of the window: the lane maxima, then the candidates (keys at or above `lo`) + compacted in block order into a shared list of up to `KPT` a thread (a window of ~21k + relu-sum scores keeps about a thousand). The narrow kernel's radix passes, pairwise rank + of the threshold's bin and block-order scan then run once over that list, `KPT` + consecutive entries a thread. A window with more candidates than the list holds (a + flat or degenerate row) takes the same steps over the whole window instead, slab by + slab, the threshold bin's `need` best kept as a cut (the worst of them, key and block) + and the block-order scan's counts carried from slab to slab.""" + + def __init__( + self, + *, + width, + key_stride, + block_topk, + compress, + token_topk, + pdl, + threads=_WIDE_THREADS, + ctas=1, + prefilter=True, + ): + self.width, self.key_stride, self.k, self.cr, self.token_topk, self.pdl = ( + width, + key_stride, + block_topk, + compress, + token_topk, + pdl, + ) + assert threads % 32 == 0 and _SCAN_THREADS <= threads <= 1024, threads + self.final = token_topk + compress - 1 + self.threads = threads + self.warps = threads // 32 + self.slab = KPT * threads + self.ctas = ctas + self.cache_slabs = _cache_slabs(width, threads, ctas) + self.per_warp = -(-block_topk // self.warps) + self.prefilter = bool(prefilter) and self.per_warp <= 32 + + def smem_bytes(self) -> int: + return (self.cache_slabs * self.slab + 2 * self.slab + _WIDE_FIXED_WORDS) * 4 + _WIDE_PAD + + @cute.jit + def __call__( + self, + keys_t: cute.Tensor, + end_ptr: cute.Pointer, + pos_ptr: cute.Pointer, + len_ptr: cute.Pointer, + topk_ptr: cute.Pointer, + sel_ptr: cute.Pointer, + rows: cutlass.Int32, + stream, + ): + mKeys = cute.make_tensor( + keys_t.iterator, cute.make_layout((rows, self.width), stride=(self.key_stride, 1)) + ) + per_row = cute.make_layout((rows,), stride=(1,)) + mEnd = cute.make_tensor(end_ptr, per_row) + mPos = cute.make_tensor(pos_ptr, per_row) + mLen = cute.make_tensor(len_ptr, per_row) + mTopk = cute.make_tensor(topk_ptr, cute.make_layout((rows, self.k), stride=(self.k, 1))) + mSel = cute.make_tensor( + sel_ptr, cute.make_layout((rows, self.final), stride=(self.final, 1)) + ) + ld_keys = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), cutlass.Int32, num_bits_per_copy=32 * KPT + ) + g2s = cute.make_copy_atom(cpasync.CopyG2SOp(), cutlass.Int32, num_bits_per_copy=32 * KPT) + self.kernel(mKeys, mEnd, mPos, mLen, mTopk, mSel, ld_keys, g2s).launch( + grid=[rows, 1, 1], + block=[self.threads, 1, 1], + smem=self.smem_bytes(), + stream=stream, + use_pdl=self.pdl, + min_blocks_per_mp=self.ctas, + ) + + @cute.jit + def _slab(self, keys, gK, sK, row, s, tidx, n, lo, ld_keys: cute.CopyAtom): + """Slab `s`'s `KPT` keys of thread `tidx` into `keys`: a negative (-0.0) as zero, -1 + past the window and below `lo`.""" + col4 = s * self.threads + tidx + e0 = col4 * KPT + if const_expr(self.cache_slabs > 0): + if s < self.cache_slabs: + cute.autovec_copy(sK[(None, col4)], keys) + else: + keys.fill(-1) + if e0 < n: + cute.copy(ld_keys, gK[row, (None, col4)], keys) + else: + keys.fill(-1) + if e0 < n: + cute.copy(ld_keys, gK[row, (None, col4)], keys) + for c in cutlass.range_constexpr(KPT): + b = cutlass.max(keys[c], cutlass.Int32(0)) + if (b < lo) | (e0 + c >= n): + b = cutlass.Int32(-1) + keys[c] = b + + @cute.jit + def _passes(self, keys, tidx, warp, lane, hist, wsum, state): + """The narrow kernel's radix passes over `keys` (this thread's share; -1 is none): + from the state's prefix, mask and shift, 11 bits a pass, until the threshold's bin + holds at most CAP keys or no bits are left.""" + for _p in cutlass.range_constexpr(_PASSES): + if state[4] == 0: + prefix = state[0] + mask = state[1] + need = state[2] + shift = state[6] + for c in cutlass.range_constexpr(KPT): + if keys[c] >= 0: + if (keys[c] & mask) == prefix: + digit = (keys[c] >> shift) & (_HIST - 1) + cute.arch.atomic_add( + hist.iterator + digit, cutlass.Int32(1), scope="cta" + ) + cute.arch.sync_threads() + self._narrow(tidx, warp, lane, hist, wsum, state, prefix, mask, need, shift) + + @cute.jit + def _narrow(self, tidx, warp, lane, hist, wsum, state, prefix, mask, need, shift): + """One pass's histogram read (and cleared) from the top bin down: the bin holding + the `need`-th key becomes the threshold's.""" + h = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), cutlass.Int32) + h.fill(0) + own = cutlass.Int32(0) + if tidx < _SCAN_THREADS: + for j in cutlass.range_constexpr(4): + h[j] = hist[_HIST - 1 - 4 * tidx - j] + hist[_HIST - 1 - 4 * tidx - j] = cutlass.Int32(0) + own = own + h[j] + incl = own + for off in (1, 2, 4, 8, 16): + up = cute.arch.shuffle_sync_up(incl, offset=off) + if lane >= off: + incl = incl + up + if lane == 31: + wsum[warp] = incl + cute.arch.sync_threads() + above = cutlass.Int32(0) + if lane < warp: + above = wsum[lane] + before = cute.arch.warp_reduction_sum(above) + incl - own + if tidx < _SCAN_THREADS: + if before < need: + if need <= before + own: + run = before + for j in cutlass.range_constexpr(4): + if run < need: + if need <= run + h[j]: + state[0] = prefix | ((_HIST - 1 - 4 * tidx - j) << shift) + state[1] = mask | ((_HIST - 1) << shift) + state[2] = need - run + state[3] = h[j] + state[4] = cutlass.Int32( + cutlass.select_((h[j] <= CAP) | (shift == 0), 1, 0) + ) + state[6] = cutlass.max(shift - 11, 0) + run = run + h[j] + cute.arch.sync_threads() + + @cute.jit + def _cut(self, tidx, state, cand_key, cand_idx, at_thr, need): + """The threshold bin's `at_thr` <= CAP keys (gathered into `cand_key`/`cand_idx`) + ranked pairwise, ties to the higher block: the `need`-th best is the cut, kept in + `state[7]` (key) and `state[8]` (block).""" + ci = tidx // _RANK_LANES + sub = tidx % _RANK_LANES + rank = cutlass.Int32(0) + kt = cutlass.Int32(0) + it = cutlass.Int32(0) + if ci < at_thr: + kt = cand_key[ci] + it = cand_idx[ci] + for j in cutlass.range(sub, at_thr, _RANK_LANES, unroll=4): + kj = cand_key[j] + ij = cand_idx[j] + better = (kj > kt) | ((kj == kt) & (ij > it)) + rank = rank + cutlass.Int32(cutlass.select_(better, 1, 0)) + for off in cutlass.range_constexpr(_RANK_LANES.bit_length() - 1): + rank = rank + cute.arch.shuffle_sync_bfly(rank, offset=1 << off) + if ci < at_thr: + if sub == 0: + if rank == need - 1: + state[7] = kt + state[8] = it + cute.arch.sync_threads() + + @cute.jit + def _mark(self, keys, idx, take, tied, everything, prefix, mask, at_thr, cut_key, cut_idx): + """Each of this thread's keys taken (above the threshold's bin, or in it at or above + the cut) or tied (in it, when the passes ran out of bits with more than CAP keys); + `everything` takes every key. Returns the packed counts, taken + (tied << 16).""" + own = cutlass.Int32(0) + for c in cutlass.range_constexpr(KPT): + t_c = cutlass.Int32(0) + q_c = cutlass.Int32(0) + if keys[c] >= 0: + if everything: + t_c = cutlass.Int32(1) + else: + mk = keys[c] & mask + if mk > prefix: + t_c = cutlass.Int32(1) + if mk == prefix: + if at_thr > CAP: + q_c = cutlass.Int32(1) + elif (keys[c] > cut_key) | ((keys[c] == cut_key) & (idx[c] >= cut_idx)): + t_c = cutlass.Int32(1) + take[c] = t_c + tied[c] = q_c + own = own + t_c + (q_c << 16) + return own + + @cute.jit + def _emit(self, row, idx, take, tied, taken_before, tied_before, skip, seqlen, mTopk, mSel): + """This thread's selected blocks at their block-order positions, and their tokens + below the sequence length (-1 past it). `tied_before - skip` of the tied keys before + these are taken (the lowest-index `skip` of the tie class are left).""" + for c in cutlass.range_constexpr(KPT): + e = idx[c] + sel = cutlass.Int32(0) + pos = cutlass.Int32(0) + if tied[c] == 1: + if tied_before >= skip: + sel = cutlass.Int32(1) + pos = taken_before + tied_before - skip + tied_before = tied_before + 1 + if take[c] == 1: + sel = cutlass.Int32(1) + pos = taken_before + cutlass.max(tied_before - skip, 0) + taken_before = taken_before + 1 + if sel == 1: + mTopk[row, pos] = e + for j in cutlass.range_constexpr(self.cr): + tok = e * self.cr + j + val = cutlass.Int32(-1) + if tok < seqlen: + val = tok + mSel[row, pos * self.cr + j] = val + + @cute.jit + def _select_listed( + self, + ckey, + cidx, + keys, + row, + tidx, + warp, + lane, + m, + seqlen, + hist, + cand_key, + cand_idx, + wsum, + state, + mTopk, + mSel, + ): + """The narrow kernel's selection over the candidate list: `KPT` consecutive entries + a thread, so one block-wide scan puts them in block order.""" + idx = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32) + take = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32) + tied = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32) + cute.autovec_copy(ckey[(None, tidx)], keys) + cute.autovec_copy(cidx[(None, tidx)], idx) + for c in cutlass.range_constexpr(KPT): + if tidx * KPT + c >= m: + keys[c] = cutlass.Int32(-1) + self._passes(keys, tidx, warp, lane, hist, wsum, state) + prefix = state[0] + mask = state[1] + need = state[2] + at_thr = state[3] + if m > self.k: + if at_thr <= CAP: + for c in cutlass.range_constexpr(KPT): + if keys[c] >= 0: + if (keys[c] & mask) == prefix: + slot = cute.arch.atomic_add( + state.iterator + 5, cutlass.Int32(1), scope="cta" + ) + cand_key[slot] = keys[c] + cand_idx[slot] = idx[c] + cute.arch.sync_threads() + self._cut(tidx, state, cand_key, cand_idx, at_thr, need) + skip = cutlass.max(at_thr - need, 0) + own = self._mark( + keys, idx, take, tied, m <= self.k, prefix, mask, at_thr, state[7], state[8] + ) + incl = own + for off in (1, 2, 4, 8, 16): + up = cute.arch.shuffle_sync_up(incl, offset=off) + if lane >= off: + incl = incl + up + if lane == 31: + wsum[warp] = incl + cute.arch.sync_threads() + above = cutlass.Int32(0) + if lane < warp: + above = wsum[lane] + pre = cute.arch.warp_reduction_sum(above) + incl - own + self._emit(row, idx, take, tied, pre & 0xFFFF, pre >> 16, skip, seqlen, mTopk, mSel) + + @cute.jit + def _select_window( + self, + keys, + gK, + sK, + row, + tidx, + warp, + lane, + n, + nslabs, + lo, + m, + seqlen, + hist, + cand_key, + cand_idx, + wsum, + wsc, + state, + mTopk, + mSel, + ld_keys: cute.CopyAtom, + ): + """The same selection when the list cannot hold every candidate: each step walks + the window slab by slab, the block-order scan's counts carried from slab to slab (a + slab's fit 16 bits each; the warp totals alternate between two halves of `wsc`).""" + nt = self.threads + nw = self.warps + for _p in cutlass.range_constexpr(_PASSES): + if state[4] == 0: + prefix = state[0] + mask = state[1] + need = state[2] + shift = state[6] + for s in cutlass.range(nslabs): + self._slab(keys, gK, sK, row, s, tidx, n, lo, ld_keys) + for c in cutlass.range_constexpr(KPT): + if keys[c] >= 0: + if (keys[c] & mask) == prefix: + digit = (keys[c] >> shift) & (_HIST - 1) + cute.arch.atomic_add( + hist.iterator + digit, cutlass.Int32(1), scope="cta" + ) + cute.arch.sync_threads() + self._narrow(tidx, warp, lane, hist, wsum, state, prefix, mask, need, shift) + prefix = state[0] + mask = state[1] + need = state[2] + at_thr = state[3] + if m > self.k: + if at_thr <= CAP: + for s in cutlass.range(nslabs): + self._slab(keys, gK, sK, row, s, tidx, n, lo, ld_keys) + e0 = (s * nt + tidx) * KPT + for c in cutlass.range_constexpr(KPT): + if keys[c] >= 0: + if (keys[c] & mask) == prefix: + slot = cute.arch.atomic_add( + state.iterator + 5, cutlass.Int32(1), scope="cta" + ) + cand_key[slot] = keys[c] + cand_idx[slot] = e0 + c + cute.arch.sync_threads() + self._cut(tidx, state, cand_key, cand_idx, at_thr, need) + skip = cutlass.max(at_thr - need, 0) + cut_key = state[7] + cut_idx = state[8] + idx = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32) + take = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32) + tied = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32) + carry_taken = cutlass.Int32(0) + carry_tied = cutlass.Int32(0) + for s in cutlass.range(nslabs): + self._slab(keys, gK, sK, row, s, tidx, n, lo, ld_keys) + e0 = (s * nt + tidx) * KPT + for c in cutlass.range_constexpr(KPT): + idx[c] = e0 + c + own = self._mark( + keys, idx, take, tied, m <= self.k, prefix, mask, at_thr, cut_key, cut_idx + ) + incl = own + for off in (1, 2, 4, 8, 16): + up = cute.arch.shuffle_sync_up(incl, offset=off) + if lane >= off: + incl = incl + up + half = (s & 1) * 32 + if lane == 31: + wsc[half + warp] = incl + cute.arch.sync_threads() + above = cutlass.Int32(0) + if lane < warp: + above = wsc[half + lane] + pre = cute.arch.warp_reduction_sum(above) + incl - own + slab_total = cutlass.Int32(0) + if lane < nw: + slab_total = wsc[half + lane] + slab_total = cute.arch.warp_reduction_sum(slab_total) + self._emit( + row, + idx, + take, + tied, + carry_taken + (pre & 0xFFFF), + carry_tied + (pre >> 16), + skip, + seqlen, + mTopk, + mSel, + ) + carry_taken = carry_taken + (slab_total & 0xFFFF) + carry_tied = carry_tied + (slab_total >> 16) + + @cute.kernel + def kernel( + self, mKeys, mEnd, mPos, mLen, mTopk, mSel, ld_keys: cute.CopyAtom, g2s: cute.CopyAtom + ): + tidx, _, _ = cute.arch.thread_idx() + row, _, _ = cute.arch.block_idx() + warp = cute.arch.warp_idx() + lane = cute.arch.lane_idx() + nt: cutlass.Constexpr = self.threads + nw: cutlass.Constexpr = self.warps + slab: cutlass.Constexpr = self.slab + k: cutlass.Constexpr = self.k + cr: cutlass.Constexpr = self.cr + smem = cutlass.utils.SmemAllocator() + sK = None + if const_expr(self.cache_slabs > 0): + # `KPT` consecutive keys (mode 0) of each thread-slab column (mode 1). + sK = smem.allocate_tensor( + cutlass.Int32, + cute.make_layout((KPT, self.cache_slabs * slab // KPT), stride=(1, KPT)), + byte_alignment=128, + ) + hist = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((_HIST,), stride=(1,)), byte_alignment=128 + ) + cand_key = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((CAP,), stride=(1,)), byte_alignment=128 + ) + cand_idx = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((CAP,), stride=(1,)), byte_alignment=128 + ) + wa = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((32,), stride=(1,)), byte_alignment=128 + ) + wb = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((32,), stride=(1,)), byte_alignment=128 + ) + wc = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((32,), stride=(1,)), byte_alignment=128 + ) + wsc = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((64,), stride=(1,)), byte_alignment=128 + ) + state = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((16,), stride=(1,)), byte_alignment=128 + ) + # The candidates in block order: keys and blocks, `KPT` consecutive a thread. + ckey = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((KPT, nt), stride=(1, KPT)), byte_alignment=128 + ) + cidx = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((KPT, nt), stride=(1, KPT)), byte_alignment=128 + ) + + # 0. Shared state and the row's host-filled metadata, before the wait. + for b in cutlass.range_constexpr(-(-_HIST // nt)): + if tidx + b * nt < _HIST: + hist[tidx + b * nt] = cutlass.Int32(0) + if tidx == 0: + state[5] = cutlass.Int32(0) # candidates compacted + n = mEnd[row] + seqlen = mLen[row] + visible = cutlass.Int32(mPos[row]) + 1 + if const_expr(self.pdl): + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + + # 1. The window's first slabs into shared memory, every load in flight at once. + gK = cute.logical_divide(mKeys, (None, KPT)) + nslabs = (n + (slab - 1)) // slab + if const_expr(self.cache_slabs > 0): + ncached = cutlass.min(nslabs, cutlass.Int32(self.cache_slabs)) + for s in cutlass.range(ncached): + col4 = s * nt + tidx + if col4 * KPT < n: + cute.copy(g2s, gK[row, (None, col4)], sK[(None, col4)]) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + + # 2. The window's largest key, and the prefilter's bound `lo`: each warp's + # `per_warp`-th largest lane maximum (ranked by every lane against all 32, ties to + # the lower lane, so exactly one lane holds each rank), then their minimum. + keys = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32) + tmax = cutlass.Int32(-1) + for s in cutlass.range(nslabs): + self._slab(keys, gK, sK, row, s, tidx, n, cutlass.Int32(0), ld_keys) + for c in cutlass.range_constexpr(KPT): + tmax = cutlass.max(tmax, keys[c]) + kmax = cute.arch.warp_reduction_max(tmax) + if lane == 0: + wa[warp] = kmax + if const_expr(self.prefilter): + ahead = cutlass.Int32(0) + for i in cutlass.range_constexpr(32): + v = cute.arch.shuffle_sync(tmax, offset=i) + if (v > tmax) | ((v == tmax) & (lane > i)): + ahead = ahead + 1 + if ahead == self.per_warp - 1: + wc[warp] = tmax + cute.arch.sync_threads() + gmax = cutlass.Int32(-1) + if lane < nw: + gmax = wa[lane] + gmax = cute.arch.warp_reduction_max(gmax) + # Keys below `lo` are no candidates; all are when the window holds at most k. + lo = cutlass.Int32(0) + if const_expr(self.prefilter): + if n > k: + bound = cutlass.Int32(_I32_MAX) + if lane < nw: + bound = wc[lane] + bound = cute.arch.warp_reduction(bound, lambda x, y: cutlass.min(x, y)) + lo = cutlass.max(bound, cutlass.Int32(0)) + + # 3. The candidates, compacted in block order: a scan of each slab's counts (warp + # totals alternating between two halves of `wsc`, one barrier a slab), carried + # from slab to slab. `m` counts them all, kept or not. + cflat_key = cute.make_tensor(ckey.iterator, cute.make_layout((KPT * nt,), stride=(1,))) + cflat_idx = cute.make_tensor(cidx.iterator, cute.make_layout((KPT * nt,), stride=(1,))) + carry = cutlass.Int32(0) + for s in cutlass.range(nslabs): + self._slab(keys, gK, sK, row, s, tidx, n, lo, ld_keys) + e0 = (s * nt + tidx) * KPT + cnt = cutlass.Int32(0) + for c in cutlass.range_constexpr(KPT): + if keys[c] >= 0: + cnt = cnt + 1 + incl = cnt + for off in (1, 2, 4, 8, 16): + up = cute.arch.shuffle_sync_up(incl, offset=off) + if lane >= off: + incl = incl + up + half = (s & 1) * 32 + if lane == 31: + wsc[half + warp] = incl + cute.arch.sync_threads() + above = cutlass.Int32(0) + if lane < warp: + above = wsc[half + lane] + pos = carry + cute.arch.warp_reduction_sum(above) + incl - cnt + slab_total = cutlass.Int32(0) + if lane < nw: + slab_total = wsc[half + lane] + slab_total = cute.arch.warp_reduction_sum(slab_total) + for c in cutlass.range_constexpr(KPT): + if keys[c] >= 0: + if pos < KPT * nt: + cflat_key[pos] = keys[c] + cflat_idx[pos] = e0 + c + pos = pos + 1 + carry = carry + slab_total + m = carry + cute.arch.sync_threads() + + # 4. The narrow kernel's selection over the list, or over the window when the list + # cannot hold every candidate. + diff = gmax ^ lo + hb = cutlass.Int32(-1) + if diff != 0: + hb = 31 - cutlass.Int32(cute.arch.clz(diff)) + low = (cutlass.Int64(1) << cutlass.Int64(hb + 1)) - 1 + mask0 = cutlass.Int32(cutlass.Int64(_I32_MAX) & (cutlass.Int64(-1) ^ low)) + if tidx == 0: + state[0] = gmax & mask0 # threshold prefix + state[1] = mask0 # its mask + state[2] = cutlass.Int32(k) # still needed at the threshold + state[3] = m # keys in the threshold's bin + state[4] = cutlass.Int32(1) # narrowing done + state[6] = cutlass.max(hb - 10, 0) # the next pass's shift + if m > k: + if hb >= 0: + state[4] = cutlass.Int32(0) + cute.arch.sync_threads() + if m <= KPT * nt: + self._select_listed( + ckey, + cidx, + keys, + row, + tidx, + warp, + lane, + m, + seqlen, + hist, + cand_key, + cand_idx, + wb, + state, + mTopk, + mSel, + ) + else: + self._select_window( + keys, + gK, + sK, + row, + tidx, + warp, + lane, + n, + nslabs, + lo, + m, + seqlen, + hist, + cand_key, + cand_idx, + wa, + wsc, + state, + mTopk, + mSel, + ld_keys, + ) + + # 5. Padding of the block list; the causal tail after the selected tokens, then -1. + total = n + if n > k: + total = cutlass.Int32(k) + for c in cutlass.range_constexpr(-(-k // nt)): + slot = tidx + c * nt + if slot < k: + if slot >= total: + mTopk[row, slot] = cutlass.Int32(-1) + tail_start = (visible // cr) * cr + tail_count = visible - tail_start + body = total * cr + if body > self.token_topk: + body = cutlass.Int32(self.token_topk) + for c in cutlass.range_constexpr(-(-self.final // nt)): + col = tidx + c * nt + if col < self.final: + if col >= body: + off = col - body + val = cutlass.Int32(-1) + if off < cr - 1: + if off < tail_count: + if tail_start + off < seqlen: + val = tail_start + off + mSel[row, col] = val + + +#: One compiled launch per score-table geometry, k, compress ratio, token budget, index +#: dtypes, PDL and wide-kernel configuration; the row count is a runtime argument. +_COMPILED: dict = {} + + +#: SMs of each device, read once. +_SMS: dict = {} + + +def _wide_config(rows: int, device) -> tuple: + """The wide kernel's `(threads, CTAs an SM, prefilter)` for a call of `rows` rows: one + 1024-thread CTA an SM (the largest window cache) while the rows fit one wave, two once + they do not, so one CTA's walks overlap the other's loads.""" + sms = _SMS.get(device.index) + if sms is None: + sms = _SMS[device.index] = torch.cuda.get_device_properties(device).multi_processor_count + return (_WIDE_THREADS, 1 if rows <= sms else 2, True) + + +_INDEX_TYPES = {torch.int32: cutlass.Int32, torch.int64: cutlass.Int64} + + +def launch( + scores, + row_ends, + query_positions, + sequence_lengths, + topk_indices, + compress_ratio, + token_topk, + pdl=True, + out=None, +): + """The launch behind `QsaTopkExpandDecode.__call__`: writes `topk_indices` and returns the + token selection `[rows, token_topk + compress_ratio - 1]` int32 (`out` when given).""" + rows, width = scores.shape + block_topk = topk_indices.shape[1] + sel = out + if sel is None: + sel = torch.empty( + (rows, token_topk + compress_ratio - 1), dtype=torch.int32, device=scores.device + ) + if rows == 0: + return sel + pdl = bool(pdl and get_env_enable_pdl()) + keys = scores.view(torch.int32) + key = ( + width, + scores.stride(0), + block_topk, + compress_ratio, + token_topk, + query_positions.dtype, + sequence_lengths.dtype, + pdl, + None if width <= MAX_WIDTH else _wide_config(rows, scores.device), + ) + ptrs = ( + row_ends.data_ptr(), + query_positions.data_ptr(), + sequence_lengths.data_ptr(), + topk_indices.data_ptr(), + sel.data_ptr(), + ) + fn = _COMPILED.get(key) + if fn is None: + gm = cute.AddressSpace.gmem + geometry = dict( + width=width, + key_stride=scores.stride(0), + block_topk=block_topk, + compress=compress_ratio, + token_topk=token_topk, + pdl=pdl, + ) + if width <= MAX_WIDTH: + kernel = _TopkExpandKernel(**geometry) + else: + threads, ctas, prefilter = key[-1] + kernel = _TopkExpandWideKernel( + **geometry, threads=threads, ctas=ctas, prefilter=prefilter + ) + fn = _COMPILED[key] = cute.compile( + kernel, + from_dlpack(keys, assumed_align=16).mark_layout_dynamic(leading_dim=1), + make_ptr(cutlass.Int32, ptrs[0], gm, assumed_align=4), + make_ptr( + _INDEX_TYPES[query_positions.dtype], + ptrs[1], + gm, + assumed_align=query_positions.element_size(), + ), + make_ptr( + _INDEX_TYPES[sequence_lengths.dtype], + ptrs[2], + gm, + assumed_align=sequence_lengths.element_size(), + ), + make_ptr(cutlass.Int32, ptrs[3], gm, assumed_align=4), + make_ptr(cutlass.Int32, ptrs[4], gm, assumed_align=4), + cutlass.Int32(rows), + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + fn(keys, *ptrs, rows) + return sel + + +def supports(scores) -> bool: + """Whether a score table this wide is one this entry serves (see `is_valid`).""" + return scores.shape[1] % KPT == 0 and scores.shape[1] <= MAX_WIDE_WIDTH + + +#: The default target's score-table widths, `block_table.shape[1] * 32 // 4`: core's `max_seq_len` +#: 9216 (288 pages) and MTP's 16384 (512 pages, `MAX_WIDTH` itself), served by the narrow +#: kernel; the tune's `max_seq_len` 262144 with MTP's extra pages (8196 pages, 65568 columns) +#: and without (65536), served by the wide one. +_CORE_WIDTH = 2304 +_MTP_WIDTH = 4096 +_TUNE_WIDTH = 65568 +_TUNE_WIDTH_NO_MTP = 65536 + + +def _spread(edges: tuple, count: int, first: int, last: int) -> tuple: + """`edges`, then `count - len(edges)` values spread evenly over `[first, last]`.""" + fill = count - len(edges) + return tuple(edges) + tuple(first + (r * (last - first)) // (fill - 1) for r in range(fill)) + + +#: 128 one-row decode requests at core's width. The edges are windows of 0 and 1 blocks, +#: k - 1, k and k + 1 blocks (512, the `n <= k` copy-through and the first narrowing), +#: and the full 2304-block window. +_CORE_DECODE_128 = _spread( + (1, 2, 3, 6, 7, 2046, 2047, 2048, 2051, 2052, 9212, 9214, 9215), 128, 8, 9215 +) +#: MTP7 verify at MTP's width: 128 requests x 8 consecutive rows, one starting where its +#: rows cross k blocks and one ending at the full 4096-block window. +_MTP7_VERIFY_1024 = tuple((start, 8) for start in _spread((1, 2041, 2044, 16376), 128, 9, 16375)) +#: Decode rows at the tune's 65536-wide table, by window: 0 and 1 blocks, k - 1, k and k + 1, +#: one slab (4096 columns) - 1, itself and + 1, the 1024-thread CTA's window cache (11 slabs, +#: 45056 blocks) - 1, itself and + 1, the whole table, and four between. +_WIDE_EDGES_16 = ( + 2, + 3, + 2043, + 2047, + 2051, + 16380, + 16384, + 16388, + 180220, + 180224, + 180228, + 262143, + 50000, + 99999, + 150001, + 220000, +) + + +class QsaTopkExpandDecode(OpWrapper): + """`qsa_topk_select` (zero row starts, `row_ends` the visible block counts), + `qsa_canonical_selection` and `qsa_expand_block_indices`, in one launch: one CTA a row, for + decode rows and a context chunk's rows alike. + + Per row: the `block_topk` largest of `scores[row, :row_ends[row]]` (non-negative fp32), + ties at the threshold to the higher block, written to `topk_indices[row]` in block order + and padded with -1 -- the composition's selection after `qsa_canonical_selection`; the + returned token selection is the expansion entry's: each selected group's members below + `sequence_lengths[row]`, packed at the front, then the incomplete group at the causal + frontier (`query_positions[row] + 1` visible tokens), then -1. `topk_indices` is written + in place, and the selection into `out` when given (a contiguous int32 + `[rows, token_topk + compress_ratio - 1]`, e.g. a row range of the step's whole + selection); nothing else is written. + + Fusion boundary: the selection and the expansion. Scoring is the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # `width` is the score table's column count (and row stride): a compile key, and the + # thread count (`_threads`: 640 at 2560, 576 at 2304, 1024 at 4096). + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode, 4 rows across the budget, one below it", + spec=dict(rows=4, pos=(8700, 4095, 1500, 9900), width=2560), + ), + # The target writes a context chunk's rows through `out` into its rows of the + # step's selection; relu-sparse scores put the zero class (negative zeros among it) + # across the budget, the ties the canonical rule decides. + Cell( + why="context chunk: 2048 rows of one request (positions 6144..8191), relu-sparse " + "scores, written through `out` into rows of a step-sized selection", + spec=dict(rows=2048, chunk=(6144, 2048), relu_sparse=True, out_margin=64, width=2560), + ), + Cell( + why="decode, 1 row of scores within 5% of each other", + spec=dict(rows=1, pos=(8700,), narrow=True, width=2560), + ), + # Serving's decode rows at 8K: positive, log-normally spread over ~2 octaves, a + # zero or two; the first pass leaves tens to ~CAP candidates at the threshold. + Cell( + why="decode, 16 rows of log-normal scores with a few zeros: threshold bins of " + "up to CAP candidates ranked pairwise, ties included", + spec=dict( + rows=16, pos=tuple(8192 + 37 * r for r in range(16)), lognormal=True, width=2560 + ), + ), + # The default target's call sites, at its own widths (`_CORE_WIDTH`, `_MTP_WIDTH`). + Cell( + why="core decode, 1 row at the 2304-wide table, its full window, relu-sparse", + spec=dict(rows=1, pos=(9215,), relu_sparse=True, width=_CORE_WIDTH), + ), + Cell( + why="core decode batch, 128 rows at the 2304-wide table (max_batch_size): windows " + "of 0, 1, k-1, k, k+1 and 2304 blocks and 115 spread between, relu-sparse", + spec=dict(rows=128, pos=_CORE_DECODE_128, relu_sparse=True, width=_CORE_WIDTH), + ), + Cell( + why="MTP7 verify, 1024 rows (128 requests x 8) at the 4096-wide table, up to its " + "full window, relu-sparse", + spec=dict(rows=1024, requests=_MTP7_VERIFY_1024, relu_sparse=True, width=_MTP_WIDTH), + ), + Cell( + why="core context chunk, 2048 rows (positions 6144..8191) at the 2304-wide table, " + "relu-sparse, written through `out`", + spec=dict( + rows=2048, chunk=(6144, 2048), relu_sparse=True, out_margin=64, width=_CORE_WIDTH + ), + ), + Cell( + why="envelope's largest row: an 8192-row chunk (max_num_tokens, positions " + "8192..16383) at the 4096-wide table, relu-sparse, written through `out`", + spec=dict( + rows=8192, chunk=(8192, 8192), relu_sparse=True, out_margin=64, width=_MTP_WIDTH + ), + ), + # The tune's tables (`max_seq_len` 262144): 65568 columns with MTP's extra pages, + # 65536 without, past `MAX_WIDTH`, so the wide kernel's. Relu-sum scores are the + # indexer's form (four heads' ReLUs summed): the prefilter keeps ~5% of a window. + Cell( + why="wide: MTP3 verify, 4 rows (positions 84000..84003) at the 65568-wide table, " + "relu-sum scores", + spec=dict(rows=4, requests=((84000, 4),), relu_sum=True, width=_TUNE_WIDTH, graph=True), + ), + Cell( + why="wide: 16 decode rows at the 65536-wide table, windows at the edges of k, a " + "slab, the window cache and the table, relu-sum scores", + spec=dict( + rows=16, pos=_WIDE_EDGES_16, relu_sum=True, width=_TUNE_WIDTH_NO_MTP, graph=True + ), + ), + Cell( + why="wide: 4 rows of log-normal scores at 65568: threshold bins ranked pairwise, " + "ties included", + spec=dict(rows=4, pos=(100000, 30000, 4095, 2047), lognormal=True, width=_TUNE_WIDTH), + ), + Cell( + why="wide: 2 rows of relu-sparse scores at 65536 (~1.4% positive): the prefilter " + "keeps the zeros, more candidates than the list holds, the whole-window path; at " + "30000 blocks the zero class (negative zeros among it) straddles the budget", + spec=dict(rows=2, pos=(262143, 120000), relu_sparse=True, width=_TUNE_WIDTH_NO_MTP), + ), + Cell( + why="wide: 3 rows of equal scores at 65568 (windows 25000, 513 and 512 blocks): " + "the tie class is the window, the highest blocks taken", + spec=dict(rows=3, pos=(100000, 2051, 2047), equal=True, width=_TUNE_WIDTH), + ), + Cell( + why="wide: ht_dep4's context chunk, 8160 rows (positions 57000..65159) at the " + "65536-wide table, relu-sum, written through `out`", + spec=dict( + rows=8160, + chunk=(57000, 8160), + relu_sum=True, + out_margin=64, + width=_TUNE_WIDTH_NO_MTP, + ), + ), + ) + + note = """ + Metadata consumed: `qsa_visible_blocks` (`row_ends`), the rows' logical positions + and sequence lengths. Preconditions: zero row starts; `scores` fp32 `[rows, W]`, unit + column stride, 16-byte aligned rows, `W % 4 == 0` and `W <= 131072` (up to 4096 the + narrow kernel, four keys a thread and at most 1024 threads; wider tables the wide + kernel, `_TopkExpandWideKernel`); `topk_indices` int32 `[rows, K]` contiguous; `row_ends` int32, + the positions and lengths int32 or int64, each a contiguous `[rows]`. Ties: the higher + block at the threshold, as `qsa_canonical_selection` rewrites the CUDA top-k's pick. + + Compiled once per score-table width and row stride (the block table's, fixed for an + engine), k, compress ratio, token budget, index dtypes and PDL: the row count is a + runtime argument, so an eager step's context rows, any count, reuse the kernel (read + from the launch, not pinned by a test). + + JIT-compiled: `cute.compile` (`--enable-tvm-ffi`) on a key's first call, for the device + that makes it. The cells ran that path on sm_103 (GB300) at widths 2560, 2304 and 4096. + + Certified, every cell bitwise against `reference` (block list and token selection): + - tp1's cells at 2560 columns: decode 1 row of near-equal scores, 4 rows, 16 rows of + log-normal scores (threshold bins ranked pairwise), and a 2048-row context chunk. + - The default target's cells at its own widths (`_CORE_WIDTH` 2304, `_MTP_WIDTH` 4096), with + relu-sparse scores: the zero class, negative zeros among it, straddles the 512-block + budget and is resolved by the tie rule. They are: + - decode, 1 and 128 rows at 2304, with windows of 0, 1, 511, 512, 513 and 2304 blocks among them; + - MTP7 verify, 1024 rows (128 requests x 8) at 4096; + - a 2048-row context chunk at 2304, through `out`; + - an 8192-row chunk (`max_num_tokens`) at 4096, through `out`. + - The tune's tables (`max_seq_len` 262144: 65568 columns with MTP's pages, 65536 + without), the wide kernel: a 4-row MTP3 verify at ~84k tokens and 16 decode rows at + the edges of k, a 4096-column slab, the 45056-block window cache and the whole table + (relu-sum scores); 4 rows of log-normal scores (threshold bins ranked pairwise); 2 + relu-sparse rows (more candidates than the list holds: the whole-window path, the + zero class straddling the budget); 3 rows of equal scores; and ht_dep4's 8160-row + context chunk through `out` (two CTAs an SM past one wave of rows). + - A CUDA-graph replay at 1x2560, 1x2304, 128x2304, 1024x4096, 4x65568, 16x65536 and + 2x65536. The scores are refilled before each replay, and the selection, read after + the wait, must equal the reference. + The widths are derived from `max_seq_len` (288, 512 and 8192 (+4) pages of 32 tokens), + not read off an engine. 4096 is `MAX_WIDTH`: 1024 threads, the largest CTA; wider + tables take the wide kernel, compiled per row-count class (one CTA an SM up to the + SM count, two past it). + + `is_valid` refuses `token_topk != k * compress_ratio`. It sizes the selection's rows + while `k * compress_ratio` places the tokens, so a smaller budget writes past each row. + + Not certified: + - threshold bins ranked pairwise at 2304 and 4096 columns (the log-normal cell is at 2560); + - decode batches other than 1, 4, 16 and 128 rows; + - a NaN inside a window (it ranks above every finite score: the caller keeps the window finite); + - index dtypes other than int64 positions and int32 lengths. + + Index limits, read from the source, not probed. The row count is an Int32 argument and + the row layouts are built from it, so these products must stay below 2^31: + - rows x row stride: 524288 rows at 4096 columns, 16383 at 131072 (`MAX_WIDE_WIDTH`); + - rows x (token_topk + compress_ratio - 1): about 1.05 M rows at 2051. + Positions are read as Int32. The envelope's largest call is 8192 rows. + """ + + def __call__( + self, + scores, + row_ends, + query_positions, + sequence_lengths, + topk_indices, + compress_ratio, + token_topk, + out=None, + ): + return launch( + scores, + row_ends, + query_positions, + sequence_lengths, + topk_indices, + compress_ratio, + token_topk, + out=out, + ) + + def reference( + self, + scores, + row_ends, + query_positions, + sequence_lengths, + topk_indices, + compress_ratio, + token_topk, + out=None, + ): + """`(topk_indices_after, selection)`: a stable descending sort per row (ties to the + higher block), the k first in block order, then the expansion entry's reference.""" + rows, _ = scores.shape + k = topk_indices.shape[1] + out = torch.full_like(topk_indices, -1) + for r in range(rows): + n = int(row_ends[r]) + s = scores[r, :n].float().clamp(min=0.0) + # A stable sort of the reversed row puts the higher block first among equals. + order = (n - 1) - torch.sort(torch.flip(s, [0]), descending=True, stable=True).indices[ + :k + ] + chosen = torch.sort(order).values + out[r, : chosen.numel()] = chosen.to(torch.int32) + sel = _expand.reference( + out, + query_positions, + sequence_lengths, + compress_ratio=compress_ratio, + token_topk=token_topk, + ) + return out, sel + + def is_valid( + self, + scores, + row_ends, + query_positions, + sequence_lengths, + topk_indices, + compress_ratio, + token_topk, + out=None, + ): + assert scores.dtype == torch.float32 and scores.stride(1) == 1, ( + "fp32 scores, unit column stride" + ) + assert scores.shape[1] % KPT == 0 and scores.shape[1] <= MAX_WIDE_WIDTH, ( + "W % 4 == 0, W <= 131072" + ) + assert scores.stride(0) % KPT == 0 and scores.data_ptr() % 16 == 0, "16-byte aligned rows" + assert topk_indices.dtype == torch.int32 and topk_indices.is_contiguous(), ( + "int32 topk_indices" + ) + rows = scores.shape[0] + assert row_ends.dtype == torch.int32 and row_ends.shape == (rows,), "int32 row_ends [rows]" + for t in (query_positions, sequence_lengths): + assert t.dtype in _INDEX_TYPES and t.shape == (rows,), "int32/int64 [rows]" + assert all(t.is_contiguous() for t in (row_ends, query_positions, sequence_lengths)), ( + "contiguous row_ends, positions and lengths" + ) + assert token_topk == topk_indices.shape[1] * compress_ratio, ( + "token_topk = k * compress_ratio" + ) + assert out is None or ( + out.dtype == torch.int32 + and out.is_contiguous() + and out.shape == (scores.shape[0], token_topk + compress_ratio - 1) + ), "out: contiguous int32 [rows, token_topk + compress_ratio - 1]" + + +qsa_topk_expand_decode = QsaTopkExpandDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_topk_select.py similarity index 58% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_topk_select.py index fa7e00d6e456..7abaa54f8115 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qsa_topk_select.py @@ -17,10 +17,29 @@ #: Score-row width the target hands over: a 260-page block table at 32 tokens #: per block, compressed 4:1. _SCORE_COLUMNS = 260 * 32 // 4 +#: The default target's widest score row: the tune envelope's `max_seq_len` +#: 262144 is an 8192-page block table at 32 tokens per block, compressed 4:1. +#: Core's 9216 (288 pages, 2304 columns) and MTP's 16384 (512 pages, 4096 +#: columns) are narrower; the kernel picks its path by each row's window, so +#: the windows below span theirs. +_ENVELOPE_SEQ_LEN = 262144 +_ENVELOPE_COLUMNS = _ENVELOPE_SEQ_LEN // 4 +#: 128 decode requests: the window edges first -- empty (position 2), one +#: block, 511 / 512 / 513 blocks around the budget (512 is the kernel's +#: copy-through shortcut, 513 its first histogram row), core's and MTP's widest +#: rows, the envelope's less one column -- then 120 spread to the envelope. +_DECODE_POSITIONS = (2, 3, 2046, 2047, 2051, 9215, 16383, _ENVELOPE_SEQ_LEN - 2) + tuple( + 4096 + (_ENVELOPE_SEQ_LEN - 1 - 4096) * i // 119 for i in range(120) +) +#: 128 verify requests at max_draft_len 7, 8 consecutive positions each: one +#: from position 0, two straddling the budget, core's and MTP's last rows, the +#: envelope's last (a full 65536-column window), then 122 spread between. +_VERIFY_STARTS = (0, 2040, 2044, 9208, 16376, _ENVELOPE_SEQ_LEN - 8) + tuple( + 4096 + (_ENVELOPE_SEQ_LEN - 16 - 4096) * i // 121 for i in range(122) +) #: Rows from which a contiguous launch runs the CuTe DSL sibling: the CUDA op -#: serves rows at or past 12288 with its radix variant, which is the slower of -#: the two (16384 rows: 310 vs 135 us at k=512 over 2592 columns, GB300); -#: below it the two run alike and the CUDA op needs no JIT compilation. +#: serves rows at or past 12288 with its slower radix variant; below it the CUDA +#: op needs no JIT compilation. _CUTE_DSL_MIN_ROWS = 12288 @@ -57,9 +76,7 @@ class QsaTopkSelect(OpWrapper): Exact, not approximate: the kernel resolves the threshold over the full 32 bits of the fp32 pattern. The **order of the offsets is unspecified** and not stable between identical launches (candidates are appended through a - shared-memory atomic; measured at `[8, 2048]`, `k=512`: four launches, four - permutations of one selected set). Ties are broken arbitrarily within the - tied value. + shared-memory atomic). Ties are broken arbitrarily within the tied value. Fusion boundary: selection only. Scoring, the causal bound the caller encodes in `row_ends`, and the block-to-token expansion are outside. @@ -110,6 +127,52 @@ class QsaTopkSelect(OpWrapper): tail="nan", ), ), + # The default target (tp4 / tep4 / dep4): the selection is per row and + # replicated across ranks, so the member changes only the row count. + Cell( + why="decode: one row at the 262144-token envelope, a 65535-block window, NaN tail", + spec=dict( + rows=1, + cols=_ENVELOPE_COLUMNS, + index_topk=_BLOCK_TOPK, + windows=("positions", (_ENVELOPE_SEQ_LEN - 2,)), + logits="relu", + tail="nan", + ), + ), + Cell( + why="decode batch 128: windows empty, 511 / 512 / 513 blocks and up to the envelope", + spec=dict( + rows=128, + cols=_ENVELOPE_COLUMNS, + index_topk=_BLOCK_TOPK, + windows=("positions", _DECODE_POSITIONS), + logits="relu", + tail="nan", + ), + ), + Cell( + why="verify MTP7 at 128 requests: 1024 rows, windows up to the envelope's 65536", + spec=dict( + rows=1024, + cols=_ENVELOPE_COLUMNS, + index_topk=_BLOCK_TOPK, + windows=("positions", tuple(s + j for s in _VERIFY_STARTS for j in range(8))), + logits="relu", + tail="nan", + ), + ), + Cell( + why="context chunk: max_num_tokens 8192 rows ending a 262144-token prompt", + spec=dict( + rows=8192, + cols=_ENVELOPE_COLUMNS, + index_topk=_BLOCK_TOPK, + windows=("chunk", _ENVELOPE_SEQ_LEN - 8192), + logits="relu", + tail="nan", + ), + ), ) note = """ @@ -120,10 +183,9 @@ class QsaTopkSelect(OpWrapper): is legal and yields an all-`-1` row; a `row_end` past the last column reads out of bounds (not probed). - Logits inside the window are finite. `-inf` is never selected, but a - NaN ranks above every finite value and IS selected: NaN in columns 0-3 - of a 32-wide window with `k=8` returned all four (measured, not pinned - by a test). Outside the window anything goes -- the cells poison it - with NaN. + NaN ranks above every finite value and IS selected (measured, not + pinned by a test). Outside the window anything goes -- the cells + poison it with NaN. - `index_topk` is the BLOCK budget (512 here). The default, 2048, is the token budget, and with a 512-wide buffer it is the layout fault `is_valid` refuses. @@ -136,21 +198,45 @@ class QsaTopkSelect(OpWrapper): The CUDA op dispatches by row index, not column count: rows below 12288 run an insertion-sort variant, rows at or above a radix variant, in one launch. A launch of `_CUTE_DSL_MIN_ROWS` (12288) rows or more with contiguous - logits runs the CuTe DSL sibling instead -- QSA's stock indexer's choice on - this GPU, measured to select identical sets (never identical permutations; - ties, as in the CUDA op, broken arbitrarily). The tuned config's + logits runs the CuTe DSL sibling instead, as QSA's stock indexer does; it + selects identical sets in a different order, ties broken arbitrarily. `max_num_tokens` 16384 reaches it at prefill-heavy mixed steps; smaller launches, every decode one among them, keep the compiled CUDA op, so a - first captured decode JIT-compiles nothing. `torch.ops.trtllm.indexer_topk_decode` is the request-grouped - sibling; QSA does not use it. + first captured decode JIT-compiles nothing. `torch.ops.trtllm.indexer_topk_decode` + is the request-grouped sibling; QSA does not use it. + + The CUDA op picks its path per row by the window's length, never by the + row width: a window of at most `index_topk` copies offsets 0..len-1 + through (unsorted, no histogram); a longer one runs the histogram passes + over the fp32 bits into a 2048-item final buffer. The width enters only as + the row stride, and a row that does not start 16-byte aligned is read with + a scalar head and tail. So the default target's cells, at the 262144-token + envelope's 65536-column rows, take windows from empty through 1, 511, 512, + 513, core's 2304 (`max_seq_len` 9216) and MTP's 4096 (16384) to 65535 and + 65536 -- rows 1, 128 (decode), 1024 (MTP7 verify at 128 requests) and 8192 + (a `max_num_tokens` context chunk). All four run the compiled-in CUDA op + (below 12288 rows); the JIT-compiled CuTe DSL sibling runs at the 16384-row + cell, on sm_103 in the same test file. Attention DP and TP do not change the + op's shape, only its row count. + + Measured, not a cell: a tied class far wider than the final buffer (this + checkpoint's relu scores leave rows with fewer positive blocks than the + budget, the rest exactly 0) still returns 512 distinct in-window offsets + whose values are the window's top 512, at 0 to 3000 positive scores in a + 65536-column window. - CUDA graphs: captured at `[256, 2048]` (32 requests at max_draft_len 7), - `k=512`, and replayed against new logits and a new `row_ends` through the - captured addresses. Bitwise comparison of two - replays is not a valid check (order is unspecified). + CUDA graphs: captured at `[256, 2048]` (32 requests at max_draft_len 7) and + at 1, 128 and 1024 rows of 65536 columns, `k=512`, each over full windows, + and replayed against new logits and a new `row_ends` through the captured + addresses (the 128-row replay at 511 blocks, the copy-through path). + Bitwise comparison of two replays is not a valid check (order is + unspecified). - Hard limits: not probed. `indices` is addressed as `row * index_topk + - slot`; the largest certified extent is 16384 x 512. + Hard limits: `indices` is addressed as `row * index_topk + slot` and the + logits row as `row * stride(0)`, both in 64-bit; `num_columns` and + `stride(0)` pass through int32, so a row past 2^31 - 1 elements is out of + reach (the envelope's is 65536). The largest certified extents are logits + 8192 x 65536 (2^29 elements) and `indices` 16384 x 512. """ def __call__( @@ -200,9 +286,9 @@ def is_valid( index_topk: int = 2048, ) -> None: # The op checks only `indices.size(1) >= index_topk`, then writes a flat - # run of `index_topk`-wide rows and ignores the real row stride. Measured - # with a [4, 12] buffer at k=8: row 1's answer landed across row 0's - # columns 8-11 and its own 0-3, and the last row was never written. + # run of `index_topk`-wide rows and ignores the real row stride, so in a + # wider buffer rows land across their neighbours and the last is never + # written. assert indices.dim() == 2 and indices.shape[1] == index_topk, ( f"indices must be [num_rows, index_topk]; got {tuple(indices.shape)} " f"for index_topk={index_topk}" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qwen4_exp_attn_decode_prologue.py similarity index 97% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qwen4_exp_attn_decode_prologue.py index 4155e8a99b40..a3b705eb80d5 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qwen4_exp_attn_decode_prologue.py @@ -47,14 +47,28 @@ import triton.language as tl from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_index_scores import ( + qsa_paged_index_scores as _index_scores, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_index_store import ( + qsa_paged_index_store as _index_store, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_kv_store import ( + qsa_paged_kv_store as _kv_store, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_kv_store_nvfp4 import ( + _e2m1x2, + _quantize_row, + _rcp_approx, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_kv_store_nvfp4 import ( + qsa_paged_kv_store_nvfp4 as _kv_store_nvfp4, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_prefill_compress import ( + qsa_prefill_compress as _compress, +) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl -from .qsa_paged_index_scores import qsa_paged_index_scores as _index_scores -from .qsa_paged_index_store import qsa_paged_index_store as _index_store -from .qsa_paged_kv_store import qsa_paged_kv_store as _kv_store -from .qsa_paged_kv_store_nvfp4 import _e2m1x2, _quantize_row, _rcp_approx -from .qsa_paged_kv_store_nvfp4 import qsa_paged_kv_store_nvfp4 as _kv_store_nvfp4 -from .qsa_prefill_compress import qsa_prefill_compress as _compress from .qwen4_exp_qk_norm_rope import _launch as _qk_launch from .qwen4_exp_qk_norm_rope import qwen4_exp_qk_norm_rope as _qk_norm_rope diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qwen4_exp_qk_norm_rope.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_qk_norm_rope.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/qwen4_exp_qk_norm_rope.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/spec_rollback_commit.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/spec_rollback_commit.py similarity index 62% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/spec_rollback_commit.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/spec_rollback_commit.py index 098506078e6f..e90cfc7c171a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/spec_rollback_commit.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/spec_rollback_commit.py @@ -2,8 +2,10 @@ # SPDX-License-Identifier: Apache-2.0 """`SpeculativeRollback`'s staging and commit of a verify step in one launch each. -Tier 5 (Triton written for this target). A verify step runs `p` rows per generation request -through the trunk; the PLE pools and the QSA planes are written for every row. The staging +Tier 5 (Triton, called by the qwen4_exp target, `qwen3_8_flash_next_nvfp4__sm_103__default`, +at every layout). Each `__call__` is one launch of a Triton kernel of this file, +JIT-compiled on first use for the device it runs on. A verify step runs `p` rows per generation +request through the trunk; the PLE pools and the QSA planes are written for every row. The staging launch copies what the commit rewrites from: each request's pre-step n-gram id history followed by its rows' ids (the `stream` whose length-`context` windows are the histories after each row), and its pre-step PLE conv state. After the draft loop the commit launch @@ -27,12 +29,22 @@ import triton.language as tl from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -from .qsa_rollback_planes import ROW +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rollback_planes import ( + ROW, +) #: Conv channels a program covers. _BLOCK_C = 256 +#: A tep4 rank (the MTP member) as the stock `Qwen4ExpHybridCacheManager` lays its pools out at +#: core's runtime: the PLE conv state `[slots, hc_count * hidden, (ple_kernel - 1) * ngram_size]` +#: = `[1874, 10240, 9]` bf16 and the n-gram history `[1874, 2]` int64 (replicated, 1874 the final +#: engine's state slots), and every attention layer's index-K plane a `[pages, 32, 1, 128]` bf16 +#: view of one coalesced pool at page stride 49152 elements, 257109 attention pages, beside the +#: `[pages, 32, 3]` int32 position plane. +_TEP4_PLE = dict(state_slots=1874, context=2, channels=10240, state_len=9) +_TEP4 = dict(planes=12, layout="coalesced", pool_pages=257109, **_TEP4_PLE) + @triton.jit def _stage_kernel( @@ -191,10 +203,22 @@ class SpecRollbackStage(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell(why="MTP3, one request", spec=dict(n=1, p=4)), Cell(why="MTP3, 128 requests", spec=dict(n=128, p=4)), + Cell( + why="tep4 MTP1 verify: 128 requests x 2 rows, the 10240 x 9 PLE conv over 1874 slots", + spec=dict(n=128, p=2, **_TEP4_PLE), + ), + Cell( + why="tep4 MTP7 verify: 128 requests x 8 rows, the 10240 x 9 PLE conv over 1874 slots", + spec=dict(n=128, p=8, **_TEP4_PLE), + ), ) note = """ Metadata consumed: none. Call site: `SpeculativeRollback.stage_ple_verify`, before `ple_ngram_embed` publishes the id history. Grid (requests, conv channel blocks). + + The tp1 cells run a 640-channel, 3-column conv state; the default target's (`_TEP4_PLE`) the + checkpoint's 10240 x 9 state at a tep4 rank's measured slot count, `ids` int32 and `slots` + int32 as the call site hands them, slots 0 and 1873 among those read. """ def __call__(self, ngram_pool, conv_pool, slots, ids, p: int): @@ -260,12 +284,34 @@ class SpecRollbackCommit(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell(why="MTP3, one request, 12 planes", spec=dict(n=1, p=4, planes=12)), Cell(why="MTP3, 128 requests, 12 planes", spec=dict(n=128, p=4, planes=12)), + Cell( + why="tep4 MTP1 verify: 128 requests x 2 rows (256), 12 planes of the coalesced pool", + spec=dict(n=128, p=2, **_TEP4), + ), + Cell( + why="tep4 MTP7 verify: 128 requests x 8 rows (1024, the largest verify step), 12 planes", + spec=dict(n=128, p=8, **_TEP4), + ), ) note = """ Metadata consumed: none. Call site: `SpeculativeRollback.commit_speculative_states`, after the draft loop. Grid (requests, max(conv channel blocks, planes)). Slots or plane slots shared by two requests (CUDA-graph padding dummies) are written by both; which one - lands is unordered, as with the indexed writes this replaces. + lands is unordered. + + Two grid regimes. The tp1 cells' 640 conv channels are 3 blocks under 12 planes, so blocks + 3..11 restore planes only. The default target's 10240 channels (`_TEP4`) are 40 blocks over 12 + planes, so blocks 12..39 write conv channels only. The tep4 cells run that layout: + - `accepted` and `slots` are int32, as the MTP worker hands them. + - The planes are the 12 trunk layers' views of the coalesced pool at the measured tep4 + page count and stride, addressed through `plane_table`. + - One request's rows straddle the pool's first and last page, and its rejected row is + restored on the last page, 257108 x 49152 elements (1.26e10) past the plane's base. + - Slots 0 and 1873 are among those written. + + The commit is pure data movement, held bitwise to `reference`. The test also holds that + nothing outside the named slots of any pool is written, and that a CUDA-graph replay at + 1024 rows re-reads its inputs. """ def __call__( @@ -330,8 +376,53 @@ def __call__( num_warps=4, ) - def reference(self, *args) -> None: - raise NotImplementedError("test_modeling_v2_spec_rollback drives the handler end to end") + def reference( + self, + accepted, + slots, + ngram_pool, + stream, + conv_pool, + conv_prev, + conv_x, + position, + position_old, + pages, + within, + planes, + index_old, + ): + """What every slot the commit names holds after it. + + `(ngram [n, context], conv [n, C, state_len], position [n * p, 3], index [planes, + n * p, ROW])`: the n-gram and conv pool rows at `slots`, and the position plane and + each index-K plane at every row's `(pages, within)`. A row the request accepted keeps + what the plane holds now; a rejected one gets `position_old` / `index_old`. + + Takes the planes themselves (`planes`, the views `plane_table` built `plane_ptrs` and + `plane_strides` from) in place of their address table. `slots`, `ngram_pool` and + `conv_pool` only say which rows those are and are not read. Data movement only: the + values are the inputs' own, so it is the exact answer. + """ + n, p, _ = conv_x.shape + context, state_len = ngram_pool.shape[1], conv_pool.shape[2] + a = accepted.to(torch.long) + req = torch.arange(n, device=a.device) + # Window k of the unfold is the id history after k rows. + ngram = stream.unfold(1, context, 1)[req, a] + # [n, C, state_len + p]; window k of the unfold is the conv state after k tokens. + history = torch.cat([conv_prev, conv_x.transpose(1, 2)], 2) + conv = history.unfold(2, state_len, 1)[req, :, a] + rejected = (torch.arange(p, device=a.device)[None, :] >= a[:, None]).reshape(-1) + pg, tok = pages.to(torch.long), within.to(torch.long) + pos = torch.where(rejected[:, None], position_old, position[pg, tok]) + index = torch.stack( + [ + torch.where(rejected[:, None], index_old[k], plane[pg, tok].reshape(-1, ROW)) + for k, plane in enumerate(planes) + ] + ) + return ngram, conv, pos, index def is_valid( self, @@ -350,18 +441,46 @@ def is_valid( plane_strides, index_old, ) -> None: + # The kernel reads every staged tensor by raw offset in the layout its producer writes, + # and the pools by their strides. A wrong shape is read short or past its end, and a + # strided staged tensor at the wrong elements, without an error. n, p, channels = conv_x.shape - assert accepted.shape == (n,) and slots.shape == (n,) - assert ngram_pool.dtype == stream.dtype == torch.int64 and ngram_pool.stride(1) == 1 - assert stream.shape == (n, ngram_pool.shape[1] + p) and stream.is_contiguous() - assert conv_pool.shape[1] == channels and conv_prev.is_contiguous() - assert conv_prev.shape == (n, channels, conv_pool.shape[2]) - assert position.dim() == 3 and position.shape[2] == 3 - assert position_old.shape == (n * p, 3) and position_old.is_contiguous() - assert pages.shape == within.shape == (n * p,) - assert pages.dtype == within.dtype == torch.int64 - assert index_old.is_contiguous() and index_old.shape[1:] == (n * p, ROW) - assert plane_ptrs.numel() == plane_strides.numel() == index_old.shape[0] + assert accepted.shape == (n,) and slots.shape == (n,), ( + f"accepted {tuple(accepted.shape)} and slots {tuple(slots.shape)} must hold one " + f"entry per generation request ({n})" + ) + assert ngram_pool.dtype == stream.dtype == torch.int64 and ngram_pool.stride(1) == 1, ( + "ngram_pool and stream must be int64, the n-gram history's last dim unit-strided" + ) + assert stream.shape == (n, ngram_pool.shape[1] + p) and stream.is_contiguous(), ( + f"stream must be the contiguous [n, context + p] = [{n}, {ngram_pool.shape[1] + p}] " + f"spec_rollback_stage wrote; got {tuple(stream.shape)}" + ) + assert conv_pool.shape[1] == channels and conv_prev.is_contiguous(), ( + "conv_prev must be contiguous (read as packed [n, C, state_len] rows), and conv_pool " + "as wide as conv_x" + ) + assert conv_prev.shape == (n, channels, conv_pool.shape[2]), ( + f"conv_prev is {tuple(conv_prev.shape)}, expected {(n, channels, conv_pool.shape[2])}" + ) + assert position.dim() == 3 and position.shape[2] == 3, ( + f"position must be the [pages, tokens_per_block, 3] plane; got {tuple(position.shape)}" + ) + assert position_old.shape == (n * p, 3) and position_old.is_contiguous(), ( + f"position_old must be the contiguous [n * p, 3] = [{n * p}, 3] gather" + ) + assert pages.shape == within.shape == (n * p,), ( + f"pages and within must hold one entry per verify row ({n * p})" + ) + assert pages.dtype == within.dtype == torch.int64, "pages and within must be int64" + assert index_old.is_contiguous() and index_old.shape[1:] == (n * p, ROW), ( + f"index_old must be the contiguous [planes, n * p, {ROW}] qsa_rollback_gather " + f"returns; got {tuple(index_old.shape)}, contiguous {index_old.is_contiguous()}" + ) + assert plane_ptrs.numel() == plane_strides.numel() == index_old.shape[0], ( + f"the plane table holds {plane_ptrs.numel()} planes and index_old " + f"{index_old.shape[0]}; the kernel restores index_old's count" + ) spec_rollback_stage = SpecRollbackStage() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/verify_step_setup.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/verify_step_setup.py similarity index 98% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/verify_step_setup.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/verify_step_setup.py index bdd150c32d71..edce2a631e91 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/verify_step_setup.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/attention/verify_step_setup.py @@ -25,8 +25,9 @@ import triton.language as tl from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -from .qsa_rollback_planes import ROW +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rollback_planes import ( + ROW, +) #: Rows a program covers. _BLOCK_T = 16 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/embedding/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/embedding/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/ngram_hash.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/embedding/ngram_hash.py similarity index 76% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/ngram_hash.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/embedding/ngram_hash.py index e3841c25f537..420b5f4a9499 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/ngram_hash.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/embedding/ngram_hash.py @@ -2,16 +2,16 @@ # SPDX-License-Identifier: Apache-2.0 """PLE n-gram hash: a token-id history becomes one global embedding row per head. -Tier-3 composed from native torch integer ops (~30 small launches); exact, so -gated bitwise. Nothing registered under ``torch.ops.trtllm`` hashes n-grams. +A macro entry: composed from native torch integer ops (~30 small launches); +exact, so gated bitwise. It serves the general form -- any history length and +`num_query_tokens`, a chunk's tail included. The Triton op +``torch.ops.trtllm.qwen4_exp_ple_ngram_hash`` covers the ``[T, 3]`` window, +last-position form only, and is not wrapped here; see `note`. ``EngramHashProvider.compute_hashes`` in -``tensorrt_llm/_torch/modules/engram/engram.py`` is the same family and a -different function on three axes: it hashes *compressed* ids (a tokenizer -table this checkpoint does not ship), derives its multipliers from a seed -(six ``(seed, layer)`` combinations reproduced none of this checkpoint's -stored ``layer_multipliers``), and pads instead of resetting at every EOS. Its -moduli rule does agree. The built-in Qwen4-Exp implementation has its own, -not wrapped here. +``tensorrt_llm/_torch/modules/engram/engram.py`` is a different function: it +hashes *compressed* ids (a tokenizer table this checkpoint does not ship), +derives its multipliers from a seed, and pads instead of resetting at every +EOS. """ from __future__ import annotations @@ -36,8 +36,8 @@ CKPT_CONTEXT_LEN = len(CKPT_MULTIPLIERS) - 1 -def _rows(batch: int, length: int, query: int, why: str, **spec) -> Cell: - return Cell(why=why, spec=dict(batch=batch, length=length, num_query_tokens=query, **spec)) +def _rows(batch: int, length: int, query: int, **spec) -> dict: + return dict(batch=batch, length=length, num_query_tokens=query, **spec) class NgramHash(OpWrapper): @@ -72,15 +72,35 @@ class NgramHash(OpWrapper): # (default True), `eos_positions` forced to EOS, and `eos_density` of the # remaining ids replaced by EOS. # - # Every cell is the call site's: one `[T, ngram_size]` window per packed - # token from `ngram_token_history`, num_query_tokens 1. + # The call site's form is one `[T, ngram_size]` window per packed token + # from `ngram_token_history`, num_query_tokens 1: T is the step's packed + # row count (requests x rows per request). The chunk-tail cell is the same + # rows hashed per request instead: `ngram_size - 1` carried ids ahead of + # the request's query tokens. CELLS: tuple[Cell, ...] = ( - _rows(4, 3, 1, "decode row whose cached context is EOS", eos_positions=(0, 1)), - _rows( - 32, 3, 1, "call site: 32 windows, decode / MTP verify", lead_eos=False, eos_density=0.2 + Cell( + why="decode row whose cached context is EOS", + spec=_rows(4, 3, 1, eos_positions=(0, 1)), ), - _rows( - 8192, 3, 1, "call site: 8192 windows, largest prefill", lead_eos=False, eos_density=0.05 + Cell( + why="call site: 32 windows, decode / MTP verify", + spec=_rows(32, 3, 1, lead_eos=False, eos_density=0.2), + ), + Cell( + why="call site: 8192 windows, largest prefill", + spec=_rows(8192, 3, 1, lead_eos=False, eos_density=0.05), + ), + Cell( + why="call site: 128 windows, a 128-request decode batch", + spec=_rows(128, 3, 1, lead_eos=False, eos_density=0.05), + ), + Cell( + why="call site: 1024 windows, 128 requests x 8 rows of an MTP-7 verify", + spec=_rows(1024, 3, 1, lead_eos=False, eos_density=0.05), + ), + Cell( + why="chunk tail: 128 requests, 2 carried ids then 8 query tokens each", + spec=_rows(128, CKPT_CONTEXT_LEN + 8, 8, lead_eos=False, eos_density=0.05), ), ) @@ -95,8 +115,11 @@ class NgramHash(OpWrapper): device (raises otherwise; a cpu `layer_multipliers` is accepted and gives the identical result). Certified on an int64 `token_history`, the window `ngram_token_history` returns. Negative ids are accepted and hashed - as ordinary ids, not padding. Longer histories and `num_query_tokens > 1` - (a chunk's tail) are not certified. + as ordinary ids, not padding. A chunk's tail (`num_query_tokens > 1`) is + certified at 8 query tokens behind the 2 carried ids, 128 requests; + longer histories are not certified. `num_query_tokens` above the history + length returns a silently shorter `[batch, history_len, heads]`, so a + `(length 3, num_query_tokens 8)` call is `is_valid`'s to refuse. A single multiplier raises `ZeroDivisionError`. The four silent mistakes are `is_valid`'s. @@ -106,8 +129,19 @@ class NgramHash(OpWrapper): (bitwise equal to slicing after), and chunked prefill with the `ngram_size - 1` carried ids reproduces one-shot bitwise. - CUDA-graph capture certified, bitwise equal to eager. Cost, not gated: - 0.387 ms for 4096 x 16, 0.370 ms for an 8-row decode -- launch-bound. + CUDA-graph capture certified at 128 and 1024 windows (the captured decode + and MTP-7 verify widths), bitwise equal to the reference after refilling + the history. + + Not wrapped: `torch.ops.trtllm.qwen4_exp_ple_ngram_hash` (Triton, + `tensorrt_llm/_torch/modules/qwen4_exp/ple_kernels.py`) on + `(window [T, 3], multipliers, head_vocab_sizes, head_offsets, eos)` -> + `[T, 16]` is bitwise equal to `reference` with num_query_tokens 1 at 1, + 128, 1024 and 8192 windows. It is that form only: a longer history raises + `ValueError`, it takes no `num_query_tokens`, and a negative id or a + product past 2^63 gives a negative row (C remainder) where this + composition's stays inside the table. The target's PLE hashes + inside `ple_ngram_embed`'s own launch and reads only `reference` here. Hard limits token id max_id * max(layer_multipliers) < 2^63 boundary 389114 for diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/bf16_gemv_pdl.py similarity index 94% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/bf16_gemv_pdl.py index c4b93e549701..75652e898bae 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/bf16_gemv_pdl.py @@ -52,8 +52,9 @@ MAX_ROWS = 4 #: (threads, outputs per CTA, elements per vector load) per K. Swept on GB300 at one row, #: graph replay over per-layer weights: K=2560 64 threads x 2 outputs; K=640 (80 16-byte -#: vectors, which no warp multiple divides) 32 threads x 4 outputs on 8-byte vectors. -TILING = {2560: (64, 2, 8), 6144: (128, 2, 8), 640: (32, 4, 4)} +#: vectors, which no warp multiple divides) 32 threads x 4 outputs on 8-byte vectors; +#: K=1536 (a tp4/tep4 rank's o / GDN out projection) as K=2560. +TILING = {2560: (64, 2, 8), 6144: (128, 2, 8), 640: (32, 4, 4), 1536: (64, 2, 8)} #: Outputs per CTA whose weight rows sit in shared memory (beside the `TILING` ones in #: registers). K=6144 swept behind the GDN recurrence and the sparse-attention merge #: (one row, graph replay): 2 + 2 a CTA beat 2 + 0, 2 + 6, 4 + 4 and 1 + 3. @@ -322,14 +323,17 @@ class Bf16GemvPdl(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # Every shape is a qwen3.8-flash-next-nvfp4 tp1 decode projection; the target calls - # it at one row (`_decode_projection`). Four cells: each `TILING` key is its own - # compiled kernel the decode forward reaches, plus the N=1 edge. + # Every shape is a qwen3.8-flash-next-nvfp4 decode projection the target calls it at + # (`_decode_projection`): one row at tp1, up to MAX_ROWS rows of a tp4/tep4 rank's o / + # GDN out projection (K=1536). Each `TILING` key is its own compiled kernel. CELLS: tuple[Cell, ...] = ( _cell("qwen3.8 decode M=1 GDN out / attention o, K=6144 shared-memory tier", 1, 2560, 6144), _cell("qwen3.8 decode M=1 GDN in-proj concat N=16480, L2 prefetch grid", 1, 16480, 2560), _cell("qwen3.8 decode M=1 shared down K=640", 1, 2560, 640), _cell("qwen3.8 decode M=1 shared gate N=1, outputs per CTA halved to 1", 1, 1, 2560), + _cell("tp4 o / GDN out M=4 (MTP-3 verify, 1 request), PDL on", 4, 2560, 1536), + _cell("tp4 o / GDN out M=1 (decode without MTP), PDL on", 1, 2560, 1536), + _cell("tp4 o / GDN out M=4, PDL off", 4, 2560, 1536, pdl=False), ) note = """ @@ -352,14 +356,14 @@ class Bf16GemvPdl(OpWrapper): Hard limits M rows per launch boundary MAX_ROWS=4 driven at 1 (cells) - K a TILING key 2560, 6144, 640 all driven + K a TILING key 2560, 6144, 640, 1536 all driven N any; the outputs per CTA halve until they divide N N=1 driven The test also replays a PDL chain of these GEMVs under a CUDA graph, each reading the one before it. - Not certified: M = 2..4; PDL off (`pdl=False`, or PDL disabled in the - environment), the same kernel without the attribute. + Not certified: M = 2..4 at K other than 1536; PDL disabled in the environment (the + PDL-off cell is the same kernel without the attribute). """ def __call__(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = True) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/fp8_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/fp8_gemv.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/fp8_gemv.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/fp8_gemv.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/fp8_per_tensor_mm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/fp8_per_tensor_mm.py similarity index 99% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/fp8_per_tensor_mm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/fp8_per_tensor_mm.py index 0ec92dde4f25..69106bdd5cd6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/fp8_per_tensor_mm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/fp8_per_tensor_mm.py @@ -156,7 +156,7 @@ class Fp8PerTensorMm(OpWrapper): ) note = """ - Metadata consumed: none. Call sites: the targets' GDN in/out and attention in/o + Metadata consumed: none. Call sites: the target's GDN in/out and attention in/o projections when the dense FP8 mode is on (`TRTLLM_MODELING_V2_DENSE_FP8`). Preconditions: `x` bf16 `[M, K]` with unit column stride, K a multiple of 16; the weight `[K, N]` column-major view of contiguous `[N, K]` e4m3 codes with N a multiple of 16 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/low_m_bf16_mm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/low_m_bf16_mm.py similarity index 85% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/low_m_bf16_mm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/low_m_bf16_mm.py index 967a23116a87..6a6722a355a6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/low_m_bf16_mm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/low_m_bf16_mm.py @@ -11,7 +11,8 @@ most such shapes. Which one, with which tactic, is a per-(N, K) table measured per M bucket on GB300 (cold L2, CUDA-graph replay); every shape the table does not list, and every M above it, stays on the trtllm cuBLASLt op. The rows -today are qwen3.8-flash-next-nvfp4's projections (hidden 2560, tp1). +are qwen3.8-flash-next-nvfp4's decode projections (hidden 2560) at tp1 and at a +tp4/tep4 rank. FlashInfer API: ``flashinfer.gemm.kernels.dense_bf16_gemm_direct.run_direct_dense`` and ``flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk.run_splitk_dense``. @@ -53,13 +54,22 @@ def _s(mma_m: int, mma_n: int, split_k: int, ab_stages: int) -> tuple: #: (N, K) -> ((M upper bound, choice), ...) in ascending M. An M above the last #: bound, and any (N, K) not listed, runs on cuBLAS. Every row below is a -#: qwen3.8-flash-next-nvfp4 tp1 projection; a target with other shapes adds its -#: own rows after measuring them. Each choice beat cublas_mm by >= 5% at the +#: qwen3.8-flash-next-nvfp4 decode projection at tp1 or at a tp4/tep4 rank; other +#: shapes are added after measuring them. Each choice beat cublas_mm by >= 5% at the #: bound's M in a sweep of these shapes (GB300, cold L2, graph replay, pdl off). #: M=3 uses the M=4 choice, and so on. _TABLE: dict[tuple[int, int], tuple[tuple[int, tuple], ...]] = { # GDN in_qkv: 16 q + 16 k heads x 128, 48 v heads x 128 (10240, 2560): ((1, _d(64, 1)), (2, _d(64, 2)), (8, _s(128, 8, 1, 6))), + # GDN in_proj at a tp4/tep4 rank: q|k (4 heads x 128 each), v (12 x 128), z (12 x 128), + # a, b (12 each) concatenated: 4120 columns. + (4120, 2560): ((1, _d(64, 1)), (2, _d(64, 2)), (8, _s(64, 8, 2, 10))), + # Attention input projection at a tp4/tep4 rank: q|gate (6 heads x 256 each), k, v (one + # KV head x 256 each), the replicated index q|k (5 x 128) and one pad head of 128: 4352. + (4352, 2560): ((2, _d(64, 2)), (8, _s(64, 8, 2, 10))), + # Attention o and GDN out at a tp4/tep4 rank (6 q heads x 256 = 12 v heads x 128 = 1536 + # in); M <= 4 go to `bf16_gemv_pdl`, these rows are what this entry runs at the same shape. + (2560, 1536): ((4, _d(64, 2)), (8, _s(64, 8, 3, 10))), # GDN in_z (6144, 2560): ((1, _d(64, 1)), (2, _d(64, 2))), # GDN in_a / in_b (48 v heads) @@ -144,8 +154,8 @@ class LowMBf16Mm(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # Every cell is a qwen3.8-flash-next-nvfp4 tp1 decode projection at M >= 2 (one row - # takes `bf16_gemv_pdl` in the target), one per branch; `kind` pins the branch `plan` + # Every cell is a qwen3.8-flash-next-nvfp4 decode projection, at tp1 (from M=2; one row + # takes `bf16_gemv_pdl` there) or at a tp4/tep4 rank; `kind` pins the branch `plan` # takes for it. The test's host-only plan pins hold the table's other shapes. CELLS: tuple[Cell, ...] = ( _cell( @@ -165,6 +175,21 @@ class LowMBf16Mm(OpWrapper): pdl=True, ), _cell("qwen3.8 decode M=128: cuBLAS branch above M=32", 128, 10240, 2560, "cublas"), + _cell( + "tp4 GDN in_proj M=4 (MTP3 verify, 1 request): split-K, pdl on", + 4, + 4120, + 2560, + "splitk", + True, + ), + _cell("tp4 GDN in_proj M=1: direct", 1, 4120, 2560, "direct", pdl=True), + _cell("tp4 attention in M=1: direct, 2-output tile, pdl on", 1, 4352, 2560, "direct", True), + _cell( + "tp4 o / GDN out M=8 (2 requests): split-K 3-way, pdl on", 8, 2560, 1536, "splitk", True + ), + _cell("router M=1: one decode token, direct 2-row tile", 1, 512, 2560, "direct", True), + _cell("router M=64: cuBLAS branch above M=32", 64, 512, 2560, "cublas", True), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/mid_rows_mm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/mid_rows_mm.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/mid_rows_mm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/mid_rows_mm.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/nvfp4_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/nvfp4_gemv.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/nvfp4_gemv.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/gemm/nvfp4_gemv.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_gated_residual.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_gated_residual.py similarity index 72% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_gated_residual.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_gated_residual.py index 9983a66b1c55..8fe4616fcea7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_gated_residual.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_gated_residual.py @@ -3,15 +3,14 @@ """Qwen4-Exp hyper-connection gating: low-rank stream mixer plus block-injection weights. -Tier-3 composed from native torch (three ``F.linear`` GEMMs and four elementwise -passes); it launches no TensorRT-LLM kernel. The ``mhc_*`` op family -(``mhc_gemm_sqrsum_fma``, ``mhc_big_fuse``, ``mhc_post_mapping``) computes a -structurally different hyper-connection -- a per-stream *scalar* gate from one -linear rather than a rank-320 SiLU bottleneck, a weighted *sum* rather than a -mean, the *raw* stream gated rather than the normalized one -- and its fused -members do not instantiate hidden 2560. The built-in Qwen4-Exp implementation -has fused ops for this (``trtllm::qwen4_exp_hc_*``), not wrapped here; the -semantics follow the HuggingFace ``Qwen4ExpTextGatedResidual`` minus its norm. +A macro entry composed from native torch (three ``F.linear`` GEMMs and four +elementwise passes); it launches no TensorRT-LLM kernel. The ``mhc_*`` op family +computes a different hyper-connection (a per-stream scalar gate, a weighted sum, +the raw stream gated) and does not instantiate hidden 2560. The semantics follow +the HuggingFace ``Qwen4ExpTextGatedResidual`` minus its norm. + +Called by the qwen4_exp target (`qwen3_8_flash_next_nvfp4__sm_103__default`) at every +layout. """ from __future__ import annotations @@ -26,7 +25,6 @@ #: This checkpoint's geometry: 4 streams of 2560, rank-320 mixer bottleneck. CKPT_HC_COUNT = 4 CKPT_HIDDEN = 2560 -CKPT_LOWRANK = 320 def gate_ratio(out: torch.Tensor, ref: torch.Tensor) -> float: @@ -84,6 +82,25 @@ class Qwen4ExpGatedResidual(OpWrapper): why="stack-final mixer: block_inject_weight None", spec=dict(leading=(256,), inject=False), ), + # The MTP layer's shared_head mixer (block_inject_weight None) at the + # rows DecodeTarget.mtp_forward runs: one request, the largest decode + # graph batch, and mtp7's first draft over 128 x (1 + 7) verify rows. + Cell( + why="1 row, no injection: one request's MTP shared_head", + spec=dict(leading=(1,), inject=False), + ), + Cell( + why="128 rows, no injection: MTP shared_head at the largest decode graph batch", + spec=dict(leading=(128,), inject=False), + ), + Cell( + why="1024 rows, no injection: mtp7's first draft over 128 x 8 verify rows", + spec=dict(leading=(1024,), inject=False), + ), + Cell( + why="8192 rows, no injection: the envelope's largest, a max_num_tokens prefill chunk", + spec=dict(leading=(8192,), inject=False), + ), ) note = """ @@ -93,33 +110,33 @@ class Qwen4ExpGatedResidual(OpWrapper): the checkpoint's final norm -- and three more in the MTP layer (two blocks and `shared_head`'s own mixer). - Preconditions, all loud (each measured to raise `RuntimeError` before - producing anything; the width one is tested): the width divisible by `hc_count` (unflatten), the - weights' input width equal to it and their dtype equal to the - activation's (`F.linear`). + Preconditions, all loud (`RuntimeError` before producing anything): the + width divisible by `hc_count` (unflatten), the weights' input width equal + to it and their dtype equal to the activation's (`F.linear`). float32 activations are accepted and not certified: `F.linear` on fp32 - follows torch's TF32 default here, ~10 mantissa bits. A literal bf16 - transcription of the pointwise chain is ~2x this entry's error (0.0449 vs - 0.0227 on `mixed` with the real layer-0 weights). + follows torch's TF32 default here, ~10 mantissa bits. Operand range: with the real layer-0 weights the normed stream reaches |30.1|, the logits 47.8 / 67.5, `inject` spans 0.430 .. 1.578. At up to 4096x that stream every output stays finite, `|mixed| <= max|normed|`, `inject` in `[0, 2]`; agreement with float64 degrades there because the - sigmoid saturates (a property of the function, outside the checkpoint's - range). + sigmoid saturates. Certified: bf16 2-D `[rows, 10240]` at 8 and 8192 rows with the - injection, 256 rows without it (the stack-final mixer); CUDA-graph replay - at 8 rows, bitwise equal to eager. Cost, not gated, shared device: ~235 us - at 1..512 tokens (overhead-bound), 475 us at 4096, 902 us at 8192. - - Not certified: other row counts, more than one leading dim, hc_count other - than 4, hidden other than 2560, lowrank other than 320, float16 and - float32, weight/activation dtype mixing; the operand-range bounds above - are measured. Hard limits: not probed (largest operand 8192 x 10240 = - 2^26.3 elements). + injection; without it (the stack-final mixer, the MTP layer's + shared_head) at 1, 128, 256, 1024 and 8192 rows; CUDA-graph replay at 8 + rows, bitwise equal to eager. On sm_103 the no-injection cells measure + `mixed` within 0.40 of the test's band, cosine >= 0.9999982 and relative + L2 <= 1.88e-3 against float64. + + Not certified: other row counts; `inject` at 128 and 1024 rows (the MTP + layer's unfused attention and MoE sites; `mixed` there is bitwise the + no-injection cells'); more than one leading dim, hc_count other than 4, + hidden other than 2560, lowrank other than 320, float16 and float32, + weight/activation dtype mixing; the operand-range bounds above are + measured. Hard limits: not probed; a 2^31-element operand would take + 209716 rows, past the 8192-row envelope. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_first_cute.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_first_cute.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_rows_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_rows_tc.py similarity index 98% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_rows_tc.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_rows_tc.py index d3756c320875..39c27a557c7b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_rows_tc.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_decode_rows_tc.py @@ -36,10 +36,12 @@ from cutlass.cutlass_dsl import T, dsl_user_op from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_inject_norm import ( + _asm, + _mapa, +) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl -from .qwen4_exp_hc_inject_norm import _asm, _mapa - #: This checkpoint's hyper-connection bottleneck width. RANK = 320 #: Rows a launch serves (one m16 MMA tile); the forward sends 2..MAX_ROWS-row sites here. @@ -704,6 +706,7 @@ def _compiled_first( block_stride, topk=0, n_perm=0, + gate_stride=1, ): kern = _First2( hidden=hidden, @@ -723,7 +726,8 @@ def _compiled_first( toprows = torch.empty(rows, max(topk, 1), dtype=torch.int32, device=dev) inj = torch.empty(rows, hc_count, dtype=bf, device=dev) shared = torch.empty(rows, hidden, dtype=bf, device=dev) - gate = torch.empty(rows, dtype=bf, device=dev) + # The kernel reads the gate logits through this layout: a compile per row stride. + gate = torch.empty(rows, gate_stride, dtype=bf, device=dev)[:, 0] norm = torch.empty(1, width, dtype=bf, device=dev) w = torch.empty(n_wrows, width, dtype=bf, device=dev) hc_out = torch.empty(rows, width, dtype=bf, device=dev) @@ -787,6 +791,9 @@ def first_rows( weight = first.t() assert weight.shape[0] >= NC * NCH top_k = topk_rows.shape[1] if topk_rows is not None else 0 + # One logit a row at any row stride: a column of the shared expert's fused gate|up GEMM + # output is read in place. + gate = gate_logits.reshape(rows) if has_shared else hyper_input.new_empty(rows) fn = _compiled_first( hidden, hc_count, @@ -801,11 +808,11 @@ def first_rows( block_out.stride(0) if has_inject else hidden, top_k, block_out.shape[0] if top_k else 0, + gate.stride(0), ) blk = block_out if has_inject else hyper_input.new_empty(rows, hidden) inj = inject if has_inject else hyper_input.new_empty(rows, hc_count) sh = shared if has_shared else hyper_input.new_empty(rows, hidden) - gate = gate_logits.reshape(rows) if has_shared else hyper_input.new_empty(rows) topw = topk_weights if top_k else hyper_input.new_empty(rows, 1) toprows = topk_rows if top_k else torch.empty(rows, 1, dtype=torch.int32, device=dev) stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream) @@ -1172,6 +1179,9 @@ class Qwen4ExpHcDecodeRowsFirst(OpWrapper): slot order from zero (a -1 slot skipped), rounded once -- bitwise what the runner's finalize stores (as `qwen4_exp_hc_decode_first_cute`'s fold). The fold needs the combine. + `gate_logits` is one logit a row, `[R, 1]` at any row stride: the target passes the column + of the shared expert's fused gate|up GEMM output, read in place (a compile per stride). + Fusion boundary: the routed experts' finalize (optional), the combine, injection, grouped RMS statistics, first projection and activations. """ @@ -1286,6 +1296,11 @@ def is_valid( ) assert (block_out is None) == (inject is None) assert (shared is None) == (gate_logits is None) + if gate_logits is not None: + assert gate_logits.dtype == torch.bfloat16 and gate_logits.numel() == rows, ( + "one bf16 gate logit a row" + ) + assert shared.is_contiguous() and shared.shape == (rows, width // hc_count) assert (topk_weights is None) == (topk_rows is None), "topk_weights and topk_rows together" if topk_rows is not None: assert shared is not None, "the finalize fold needs the combine" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_first_rescale_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_first_rescale_tc.py similarity index 99% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_first_rescale_tc.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_first_rescale_tc.py index 9c4a5dbeadcc..2e21b5c45ddc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_first_rescale_tc.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_first_rescale_tc.py @@ -39,13 +39,20 @@ from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_inject_norm import ( + _bf16_hi, + _bf16_lo, + _bf16x2, + _ldg_v4, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_up_gate_mean import ( + _asm, + _st_global_v4, +) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from tensorrt_llm._torch.kernels.blackwell.custom_pipeline import PipelineTmaUmma, PipelineUmmaAsync from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr -from .qwen4_exp_hc_inject_norm import _bf16_hi, _bf16_lo, _bf16x2, _ldg_v4 -from .qwen4_exp_hc_up_gate_mean import _asm, _st_global_v4 - #: Programmatic dependent launch (sm_90+), off with TRTLLM_ENABLE_PDL=0. _PDL = get_env_enable_pdl() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_inject_norm.py similarity index 84% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_inject_norm.py index 9ddaa39c7e7f..ba8d55b24a46 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_inject_norm.py @@ -6,18 +6,15 @@ block output itself, from the routed experts' output, the shared expert's output and its gate logit (the shared-expert combine). -Tier 5 (Triton written for this target). No ``torch.ops.trtllm.*`` op or -FlashInfer callable accumulates a block output into ``hc_count`` streams with -per-stream gates and normalizes each stream with its own RMS in one launch; the -two-launch composition (``qwen4_exp_post_inject`` + four -``flashinfer_gemma_rmsnorm`` calls and a concat) is what this replaces, and -with ``shared`` / ``gate_logits`` also the ``fused_sigmoid_gate_mul_add`` -launch before it -- and with ``topk_weights`` / ``topk_rows`` the MoE runner's -finalize launch as well. +Tier 5 (Triton). No ``torch.ops.trtllm.*`` op or FlashInfer callable +accumulates a block output into ``hc_count`` streams with per-stream gates and +normalizes each stream with its own RMS in one launch. It computes what +``qwen4_exp_post_inject`` followed by a per-stream ``flashinfer_gemma_rmsnorm`` +computes; with ``shared`` / ``gate_logits`` also ``fused_sigmoid_gate_mul_add`` +before it, and with ``topk_weights`` / ``topk_rows`` the MoE runner's finalize. The decode MoE sites (the finalize fold below `_FEW_WARPS_MIN_ROWS` rows) run a -CuTe DSL kernel instead (Tier 6, written for this target), bitwise the Triton -kernel's outputs there. +CuTe DSL kernel instead (Tier 6), bitwise the Triton kernel's outputs there. """ from __future__ import annotations @@ -43,8 +40,7 @@ _PDL_TRIGGER = get_env_enable_pdl() #: Rows from which a wide (`hidden >= 2048`) launch runs two warps a program -#: instead of eight: prefill-sized grids are bandwidth-bound and run faster -#: with fewer, longer threads; decode-sized ones keep eight. +#: instead of eight (bandwidth-bound prefill-sized grids). _FEW_WARPS_MIN_ROWS = 512 @@ -113,7 +109,7 @@ def _hc_inject_norm_kernel( b = tl.load(block_ptr + row * stride_block + cols, mask=mask, other=0.0).to(tl.float32) if HAS_SHARED: # The MoE block output, `routed + sigmoid(gate) * shared` in fp32, rounded - # once to the routed output's dtype -- the value that block used to store. + # once to the routed output's dtype. sg = tl.sigmoid(tl.load(gate_ptr + row * stride_gate).to(tl.float32)) sh = tl.load(shared_ptr + row * stride_shared + cols, mask=mask, other=0.0) b = (b + sg * sh.to(tl.float32)).to(block_ptr.dtype.element_ty).to(tl.float32) @@ -138,15 +134,15 @@ def _hc_inject_norm_kernel( # -------------------------------------------------------------------------------------- # The decode MoE sites in CuTe DSL. There the Triton kernel runs two programs a row, each -# gathering all `k` permuted FC2 rows and normalizing its two streams one after the other -# (6.5 us at 128 rows in the con128 decode step on GB300). Here the row's two CTAs are a -# cluster: CTA j forms the block output's columns [j * H / 2, (j + 1) * H / 2) -- the -# finalize's slot-order FMA chain and the combine, the Triton kernel's instructions -- and -# stores them into both CTAs' shared memory; after a cluster barrier each CTA reads the -# whole block row locally and accumulates and normalizes its two streams together. A CTA -# is 256 threads in the Triton kernel's layout at a 4096 block and eight warps (thread t: -# columns [8t, 8t + 8) and, below the hidden, [2048 + 8t, 2048 + 8t + 8)), and the sum of -# squares is its reduction tree: `hc_out` and `normed` are bitwise the Triton kernel's. +# gathering all `k` permuted FC2 rows and normalizing its two streams one after the +# other. Here the row's two CTAs are a cluster: CTA j forms the block output's columns +# [j * H / 2, (j + 1) * H / 2) -- the finalize's slot-order FMA chain and the combine, +# the Triton kernel's instructions -- and stores them into both CTAs' shared memory; +# after a cluster barrier each CTA reads the whole block row locally and accumulates and +# normalizes its two streams together. A CTA is 256 threads in the Triton kernel's layout +# at a 4096 block and eight warps (thread t: columns [8t, 8t + 8) and, below the hidden, +# [2048 + 8t, 2048 + 8t + 8)), and the sum of squares is its reduction tree: `hc_out` and +# `normed` are bitwise the Triton kernel's. # -------------------------------------------------------------------------------------- _CUTE_THREADS = 256 @@ -771,22 +767,25 @@ class Qwen4ExpHcInjectNorm(OpWrapper): routing weights. The fold needs the combine (`shared`, `gate_logits`). Fusion boundary: the routed experts' finalize (optional), the - shared-expert combine (optional), the stream-broadcast - gated accumulate (the old `qwen4_exp_post_inject`) and the grouped gemma - RMS norm of the next site (the old `_grouped_norm`). The routed and shared - experts, the gate projection and the mixer that consumes `normed` are not - inside. + shared-expert combine (optional), the stream-broadcast gated accumulate + (`qwen4_exp_post_inject`'s) and the grouped gemma RMS norm of the next + site. The routed and shared experts, the gate projection and the mixer + that consumes `normed` are not inside. """ ARCHS = frozenset({Arch.SM_103}) # spec: rows, inject (True), moe (False: shared + gate_logits [T, 1] too), fold - # (0: top-k of the finalize folded in), drop (False: some slots -1). The - # checkpoint's geometry throughout: hc_count 4, hidden 2560, bf16, contiguous rows. - # Four cells: the forward reaches four distinct launches -- the Triton kernel's + # (0: top-k of the finalize folded in), drop (False: some slots -1), site + # ("trunk": the norm weight a trunk site's; "stack_final": the stack-final mixer's). + # The checkpoint's geometry throughout: hc_count 4, hidden 2560, bf16, contiguous rows. + # The first four reach the target's four distinct launches at tp1 -- the Triton kernel's # eight-warp form (decode), its two-warp form (prefill, here with the combine), # the CuTe fold kernel (decode MoE sites; the -1 slots ride it) and the norm-only - # variant -- none of which another covers. + # variant -- none of which another covers. The rest are the other layouts': one + # form at every row count, its MoE sites the combine over the routed rows already + # finalized and reduced across ranks (no fold), so the Triton kernel with the + # combine at decode sizes, and every form from 1 row up to max_num_tokens. CELLS: tuple[Cell, ...] = ( Cell( why="checkpoint 4x2560 bf16, inject, decode batch 128 (eight-warp form)", @@ -804,11 +803,47 @@ class Qwen4ExpHcInjectNorm(OpWrapper): why="norm only (no inject: stack start / after PLE), decode 128", spec=dict(rows=128, inject=False), ), + Cell( + why="1 row, attention site: inject of the reduced block output, eight-warp form", + spec=dict(rows=1, moe=False, inject=True), + ), + Cell( + why="1 row, MoE site: combine over the reduced routed rows, no fold", + spec=dict(rows=1, moe=True, inject=True), + ), + Cell( + why="32 rows, MoE site: combine over the reduced routed rows, no fold", + spec=dict(rows=32, moe=True, inject=True), + ), + Cell( + why="128 rows, MoE site: combine without the fold (the con-128 decode batch)", + spec=dict(rows=128, moe=True, inject=True), + ), + Cell( + why="1024 rows, attention site: two-warp form (an MTP-7 verify of 128 requests)", + spec=dict(rows=1024, moe=False, inject=True), + ), + Cell( + why="8192 rows, MoE site: combine, two-warp form (max_num_tokens)", + spec=dict(rows=8192, moe=True, inject=True), + ), + Cell( + why="1 row, stack-final site: last MoE combine, the final mixer's norm weight", + spec=dict(rows=1, moe=True, inject=True, site="stack_final"), + ), + Cell( + why="1 row, norm only: the stack's first site at decode", + spec=dict(rows=1, moe=False, inject=False), + ), + Cell( + why="8192 rows, norm only: the stack's first site at max_num_tokens (two-warp)", + spec=dict(rows=8192, moe=False, inject=False), + ), ) note = """ Metadata consumed: none; every (row, stream) independent. Call sites: in - the trunk forward of both targets, once after each block (attention / GDN + the trunk forward of the target, once after each block (attention / GDN and MoE) of each layer, normalizing for the next site -- the layer's MoE mixer, the next layer's attention mixer, or the stack-final mixer -- plus norm-only calls at the stack's first site and after the PLE add. 97 calls @@ -818,13 +853,28 @@ class Qwen4ExpHcInjectNorm(OpWrapper): `topk_rows`: the runner's finalize folded in). The MTP layer and the PLE block keep the unfused ops. + In the default (tp4 / tep4 / dep4) target HC is replicated and runs on the + rank's own rows, one form at every row count (1 to max_num_tokens 8192, an + MTP-7 verify 1024): the attention site injects the block output after its + all-reduce (none at dep4); the MoE site passes the routed rows already + finalized and all-reduced / reduce-scattered with this rank's replicated + shared expert output and gate logit -- the combine, never the fold; the + last layer's MoE site normalizes with the stack-final mixer's weight + (`hyper_connection_mixer.hc_norm`); the stack's first site is norm only. + The MTP draft layer's sites are the trunk's too, at a draft step's rows (1 + to max_num_tokens): norm only over the trunk's residual (`pre_h_norm`, a + column view of the trunk's packed output on the first draft step) and + before the shared head's mixer; the attention site with the front end's + `x_h + x_e` as its injection at unit weight (`inject` all ones: bitwise + the bf16 add); the MoE-side site after the output projection's all-reduce. + Preconditions: 2-D operands, last-dim stride 1 (row strides free); the width divisible by `hc_count`; `inject` `[T, hc_count]`, `block_out` and `shared` `[T, hidden]` of one dtype (with the fold `block_out` is `[P, - hidden]`, any `P`, read through int32 `topk_rows` `[T, k]`); `gate_logits` `T` elements along its - first axis (`[T, 1]` or `[T]`, any row stride); `norm_weight` contiguous - `[width]`. All on one device. `shared` and `gate_logits` come together and - only with `block_out`. + hidden]`, any `P`, read through int32 `topk_rows` `[T, k]`); `gate_logits` + `T` elements along its first axis (`[T, 1]` or `[T]`, any row stride); + `norm_weight` contiguous `[width]`. All on one device. `shared` and + `gate_logits` come together and only with `block_out`. Numerics: the combine is fp32 with one rounding to `block_out`'s dtype (bitwise `fused_sigmoid_gate_mul_add`: the same `tl.sigmoid` and fma); the @@ -839,9 +889,10 @@ class Qwen4ExpHcInjectNorm(OpWrapper): Launch: one program per (row, stream), stream-minor in a 1-D grid, so a row's `block_out` / `shared` are read from HBM once and from L2 by the other `hc_count - 1` programs; with the fold, one program per two streams - of a row (the gathered rows are read and summed half as often). Every program releases programmatic - dependents at its start (`griddepcontrol.launch_dependents`; off with - TRTLLM_ENABLE_PDL=0): a consumer launched with PDL (the decode mixer's + of a row (the gathered rows are read and summed half as often). Every + program releases programmatic dependents at its start + (`griddepcontrol.launch_dependents`; off with TRTLLM_ENABLE_PDL=0): a + consumer launched with PDL (the decode mixer's `qwen4_exp_hc_decode_tc_first`) streams its weights meanwhile and waits for this grid before reading `normed`. A no-op for any other consumer. At a hidden of 2048 or more a program runs eight warps below @@ -855,19 +906,35 @@ class Qwen4ExpHcInjectNorm(OpWrapper): finalize and the combine) into both CTAs' shared memory, then each accumulating and normalizing two streams, in the Triton kernel's layout and reduction tree -- bitwise its `hc_out` and `normed` (tested at 511 rows, the - last before the switch). GB300, graph replay, - 16 rotating sites, us: 128 rows 7.00 -> 5.97, 64 rows 6.63 -> 5.25, 8 rows - 3.59 -> 3.09; it releases dependents at its start as the Triton kernel does. - - Certified: bf16, the checkpoint's 4 x 2560: the injection at 128 rows - (eight-warp form), the combine at 16384 rows (two-warp form), the combine - with the fold (three slots -1) at 128 rows (CuTe kernel), norm-only at 128 - rows; the CuTe fold bitwise the Triton kernel at 511 rows. + last before the switch). It releases dependents at its start as the Triton + kernel does. + + Certified on sm_103 (GB300), bf16, the checkpoint's 4 x 2560, PDL trigger + on (the default): the Triton kernel and the CuTe fold kernel are JIT- + compiled at first call for the running GPU, and the cells ran that path. + tp1: the injection at 128 rows (eight-warp form), the combine at 16384 + rows (two-warp form), the combine with the fold (three slots -1) at 128 + rows (CuTe kernel), norm-only at 128 rows; the CuTe fold bitwise the + Triton kernel at 511 rows. + default: the injection at 1 row (eight-warp) and 1024 rows (two-warp); + the combine without the fold at 1, 32 and 128 rows (eight-warp) and + 8192 rows (two-warp); the combine at 1 row under the stack-final + weight (RMS 3.1, -0.37 to 12.4, so `1 + w` up to 13.4); norm-only at 1 + and 8192 rows (two-warp). + Measured against the float64 reference rounded once: the decode cells (1, + 32 rows) bitwise or one ulp on 4 of 327680 `normed` elements; at 128 rows + one `hc_out` element and 5 `normed` elements of 1310720 differ; at 1024 + rows `hc_out` bitwise, 74 of 10485760 `normed` one ulp; norm-only at 8192 + rows 602 of 83886080 one ulp. The combine at 8192 / 16384 rows moves 131 / + 192 `hc_out` elements, each inside one ulp of `hc_out` plus one ulp of the + combine times `|inject|` (worst 0.95 of that band; up to 31 ulp of a + near-zero residual), and `normed` within one ulp of the norm of the + stored `hc_out`. Relative L2 at most 1.8e-5, cosine 1.0 on both outputs at + every cell. Not certified: hc_count other than 4, hidden other than 2560, float16 operands, a flat `[T]` gate, row strides other than the width, zero rows, - the injection without the combine at prefill sizes, the combine without the - fold at decode sizes, CUDA-graph replay (the kernels take them; no cell + TRTLLM_ENABLE_PDL=0, CUDA-graph replay (the kernels take them; no cell drives them). Hard limits: row offsets are int64 (`program_id(0)` widened); the grid is @@ -980,11 +1047,9 @@ def reference( topk_rows: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """float64 throughout. The finalize (with `topk_rows`) and the combine - (with `shared`) are each rounded once to `block_out`'s dtype, as the op - states; `hc_out` is the float64 - accumulate rounded once; `normed` is the float64 norm of that rounded - residual, rounded once -- the same dataflow the op states, with no - other intermediate rounding.""" + (with `shared`) are each rounded once to `block_out`'s dtype; `hc_out` + is the float64 accumulate rounded once; `normed` is the float64 norm of + that rounded residual, rounded once. No other intermediate rounding.""" dtype = hyper_input.dtype rows, width = hyper_input.shape hidden = width // hc_count diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_rstd.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_inject_rstd.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_rstd.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_inject_rstd.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_lowrank_act.py similarity index 63% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_lowrank_act.py index 493ca58301c4..bb5368c55c33 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_lowrank_act.py @@ -4,10 +4,8 @@ the rank-`R` bottleneck and the block-injection weights, from the first projection's fp32 logits, in one Triton row kernel. -Tier 5 (Triton written for this target). Used where the first projection runs on -cuBLAS (prefill, and decode batches past `qwen4_exp_hc_decode_tc_first`'s rows), -between `cublas_mm` (first projection) and the fused up projection -`qwen4_exp_hc_up_gate_mean` (or its NVFP4 variant). +Called by the qwen4_exp target (`qwen3_8_flash_next_nvfp4__sm_103__default`) at every +layout. """ from __future__ import annotations @@ -68,23 +66,59 @@ class Qwen4ExpHcLowrankAct(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # spec: rows, inject (True), hc_count (4), rank (320), logits_width (384), dtype (bf16). + # The default target's HC mixers run this form at every row count, on the fresh + # fp32 `cublas_mm` output over the first `hc_first_cols` = ceil((320 + 4) / 32) * 32 + # = 352 columns of the first projection: hence the 352-wide cells. CELLS: tuple[Cell, ...] = ( Cell(why="checkpoint rank 320 + 4, padded to 384, 128 rows", spec=dict(rows=128)), Cell(why="checkpoint rank 320 + 4, padded to 384, prefill 8192", spec=dict(rows=8192)), Cell(why="stack-final mixer: no inject, 8192 rows", spec=dict(rows=8192, inject=False)), + Cell( + why="default decode HC mixer: one token, 352-wide logits", + spec=dict(rows=1, logits_width=352), + ), + Cell( + why="default dep4 rank's largest captured decode batch: 32 rows", + spec=dict(rows=32, logits_width=352), + ), + Cell( + why="default tep4 MTP-7 verify: 128 requests x 8 rows = 1024", + spec=dict(rows=1024, logits_width=352), + ), + Cell( + why="default prefill chunk at max_num_tokens: 8192 rows, 352-wide", + spec=dict(rows=8192, logits_width=352), + ), + Cell( + why="default stack-final mixer (no inject): one decode token", + spec=dict(rows=1, inject=False, logits_width=352), + ), ) note = """ - Metadata consumed: none. Call sites: the trunk's hyper-connection sites in - `PrefillTarget` (all 97), between the first projection and the up - projection. Preconditions: `logits` fp32, unit column stride, at least - `rank (+ hc_count)` columns; `hc_count` a power of two. Graph-capturable. - - Certified: bf16 outputs from 384-wide fp32 logits, with the injection at 128 - and 8192 rows, without it at 8192 rows; graph replay. - - Not certified: zero rows, hc_count other than 4, rank other than 320, - float16 outputs. + Metadata consumed: none. Call sites: every hyper-connection mixer, between + the first projection (`cublas_mm`, fp32 out, 352 columns) and the up + projection. HC is replicated, so the rows are this rank's tokens at every + member: decode 1..128 (tp4 / tep4) or 1..32 (a dep4 rank), MTP verify up to + 128 x 8 = 1024 rows (tep4, draft 7), prefill one chunk of up to 8192 rows + (`max_num_tokens`). Preconditions: `logits` fp32, unit column stride, at + least `rank (+ hc_count)` columns; `hc_count` a power of two. + Graph-capturable. + + Source: a Triton kernel JIT-compiled at first call for the running device; + the cells ran on sm_103. One program per row with an int64 row offset; + 2^31 elements is 6.1e6 rows of 352, far past the envelope. + + Certified: bf16 outputs from fp32 logits 352 wide (the call site) at 1, 32, + 1024 and 8192 rows with the injection and at 1 row without it, and 384 wide + at 128 and 8192 rows (without the injection at 8192); graph replay. + Against the float64 reference: bitwise at 1, 32 and 128 rows; at 1024 and + 8192 rows at most 0.99 ulp on a few `low` elements, relative L2 <= 1.6e-5. + + Not certified: zero rows; rows past 8192; hc_count other than 4; rank + other than 320; float16 outputs. The kernel has no row-count switch, so row + counts between the cells take the same path, and the logits width enters + only as the row stride (128 rows is driven at 384 wide only). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_up_gate_mean.py similarity index 91% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_up_gate_mean.py index 755fc8535e08..e4a4561ecd0c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_up_gate_mean.py @@ -5,14 +5,14 @@ epilogue applies the sigmoid gate to the normed stream and averages the streams, so the `[T, C*H]` gate logits never reach global memory. -Tier 6 (CuTe DSL written for this target). A persistent warp-specialized -kernel: one TMA warp loads `low` and the up weight, one MMA warp issues -`tcgen05.mma` into a double-buffered TMEM accumulator, four epilogue warps -read the accumulator, load the matching normed elements straight from global -memory and write the mean. An MMA tile is 128 rows by `C * BH` up-projection -outputs: the `BH`-wide hidden slice `[h0, h0 + BH)` of every stream, read -through a strided view of the unmodified `[C*H, R]` weight (no copy, no -re-layout), so the stream reduction stays inside one thread. +A persistent warp-specialized CuTe DSL kernel: one TMA warp loads `low` and +the up weight, one MMA warp issues `tcgen05.mma` into a double-buffered TMEM +accumulator, four epilogue warps read the accumulator, load the matching +normed elements straight from global memory and write the mean. An MMA tile +is 128 rows by `C * BH` up-projection outputs: the `BH`-wide hidden slice +`[h0, h0 + BH)` of every stream, read through a strided view of the unmodified +`[C*H, R]` weight (no copy, no re-layout), so the stream reduction stays inside +one thread. """ from __future__ import annotations @@ -313,7 +313,7 @@ def __call__( ) # Rastered along N: consecutive CTAs take consecutive hidden slices of the # same rows, so the normed reads in flight cover whole rows rather than - # one 64-byte column block of many rows (~2% faster at 8K-16K rows). + # one 64-byte column block of many rows. sched = utils.PersistentTileSchedulerParams( ((m + 127) // 128, nt, 1), (*self.cluster_shape_mn, 1), 1, False ) @@ -811,8 +811,7 @@ def hc_up_gate_mean_fp4( return mixed, xq, xsf -# Hidden columns per tile. 32 (MMA N = 4 * 32 = 128, five A/B stages) measured -# equal to 64 (N = 256, three stages) at 8K-16K rows and faster at 32..2048 rows. +# Hidden columns per tile: MMA N = 4 * 32 = 128, five A/B stages. _BLOCK_H = 32 @@ -832,37 +831,61 @@ class Qwen4ExpHcUpGateMean(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # spec: rows, hc_count (4), hidden (2560), rank (320); bf16 only. - CELLS: tuple[Cell, ...] = tuple( - Cell(why=f"checkpoint 4x2560, rank 320, {why}", spec=dict(rows=t)) - for t, why in ( - (128, "decode 128 rows, one full tile"), - (16384, "prefill chunk 16384 (max_num_tokens)"), + CELLS: tuple[Cell, ...] = ( + tuple( + Cell(why=f"checkpoint 4x2560, rank 320, {why}", spec=dict(rows=t)) + for t, why in ( + (128, "decode 128 rows, one full tile"), + (16384, "prefill chunk 16384 (max_num_tokens)"), + ) + ) + + (Cell(why="ragged tail, 16320 + 70 rows (mixed step)", spec=dict(rows=16390)),) + # The tp4 / tep4 / dep4 target: the mixer runs on the full 4 x 2560 + # residual on every rank, so only the row count moves per member. + + tuple( + Cell(why=why, spec=dict(rows=t)) + for t, why in ( + (1, "1 row: one decode token, the smallest captured batch; 1 of a tile's 128 rows"), + (32, "32 rows: dep4's per-rank decode cap (128 / 4), a quarter of one tile"), + (1024, "1024 rows: mtp7 verify at tep4, 128 requests x 8 rows"), + (8192, "8192 rows: one prefill chunk at max_num_tokens, the envelope's largest"), + ) ) - ) + (Cell(why="ragged tail, 16320 + 70 rows (mixed step)", spec=dict(rows=16390)),) + ) note = """ - Tier 6 (CuTe DSL written for this target): a persistent tcgen05 GEMM whose - epilogue applies the gate, so the `[T, C*H]` gate logits never reach HBM; - tiers 1-5 have no GEMM with a custom multi-input epilogue (cuBLAS writes - the logits). Metadata consumed: none. Call sites: - the trunk's hyper-connection mixers (`_hc_mixer`). Preconditions (bf16): - `low` `[T, R]` contiguous, `up` the transposed view of a contiguous - `[C*H, R]` weight, `normed` `[T, C*H]` contiguous, 16-byte aligned - pointers, `R % 64 == 0`, `H % 32 == 0`, `C * 32 <= 256`. Compiles once per - `(block_h, C, H, R)` on first call (~1 s); the row count is a runtime - argument. Graph-capturable (launches on the current stream, no host sync). - Launched with programmatic dependent launch (TRTLLM_ENABLE_PDL=0 turns it - off): each CTA's first tile streams its up-weight stages before - `griddepcontrol.wait` and `low` / `normed` after it, so behind a producer - that releases dependents early (`qwen4_exp_hc_decode_tc_first`) the weight - read overlaps the producer. - - Certified: bf16 at the checkpoint's 4 x 2560, rank 320: 128 rows (one full - tile), 16384 rows, 16390 rows (a partial last tile); graph replay. - - Not certified: zero rows, fewer rows than a tile, hc_count other than 4, - hidden other than 2560, rank other than 320 (the kernel compiles for them; - no cell drives them). + Source: the CuTe DSL kernel in this module (no trtllm op or FlashInfer call + fuses this multi-input epilogue). It is JIT-compiled by `cute.compile` on + the first call per `(block_h, C, H, R)` for the running device (sm_103 in + the cells, no `CUTE_DSL_ARCH` override); the row count is a runtime + argument. Metadata consumed: none. Call sites: the hyper-connection mixers + at decode, MTP verify (4 rows per request) and prefill. Under tp4, tep4 and + dep4 every rank mixes the full 4 x 2560 residual, so only the row count + moves per member. Preconditions (bf16): `low` `[T, R]` contiguous, `up` the + transposed view of a contiguous `[C*H, R]` weight, `normed` `[T, C*H]` + contiguous, 16-byte aligned pointers, `R % 64 == 0`, `H % 32 == 0`, + `C * 32 <= 256`. Graph-capturable (launches on the current stream, no host + sync). Launched with programmatic dependent launch (TRTLLM_ENABLE_PDL=0 + turns it off): each CTA's first tile streams its up-weight stages before + `griddepcontrol.wait` and reads `low` / `normed` after it. + + Certified: bf16 at the checkpoint's 4 x 2560, rank 320, at 1 and 32 rows + (part of one tile), 128 (one full tile), 1024, 8192, 16384 and 16390 rows + (a partial last tile); graph replay. + + Measured on GB300 against the float64 reference, with the band + `rtol = eps`, `atol = eps * max|ref|` (`eps` the bf16 epsilon): worst + element at most 0.40 of the band, relative L2 <= 8.1e-5, cosine + >= 0.999999997. The band does not separate a bf16-rounded gate logit + (0.28x to 0.46x), which is a legitimate two-launch result. + + Not certified: + - rows that no cell drives: zero, 2-31, 33-127 (the same one-tile path + as 1 and 32), between the full-tile cells, and above 16390; + - hc_count other than 4, hidden other than 2560, rank other than 320 (the + kernel compiles for them); + - `hc_up_gate_mean_fp4`, which this module also defines: no cell here + drives it. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_rstd.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_up_gate_mean_rstd.py similarity index 99% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_rstd.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_up_gate_mean_rstd.py index 69fcca405ae8..a42b236a8f13 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_rstd.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_hc_up_gate_mean_rstd.py @@ -43,13 +43,20 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( fp4_quantize, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_inject_norm import ( + _bf16x2, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_up_gate_mean import ( + _asm, + _rcp_approx_ftz, + _st_global_b16, + _st_global_v4, +) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from tensorrt_llm._torch.kernels.blackwell.custom_pipeline import PipelineTmaUmma, PipelineUmmaAsync from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr from .qwen4_exp_hc_first_rescale_tc import _lds_v4, _sts_v4 -from .qwen4_exp_hc_inject_norm import _bf16x2 -from .qwen4_exp_hc_up_gate_mean import _asm, _rcp_approx_ftz, _st_global_b16, _st_global_v4 #: Programmatic dependent launch (sm_90+), off with TRTLLM_ENABLE_PDL=0. _PDL = get_env_enable_pdl() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_post_inject.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_post_inject.py similarity index 71% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_post_inject.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_post_inject.py index d51ddde1e610..eadadb7e44ba 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_post_inject.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/hc/qwen4_exp_post_inject.py @@ -3,12 +3,11 @@ """Hyper-connection post-block injection: a per-stream gated accumulate of one block output into the multi-stream residual. -Tier-3 composed from native torch: one kernel-launching call -(``torch.addcmul``) plus view glue. No ``torch.ops.trtllm.*`` op computes this -accumulate with an *identity* cross-stream mix: ``mhc_post_mapping`` applies a -Sinkhorn doubly-stochastic mix across streams, fixes the stream count at 4 and -instantiates its MMA paths for hidden sizes this checkpoint does not use. The -built-in Qwen4-Exp implementation has its own fused op for it, not wrapped here. +Native torch (sourcing tier 4): one kernel-launching call (``torch.addcmul``) +plus view glue. No ``torch.ops.trtllm.*`` op computes this accumulate from the +gates with an *identity* cross-stream mix: ``mhc_post_mapping`` applies a +Sinkhorn mix across streams, and ``trtllm::qwen4_exp_hc_combine`` takes the +injection *logits* rather than the gates ``qwen4_exp_gated_residual`` returns. """ from __future__ import annotations @@ -59,6 +58,23 @@ class Qwen4ExpPostInject(OpWrapper): spec=dict(leading=(8,), hc_count=CKPT_HC_COUNT, hidden=CKPT_HIDDEN, layout="col_slice"), ), ) + # The tp4/tep4/dep4 layouts' sites, each with `hyper_input` a column + # slice of a wider row buffer (row stride twice the packed width). The op + # sees a rank's full 4 x 2560 residual at every member; only the row count + # moves: one decode request, max_batch_size, a verify step's trunk and a + # max_num_tokens prefill chunk. + CELLS += tuple( + Cell( + why=f"col_slice {t} rows: {why}", + spec=dict(leading=(t,), hc_count=CKPT_HC_COUNT, hidden=CKPT_HIDDEN, layout="col_slice"), + ) + for t, why in ( + (1, "one decode request; the slice reads as contiguous at one row"), + (128, "a max_batch_size decode step"), + (1024, "a verify step's trunk, 128 requests x verify width 8"), + (8192, "a max_num_tokens prefill chunk, the envelope's largest row"), + ) + ) note = """ Metadata consumed: none. Call sites: twice per decoder layer (after the @@ -69,32 +85,31 @@ class Qwen4ExpPostInject(OpWrapper): Preconditions, all loud except the guarded one: the packed width must equal `inject.shape[-1] * block_out.shape[-1]` (the unflatten raises); all three on one CUDA device (raises). The leading dims must be **equal**, not merely - broadcastable -- a 1-D `inject` or `block_out` was measured to be accepted - and to broadcast one row over every token (`is_valid`). + broadcastable -- a 1-D `inject` or `block_out` is accepted and broadcasts + one row over every token (`is_valid`). Mixed dtypes are not rejected, they promote: an fp32 `inject` beside a bf16 stream returns an fp32 result -- correct values, twice the bytes. Recorded, not guarded. Rounding: `addcmul` forms product and sum in fp32 opmath and rounds once - on the store. Measured bit-exact against an fp32 reference rounded once at - 1..8192 tokens; the bf16-multiply-then-bf16-add form differs by 0.0625. - Wrong operators at `[64, 10240]` bf16, max deviation from the bit-exact - result: stream-minor packing 10.34, no gate 3.77, gate on the residual 5.89. + on the store; bit-exact against an fp32 reference rounded once at 1..8192 + tokens. CUDA-graph capture certified; the output is a fresh allocation, so under - capture it comes from the graph's private pool. Cost, not gated: 17 us at - 64 tokens, 53 us at 2048, 213 us at 8192 (~1.78 TB/s) -- the stride-0 - broadcast operands take a non-vectorized path. An in-place form was not - measured. + capture it comes from the graph's private pool. Certified: bf16 2-D operands at the checkpoint's 4 x 2560, 64 and 8192 - rows contiguous, 8 rows with `hyper_input` a column slice; graph replay. + rows contiguous; `hyper_input` a column slice (`buf[:, :width]` of a + 2 x width row buffer) at 1, 8, 128, 1024 and 8192 rows; graph replay. At + one row the slice is contiguous by torch's definition (a size-1 dim's + stride is ignored), so that cell is the call site's layout, not a strided + read. Not certified: other row counts, a row-slice `hyper_input`, hc_count other than 4, hidden other than 2560, float16, float32 and float8 operands, - leading dims beyond one. Hard limits: not - probed (the largest measured operand is 8192 x 10240 = 2^26.3 elements). + leading dims beyond one. Hard limits: not probed; 2^31 elements of the + packed residual is 209,715 rows, past the target's max_num_tokens 8192. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/cute_dsl_nvfp4_grouped_gemm.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/cute_dsl_nvfp4_grouped_gemm.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_permuted.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_permuted.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_routed.py similarity index 70% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_routed.py index 0377db559291..f6257976bc31 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_block_scale_moe_renorm_routed.py @@ -3,7 +3,13 @@ """NVFP4 MoE layer entered with router logits: the trtllm-gen runner's own Renormalize routing (top-k over the raw logits, softmax over those k) + grouped FC1 GEMM + SwiGLU + NVFP4 requantization + grouped FC2 GEMM + routing-weighted -combine, in one trtllm-gen call.""" +combine, in one trtllm-gen call. + +It runs over an expert window (`local_expert_offset`, `local_num_experts`; the whole routing +space by default) at the intermediate width a rank holds: the window of a tep4 / dep4 / dep2 +rank, every expert at a tp4 rank's padded slice, every expert at tp1. A decode call over the +whole 512-expert layer at I 640 runs its pinned tactic (`DECODE_TACTICS`), every other call the +AutoTuner's.""" from __future__ import annotations @@ -16,13 +22,14 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( # noqa: E501 fp4_block_scale_moe_runner as _pre_routed, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.renorm_moe_routing_op import ( + renorm_moe_routing_op as _renorm, +) from tensorrt_llm._torch.autotuner import AutoTuner from tensorrt_llm._torch.custom_ops.trtllm_gen_custom_ops import FP4BlockScaleMoERunner from tensorrt_llm._torch.utils import last_positive_power_of_2 from tensorrt_llm.logger import logger -from .renorm_moe_routing_op import renorm_moe_routing_op as _renorm - #: `RoutingMethodType::Renormalize`: NoOp preprocess, top-k, softmax over the k. ROUTING_RENORMALIZE = 1 #: `act_type` SwiGlu. @@ -131,15 +138,30 @@ def run_renorm_moe( top_k: int, intermediate_size: int, do_finalize: bool, + local_expert_offset: int = 0, + local_num_experts: int | None = None, ) -> list[torch.Tensor]: - """One logits-routed trtllm-gen FP4 MoE call (Renormalize, SwiGlu, no expert - parallelism): `[out]` with `do_finalize`, else `[rows, topk_weights, topk_rows]`. A - `DECODE_TACTICS` bucket runs its pinned tactic; any other call is the AutoTuner's.""" - tactic = _pinned_tactic( - hidden_states.shape[0], num_experts, top_k, hidden_states.shape[1] * 2, intermediate_size + """One logits-routed trtllm-gen FP4 MoE call (Renormalize, SwiGlu) over the expert window + `[local_expert_offset, local_expert_offset + local_num_experts)` (all experts when + `local_num_experts` is None): `[out]` with `do_finalize`, else `[rows, topk_weights, + topk_rows]`. A `DECODE_TACTICS` bucket of the whole layer runs its pinned tactic; any other + call is the AutoTuner's.""" + if local_num_experts is None: + local_num_experts = num_experts + whole = local_expert_offset == 0 and local_num_experts == num_experts + tactic = ( + _pinned_tactic( + hidden_states.shape[0], + num_experts, + top_k, + hidden_states.shape[1] * 2, + intermediate_size, + ) + if whole + else None ) if tactic is None: - return torch.ops.trtllm.fp4_block_scale_moe_runner( + return _pre_routed( routing_logits, None, hidden_states, @@ -161,8 +183,8 @@ def run_renorm_moe( None, None, intermediate_size, - 0, - num_experts, + local_expert_offset, + local_num_experts, None, ROUTING_RENORMALIZE, do_finalize, @@ -215,20 +237,36 @@ def _ckpt(tokens: int, dist: str = "randn") -> dict: ) +def _layer(tokens: int, off: int, local: int, inter: int = 640) -> dict: + """qwen4_exp's routed geometry (512 experts, top-10, H 2560, bf16 logits) at one rank's + window `[off, off + local)` and intermediate width `inter`.""" + return dict( + tokens=tokens, + hidden=2560, + intermediate=inter, + num_experts=512, + top_k=10, + local_expert_offset=off, + local_num_experts=local, + logits_dtype=torch.bfloat16, + dist="randn", + ) + + class Fp4BlockScaleMoeRenormRouted(OpWrapper): - """One NVFP4 MoE layer whose routing happens inside the call. - - `routing_logits` is the router GEMM's `[T, num_experts]` output. Per token - the kernel picks the `top_k` largest logits (ties to the LOWER expert index) - and weights them by an fp32 softmax over just those `top_k`, rounded to - bf16 -- the same function as `ops/moe/renorm_moe_routing_op` followed by the - pre-routed `catalog/moe/fp4_block_scale_moe_runner`, and bitwise equal to - that composition at every cell. Returns the combined `[T, H]` bf16 rows. - - Fixed, not arguments: no routing bias, Renormalize routing, SwiGlu, no - expert parallelism (`local_num_experts == num_experts`), `do_finalize=True`. - The weight layout (row shuffle + 128x4 scale swizzle) and the three - per-expert scalars are exactly the pre-routed entry's. + """One NVFP4 MoE layer whose routing happens inside the call, over an expert window. + + `routing_logits` is the router GEMM's `[T, num_experts]` output over the whole routing + space. Per token the kernel picks the `top_k` largest logits (ties to the LOWER expert + index) and weights them by an fp32 softmax over just those `top_k`, rounded to bf16; + slots routed outside `[local_expert_offset, local_expert_offset + local_num_experts)` add + nothing. The same function as `ops/moe/renorm_moe_routing_op` followed by the pre-routed + `catalog/moe/fp4_block_scale_moe_runner` over the same window, and bitwise equal to that + composition at every cell. Returns the combined `[T, H]` bf16 rows (the window's partial). + + Fixed, not arguments: no routing bias, Renormalize routing, SwiGlu, `do_finalize=True`. + The weight layout (row shuffle + 128x4 scale swizzle), the window's weight stacks and the + three per-expert scalars are exactly the pre-routed entry's. """ ARCHS = frozenset({Arch.SM_103}) @@ -242,6 +280,32 @@ class Fp4BlockScaleMoeRenormRouted(OpWrapper): why="16384 rows: max_num_tokens of the tuned config; past the 8192 tuning bucket", spec=_ckpt(16384), ), + # The windows: tep4 / dep4 a 128-expert window at I 640, dep2 a 256-expert window, tp4 + # all 512 experts at its 160-wide slice padded to 256. Rows: 1 (decode), 4 (an MTP-3 + # verify step), 5 (past the block routing kernel), 128 (dep4's 4 x 32 gathered rows), + # 1024 (MTP-7 verify of 128 requests) and 8192 (a prefill chunk); offsets cover every + # rank. + Cell(why="tep4 decode: one token, rank 0's 128-expert window", spec=_layer(1, 0, 128)), + Cell(why="tep4 MTP-3 verify at con1: 4 rows, rank 1's window", spec=_layer(4, 128, 128)), + Cell(why="ep4 at 5 rows, rank 2's window: past the block kernel", spec=_layer(5, 256, 128)), + Cell( + why="dep4 gathered decode 4 x 32 = 128 rows, rank 3's window", + spec=_layer(128, 384, 128), + ), + Cell(why="tep4 MTP-7 verify of 128 requests: 1024 rows", spec=_layer(1024, 0, 128)), + Cell(why="ep4 prefill chunk at 8192 rows, rank 3's window", spec=_layer(8192, 384, 128)), + Cell(why="tp4 MTP-3 verify: 4 rows, 512 experts at I 256", spec=_layer(4, 0, 512, 256)), + Cell(why="tp4 at 128 rows, 512 experts at I 256", spec=_layer(128, 0, 512, 256)), + Cell( + why="tp4 prefill chunk 8192 rows, 512 experts at I 256", spec=_layer(8192, 0, 512, 256) + ), + Cell( + why="dep2 gathered decode: 2 rows, rank 1's 256-expert window", spec=_layer(2, 256, 256) + ), + Cell( + why="dep2 prefill chunk at 8192 rows, rank 0's 256-expert window", + spec=_layer(8192, 0, 256), + ), ) note = """ @@ -274,7 +338,14 @@ class Fp4BlockScaleMoeRenormRouted(OpWrapper): warm in L2 under uniform routing, takes the rest and the warmup's tuning pass. The test holds each pinned tactic bitwise to the AutoTuner path. + Windows: a window's weight stacks are its own experts only (the op compares the scales' + expert dimension with `local_num_experts`; the whole stack beside a window's scales would be + read as its first experts, refused by `is_valid`). Not certified: windows of the padded + I 256 stack (tp4 does not split experts); window rows past 8192 (the caller chunks). + Hard limits + window the whole 512 at I 640 (tp1), 128 of 512 at I 640 (every offset), 256 of + 512 at I 640, 512 at I 256 T certified at 1, 128 and 16384 (max_num_tokens of the tuned config). Autotuner buckets stop at 8192, the op's `tune_max_num_tokens` default, which this entry does not @@ -297,6 +368,8 @@ def __call__( num_experts: int, top_k: int, intermediate_size: int, + local_expert_offset: int = 0, + local_num_experts: int | None = None, ) -> torch.Tensor: return run_renorm_moe( routing_logits, @@ -313,6 +386,8 @@ def __call__( top_k, intermediate_size, True, + local_expert_offset, + local_num_experts, )[0] def reference( @@ -324,8 +399,10 @@ def reference( gate: tuple, down: tuple, fc2_global_scale: float, + local_expert_offset: int = 0, + local_num_experts: int | None = None, ) -> torch.Tensor: - """The routing, then the layer, both by their own entries' references. + """The routing, then the window's layer, both by their own entries' references. Routing: `renorm_moe_routing_op.reference` (stable descending sort in the logits' dtype, fp32 softmax over the k, bf16 weights -- the dtype the @@ -335,7 +412,17 @@ def reference( NVFP4 activation). Returns fp32 `[T, H]`. """ ids, weights = _renorm.reference(routing_logits, top_k, torch.bfloat16) - return _pre_routed.reference(hidden_states, ids, weights, up, gate, down, fc2_global_scale) + return _pre_routed.reference( + hidden_states, + ids, + weights, + up, + gate, + down, + fc2_global_scale, + local_expert_offset=local_expert_offset, + local_num_experts=local_num_experts, + ) def is_valid( self, @@ -352,6 +439,8 @@ def is_valid( num_experts: int, top_k: int, intermediate_size: int, + local_expert_offset: int = 0, + local_num_experts: int | None = None, ) -> None: # Raw data pointers, dense row-major assumed: a strided view is read as # if dense. A router GEMM output written into a wider buffer is the @@ -383,6 +472,13 @@ def is_valid( assert gemm1_weights.shape[-1] * 2 % 256 == 0, ( "hidden size (gemm1_weights.shape[-1] * 2) must be a multiple of 256" ) + # The op checks the scales' expert dimension, not the weights'. + local = num_experts if local_num_experts is None else local_num_experts + for name, tensor in (("gemm1_weights", gemm1_weights), ("gemm2_weights", gemm2_weights)): + assert tensor.shape[0] == local, ( + f"{name} holds {tensor.shape[0]} experts but the window has {local}; pass the " + "window's own weights" + ) fp4_block_scale_moe_renorm_routed = Fp4BlockScaleMoeRenormRouted() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_moe_renorm_cute_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_moe_renorm_cute_dsl.py similarity index 97% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_moe_renorm_cute_dsl.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_moe_renorm_cute_dsl.py index f7bc335a4e99..1d7cf6041e65 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_moe_renorm_cute_dsl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp4_moe_renorm_cute_dsl.py @@ -15,13 +15,18 @@ import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* import tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops as _cdo from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_block_scale_moe_renorm_routed import ( + CERTIFIED_PAIRS, + _ckpt, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_block_scale_moe_renorm_routed import ( + fp4_block_scale_moe_renorm_routed as _routed, +) from tensorrt_llm._torch.utils import ActivationType from ...kernels.moe.blockscaled_contiguous_grouped_gemm_finalize_fusion import ( Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel, ) -from .fp4_block_scale_moe_renorm_routed import CERTIFIED_PAIRS, _ckpt -from .fp4_block_scale_moe_renorm_routed import fp4_block_scale_moe_renorm_routed as _routed #: Token tile of the grouped GEMMs (the kernels' M tile). TILE_TOKENS = 128 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp8_block_scale_moe_runner.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp8_block_scale_moe_runner.py similarity index 72% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp8_block_scale_moe_runner.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp8_block_scale_moe_runner.py index 2bb80ff63ecb..392908f43983 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp8_block_scale_moe_runner.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fp8_block_scale_moe_runner.py @@ -2,7 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 """DeepSeek-style FP8 block-scale mixture-of-experts layer: given top-k + grouped FC1 GEMM + SwiGLU + FP8 requantization + grouped FC2 GEMM + -routing-weighted combine, in one trtllm-gen call.""" +routing-weighted combine, in one trtllm-gen call. + +Called by the qwen4_exp target's MTP layer at every layout.""" from __future__ import annotations @@ -67,6 +69,12 @@ def _geom(hidden, inter, num_experts, top_k, offset=0, local=None): _CKPT = _geom(2560, 640, 512, 10) +def _window(offset: int) -> dict: + """One rank's 128-expert window of the 512 (the qwen4 default target at tp4, + tep4 and dep4): the router still ranks all 512, so `num_experts` stays 512.""" + return _geom(2560, 640, 512, 10, offset=offset, local=128) + + class Fp8BlockScaleMoeRunner(OpWrapper): """One FP8 block-scale MoE layer over e4m3 activations, entered pre-routed. @@ -115,13 +123,41 @@ class Fp8BlockScaleMoeRunner(OpWrapper): why="16384 rows: max_num_tokens of the tuned config; past the last tuning bucket", spec=dict(tokens=16384, **_CKPT), ), + # The qwen4 default target: every member runs the MTP experts as four + # 128-wide windows, one per rank -- at tp4 too, where 640 / 4 = 160 is not + # a multiple of 128 and the intermediate cannot be split. One cell per + # window offset, each at a row count the MTP layer reaches. + Cell( + why="default MTP drafter decode: one token, rank 0's 128-expert window", + spec=dict(tokens=1, **_window(0)), + ), + Cell( + why="default: 128 drafter rows, the largest captured batch; rank 1's window", + spec=dict(tokens=128, **_window(128)), + ), + Cell( + why="default: 1024 rows, the first draft over mtp7's 128 x 8 verify rows; window 2", + spec=dict(tokens=1024, **_window(256)), + ), + Cell( + why="default: 8192 rows, a max_num_tokens context chunk; the last window (384)", + spec=dict(tokens=8192, **_window(384)), + ), ) note = """ Metadata consumed: none. Two process-global caches sit behind the call -- the autotuner's profiling cache and the C++ runner cache -- and a warm tuned tactic is bitwise equal to the cold fallback (the test holds this at - 256 tokens). + 256 tokens, for the whole stack and for a 128-expert window, which tunes + under its own `local_num_experts` key). + + The op is a Python `torch.library.custom_op` + (`trtllm::fp8_block_scale_moe_runner` in + `_torch/custom_ops/trtllm_gen_custom_ops.py`) over the compiled-in C++ class + `torch.classes.trtllm.FP8BlockScaleMoERunner` + (`cpp/tensorrt_llm/thop/moe/trtllmGen/fp8BlockScaleMoe.cpp`), which launches + trtllm-gen's precompiled batched-GEMM cubins: nothing is JIT-compiled. Preconditions. `hidden_states` is `[T, H]` e4m3 and `hidden_states_scale` `[H/128, T]` fp32 (see the class docstring); weights are the checkpoint's @@ -135,13 +171,31 @@ class Fp8BlockScaleMoeRunner(OpWrapper): Supply distinct ids per row: a repeated id counted once at 16 tokens and twice at 256 (the token count picks the routing kernel). - The window: local slot i answers for global id `local_expert_offset + i`. - Nothing checks the window against `num_experts`; a token with no slot in it - comes back bitwise zero, and a window starting at `num_experts` returns all - zeros. Ids outside the routing space (negative, `>= num_experts`) were - measured to be dropped; the router never emits one, and no cell certifies it. - Four 128-wide windows of 512 summed match the full call inside the test's band - but not bitwise (each rounds its partial to bf16). + The window: local slot i answers for global id `local_expert_offset + i`, + and `num_experts` stays the routing space (512) on every window call. The + weights and all four scale tensors hold the window's experts only: the op + compares each one's dim 0 with `local_num_experts` and raises on the whole + stack (the test holds all four). A token with no slot in the window comes + back bitwise zero (every window cell holds this). + + The qwen4 default target's four 128-wide windows of 512 are certified at the + row counts its MTP layer reaches, one window offset per cell, on sm_103: 1 + to 8192 rows, at most 7.95 bf16 ulp per element of the row max and 0.227 + ulp RMS, inside the test's 16 / 0.75 gate. Each cell asserts some slot routes + into its window. The four windows' bf16 outputs added in fp32 -- what the + target's all-reduce or reduce-scatter adds -- sit at most 4.41 / 0.297 ulp + from the 512-expert reference at 1, 128 and 1024 rows, and the sum with any + one window left out fails the gate. Not bitwise to the 512-expert call (each window rounds its + partial to bf16): at most 1.97 / 0.187 ulp from it. + + `num_experts` passed as the window's own count (128) puts windows 1..3 past + the routing space: at small row counts the op answers all zeros with no + error, and at 16 rows it reads out of bounds (an illegal memory access) -- + ids `>= num_experts` are dropped only at small row counts. `is_valid` + refuses a window that starts outside `[0, num_experts)`; window 0 of that + mistake still passes it, and ids past `num_experts` cannot be seen without + reading the ids. A window that starts inside the routing space and runs + past its end ([448, 576) of 512) computes its in-range experts correctly. Inert on the certified path, bitwise: `routing_method_type` 0/1/2/4 (the pre-routed path forces the routing postprocess), `n_group`/`topk_group` @@ -150,10 +204,11 @@ class Fp8BlockScaleMoeRunner(OpWrapper): applied after the FC1 requantization; None is bitwise equal to omitting it. Silent on the raw op, refused by `is_valid` (the test holds the strided - `hidden_states` refusal): any non-contiguous tensor (read as dense, 55-115 - ulp RMS off); `act_type != 0` (the Python runner forces - SwiGlu, so 1 and 2 are bitwise identical to 0); `routing_bias` alongside - `topk_ids` (ignored; a +-1e3 bias leaves the output bitwise unchanged). + `hidden_states` and the outside-the-routing-space refusals): any + non-contiguous tensor (read as dense, 55-115 ulp RMS off); `act_type != 0` + (the Python runner forces SwiGlu, so 1 and 2 are bitwise identical to 0); + `routing_bias` alongside `topk_ids` (ignored; a +-1e3 bias leaves the output + bitwise unchanged); a `local_expert_offset` outside `[0, num_experts)`. `num_fused_shared_experts > 0` is rejected loudly on the pre-routed path, which is the only path the C++ runner's small-tile tactic exclusion guards. @@ -161,21 +216,26 @@ class Fp8BlockScaleMoeRunner(OpWrapper): states): an upstream defect, eager is unaffected. Hard limits - T certified at 1, 20 and 16384 (max_num_tokens of the tuned config); - chunk above it, not probed. Autotuner buckets stop at + T whole stack certified at 1, 20 and 16384 (max_num_tokens of the + tuned config); 128-expert windows at 1, 128, 1024 and 8192 + (the default target's max_num_tokens chunk). Chunk above + them, not probed. Autotuner buckets stop at `tune_max_num_tokens` (default 8192, which the target leaves): past it a call runs the 8192 bucket's tactic; a larger value makes the op append a bucket at that size top_k <= 32 (raises past it); certified 10, the checkpoint's E, H, I certified at the checkpoint's 512, 2560, 640 - - Not certified (the statements above were measured at bring-up and are not - re-driven by the test): expert windows (`local_expert_offset` / - `local_num_experts` other than the full stack), repeated or out-of-range - ids, `gemm1_clamp_limit`, the inert arguments, `output=`, the registered - fake, the op's loud rejections. Nor: the in-kernel routing path - (`routing_logits` given), multi-rank execution, geometries other than the - checkpoint's, sm_100. + window the whole stack, and 128 wide at offsets 0, 128, 256 and 384 of 512 + + Not certified (measured, but not re-driven by the test): windows of other + widths (a 256-wide dep2 window, say) or offsets, a window past the routing + space's end, repeated or + out-of-range ids, `gemm1_clamp_limit`, the inert arguments, `output=`, the + registered fake, the op's other loud rejections. Nor: CUDA-graph capture + of the call (the cells run eager), the in-kernel routing path + (`routing_logits` given), multi-rank execution (every window here runs on + one device; the cross-rank sum is the collective's), geometries other than + the checkpoint's, sm_100. """ def __call__( @@ -325,9 +385,9 @@ def is_valid( tune_max_num_tokens: int = 8192, use_dp: bool = False, ) -> None: - """Three inputs the op accepts and answers wrongly, each measured. + """Four inputs the op accepts and answers wrongly, each measured. - The transposed activation-scale buffer at `T == H/128` is the fourth, + The transposed activation-scale buffer at `T == H/128` is the fifth, and no metadata check can see it. """ # The kernel takes raw data pointers: a strided view is read as dense. @@ -360,6 +420,15 @@ def is_valid( "routing_bias is ignored when topk_ids/topk_weights are given " "(routing has already happened); fold it into the router logits" ) + # `num_experts` is the routing space, not the window: a window that + # starts at or past it holds no routable id. Passing the window's own + # count as `num_experts` puts ranks 1..3 there: the op answers all zeros + # with no error at small row counts and reads out of bounds at 16 rows. + assert 0 <= local_expert_offset < num_experts, ( + f"local_expert_offset={local_expert_offset} is outside the routing space " + f"[0, num_experts={num_experts}); num_experts is every routed expert " + "(512 for a 128-expert window of 512), not the window's count" + ) fp8_block_scale_moe_runner = Fp8BlockScaleMoeRunner() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fused_sigmoid_gate_mul_add.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fused_sigmoid_gate_mul_add.py similarity index 58% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fused_sigmoid_gate_mul_add.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fused_sigmoid_gate_mul_add.py index 367a6139f8a8..c6ac057e9432 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fused_sigmoid_gate_mul_add.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/fused_sigmoid_gate_mul_add.py @@ -10,8 +10,10 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -# No torch.ops boundary exists around this kernel: the module callable is the -# stable surface, and its body is one Triton launch after argument checks. +# The module callable is the surface: outside torch.compile its body is one Triton +# launch after argument checks. `torch.ops.trtllm.fused_sigmoid_gate_mul_add_` (same +# module) is only its torch.compile boundary, which the callable takes itself when +# compiling, with a different signature (mutated `output`, separate addend). from tensorrt_llm._torch.moe.fused_shared_expert import ( fused_sigmoid_gate_mul_add as _fused_sigmoid_gate_mul_add, ) @@ -47,9 +49,8 @@ class FusedSigmoidGateMulAdd(OpWrapper): Fusion boundary: sigmoid, broadcast, multiply and add. The routed MoE, the shared expert and the `[hidden, 1]` gate projection are the caller's. The - single rounding is the point: a sigmoid-mul then add rounds the product to - bf16 first and moves the output by up to 0.0312 at [17, 2560] (51.6x the - bf16 tolerance). + single rounding matters: rounding the product to bf16 before the add moves + the output by up to 0.0312 at [17, 2560], outside the bf16 tolerance. """ ARCHS = frozenset({Arch.SM_103}) @@ -65,11 +66,40 @@ class FusedSigmoidGateMulAdd(OpWrapper): "GEMM leaves it", spec=_ckpt(33, gate_layout="strided"), ), + # The default target's call sites, the decode PLE-layer combine and the MTP layer's + # MoE, on the gate the shared expert's fused GEMM leaves: column 1280 of the + # replicated [T, 1288] `gug` output. + Cell( + why="default decode PLE-layer / MTP MoE combine: one token, strided gate", + spec=_ckpt(1, gate_layout="strided"), + ), + Cell( + why="default tp4/tep4 decode batch at the largest captured size: 128 rows", + spec=_ckpt(128, gate_layout="strided"), + ), + Cell( + why="default tep4 MTP-7 at batch 128: 128 x 8 = 1024 rows, strided gate", + spec=_ckpt(1024, gate_layout="strided"), + ), + Cell( + why="the envelope's largest row: max_num_tokens 8192, strided gate", + spec=_ckpt(8192, gate_layout="strided"), + ), ) note = """ Metadata consumed: none; stateless, no autotuner, no architecture branch. - The Triton kernel JIT-compiles on first call. + Source: a Triton kernel (`_sigmoid_gate_mul_add_kernel`) JIT-compiled on first + call for the running device; the cells ran on sm_103. + + Call sites in the default target: the decode PLE-layer combine + (`DecodeTarget.forward`) and the MTP layer's MoE (`DecodeTarget.mtp_forward`), + in place on the routed output. The shared expert is replicated at every member, + so the gate is column 1280 of this rank's `[T, 1288]` `gug` output (gate|up 2 x + 640 plus 8 gate-logit rows) and hidden is 2560 everywhere. Rows are this rank's + tokens: decode 1..128 at tp4 / tep4 (1..32 a dep4 rank's own), MTP-7 at batch + 128 up to 128 x 8 = 1024. Every row count takes one path: hidden 2560 gives + one 4096-wide block and 8 warps per row, grid `(T, 1)`. Preconditions, all enforced loudly by the callable (AssertionError before any launch; the test holds the wrong-length gate): `final_hidden_states` @@ -84,18 +114,27 @@ class FusedSigmoidGateMulAdd(OpWrapper): calls it with; fp16 and fp32 run the same kernel and are not certified. bf16 at hidden 2560 is bitwise equal at 1-2 tokens and >= 99.993% bitwise at 17-8192 tokens, one ulp otherwise -- `tl.sigmoid`'s fast exponential. - Wrong variants land 29,000-152,000x outside the bf16 band: no sigmoid, - multiply instead of add, gate on the routed half, gate dropped, one gate - for the whole batch. + The bf16 band refuses the wrong variants: no sigmoid, multiply instead of + add, gate on the routed half, gate dropped, one gate for the whole batch. + At the default cells, wherever |out| >= 2^-6 each difference from + `reference` is exactly one bf16 ulp of the output; the rest sit where the + sum cancels toward zero, the largest 1.6e-2 absolute at 8192 rows (one ulp + of a [2, 4) value). Against a float64 reference the relative L2 is 1.66e-3 + at every cell, equal to `reference`'s own (bf16 output rounding). A strided + `gug` column and the same logits as a dense `[T, 1]` are bitwise one call + (held at 1024 rows). Graph-capturable: one capture at the 128-row site, + replayed on refilled operands, matches `reference` (held). Hard limits - num_tokens row offset widened to int64 in the kernel; largest run 8192 + num_tokens largest run 8192 (max_num_tokens). The row offset is int64 in + the kernel and column offsets are int32 within one row, so no + element index overflows int32. Not certified: hidden sizes other than 2560 (the column block and warp count step with it), gate layouts other than `[T, 1]` and the strided column, an `output` buffer (including one aliasing `shared_expert_output`), row-padded - operands, host operands, CUDA-graph capture, torch.compile, operands on two - devices, float8 activations. + operands, host operands, torch.compile, operands on two devices, float8 + activations, gate dtypes other than bf16, rows past 8192. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_sort.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/moe_sort.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_sort.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/moe_sort.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_unpermute_inplace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/moe_unpermute_inplace.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_unpermute_inplace.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/moe_unpermute_inplace.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/renorm_moe_routing_op.py similarity index 64% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/renorm_moe_routing_op.py index c97fd51ca4e9..ee9c4e5252c8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/moe/renorm_moe_routing_op.py @@ -1,7 +1,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Renormalized MoE routing: top-k over the raw router logits, then a softmax -over just those k, in one trtllm call.""" +over just those k, in one trtllm call. + +Called by the qwen4_exp target (`qwen3_8_flash_next_nvfp4__sm_103__default`) at +every layout.""" from __future__ import annotations @@ -47,8 +50,33 @@ class RenormMoeRoutingOp(OpWrapper): ARCHS = frozenset({Arch.SM_103}) + # The trunk pre-routes the trtllm-gen FP4 runner with bf16 weights at every + # member: tp4/tep4 route every row the rank holds; dep4 routes the + # all-gathered rows of the 4 ranks, then runs its expert window in chunks of + # at most 8192 rows. CELLS: tuple[Cell, ...] = ( - Cell(why="MTPLayer._moe decode: one token, bf16 weights", spec=_ckpt(1)), + Cell( + why="MTPLayer._moe decode: one token, bf16 weights; the trunk's decode at batch 1", + spec=_ckpt(1), + ), + Cell( + why="the trunk's decode at the largest captured batch: 128 rows at tp4/tep4, " + "4 x 32 gathered at dep4, bf16 weights", + spec=_ckpt(128), + ), + Cell( + why="the trunk's MTP7 verify at tep4: 128 requests x 8 rows, bf16 weights", + spec=_ckpt(1024), + ), + Cell( + why="the trunk's prefill chunk at max_num_tokens 8192, bf16 weights", + spec=_ckpt(8192), + ), + Cell( + why="dep4 prefill over the gathered rows, 4 ranks x max_num_tokens 8192: " + "the envelope's largest call, bf16 weights", + spec=_ckpt(32768), + ), Cell( why="_moe_block at 16384 rows (max_num_tokens of the tuned config): fp32 weights, " "as the CuteDSL experts take them; thousands of rows carry an exact bf16 tie", @@ -61,12 +89,30 @@ class RenormMoeRoutingOp(OpWrapper): ) note = """ + Source: compiled-in C++, no JIT path. `torch_ext::renorm_moe_routing_op` + in `cpp/tensorrt_llm/thop/moe/routing/customMoeRoutingOp.cpp` (a + `TORCH_LIBRARY_FRAGMENT(trtllm)` schema) launches `customMoeRoutingKernel` + from `cpp/tensorrt_llm/kernels/moe/routing/customMoeRoutingKernels.cu`. + Metadata consumed: none; stateless, deterministic, runs on the ambient stream. The launch picks its kernel instance by `nextPowerOfTwo(num_experts)` and `nextPowerOfTwo(topk)` -- 512 / top-10 is the (512, 16) instance with a 256-thread block. PDL (`TRTLLM_ENABLE_PDL`) was measured bitwise inert. + One warp per row, `ceil(T / 8)` blocks capped at 8192 + (`invokeCustomMoeRouting`), so the row count changes only the grid: every + certified count is under the cap (32768 rows = 4096 blocks), and the + grid-stride loop past 65536 rows is not exercised. + + The trunk's bf16-weight cells, measured on GB300: ids equal and weights + bit-exact at 1 / 128 / 1024 / 8192 / 32768 rows. Their randn logits tie + at the 10th / 11th place, so the ties-to-lower-index rule is exercised at + every multi-row cell. The bf16 weights of a row sum to 1 within 2.44e-3 + (the rounding of ten weights). + CUDA-graph replay at 128 rows is bit-exact against the reference on + refilled logits (the test holds it). + Preconditions, enforced loudly by the op: `router_logits` 2-D CUDA bf16 / fp16 / fp32; `output_dtype` fp32, bf16 or None (fp16 raises even for fp16 logits); `topk <= 16`; `num_experts <= 512` -- this checkpoint sits on @@ -83,21 +129,29 @@ class RenormMoeRoutingOp(OpWrapper): 512 / 16) is not bit-compatible: it uses `torch.topk` and casts before the softmax. - Bit-exact against `reference` over 2040 configurations at bring-up (in - dtype x out dtype x num_experts {1..512} x topk {0..16} x tokens - {0..8192}); the cells hold the target's router (512 / top-10, bf16 logits, + Measured bit-exact against `reference` over 2040 configurations (in dtype + x out dtype x num_experts {1..512} x topk {0..16} x tokens {0..8192}); the + cells hold the target's router (512 / top-10, bf16 logits, bf16 and fp32 weights), and the test holds the weights to HF's `norm_topk_prob` formulation at bf16 logits. Hard limits num_experts <= 512 (raises past it) topk <= 16 (raises past it); > num_experts silent, guarded - T certified 1..16384 (max_num_tokens of the tuned config) + T certified 1..32768 with bf16 weights (dep4's gathered + prefill, 4 x max_num_tokens 8192), 16384 with fp32 weights. + From the source, not probed: `numTokens` is narrowed to + int32 and the row offset `tokenId * numExperts` is uint32, + which wraps at 2^32 / 512 = 8388608 rows. Not certified: other num_experts / topk / dtype combinations, `topk == 0` and `num_tokens == 0` (both return empties), non-finite logits (`+inf`, NaN and an all `-inf` row poison the row to NaN; a `-inf` among finite - logits is a soft mask). + logits is a soft mask). Row counts between the cells -- dep4's smallest + gathered decode (4 rows), the MTP3 verify (128 x 4) -- run the same kernel + instance on a different grid size and are not cells. Router logits from + the checkpoint's own router GEMM were not driven; the cells use randn and + coarse logits. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/flashinfer_gemma_rmsnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/flashinfer_gemma_rmsnorm.py similarity index 59% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/flashinfer_gemma_rmsnorm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/flashinfer_gemma_rmsnorm.py index 4ef212a17f9b..5d4afe0c7dc4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/flashinfer_gemma_rmsnorm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/flashinfer_gemma_rmsnorm.py @@ -44,19 +44,54 @@ class FlashinferGemmaRmsnorm(OpWrapper): why="MTP pre_e_norm: the contiguous [T, 2560] embedding", spec=dict(shape=(4, 2560), dtype=torch.bfloat16, eps=1e-6), ), + # The default target's grouped-norm sites, per rank: every tp4/tep4 rank + # holds all of a step's tokens, so its rows are the step's rows. + Cell( + why="decode, one token: a single stream-1 row, which torch reports contiguous", + spec=dict(shape=(1, 2560), col_offset=2560, **_STREAM), + ), + Cell( + why="decode at batch 128, the largest tp4/tep4 capture: stream 1", + spec=dict(shape=(128, 2560), col_offset=2560, **_STREAM), + ), + Cell( + why="MTP7 verify, 128 requests x 8 rows = 1024: stream 2", + spec=dict(shape=(1024, 2560), col_offset=5120, **_STREAM), + ), + Cell( + why="prefill chunk at max_num_tokens 8192, the envelope's largest row: stream 2", + spec=dict(shape=(8192, 2560), col_offset=5120, **_STREAM), + ), ) note = """ Metadata consumed: none. Registered only when flashinfer is importable. + `torch.ops.trtllm.flashinfer_gemma_rmsnorm` is a Python `custom_op` + (`custom_ops/flashinfer_custom_ops.py`) over `flashinfer.norm.gemma_rmsnorm`, + which JIT-compiles a CuTe DSL kernel per (dtype, H, weight_bias, PDL, SM, + contiguous) through TVM-FFI; on GB300 the DSL path ran + (`flashinfer.norm._use_cuda_norm` is False there). Certified path: flashinfer's CuTe DSL `rmsnorm_cute` kernel for 2-D bf16 `x` (`weight_bias=1.0`) at hidden 2560: a stream slice at the residual's - 10240 row stride (33 and 2048 rows; a strided row selects a strided kernel - variant) and a contiguous `[4, 2560]` (MTP `pre_e_norm`). Not certified: a - single strided row, other widths, other row counts, the 3-D per-head kernel - (`qk_rmsnorm_cute`) and fp16/fp32 operands. The CUDA-JIT fallback (`FLASHINFER_USE_CUDA_NORM=1`, + 10240 row stride (33, 128, 1024, 2048 and 8192 rows; a strided view selects + the strided kernel variant), one stream row, and a contiguous `[4, 2560]` + (MTP `pre_e_norm`). + + One row of a stream is contiguous to torch (a size-1 dim ignores its + stride), so it takes the compact kernel, compiled for 128-byte aligned + rows: every stream starts a multiple of 5120 bytes into the residual row, + so the slice is aligned. Its output keeps stride `(10240, 1)` (`empty_like` + preserves a dense view's strides); with one row nothing reads past it. + Streams 0 and 3 (column offsets 0 and 7680) run the same strided kernel at + another base pointer; probed once at 128 rows within the default + tolerance, they are not cells. + + Not certified: other widths, the 3-D per-head kernel (`qk_rmsnorm_cute`) + and fp16/fp32 operands. The CUDA-JIT fallback (`FLASHINFER_USE_CUDA_NORM=1`, or a DSL that cannot target the device) is a different implementation and is not certified. PDL (`TRTLLM_ENABLE_PDL`, default on) changes scheduling, - not results. + not results. A CUDA graph captured at the 128-row stream slice replays + bitwise equal to the eager call (the test holds this). Loud (measured), so not guarded: 1-D and 4-D `x`, `x.stride(-1) != 1`, a weight of the wrong length, dtype or device, and a CPU `x` raise ValueError; @@ -70,12 +105,19 @@ class FlashinferGemmaRmsnorm(OpWrapper): Numerics: default tolerance holds for bf16, nothing loosened; the largest deviation measured is one bf16 ulp (0.03125) at 2048 and 8192 rows - of 5120. The `+ 1` is measured: against a plain-`w` reference the output is - ~4 away, so the cells tell the two apart. + of 5120. At the stream-slice cells on GB300 the worst element is one bf16 + ulp of its own magnitude at every row count (1 to 8192), on at most 1 in + 2560 elements; relative L2 at most 2.5e-5. + The `+ 1` is measured: against a plain-`w` reference the output is + 3.35 to 5.44 away, so the cells tell the two apart. Hard limits T * H contiguous inputs above 2^31 - 1 elements route to the strided (int64-offset) kernel boundary 2^31 read from the dispatch, not probed + rows a stream slice at row stride 10240 first passes 2^31 elements of + offset at 209716 rows; the strided kernel's row stride and row count + are int64. The envelope's largest is 8192 rows (8.4e7 elements of + offset), so that boundary is not reached and was not probed """ def __call__(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/layer_norm.py similarity index 67% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/layer_norm.py index 35201417be8a..f5d9025f186d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/layer_norm.py @@ -57,6 +57,13 @@ class LayerNorm(OpWrapper): ) note = """ + Source: the in-tree Python callable `layernorm_fn` + (`tensorrt_llm/_torch/modules/fla/layernorm_gated.py`: `LayerNormFn.apply` + -> `_layer_norm_fwd` -> `_layer_norm_fwd_1pass_kernel`). A Triton kernel, + JIT-compiled on first use, not in the wheel: BLOCK_N 2048 and 8 warps for + width 1152; `M` and both row strides are passed as `i32`. One program per + row. + Metadata consumed: none. Inputs are not mutated; two scratch fp32 [M] tensors (mean, rstd) are allocated per call and discarded. @@ -74,17 +81,39 @@ class LayerNorm(OpWrapper): 64KB` raises "This layer norm doesn't support feature dim >= 64KB" (N = 65536 in bf16); CPU tensors fail before or at the Triton launch. - Accuracy, measured beyond the cells' default gate and not certified: at - the checkpoint's affine (weights to 17.5, biases to -25.6) the bf16 result - is the correctly rounded fp64 value on >= 99.9989% of elements; where it - differs, by <= 2^-7 of the summed term magnitude. A constant row returns - `bias` bitwise. `eps` sits inside the sqrt. + Hard limit: the row offsets `row * stride_x_row` and `row * stride_y_row` + are `i32` products. At width 1152 the first row whose offset reaches 2^31 + is 1864136, so an `x` of >= 1864137 rows wraps. In a fresh process this + was measured as a CUDA illegal memory access. With other allocations + mapped below `x`, the wrapped read could land in them and be silent; that + was not measured and is not guarded. The envelope never reaches it: + `VisionTower` cuts a request's frames into calls of at most 32768 rows. + + Accuracy at the cells against an fp64 reference, measured on GB300 beyond + the default gate: + rows differs from the correctly rounded fp64 worst |err| / row max + 1 0 / 1152 0.907 bf16 ulp + 1024 37 / 1179648 (3.1e-5) 0.949 bf16 ulp + 32768 1096 / 37748736 (2.9e-5) 0.995 bf16 ulp + Cosine 0.9999986 and relative L2 1.65e-3 at every cell (the bf16 output + rounding itself). The elements that differ are outputs where `x_hat * w` + and `b` cancel toward zero: up to 205 ulp of their own tiny magnitude, + and each within torch's bf16 tolerance. A constant row returns `bias` + bitwise. `eps` sits inside the sqrt. + + What the default bf16 gate tells apart: an unbiased (N - 1) variance, a + missing mean subtraction and a missing bias each fail it at every cell + (`test_the_gate_discriminates_a_wrong_normalization` holds these). An eps + of 1e-5, or eps outside the sqrt, fails it only at some cells, so the gate + does not certify the eps placement. Accepted but not certified: an fp8-e4m3 `x` returns fp8-e4m3 (output is `empty_like(x)`), within one e4m3 ulp of the exact value; fp64 `x`. Forward only: the autograd Function has no backward. Measured, not certified: one kernel launch, runs unchanged under `inference_mode`, - CUDA-graph capturable after a warm-up call compiles the Triton kernel. + CUDA-graph capturable after a warm-up call compiles the Triton kernel (a + 1024-row capture replayed on refilled input is bitwise the eager call). + The call site is eager-only. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/rms_norm_gated_token_major.py similarity index 69% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/rms_norm_gated_token_major.py index fb79cf780dfd..b9ccf66104b1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/norm/rms_norm_gated_token_major.py @@ -18,10 +18,13 @@ #: The checkpoint's GDN epilogue: 48 value heads of 128, `output_gate_type: sigmoid`. _CKPT = dict(heads=48, width=128, dtype=torch.bfloat16, eps=1e-6, gate_activation="sigmoid") +#: One tp4/tep4 rank's share: 12 of the 48 value heads, `z` the column view of that rank's +#: 4120-wide input projection `[q k v | z | a | b]` = 2560 | 1536 | 12 | 12. +_TP4 = dict(_CKPT, heads=12, z="in_proj", z_row_stride=4120) +#: A dep4 rank: all 48 heads, `z` in the full 16480-wide projection. +_DEP4 = dict(_CKPT, z="in_proj", z_row_stride=16480) #: Prefill-sized launches (at least this many rows, i.e. 683 tokens of 48 heads) run -#: the multi-row kernel at `_PREFILL_ROWS` rows a program: at 16384 tokens 94 vs 143 -#: us on GB300 (traffic roofline ~81 us), where 4 rows a program leave 196608 small -#: programs. +#: the multi-row kernel at `_PREFILL_ROWS` rows a program. _PREFILL_MIN_ROWS = 32768 # capacity: fixed: a dispatch threshold, not a buffer size _PREFILL_ROWS = 8 @@ -68,8 +71,11 @@ class RmsNormGatedTokenMajor(OpWrapper): `x.dtype`, multiplied by `1 / fp8_scale[0]` and saturated at +-448 -- byte-identical to the unfused norm -> static-quantize pair. - Call site: `z` is the `[T, 6144]` column view of the GDN layer's one - `[T, 16480]` input-projection output, reshaped `[T, 48, 128]`. + Call sites (`PrefillTarget._gdn_epilogue`, and `DecodeTarget._gdn_epilogue` + at an MTP verify): `z` is the column view of the GDN layer's one + input-projection output, `[T, 6144]` of `[T, 16480]` reshaped `[T, 48, 128]` + on a dep4 rank, `[T, 1536]` of `[T, 4120]` reshaped `[T, 12, 128]` on a + tp4/tep4 rank, which holds 12 of the 48 value heads. Fusion boundary: norm, weight, gate and the optional quantize. The caller owns both projections producing `x` and `z`, and every reshape. @@ -90,6 +96,30 @@ class RmsNormGatedTokenMajor(OpWrapper): why="683 tokens: the first launch at 8 rows a program", spec=dict(tokens=683, z="in_proj", **_CKPT), ), + Cell( + why="tp4/tep4, one token: 12 heads, z at token stride 4120; 4-row kernel, PDL", + spec=dict(tokens=1, **_TP4), + ), + Cell( + why="tp4/tep4, 4 tokens (an MTP-3 verify of one request): 4-row kernel, PDL", + spec=dict(tokens=4, **_TP4), + ), + Cell( + why="tp4/tep4, 128 tokens (the con-128 decode batch): grid past the SM count", + spec=dict(tokens=128, **_TP4), + ), + Cell( + why="tp4/tep4, 1024 tokens: the MTP-7 verify of 128 requests x 8 rows", + spec=dict(tokens=1024, **_TP4), + ), + Cell( + why="tp4/tep4 prefill at max_num_tokens 8192: 98304 rows on the 8-row kernel", + spec=dict(tokens=8192, **_TP4), + ), + Cell( + why="dep4 decode: 48 heads at token stride 16480, a rank's 32 tokens", + spec=dict(tokens=32, **_DEP4), + ), ) note = """ @@ -113,19 +143,39 @@ class RmsNormGatedTokenMajor(OpWrapper): `weight` and `z` are widened to fp32 and need not match `x.dtype`. The output is freshly allocated and contiguous; no input is mutated. - Certified: bf16 `x` with the sigmoid gate over the in_proj column view of - `z` (48 heads of N = 128): 128 tokens on the 4-row kernel, 683 tokens (the - first 8-row launch) and 16384 tokens on the 8-row kernel. Not certified: - fp8 output, the silu gate, the generic fallback kernel, fp16/fp32 - operands, a strided `x`. + Both kernels are Triton, JIT-compiled at first call for the running GPU; + the cells ran that JIT path on sm_103 (GB300, 152 SMs). - Numerics: bf16 inside the default tolerance against the fp32 reference, - with one-ulp differences on ~1e-5 of elements at prefill row counts. fp8 - output byte-identical (measured). Nothing loosened. + Certified: bf16 `x` with the sigmoid gate over the in_proj column view of + `z`, N = 128. + 48 heads, token stride 16480 (tp1 and dep4): 128 and 32 tokens on the + 4-row kernel without PDL; 683 tokens (the first 8-row launch) and + 16384 tokens on the 8-row kernel. + 12 heads, token stride 4120 (a tp4/tep4 rank): 1 and 4 tokens on the + 4-row kernel with PDL (grids of 3 and 12 programs); 128 and 1024 tokens + on it without; 8192 tokens (98304 rows) on the 8-row kernel. + A one-token `z` view reports the dense token stride (1536 at 12 heads); + only token 0 is read, so it computes the same. + Not certified: fp8 output, the silu gate, the generic fallback kernel, + fp16/fp32 operands, a strided `x`. At 12 heads, the 8-row kernel's first + launch (2731 tokens) and its masked last program, which an odd token count + from 2731 up leaves half full (no 48-head row count does, and the 4-row + kernel never has one at either width): a prompt's last prefill chunk. + + Numerics: bf16 inside the default tolerance against the fp32 reference. + Measured on sm_103: at most one bf16 ulp from the bf16-rounded reference, + on about 1 in 10^5 elements or fewer (none at 1, 4 and 32 tokens). + Against the unrounded fp32 value the output's relative L2 (1.56e-3 to + 1.66e-3) equals the rounded reference's, and two calls are bitwise equal. + fp8 output byte-identical (measured). Nothing loosened. Hard limits N * itemsize fallback holds a row in 64 KB boundary 65536 B probed at bf16 N = 32768 (returns) and 40960 (RuntimeError) + offsets the multi-row kernel (4 and 8 rows a program) indexes x, z + and out in int64; the widest call the envelope makes (8192 + tokens x 48 heads x 128, z at stride 16480) stays below + 2^31 elements on every operand anyway """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ngram_token_history.py similarity index 65% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ngram_token_history.py index a6f8c97dbd5e..f0d557a20929 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ngram_token_history.py @@ -4,14 +4,14 @@ per token, its out-of-slice predecessors read from -- and its tail published back to -- the per-slot n-gram context pool. -Tier-3 composed from native torch integer ops (~20 small launches); exact, so -gated bitwise. Nothing registered under ``torch.ops.trtllm`` assembles an -n-gram id window, and the one kernel-plane callable in the family -(``EngramHashProvider.compute_hashes`` / ``NgramHashMapping._get_ngram_hashes`` -in ``tensorrt_llm/_torch/modules/engram/engram.py``) has no cross-step state: -it left-pads every chunk from a pad id, so a decode step there loses the -sequence's whole context. The built-in Qwen4-Exp implementation keeps its own, -not wrapped here. +A macro entry, composed from native torch integer ops (~20 small launches); +exact, so gated bitwise. Called by the qwen4_exp target at every layout; its +per-rank slot pools are among the cells. Nothing registered +under ``torch.ops.trtllm`` assembles an n-gram id window, and +``EngramHashProvider.compute_hashes`` / ``NgramHashMapping._get_ngram_hashes`` +(``tensorrt_llm/_torch/modules/engram/engram.py``) has no cross-step state: it +left-pads every chunk from a pad id, so a decode step there loses the +sequence's whole context. """ from __future__ import annotations @@ -26,10 +26,16 @@ CKPT_EOS_TOKEN_ID = 248044 CKPT_VOCAB_SIZE = 248320 #: A `ple_ngram` slot count of a full-size engine run. The op addresses a slot by its index -#: alone, so a pool of another size (the pool shares the recurrent pool's slot space, and a -#: bf16 recurrent pool holds more slots in the same memory) takes the same path. The target -#: hands the whole pool over unsliced. +#: alone, so a pool of another size takes the same path. The target hands the whole pool +#: over unsliced. ENGINE_SLOTS = 461 +#: The per-rank `ple_ngram` slot counts of the default target's members, from the stock +#: `Qwen4ExpHybridCacheManager` (`max_batch_size` 128, `max_seq_len` 9216, `max_num_tokens` +#: 8192, n-gram table host-resident), the same on all four ranks. The pool shares the GDN +#: recurrent pools' slot space, which the V2 manager sizes from the rank's memory budget, +#: not from `max_batch_size`; a dep4 rank, which keeps every GDN head, holds fewer. +TEP4_SLOTS = 1874 +DEP4_SLOTS = 850 def _scattered(num_seqs: int) -> tuple[int, ...]: @@ -39,16 +45,22 @@ def _scattered(num_seqs: int) -> tuple[int, ...]: return slots -def _cell(why: str, lengths, slots, carries, pool_slots: int, **spec) -> Cell: - return Cell( - why=why, - spec=dict( - lengths=tuple(lengths), - slots=tuple(slots), - carries=tuple(carries), - pool_slots=pool_slots, - **spec, - ), +def _spread(num_seqs: int, pool_slots: int) -> tuple[int, ...]: + """`num_seqs` distinct slots evenly over a `pool_slots` pool, both ends in, out of + order: batch row `s` is never slot `s`, as the scheduler's slot vector allows.""" + even = sorted({round(i * (pool_slots - 1) / (num_seqs - 1)) for i in range(num_seqs)}) + slots = tuple(even[1::2][::-1]) + tuple(even[0::2]) + assert len(set(slots)) == num_seqs and {0, pool_slots - 1} <= set(slots) + return slots + + +def _spec(lengths, slots, carries, pool_slots: int, **spec) -> dict: + return dict( + lengths=tuple(lengths), + slots=tuple(slots), + carries=tuple(carries), + pool_slots=pool_slots, + **spec, ) @@ -93,31 +105,71 @@ class NgramTokenHistory(OpWrapper): # "literal" | "seq_lens" (zeros + int64 cumsum of an int32 seq_lens prefix -- # `cu_seqlens_from_seq_lens`' form). `ngram_size` is the checkpoint's 3. CELLS: tuple[Cell, ...] = ( - _cell( - "MTP verify p=8: 32 decodes x 8 rows, qsl rebuilt", - (8,) * 32, - _scattered(32), - (True,) * 32, - ENGINE_SLOTS, - qsl_from="seq_lens", - **_CALL_SITE, + Cell( + why="MTP verify p=8: 32 decodes x 8 rows, qsl rebuilt", + spec=_spec( + (8,) * 32, + _scattered(32), + (True,) * 32, + ENGINE_SLOTS, + qsl_from="seq_lens", + **_CALL_SITE, + ), + ), + Cell( + why="MTP verify mixed: 37-token context + 8 x 3 rows", + spec=_spec( + (37,) + (3,) * 8, + _scattered(9), + (False,) + (True,) * 8, + ENGINE_SLOTS, + qsl_from="seq_lens", + eos_density=0.05, + **_CALL_SITE, + ), + ), + Cell( + why="zero-token rows: carry kept, fresh reseeded", + spec=_spec((2, 0, 0), (0, 1, 2), (True, True, False), 3), + ), + Cell( + why=f"tep4 decode: 128 x 1 row at the rank's {TEP4_SLOTS}-slot pool", + spec=_spec( + (1,) * 128, + _spread(128, TEP4_SLOTS), + (True,) * 128, + TEP4_SLOTS, + qsl_from="seq_lens", + **_CALL_SITE, + ), + ), + Cell( + why=f"dep4 decode: a rank's 32 x 1 row at its {DEP4_SLOTS}-slot pool", + spec=_spec( + (1,) * 32, + _spread(32, DEP4_SLOTS), + (True,) * 32, + DEP4_SLOTS, + qsl_from="seq_lens", + **_CALL_SITE, + ), ), - _cell( - "MTP verify mixed: 37-token context + 8 x 3 rows", - (37,) + (3,) * 8, - _scattered(9), - (False,) + (True,) * 8, - ENGINE_SLOTS, - qsl_from="seq_lens", - eos_density=0.05, - **_CALL_SITE, + Cell( + why=f"tep4 MTP7 verify: 128 x 8 rows at the {TEP4_SLOTS}-slot pool", + spec=_spec( + (8,) * 128, + _spread(128, TEP4_SLOTS), + (True,) * 128, + TEP4_SLOTS, + qsl_from="seq_lens", + **_CALL_SITE, + ), ), - _cell( - "zero-token rows: carry kept, fresh reseeded", - (2, 0, 0), - (0, 1, 2), - (True, True, False), - 3, + Cell( + why="tep4 prefill: one fresh 8192-token context, max_num_tokens", + spec=_spec( + (8192,), (TEP4_SLOTS - 1,), (False,), TEP4_SLOTS, qsl_from="seq_lens", **_CALL_SITE + ), ), ) @@ -136,7 +188,7 @@ class NgramTokenHistory(OpWrapper): Preconditions: - Every slot a real one, `0 <= slot < slots`. The carry read bounds-checks, so -1 or an index past the pool is a device-side assert - that poisons the CUDA context (observed in its own process). **A -1 + that poisons the CUDA context. **A -1 pad sentinel is not supported**; give padded rows real slots. - Slots distinct: two sequences on one slot is accepted and the publish is unordered (data; not guarded). @@ -151,23 +203,38 @@ class NgramTokenHistory(OpWrapper): A wrong window is a window of valid ids: newest-first order, a zero seed, ignored carry flags, packing across sequence boundaries, a rotated carry read and publishing the head all hash to valid rows. Chunked - prefill (96 tokens split at 41) is bitwise equal to one-shot, window and - pool. At 461 slots the result is bitwise identical to an 11-slot pool - with the same slot content. + prefill is bitwise equal to one-shot, window and pool, and the pool size + does not change the result (461 and 11 slots bitwise identical). - Certified at the forward's two verify-step shapes, both on the + Certified at the tp1 forward's two verify-step shapes, both on the `cu_seqlens_from_seq_lens` rebuild at the 461-slot pool (32 x 8 rows; a 37-token fresh context with in-stream EOS beside 8 x 3 rows), and at - zero-token rows. Not certified: a plain decode (one row per sequence) or - a context-only batch -- outside a verify step the forward takes the - fused `ple_ngram_embed`, not this op. + zero-token rows. At the default target's per-rank pools (`TEP4_SLOTS` 1874, + `DEP4_SLOTS` 850; every cell on the rebuild, with slots spread over the + whole pool, both ends in): a plain decode of 128 x 1 rows at tep4 and of + a dep4 rank's 32 x 1 rows, the MTP7 verify's 128 x 8 rows at tep4, and + one fresh 8192-token context at tep4 (`max_num_tokens`, the envelope's + largest row). + + The pool the manager hands over is a contiguous `[slots, 2]` int64 + (stride (2, 1), offset 0) at tp4, tep4 and dep4, with the GDN pools' + slot count. The call site is `ple_ngram_embed`, whose `reference` + composes this entry's `reference`; the tp1 forward launches the fused + entry, not this op. + + Not certified: tp4's pool (pure tp4 is not served), the MTP and + multimodal members' pools, and the dep4 variants' pools (kv_fp8, + kv_nvfp4, chunked, reuse with its recurrent-state snapshots). The op + addresses a slot by index alone, so these differ from the cells in pool + length only. CUDA-graph capture certified (no `repeat_interleave` on device lengths, - so no data-dependent shape). Cost, not gated: ~0.30 ms for a 4096-token - prefill, a 32-row decode or a mixed 2 x 1024 + 30 x 1 batch alike -- - launch-bound. + so no data-dependent shape). - Hard limits: not probed (`num_tokens` measured to 4096 in timing). + Hard limits: none reached. `num_tokens` is certified to 8192; every index + is int64 and the largest tensor is the `[num_tokens, ngram_size]` window, + far from 2^31 elements. Pure torch: no JIT path and no compiled kernel of + its own, so nothing in it is architecture-specific. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_gate_conv_inject_norm.py similarity index 82% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_gate_conv_inject_norm.py index bb2814cac29b..f1eebf8224e8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_gate_conv_inject_norm.py @@ -6,12 +6,12 @@ residual add into the hyper-connection stream, and the next site's grouped norm -- three Triton launches behind one call. -Tier 5 (Triton written for this target). It replaces, on the decode path, -three `_grouped_norm` calls (12 ``flashinfer_gemma_rmsnorm`` + 3 cats), -``gated_ngram_injection`` (~10 launches), ``dilated_short_conv_paged`` (~15 -launches around ``F.conv1d``), two adds and the norm-only -``qwen4_exp_hc_inject_norm`` after the PLE add. No ``torch.ops.trtllm`` op or -FlashInfer callable takes a dilation (see ``dilated_short_conv_paged``). +Called by the qwen4_exp target at every layout; uses the row-to-sequence +search of ``ple_ngram_embed``. It computes the composition of three grouped +norms, ``gated_ngram_injection``, ``dilated_short_conv_paged``, two adds and +the norm-only ``qwen4_exp_hc_inject_norm`` after the PLE add. No +``torch.ops.trtllm`` op or FlashInfer callable takes a dilation (see +``dilated_short_conv_paged``). The three launches split where the parallelism changes shape. The norms and the gate are row reductions, one program per (row, stream). The conv and the @@ -19,8 +19,6 @@ `[channels, 9]` rows: one program per (sequence, block of channels), so the state moves as contiguous spans across many SMs instead of strided through four. The residual add and the next site's norm are row reductions again. -Folding all three into one (row, stream) program was measured 2.5-3x slower -at 128 rows: four programs per sequence cannot move the state fast enough. """ from __future__ import annotations @@ -33,8 +31,14 @@ import triton.language as tl from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -from .ple_ngram_embed import PER_SEQ_MAX_ROWS, _seq_of_rows +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ngram_token_history import ( + DEP4_SLOTS, + TEP4_SLOTS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ple_ngram_embed import ( + PER_SEQ_MAX_ROWS, + _seq_of_rows, +) @triton.jit @@ -298,19 +302,25 @@ def _state_head(state_len: int) -> int: CKPT_HIDDEN = 2560 CKPT_DILATION = 3 CKPT_CLAMP_MIN = 1e-6 -#: Largest PLE-conv pool seen for this checkpoint (see `dilated_short_conv_paged`). +#: Largest PLE-conv pool seen for this checkpoint at tp1 (see `dilated_short_conv_paged`). ENGINE_SLOTS = 461 +#: The default target's per-rank PLE-conv pools: the conv state and the n-gram context are +#: bound over one slot space (the stock manager asserts their slot counts equal), so the +#: counts are the n-gram pool's (`TEP4_SLOTS`, `DEP4_SLOTS`). Every rank holds the whole +#: `[hc_count * hidden, 9]` conv state at every member. A tp4 rank holds the same per-rank +#: PLE and GDN state as a tep4 rank, so tp4 is bounded by tep4's. +TP4_SLOTS = TEP4_SLOTS -def _cell(why: str, lengths, **spec) -> Cell: - return Cell(why=why, spec=dict(lengths=tuple(lengths), **spec)) +def _spec(lengths, **spec) -> dict: + return dict(lengths=tuple(lengths), **spec) class PleGateConvInjectNorm(OpWrapper): """With `C = hc_count`, `D = value.shape[-1]`, every grouped norm a gemma RMS norm over each `D`-wide stream (`n(x, w) = x * rsqrt(mean(x^2) + eps) * (1 + w)`), and each step rounded to the activation dtype where the - composition it replaces rounds -- for row `t` of sequence `s`, stream `c`:: + unfused composition rounds -- for row `t` of sequence `s`, stream `c`:: kn = round(n(key[t, c], key_norm[c])) qn = round(n(hc[t, c], query_norm[c])) @@ -324,8 +334,9 @@ class PleGateConvInjectNorm(OpWrapper): where `window` is `[conv_states[slot] | xc of the sequence's rows]` and `t'` the row's offset in its sequence; the state part is zero for a - sequence whose `has_initial_states` flag is false (None: all carry). Then `conv_states[slot]` becomes - `window[:, L : L + state_len]`, the raw conv-input tail. Every sequence + sequence whose `has_initial_states` flag is false (None: all carry). Then + `conv_states[slot]` becomes `window[:, L : L + state_len]`, the raw + conv-input tail. Every sequence carries its state (a decode batch); a `pad_slot_id` sequence reads a zero state and publishes nothing. @@ -348,18 +359,63 @@ class PleGateConvInjectNorm(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # spec: lengths (rows per sequence), pool_slots (default num_seqs + 2), - # max_rows (max(lengths)), carries (None -> the argument is None). The + # max_rows (max(lengths)), carries (None -> the argument is None), + # conv_input (True -> a staged conv-input buffer is passed). The # checkpoint's 4 x 2560 bf16 geometry throughout. CELLS: tuple[Cell, ...] = ( - _cell("decode 128 x 1 row at the 461-slot pool", (1,) * 128, pool_slots=ENGINE_SLOTS), - _cell("MTP3 verify: 64 sequences x 4 rows", (4,) * 64, pool_slots=ENGINE_SLOTS), + Cell( + why="decode 128 x 1 row at the 461-slot pool", + spec=_spec((1,) * 128, pool_slots=ENGINE_SLOTS), + ), + Cell( + why="MTP3 verify: 64 sequences x 4 rows", + spec=_spec((4,) * 64, pool_slots=ENGINE_SLOTS), + ), # Context-shaped (prefill / mixed) steps: the packed-row conv. - _cell("prefill 8192 rows, fresh", (8192,), carries=(False,), pool_slots=ENGINE_SLOTS), - _cell( - "mixed step: 2 contexts + 64 decode rows", - (3001, 5000) + (1,) * 64, - carries=(False, True) + (True,) * 64, - pool_slots=ENGINE_SLOTS, + Cell( + why="prefill 8192 rows, fresh", + spec=_spec((8192,), carries=(False,), pool_slots=ENGINE_SLOTS), + ), + Cell( + why="mixed step: 2 contexts + 64 decode rows", + spec=_spec( + (3001, 5000) + (1,) * 64, + carries=(False, True) + (True,) * 64, + pool_slots=ENGINE_SLOTS, + ), + ), + # The default target's members, each at the per-rank conv pool its rank binds, as + # the forward calls it: decode and verify with `has_initial_states` None and + # `max_rows` the rows per sequence (a verify step also stages the conv + # input), prefill with one flag per sequence and `max_rows` the step's rows. + Cell( + why=f"tep4 decode: 128 x 1 row, {TEP4_SLOTS} slots", + spec=_spec((1,) * 128, pool_slots=TEP4_SLOTS), + ), + Cell( + why=f"dep4 decode: a rank's 32 x 1 row, {DEP4_SLOTS} slots", + spec=_spec((1,) * 32, pool_slots=DEP4_SLOTS), + ), + Cell( + why=f"tep4 MTP1 verify: 128 sequences x 2 rows, {TEP4_SLOTS} slots", + spec=_spec((2,) * 128, pool_slots=TEP4_SLOTS, conv_input=True), + ), + Cell( + why=f"tep4 MTP7 verify: 128 sequences x 8 rows (packed-row conv), {TEP4_SLOTS} slots", + spec=_spec((8,) * 128, pool_slots=TEP4_SLOTS, conv_input=True), + ), + Cell( + why=f"tep4 prefill: one fresh 8192-token context, {TEP4_SLOTS} slots", + spec=_spec((8192,), carries=(False,), pool_slots=TEP4_SLOTS), + ), + Cell( + why=f"envelope: 8192 rows over 128 sequences (1 context + 127 decode), {TEP4_SLOTS} slots", + spec=_spec( + (8192 - 127,) + (1,) * 127, + carries=(False,) + (True,) * 127, + max_rows=8192, + pool_slots=TEP4_SLOTS, + ), ), ) @@ -377,9 +433,31 @@ class PleGateConvInjectNorm(OpWrapper): column stride; weights contiguous. Certified at the checkpoint's 4 x 2560 bf16 geometry with every slot live (the target never passes `pad_slot_id`), at a 128 x 1-row decode, a 64 x 4-row MTP3 verify, an 8192-row fresh - prefill and a mixed 2-context + 64-decode step, all on the 461-slot pool. + prefill and a mixed 2-context + 64-decode step, all on the 461-slot pool; + and at the default target's per-rank pools (`TEP4_SLOTS` 1874, `DEP4_SLOTS` 850, + slots spread over the whole pool with both ends taken): a 128 x 1 tep4 and + a 32 x 1 dep4 decode, 128 x 2 (MTP1) and 128 x 8 (MTP7, the packed-row conv + with no carry flags) verifies staging `conv_input`, an 8192-row fresh + prefill, and the envelope's largest step, 8192 rows over 128 sequences (one + 8065-row context, 127 carrying decode rows). A tp4 rank binds the same + conv state as a tep4 rank, so the tep4 cells cover it by bound. Not certified: sequences shorter than `max_rows` on the per-sequence path - (a verify step gives every sequence `max_rows` rows). + (a verify step gives every sequence `max_rows` rows); verify widths 3, 5, + 6 and 7 rows at the per-rank pools (4 rows only at 64 x 4 over 461 slots); + the MTP member's own pool (max_seq_len 16384; the verify cells run at the + tep4 count); a dep4 prefill or verify step; a chunked continuation at the + per-rank pools (a carrying context chunk is certified only by the 461-slot + mixed cell). Every pool count is the stock engine's with the n-gram table + host-resident. + + Index range: slots and rows are widened to int64 before every offset, so + neither the pool (2^31 elements at 23301 slots of `[10240, 9]`) nor the + rows (2^31 elements at 209715 rows of 10240) is near a 32-bit limit at + the envelope (1874 slots, 8192 rows); nothing is indexed in fp32. + + Source: the four Triton kernels in this file (three launches a call, the + packed-row conv a fourth on context-shaped steps), JIT-compiled for the + running device; every cell ran the JIT path on sm_103 (GB300). Numerics: every rounding point of the composition is kept, so the only differences are fp32 reduction order (the three norms' sums of squares @@ -552,7 +630,7 @@ def reference( ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """`(hc_out, normed, conv_states)` as the call leaves them: float64 throughout, rounded to the activation dtype exactly where the op (and - the composition it replaces) rounds; the conv one sequence at a time as + the unfused composition) rounds; the conv one sequence at a time as an explicit tap sum. Mutates nothing.""" dtype = hc.dtype rows, width = hc.shape diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_ngram_embed.py similarity index 79% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_ngram_embed.py index f8c2236854cf..f727731c8386 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ple/ple_ngram_embed.py @@ -4,12 +4,12 @@ table gather and per-tensor fp8 dequantization in one Triton launch, with the per-slot n-gram context published back in the same pass. -Tier 5 (Triton written for this target). It replaces, on the decode path, -the composition ``ngram_token_history`` (~20 launches) + ``ngram_hash`` (~30 -launches) + ``nn.functional.embedding`` + ``dequantize_e4m3_per_tensor`` -- -nothing registered under ``torch.ops.trtllm`` or in the installed FlashInfer hashes -n-grams or gathers-and-dequantizes, and each of those entries is launch-bound -at decode (0.3-0.4 ms eager each). +Called by the qwen4_exp target at every layout. It computes the composition +``ngram_token_history`` + ``ngram_hash`` + ``nn.functional.embedding`` + +``dequantize_e4m3_per_tensor``, whose references it reads. Nothing registered +under ``torch.ops.trtllm`` or in the installed FlashInfer gathers and +dequantizes n-gram rows (``torch.ops.trtllm.qwen4_exp_ple_ngram_hash`` hashes a +window only). One program per **sequence**, which loops over that sequence's rows: the carry read and the publish of one slot happen in one program, so no other program can @@ -30,9 +30,14 @@ import triton.language as tl from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -from ..embedding.ngram_hash import ngram_hash -from .ngram_token_history import ngram_token_history +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.embedding.ngram_hash import ( + ngram_hash, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ngram_token_history import ( + DEP4_SLOTS, + TEP4_SLOTS, + ngram_token_history, +) #: This checkpoint's PLE geometry: 3-gram, 16 heads of 160 fp8 columns. CKPT_NGRAM_SIZE = 3 @@ -40,8 +45,11 @@ CKPT_TABLE_ROWS = 320001536 CKPT_EOS_TOKEN_ID = 248044 CKPT_SCALE = 0.00019931793212890625 -#: Largest `ple_ngram` pool seen for this checkpoint (see `ngram_token_history`). +#: Largest `ple_ngram` pool seen for this checkpoint at tp1 (see `ngram_token_history`). ENGINE_SLOTS = 461 +#: A tp4 rank's pool, not counted on an engine. A tp4 rank holds the same per-rank PLE +#: and GDN state as a tep4 rank, so the pool is bounded by tep4's (`TEP4_SLOTS`). +TP4_SLOTS = TEP4_SLOTS #: Above this many rows per sequence (MTP verify carries at most 4) the call #: takes the context-shaped path: packed row tiles plus a separate publish. PER_SEQ_MAX_ROWS = 4 @@ -279,8 +287,8 @@ def _ple_ngram_publish_kernel( tl.store(pool_row + c, published, mask=cmask) -def _cell(why: str, lengths, **spec) -> Cell: - return Cell(why=why, spec=dict(lengths=tuple(lengths), **spec)) +def _spec(lengths, **spec) -> dict: + return dict(lengths=tuple(lengths), **spec) class PleNgramEmbed(OpWrapper): @@ -315,33 +323,68 @@ class PleNgramEmbed(OpWrapper): # 320001536-row device table, rows filled where touched). The checkpoint's # 3-gram of 8 heads per order, int32 ids. CELLS: tuple[Cell, ...] = ( - _cell( - "decode 128 x 1 row, host-resident table, 461 slots", - (1,) * 128, - pool_slots=ENGINE_SLOTS, - table="host", + Cell( + why="decode 128 x 1 row, host-resident table, 461 slots", + spec=_spec((1,) * 128, pool_slots=ENGINE_SLOTS, table="host"), ), - _cell( - "decode 128 x 1 row, the full checkpoint-height table", - (1,) * 128, - pool_slots=ENGINE_SLOTS, - table="ckpt", - eos_density=0.05, + Cell( + why="decode 128 x 1 row, the full checkpoint-height table", + spec=_spec((1,) * 128, pool_slots=ENGINE_SLOTS, table="ckpt", eos_density=0.05), ), - _cell( - "MTP3 verify: 64 sequences x 4 rows, 461 slots", - (4,) * 64, - pool_slots=ENGINE_SLOTS, - table="host", - eos_density=0.05, + Cell( + why="MTP3 verify: 64 sequences x 4 rows, 461 slots", + spec=_spec((4,) * 64, pool_slots=ENGINE_SLOTS, table="host", eos_density=0.05), ), # Context-shaped (prefill / mixed) steps: the packed-row form. - _cell( - "mixed step: 2 contexts + 64 decode rows", - (3001, 5000) + (1,) * 64, - carries=(False, True) + (True,) * 64, - pool_slots=ENGINE_SLOTS, - eos_density=0.05, + Cell( + why="mixed step: 2 contexts + 64 decode rows", + spec=_spec( + (3001, 5000) + (1,) * 64, + carries=(False, True) + (True,) * 64, + pool_slots=ENGINE_SLOTS, + eos_density=0.05, + ), + ), + # The default target's members, each at the per-rank n-gram pool its rank binds, + # as the forward calls it: decode and verify with `has_initial_states` + # None and `max_rows` the rows per sequence, prefill with one flag per + # sequence and `max_rows` the step's row count. + Cell( + why=f"tep4 decode: 128 x 1 row, host-resident table, {TEP4_SLOTS} slots", + spec=_spec((1,) * 128, pool_slots=TEP4_SLOTS, table="host", eos_density=0.05), + ), + Cell( + why=f"dep4 decode: a rank's 32 x 1 row, host-resident table, {DEP4_SLOTS} slots", + spec=_spec((1,) * 32, pool_slots=DEP4_SLOTS, table="host", eos_density=0.05), + ), + Cell( + why=f"tp4 decode: 128 x 1 row, the full checkpoint-height table, {TP4_SLOTS} slots", + spec=_spec((1,) * 128, pool_slots=TP4_SLOTS, table="ckpt", eos_density=0.05), + ), + Cell( + why=f"tep4 MTP1 verify: 128 sequences x 2 rows, {TEP4_SLOTS} slots", + spec=_spec((2,) * 128, pool_slots=TEP4_SLOTS, table="host", eos_density=0.05), + ), + Cell( + why=f"tep4 MTP7 verify: 128 sequences x 8 rows (packed-row form), {TEP4_SLOTS} slots", + spec=_spec((8,) * 128, pool_slots=TEP4_SLOTS, table="host", eos_density=0.05), + ), + Cell( + why=f"tep4 prefill: one fresh 8192-token context, {TEP4_SLOTS} slots", + spec=_spec( + (8192,), carries=(False,), pool_slots=TEP4_SLOTS, table="host", eos_density=0.05 + ), + ), + Cell( + why=f"envelope: 8192 rows over 128 sequences (1 context + 127 decode), {TEP4_SLOTS} slots", + spec=_spec( + (8192 - 127,) + (1,) * 127, + carries=(False,) + (True,) * 127, + max_rows=8192, + pool_slots=TEP4_SLOTS, + table="host", + eos_density=0.05, + ), ), ) @@ -368,6 +411,27 @@ class PleNgramEmbed(OpWrapper): carrying) + 64 decode rows on the packed-row form. Not certified: zero-row sequences. + At the default target's per-rank n-gram pools, with slots spread over the whole + pool and both ends taken (a lone sequence on the last slot): + - tep4 (`TEP4_SLOTS` 1874): 128 x 1 decode; MTP1 verify, 128 x 2 rows, + on the per-sequence path; MTP7 verify, 128 x 8 rows, on the packed-row + form with no carry flags (the decode call site's `None`); one fresh + 8192-token context; and the envelope's largest step, 8192 rows over + 128 sequences (one fresh 8065-token context + 127 carrying decode rows). + - dep4 (`DEP4_SLOTS` 850): a rank's 32 x 1 decode. + - tp4 (`TP4_SLOTS`, bounded by tep4's pool): 128 x 1 decode on the full + checkpoint-height device table. + The pools are the stock engine's with a host-resident table; an MTP + member reserves its own pool, not counted here. CUDA-graph replay + is certified at 16 x 1 rows over 461 slots and at 128 x 1 and 128 x 8 + rows over 1874 slots. + + The kernels are Triton, JIT-compiled on first call for the running + device: every cell above ran that path on sm_103. The host-resident + cells read a pinned host table through its mapped device address, the + stock PLE host offload's placement. The checkpoint-height cells allocate + the 47.7 GiB table on the device. + Bitwise equal to `ngram_token_history` + `ngram_hash` + embedding + `dequantize_e4m3_per_tensor`, window, rows, values and pool alike. CUDA-graph capturable: no host sync, one launch, output allocated per call. @@ -489,7 +553,7 @@ def reference( max_rows: int, ) -> tuple[torch.Tensor, torch.Tensor]: """`(out, context_pool)` as the call leaves them, from the scalar - references of the two entries this replaces -- `ngram_token_history` + references of the two entries it fuses -- `ngram_token_history` (window and published pool) and `ngram_hash` (rows) -- and a float64 gather-and-scale rounded once. Mutates nothing.""" carries = ( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/quantization/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/quantization/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/bf16_to_fp32.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/quantization/bf16_to_fp32.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/bf16_to_fp32.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/quantization/bf16_to_fp32.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/quantization/fp8_quantize_1x128.py similarity index 73% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/quantization/fp8_quantize_1x128.py index 51e97e8e5444..1ac77fe895e5 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/quantization/fp8_quantize_1x128.py @@ -79,6 +79,18 @@ class Fp8Quantize1x128(OpWrapper): why="16384 rows: max_num_tokens of the tuned config", spec=dict(tokens=16384, **_CKPT), ), + Cell( + why="128 rows: the MTP drafter's decode step at max_batch_size 128", + spec=dict(tokens=128, **_CKPT), + ), + Cell( + why="1024 rows: the first draft over mtp7's verify rows, 128 requests x 8", + spec=dict(tokens=1024, **_CKPT), + ), + Cell( + why="8192 rows: the MTP layer over one max_num_tokens context chunk", + spec=dict(tokens=8192, **_CKPT), + ), ) note = """ @@ -95,18 +107,36 @@ class Fp8Quantize1x128(OpWrapper): `data` is a leading-row view of an allocation padded to a multiple of 4 rows -- the padded rows are uninitialized, so a raw-pointer reader must - stop at `T`. CUDA-graph capture and replay are certified. + stop at `T`. At a `T` that is a multiple of 4 (128, 1024, 8192) nothing is + padded and `scales` is the kernel's own buffer, not a compacting copy. + CUDA-graph capture and replay are certified at 64, 128 and 1024 rows. + + Source: compiled-in C++, `torch.ops.trtllm.fp8_quantize_1x128` from + `thop/fp8Quantize.cpp` (`fp8_1x128_cs` -> `launch_scale_1x128`). The + launcher takes `scale_1x128_vec_kernel` when `can_vectorize_1x128` holds + (H % 8 == 0, 16-byte aligned input, 8-byte aligned output), as at every + cell; the element-wise `scale_1x128_kernel` it falls back to (an input + whose storage offset breaks the 16-byte alignment) is not certified. The + grid is persistent, `kNumDeviceSMs * 8` blocks of 8 warps quantizing 2 + groups each (19456 groups a sweep on GB300); the cells run under one sweep + (128 rows), just past one (1024) and nine (8192). + + Measured on GB300 at H 2560: codes bit-exact at 1 row and on >= 99.94% of + elements at every other row count, each miss within 1.122e-07 (relative) + of its rounding midpoint; scales within the 4-ulp gate. Hard limits - T certified at 1, 20 and 16384 at H 2560 (16384: max_num_tokens of the - tuned config); above that not probed + T certified at 1, 20, 128, 1024, 8192 and 16384 at H 2560 (128: the + drafter's decode at max_batch_size 128; 1024: the first draft over + mtp7's 128 x 8 verify rows; 8192: a max_num_tokens context chunk; + 16384: max_num_tokens of the tuned config); above that not probed H multiple of 16 (raises otherwise); certified at 2560, the draft layer's width Not certified: `T == 0` (the op accepts it and returns `[0, H]` / `[NB, 0]`); a short tail block (`H % 128 != 0`); all-zero blocks (the `AMAX_FLOOR` scale), `inf` / NaN inputs and saturation; `use_ue8m0=True`; - the registered fake. + the element-wise fallback kernel; the registered fake. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_triton_update.py similarity index 65% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_triton_update.py index cbcd77719624..1bc5b6bf9d6a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_triton_update.py @@ -31,6 +31,25 @@ ENGINE_SLOTS = 461 #: The engine's `max_num_sequences`: the snapshot buffer's row count. ENGINE_SNAPSHOT_ROWS = 32 +#: Row stride of tp1's (and a dep4 rank's) input-projection output `[q|k|v|z|a|b]`. +CKPT_X_STRIDE = 16480 + +#: A tp4/tep4 rank's GDN conv width (`[q | k | v]` of 4 key and 12 value heads of 128) and +#: the row stride of its input-projection output, `[q | k | v | z | a | b]` = 2560 + 1536 + +#: 12 + 12. +RANK_DIM = 2560 +RANK_X_STRIDE = 4120 +#: Conv-pool slots of one tep4 rank's engine in the V2 hybrid manager (memory-derived; the +#: PLE and GDN pools share one slot space). +TEP4_SLOTS = 1874 +#: GDN layers whose conv states share one slot's page in the V2 manager: a layer's view +#: `[slots, dim, 3]` has slot stride 36 * dim * 3 (276480 at 2560) and layer `i` sits +#: `i * dim * 3` elements into the page. +GDN_PAGE_LAYERS = 36 +#: The stock GDN speculative state's snapshot buffer is `[36 layers, max_batch_size, +#: tokens_per_gen_step, dim, 3]`, a layer's view `[max_batch_size, p, dim, 3]`, with +#: `intermediate_state_indices = arange(max_batch_size)` (`GDNReplayState.bind`). +MAX_BATCH_SIZE = 128 def _pinned(batch: int) -> tuple[int, ...]: @@ -38,6 +57,33 @@ def _pinned(batch: int) -> tuple[int, ...]: return (0,) + tuple(i * 14 + 3 for i in range(1, batch - 1)) + (ENGINE_SLOTS - 1,) +def _spread(batch: int, pool_slots: int) -> tuple[int, ...]: + """`batch` distinct slots of a `pool_slots` pool: the top slot (the page's highest + addresses) and slot 0 first, then a stride-37 walk over the interior (37 is coprime to + `pool_slots - 2`), so batch row `b` is not slot `b`.""" + rest = tuple((i * 37 + 4) % (pool_slots - 2) + 1 for i in range(max(batch - 2, 0))) + slots = ((pool_slots - 1, 0) + rest)[:batch] + assert len(set(slots)) == batch + return slots + + +def _tep4(batch: int, seqlen: int) -> dict: + """A tep4 rank's verify call: 2560 channels of x at row stride 4120, the layer's conv + view the last of the manager's 36-layer page of 1874 slots, and its snapshots into a + layer's `[128, seqlen, 2560, 3]` view of the stock buffer.""" + return dict( + dim=RANK_DIM, + x_row_stride=RANK_X_STRIDE, + batch=batch, + seqlen=seqlen, + pool_slots=TEP4_SLOTS, + pool_layers=GDN_PAGE_LAYERS, + slots=_spread(batch, TEP4_SLOTS), + snapshots=MAX_BATCH_SIZE, + snapshot_layers=2, + ) + + class CausalConv1dTritonUpdate(OpWrapper): """Convolve each row's new tokens against its cached window, out of place, and advance the window; optionally snapshot the window after every token. @@ -68,10 +114,14 @@ class CausalConv1dTritonUpdate(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec keys: dim, batch, seqlen, pool_slots, slots, snapshots (rows of the - # snapshot buffer), snapshot_steps (the buffer's step depth, default - # seqlen). Every cell is the call site: `x` the token-major q|k|v columns - # of a [n * p, 16480] projection seen as [n, C, p], bf16, W 4, no bias, + # spec keys: dim, x_row_stride (default 16480), batch, seqlen, pool_slots, + # pool_layers (the pool the last layer's view of a page this many layers + # deep, default 1), slots, snapshots (rows of the snapshot buffer), + # snapshot_steps (the buffer's step depth, default seqlen), snapshot_layers + # (the buffer the first layer's view of one this many layers deep, default + # 1), launch_dependent_kernels (default False). Every cell is the call site: `x` the + # token-major q|k|v columns of a + # [n * p, x_row_stride] projection seen as [n, C, p], bf16, W 4, no bias, # SiLU, pad_slot_id -1, no accepted-token roll. CELLS: tuple[Cell, ...] = tuple( Cell( @@ -100,6 +150,31 @@ class CausalConv1dTritonUpdate(OpWrapper): snapshot_steps=8, ), ), + # The default target's tep4 member (the MTP unit's): a rank's 4 k / 12 v heads. + Cell( + why="tep4 verify, 1 row x 4 tok (MTP3) at 2560 / stride 4120, the 1874-slot page", + spec=_tep4(1, 4), + ), + Cell( + why="tep4 verify, 1 row x 4 tok, launch_dependent_kernels: the PDL chain into the replay", + spec=dict(_tep4(1, 4), launch_dependent_kernels=True), + ), + Cell( + why="tep4 verify, 4 rows x 4 tok, launch_dependent_kernels (largest ll graph batch)", + spec=dict(_tep4(4, 4), launch_dependent_kernels=True), + ), + Cell( + why="tep4 verify, 128 rows (max_batch_size) x 4 tok (MTP3) at 2560 / stride 4120", + spec=_tep4(128, 4), + ), + Cell( + why="tep4 verify, 128 rows x 2 tok (draft length 1) at 2560 / stride 4120", + spec=_tep4(128, 2), + ), + Cell( + why="tep4 verify, 128 rows x 8 tok (draft length 7, the envelope's largest step)", + spec=_tep4(128, 8), + ), ) note = """ @@ -116,8 +191,33 @@ class CausalConv1dTritonUpdate(OpWrapper): of the projection: 32 rows at seqlen 3 and at seqlen 8 (the widest verify step), and 8 rows at seqlen 5 into an 8-step buffer (a buffer deeper than seqlen keeps its later steps untouched); bf16, W 4, no bias, SiLU only. - Not certified: a contiguous `x`, other seqlens, and pad-sentinel rows - (the target never sends one). + + The default target's tep4 member (the MTP unit's verify, `DecodeTarget._gdn_verify`): + a rank's 4 k / 12 v heads, so `x` is the `[n, 2560, p]` view of the q|k|v + columns of the `[n * p, 4120]` in_proj output (strides `(p * 4120, 1, + 4120)`); the pool is one layer's view of the V2 manager's 36-layer conv page + at the measured 1874 slots (slot stride 276480, the last layer's view); the + snapshots go to a layer's `[128, p, 2560, 3]` view of the GDN speculative + state's `[36, max_batch_size, p, 2560, 3]` buffer with + `intermediate_state_indices = arange(n)`. Certified at 1 row x 4 tokens and + at 128 rows (max_batch_size) x 2, 4 and 8 tokens (draft lengths 1, 3 and 7; + 128 x 8 = 1024 rows is the envelope's largest verify step). The result + differs from the reference by one bf16 ulp in at most a few elements + (relative L2 <= 2.0e-5), the pool and the snapshots are bitwise the + reference, and no other layer's conv page or snapshots, nor x's buffer, + is written. + + The kernel is Triton, JIT-compiled at the first call; the cells ran that + path on sm_103. The x strides, the seqlen and the snapshot strides are + `tl.constexpr`, so each (p, row stride) the target sends is its own + compile. + + Not certified: a contiguous `x`, other seqlens, pad-sentinel rows (the + target never sends one), a dep4 rank (10240 channels at row stride 16480 + over its 850-slot page) and tp4 (no MTP member). The tep4 slot count is + memory-derived and was not measured on an MTP-config engine; the kernel + addresses a slot by its index alone, so the pool's size does not change + the path. Raise loudly, so not guarded: width 1 (triton CompilationError); a missing `intermediate_state_indices` with snapshots and slot indices given. @@ -127,14 +227,23 @@ class CausalConv1dTritonUpdate(OpWrapper): roll is uncertified; the target passes None). Out-of-range slot values: illegal access, not probed. - Accepted, uncertified: `launch_dependent_kernels` (the target passes - False). Accepted but not exposed: `cache_seqlens` (ignored by the kernel -- + `launch_dependent_kernels=True` (the target's verify step: the CuTe replay after it + launches with PDL and streams its state before `griddepcontrol.wait`) only adds + `griddepcontrol.launch_dependents` at the kernel's start: certified at the tep4 shape, + 1 and 4 rows x 4 tokens, held to the reference like every other cell. + + Accepted but not exposed: `cache_seqlens` (ignored by the kernel -- there is no ring mode), the eagle-tree mask trio, `validate_data`, `metadata`, the `_block_n`/`_num_warps`/`_num_stages` overrides. A `conv_state` dtype different from `x` is accepted and uncertified. Hard limits - none probed. + slot offset int64 in the kernel (`conv_state_indices` loaded `.to(tl.int64)`): + at the 1874-slot page the last slot's offset is 517.8 M + elements, below 2^31 either way. + x offset `row * stride_x_seq` in int32: 128 x 8 x 4120 = 4.2 M + elements at the envelope, far below 2^31. + none probed beyond the cells. """ def __call__( @@ -243,14 +352,12 @@ def is_valid( seqlen = 1 if x.dim() == 2 else x.shape[2] width = weight.shape[1] state_len = width - 1 + (seqlen - 1 if num_accepted_tokens is not None else 0) - # The tap loop is specialised for W 2..4: W 5 runs and returned an - # answer up to 47.2 off (mean 1.89). W 1 fails to compile. + # The tap loop is specialised for W 2..4; W 1 fails to compile. assert 2 <= width <= 4, ( f"weight width must be 2, 3 or 4; got {width} -- width 5 is accepted " "by the kernel and silently computes a wrong answer" ) - # Reads and writes `state_len` columns unmasked against the pool: a - # 2-column pool at W 4 wrote one element into the next slot. + # Reads and writes `state_len` columns unmasked against the pool. assert conv_state.shape[2] >= state_len, ( f"conv_state must hold at least {state_len} columns (width={width}, " f"seqlen={seqlen}, num_accepted_tokens=" @@ -258,17 +365,15 @@ def is_valid( f"{conv_state.shape[2]} -- a shorter pool is accepted and writes past " "the slot into the next one" ) - # The channel axis is masked against x's dim only: a half-width pool - # changed 384 elements of unnamed slots; a half-height weight was read - # past its end. + # The channel axis is masked against x's dim only. assert conv_state.shape[1] >= dim and weight.shape[0] >= dim, ( f"conv_state and weight must cover all {dim} channels of x; got " f"{conv_state.shape[1]} and {weight.shape[0]} -- a narrower pool or " "weight is accepted and reads/writes out of bounds" ) if intermediate_conv_window is not None: - # No bound on the step axis: a 2-step buffer at seqlen 4 wrote - # into the next row. + # No bound on the step axis: a shallower buffer is written into + # the next row. assert ( intermediate_conv_window.shape[1] >= seqlen and intermediate_conv_window.shape[2] >= dim diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_update.py similarity index 69% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_update.py index 14413d38be2b..f8f98d2a9807 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/causal_conv1d_update.py @@ -3,9 +3,9 @@ """Single-step causal depthwise conv1d against a slot-indexed conv-state pool. The decode step's shape (one token a row, SiLU, no bias, a width-4 conv over a -3-column bf16 pool, slot indices) from `_CUTE_MIN_ROWS` rows runs a CuTe DSL -kernel (Tier 6, written for this target) that is bitwise the TRT-LLM kernel it -replaces there; every other call is `torch.ops.trtllm.causal_conv1d_update`. +3-column bf16 pool, slot indices, a channel count that is a multiple of 1024) +from `_CUTE_MIN_ROWS` rows runs a CuTe DSL kernel that is bitwise the TRT-LLM +kernel; every other call is `torch.ops.trtllm.causal_conv1d_update`. """ from __future__ import annotations @@ -25,25 +25,40 @@ #: This checkpoint's GDN `[q|k|v]` width; the conv is width 4 over a 3-column pool. CKPT_DIM = 10240 -#: Largest conv-pool slot count the cache manager was seen to allocate for this -#: checkpoint (10 .. 461 over twelve runtime-value dumps, batch never above 32). -#: The decode step hands the whole pool over unsliced. +#: Largest conv-pool slot count the cache manager allocates for this checkpoint at a batch +#: of at most 32. The decode step hands the whole pool over unsliced. ENGINE_SLOTS = 461 +#: A tp4/tep4 rank's GDN conv width (`[q | k | v]` of 4 key and 12 value heads of 128) and +#: the row stride of its input-projection output, `[q | k | v | z | a | b]` = 2560 + 1536 + +#: 12 + 12; at tp1 and dep4 the projection is 16480 wide. +RANK_DIM = 2560 +RANK_X_STRIDE = 4120 +CKPT_X_STRIDE = 16480 +#: Per-rank conv-pool slot counts of the V2 hybrid manager at bf16 KV and recurrent state, +#: KV fraction 0.6, `max_batch_size` 128, `max_seq_len` 9216 (the PLE and GDN pools share +#: one slot space). The count is memory-derived, so a dep4 rank, which keeps every GDN head, +#: holds fewer than a tep4 rank. tp4 has no measured count; its rank holds tep4's conv +#: geometry, so tep4's count stands in for it. +TEP4_SLOTS = 1874 +DEP4_SLOTS = 850 +TP4_SLOTS = TEP4_SLOTS +#: GDN layers whose conv states share one slot's page in the V2 manager: a layer's view +#: `[slots, dim, 3]` has slot stride 36 * dim * 3 (276480 at 2560, 1105920 at 10240) and +#: layer `i` sits `i * dim * 3` elements into the page (at tp4, tep4 and dep4). +GDN_PAGE_LAYERS = 36 + # -------------------------------------------------------------------------------------- -# The decode shape in CuTe DSL. TRT-LLM's seqlen-1 kernel runs a thread per (row, -# channel) with 2-byte accesses -- at 128 rows 1.3 M threads, four waves of one load -# latency each (7.5 us a layer on GB300). Here a thread owns 8 channels of one row: the -# weights (64 B), its state (8 channels x 3 columns, 48 contiguous bytes) and x (16 B) -# by 16-byte loads, then the TRT-LLM kernel's arithmetic instruction for instruction +# The decode shape in CuTe DSL. A thread owns 8 channels of one row: the weights +# (64 B), its state (8 channels x 3 columns, 48 contiguous bytes) and x (16 B) by +# 16-byte loads, then the TRT-LLM kernel's arithmetic instruction for instruction # (`__fmaf_rn` over the taps from 0, `v * __frcp_rn(1 + __expf(-v))`, round to bf16), # x written in place and the state shifted by the new token. # -------------------------------------------------------------------------------------- -#: Rows from which the decode shape runs the CuTe kernel (GB300, graph replay, 8 -#: rotating layers, us: 128 rows 7.46 -> 5.41, 121 rows 7.40 -> 5.13; at 33 rows -#: and below the TRT-LLM kernel's one wave is faster: 3.10 vs 3.39). +#: Rows from which the decode shape runs the CuTe kernel; below it the TRT-LLM kernel's +#: single wave is faster. _CUTE_MIN_ROWS = 64 _CH = 8 _THREADS = 128 @@ -313,6 +328,45 @@ def _pinned(batch: int) -> tuple[int, ...]: return (0,) + tuple(i * 14 + 3 for i in range(1, batch - 1)) + (ENGINE_SLOTS - 1,) +def _spread(batch: int, pool_slots: int, pad_every: int = 0) -> tuple[int, ...]: + """`batch` distinct slots of a `pool_slots` pool: the top slot (the page's highest + addresses) and slot 0 first, then a stride-37 walk over the interior (37 is coprime to + every `pool_slots - 2` here), so batch row `b` is not slot `b`. With `pad_every`, rows + `5, 5 + pad_every, ...` take the pad slot -1 instead.""" + rest = tuple((i * 37 + 4) % (pool_slots - 2) + 1 for i in range(max(batch - 2, 0))) + slots = ((pool_slots - 1, 0) + rest)[:batch] + assert len(set(slots)) == batch + if pad_every: + slots = tuple(-1 if i % pad_every == 5 else s for i, s in enumerate(slots)) + return slots + + +def _rank(batch: int, pool_slots: int, pad_every: int = 0) -> dict: + """A tp4/tep4 rank's call: 2560 conv channels at x row stride 4120, the layer's view the + last of the manager's 36-layer conv page.""" + return dict( + dim=RANK_DIM, + x_row_stride=RANK_X_STRIDE, + batch=batch, + pool_slots=pool_slots, + slots=_spread(batch, pool_slots, pad_every), + pool_layers=GDN_PAGE_LAYERS, + ) + + +def _dep4(batch: int) -> dict: + """A dep4 rank's call: tp1's 10240 channels at row stride 16480, the layer's view the + last of the manager's 36-layer conv page.""" + return dict( + dim=CKPT_DIM, + x_row_stride=CKPT_X_STRIDE, + batch=batch, + pool_slots=DEP4_SLOTS, + slots=_spread(batch, DEP4_SLOTS), + pool_layers=GDN_PAGE_LAYERS, + ) + + class CausalConv1dUpdate(OpWrapper): """One causal depthwise conv step per batch row against that row's cached window, then the window's advance -- `x` and the pool both written in place. @@ -336,11 +390,11 @@ class CausalConv1dUpdate(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec keys: dim, batch, pool_slots, slots ("random" | tuple), pool_layers - # (the pool a one-layer view of a buffer this many layers deep, default 1). - # Every cell is the call site: x the q|k|v column view of the decode rows - # of a [rows, 16480] in_proj output, bf16, W 4, a 3-column pool, no bias, - # SiLU, no ring, pad_slot_id -1. + # spec keys: dim, x_row_stride (default 16480), batch, pool_slots, slots + # ("random" | tuple), pool_layers (the pool the last layer's view of a buffer + # this many layers deep, default 1). Every cell is the call site: x the + # q|k|v column view of the decode rows of a [rows, x_row_stride] in_proj + # output, bf16, W 4, a 3-column pool, no bias, SiLU, no ring, pad_slot_id -1. CELLS: tuple[Cell, ...] = ( Cell( why="x the q|k|v column view of the [rows, 16480] in_proj, batch 32", @@ -360,6 +414,34 @@ class CausalConv1dUpdate(OpWrapper): pool_layers=2, ), ), + # The default target's members, each over its rank's measured conv page. Every + # 2560-channel call runs the TRT-LLM kernel (2560 is not a multiple of the + # CuTe kernel's 1024 channels a block row). + Cell( + why="tep4 decode, 1 row at 2560 / stride 4120: the top slot of the 1874-slot page", + spec=_rank(1, TEP4_SLOTS), + ), + Cell( + why="tep4, 64 rows at 2560 / stride 4120 over the 1874-slot page", + spec=_rank(64, TEP4_SLOTS), + ), + Cell( + why="tep4, 128 rows (max_batch_size) at 2560 / stride 4120 over the 1874-slot page", + spec=_rank(128, TEP4_SLOTS), + ), + Cell( + why="tp4, 128 rows at 2560 / stride 4120 over tp4's page, pad rows (-1) on the " + "TRT-LLM kernel", + spec=_rank(128, TP4_SLOTS, pad_every=21), + ), + Cell( + why="dep4 decode, 32 rows at 10240 / stride 16480: the top slot of the 850-slot page", + spec=_dep4(32), + ), + Cell( + why="dep4, 128 rows (max_batch_size, the CuTe kernel) over the 850-slot page", + spec=_dep4(128), + ), ) note = """ @@ -377,11 +459,10 @@ class CausalConv1dUpdate(OpWrapper): `[dim, 3]` contiguous and 16-byte aligned rows, it runs `_DecodeConvKernel` (CuTe DSL) instead: 8 channels a thread by 16-byte loads, the TRT-LLM kernel's instructions, bitwise its result (x and pool, - pad rows untouched); 128 rows 7.46 -> 5.41 us a layer (GB300, graph - replay over 8 rotating layers). Certified at that call: 32 rows (the - TRT-LLM kernel) over the 461-slot pool, 128 rows (the CuTe kernel) over a - 256-slot pool, and 64 rows (the CuTe kernel's first row count) with pad - rows over the 461-slot pool seen as one layer of a two-layer buffer. + pad rows untouched). Certified at that call: 32 rows (the TRT-LLM kernel) + over the 461-slot pool, 128 rows (the CuTe kernel) over a 256-slot pool, + and 64 rows (the CuTe kernel's first row count) with pad rows over the + 461-slot pool seen as one layer of a two-layer buffer. The op checks metadata itself, loudly: width 2..4, `state_len >= W - 1`, one dtype across x/conv_state/weight (bias matches weight), weight shape, @@ -389,16 +470,42 @@ class CausalConv1dUpdate(OpWrapper): `cache_seqlens`, a batch-sized pool when indices are absent. So no guard. Not guarded because it is content, not metadata: two rows naming one slot - race -- the slot ended up matching neither row's expected state in 10 of - 10 runs while both outputs matched. Slot uniqueness is the caller's. + race -- the slot ends up matching neither row's expected state while both + outputs match. Slot uniqueness is the caller's. Out-of-range slot values are an illegal access, not probed. A row whose index equals `pad_slot_id` (-1 at the call site) is skipped: - its x rows and the pool stay bitwise unchanged (certified on the CuTe - kernel only). - - Hard limits - none probed. + its x rows and the pool stay bitwise unchanged (certified on both kernels: + the CuTe one at 64 rows of 10240, the TRT-LLM one at tp4's 128 rows). + + The default target's call site, per rank. tp4 / tep4: x `[num_decodes, 2560]`, + the q|k|v view (row stride 4120) of the rank's `[T, 4120]` projection; the + pool one layer's view `[slots, 2560, 3]` of the V2 manager's conv page, + 36 GDN layers a slot (slot stride 276480). dep4: tp1's 10240 channels at + stride 16480, page slot stride 1105920. Certified at the measured per-rank + slot counts (`TEP4_SLOTS` 1874, `DEP4_SLOTS` 850; tp4 has none measured and + takes tep4's, `TP4_SLOTS`), the view the page's last layer and the top slot + named, so a cell reaches the pool's highest address: tep4 / tp4 at 1, 64 + and 128 rows, dep4 at 32 and 128 rows. + + Which kernel runs: 2560 is not a multiple of `_CH * _THREADS` (1024), so + every tp4 / tep4 call runs TRT-LLM's seqlen-1 kernel + (`causal_conv1d_update_kernel_sl1`, one thread per (row, channel), 20 + blocks of 128 channels a row), whatever the row count. dep4 runs the CuTe + kernel from 64 rows, as tp1 does. + + Sources: the TRT-LLM kernel is compiled into the wheel + (`cpp/tensorrt_llm/kernels/causalConv1d/causalConv1d.cu`, registered by + `cpp/tensorrt_llm/thop/causalConv1dOp.cpp`); the CuTe kernel is JIT-compiled + by `cute.compile` on its first call for the running device, sm_103 here. + + Hard limits (from source, not probed: the measured pools stay far below) + The TRT-LLM kernel forms a slot's offset as `int slot * uint32_t + conv_state_batch_stride`, 32-bit unsigned elements: it wraps from slot + 15535 at stride 276480 (tp4 / tep4) and slot 3884 at 1105920 (dep4). + The measured pools reach 518,123,519 (tep4) and 940,031,999 (dep4) + elements, under 2^31 elements and bytes. The CuTe kernel forms slot and + row offsets in 64 bits. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/cu_seqlens_from_seq_lens.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/cu_seqlens_from_seq_lens.py similarity index 59% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/cu_seqlens_from_seq_lens.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/cu_seqlens_from_seq_lens.py index 904246678bc0..c94f9601b4d3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/cu_seqlens_from_seq_lens.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/cu_seqlens_from_seq_lens.py @@ -3,11 +3,7 @@ """Prepend-zero cumulative sum of a per-request query-row count: the ``cu_seqlens`` segmentation vector every varlen recurrent op takes. -Tier-3 composed from native torch (``torch.zeros`` + ``torch.cumsum``, three -launches). No registered op or ``modules/`` callable produces this vector: the -recurrent metadata's ``query_start_loc*`` is built inline by -``Mamba2Metadata.prepare`` and segments the generation tail one row per -sequence, which is short under speculative decoding. +A macro entry composed from native torch (``torch.zeros`` + ``torch.cumsum``). """ from __future__ import annotations @@ -22,11 +18,20 @@ #: Batch shapes this checkpoint's serving path produces where it calls this entry. #: ``p`` is the trunk's query rows per generation request, ``1 + runtime_draft_len`` -#: under MTP (2..8: `max_draft_len` 1..7); at ``p = 1`` the target reuses the -#: recurrent metadata's vector instead. +#: under MTP (2..8: `max_draft_len` 1..7); at ``p = 1`` the target at tp1 reuses the +#: recurrent metadata's vector instead, and ``decode_p1_128`` certifies the call +#: there all the same. +#: The ``_128`` batches are the serving envelope's widest step at every member of +#: the default target (``max_batch_size`` 128 requests per rank, ``max_num_tokens`` +#: 8192 rows): ``mixed_p8_128`` is one context chunk beside 127 MTP-7 verifies, +#: 7176 + 127 * 8 = 8192 rows. BATCHES = { "mixed_p2": (1024, 37, 2, 2, 2, 2), "padded_decode_p2": (2, 2, 2, 0, 0, 0), + "decode_p1_128": (1,) * 128, + "verify_p8_128": (8,) * 128, + "chunk_8192": (8192,), + "mixed_p8_128": (7176,) + (8,) * 127, } @@ -67,6 +72,42 @@ class CuSeqlensFromSeqLens(OpWrapper): lengths=BATCHES["padded_decode_p2"], in_dtype=torch.int32, out_dtype=torch.int64 ), ), + Cell( + why="decode at con 128: 128 one-row requests, prefix int32->int64", + spec=dict( + lengths=BATCHES["decode_p1_128"], + in_dtype=torch.int32, + out_dtype=torch.int64, + layout="prefix", + ), + ), + Cell( + why="MTP-7 verify at con 128: 128 x 8 rows, prefix int32->int64", + spec=dict( + lengths=BATCHES["verify_p8_128"], + in_dtype=torch.int32, + out_dtype=torch.int64, + layout="prefix", + ), + ), + Cell( + why="one 8192-row prefill chunk (max_num_tokens), prefix int32->int64", + spec=dict( + lengths=BATCHES["chunk_8192"], + in_dtype=torch.int32, + out_dtype=torch.int64, + layout="prefix", + ), + ), + Cell( + why="widest step: 7176-row chunk + 127 MTP-7 verifies = 8192 rows, 128 requests", + spec=dict( + lengths=BATCHES["mixed_p8_128"], + in_dtype=torch.int32, + out_dtype=torch.int64, + layout="prefix", + ), + ), ) note = """ @@ -78,27 +119,36 @@ class CuSeqlensFromSeqLens(OpWrapper): every generation request contributes one row. Preconditions, none checked (data, not metadata -- a check would sync): - entries non-negative (a negative length was silently wrong in three - consumer probes and a device-side assert in the fourth); a padded row is - length 0, never -1; `out[-1]` equals the packed activation's row count. + entries non-negative (a negative length is silently wrong or a + device-side assert in the consumers); a padded row is length 0, never -1; + `out[-1]` equals the packed activation's row count. `dtype` has no default: `causal_conv1d_fwd` requires int32 and `chunk_gated_delta_rule` keys its autotuner on `cu_seqlens`' dtype. Certified bitwise at the call site's prefix view of `seq_lens_cuda` (a mixed step at p 2, and the p 8 verify of 32 requests) and at a decode - batch padded with zero-length rows, int32 in and int64 out. CUDA-graph - capture is certified (only `seq_lens.shape[0]` reaches Python; a captured - graph follows lengths rewritten in place). Three launches, about 0.03 ms - at batch 8..256 -- launch bound. + batch padded with zero-length rows, int32 in and int64 out. Also at a + rank's serving envelope (`max_batch_size` 128, `max_num_tokens` 8192, the + same at tp4, tep4 and dep4): 128 one-row requests, 128 MTP-7 verifies of + 8 rows, one 8192-row chunk, and a 7176-row chunk beside 127 verifies. + CUDA-graph capture is certified at 32 requests (int32 out) and 128 + requests (int64 out): only `seq_lens.shape[0]` reaches Python, and a + captured graph follows lengths rewritten in place, up to a mixed step of + 8192 rows. + + Composition (int32 in, int64 out): four kernels -- a fill, the int32 -> + int64 cast copy, and cub's scan-init and scan, all compiled into torch (no + JIT path). Not certified: a `seq_lens` dtype other than int32 (the guard admits any - integer dtype), a non-contiguous `seq_lens` view. + integer dtype), a non-contiguous `seq_lens` view, more than 128 requests, + and a step over 8192 rows (an int64 `out` holds any reachable total). Hard limits out[-1] accumulated in `dtype` boundary 2^31 - 1 for int32 - probed at [2e9, 2e9]: int64 gives [0, 2e9, 4e9]; int32 wraps silently - to [0, 2000000000, -294967296]. Not guarded (data). + int32 wraps silently ([2e9, 2e9] gives [0, 2000000000, -294967296]); + int64 does not. Not guarded (data). """ def __call__(self, seq_lens: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: @@ -115,10 +165,9 @@ def reference(self, seq_lens: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: return torch.tensor(out, dtype=dtype, device=seq_lens.device) def is_valid(self, seq_lens: torch.Tensor, dtype: torch.dtype) -> None: - # Measured: a [1, n] row vector returned a length-2 [0, n_0]; an [n, 2] - # input a length-(n + 1) vector of nonsense (only a resize warning); a - # float32 [7.9, 1.2, 13.6] was truncated to [0, 7, 8, 21]; a bool input - # was counted as nonzero rows; a float `dtype` returned a float vector. + # Each of these answers wrongly without an error: a 2-D input gives a + # short or nonsense vector, a float input is truncated, a bool input is + # counted as nonzero rows, and a float `dtype` returns a float vector. assert seq_lens.dim() == 1, ( f"seq_lens must be a 1-D per-request row count, got shape {tuple(seq_lens.shape)}" ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/flashinfer_chunk_gated_delta_rule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/flashinfer_chunk_gated_delta_rule.py similarity index 58% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/flashinfer_chunk_gated_delta_rule.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/flashinfer_chunk_gated_delta_rule.py index c0f7de9908f2..85a0d7c477f9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/flashinfer_chunk_gated_delta_rule.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/flashinfer_chunk_gated_delta_rule.py @@ -36,16 +36,67 @@ #: qwen3.8-flash-next-nvfp4's Gated DeltaNet at tp1: (H, HV, K, V) -- 16 q/k #: heads, 48 v heads, head dims 128. QWEN38_GEOMETRY = (16, 48, 128, 128) -#: Largest slot count the cache manager was seen to allocate the recurrent pool -#: at for qwen3.8-flash-next-nvfp4 (10 .. 461 over twelve runtime-value dumps). +#: Largest recurrent-pool slot count the cache manager allocates for +#: qwen3.8-flash-next-nvfp4. ENGINE_SLOTS = 461 +#: A tp4/tep4 rank's Gated DeltaNet: tp1's 16 key / 48 value heads split four ways, the +#: key-to-value head ratio (3) unchanged. A dep4 rank keeps `QWEN38_GEOMETRY`. +RANK4_GEOMETRY = (4, 12, 128, 128) +#: Per-rank recurrent-pool slot counts of `Qwen4ExpHybridCacheManager` on this checkpoint +#: at bf16 KV and state, KV fraction 0.6, `max_batch_size` 128, `max_seq_len` 9216, TRTLLM +#: MoE and the n-gram table host-resident (the PLE and GDN pools share one slot space). The +#: V2 manager sizes the pool from the rank's memory budget, not from `max_batch_size`; a +#: dep4 rank keeps every GDN head, so it holds fewer slots than a tep4 rank. +TEP4_SLOTS = 1874 +DEP4_SLOTS = 850 +#: tp4 has no measured count. A tp4 rank's GDN state per slot is a tep4 rank's, so tep4's +#: count stands in for it. +TP4_SLOTS = TEP4_SLOTS +#: GDN layers whose states share one slot page in the stock manager (48 layers, every +#: 4th full attention): the pool one layer receives is a view whose slot stride is 36 +#: `[HV, V, K]` states, so a tep4 rank's last slot starts 1.3e10 elements past the base. +PAGE_LAYERS = 36 + + def _engine_slots(num_seqs: int) -> tuple[int, ...]: """`num_seqs` distinct slots of the engine pool, slot 0 and the last live.""" return ((0, ENGINE_SLOTS - 1) + tuple(i * 13 + 4 for i in range(num_seqs)))[:num_seqs] +def _rank_slots(num_seqs: int, slots: int) -> tuple[int, ...]: + """`num_seqs` distinct slots of a `slots`-slot pool: the last and the first, then a + stride-37 walk over the interior (37 is coprime to 1874 - 2 and 850 - 2).""" + rest = tuple((i * 37 + 4) % (slots - 2) + 1 for i in range(max(num_seqs - 2, 0))) + picked = ((slots - 1, 0) + rest)[:num_seqs] + assert len(set(picked)) == num_seqs + return picked + + +def _rank_cell(why: str, geometry, seqlens, slots: int, initial, fi_path: str) -> Cell: + """The target's call (`PrefillTarget._gdn_core`) at one rank's pool: the stock + manager's bf16 state page at the member's slot count, the layer's view of it, a + prefix view of a wider output buffer, `has_initial_state`, `linear_gate`. `fi_path` is + the FlashInfer kernel the call reaches on sm_103: `cp` (the context-parallel pipeline) + or `chunked`.""" + return Cell( + why=why, + spec=dict( + geometry=geometry, + seqlens=seqlens, + slots=slots, + page_layers=PAGE_LAYERS, + indices=_rank_slots(len(seqlens), slots), + pool_dtype=torch.bfloat16, + output="slice", + initial=initial, + linear_gate=True, + fi_path=fi_path, + ), + ) + + #: Elements of a `[V, K]` state plane a gather program copies. _GATHER_BLOCK = 4096 @@ -98,6 +149,79 @@ def _gather_carried_states( return out +#: Context tokens up to which a `has_initial_state` call runs in place over the pool: one +#: chunk (64 tokens) a sequence at most, where FlashInfer's chunked kernel is a single +#: launch and its context-parallel pipeline (five launches plus the gather and the scatter) +#: buys nothing. +INPLACE_MAX_TOKENS = 64 + + +@triton.jit +def _zero_fresh_states_kernel( + pool_ptr, + slot_ptr, + init_ptr, + stride_ps, + stride_ph, + stride_pv, + stride_pk, + K: tl.constexpr, + PLANE: tl.constexpr, + BLOCK: tl.constexpr, +): + """One program: `BLOCK` elements of sequence `n`'s head `h` plane of the pool row + `slot_ptr[n]`, zeroed when `init_ptr[n]` is false; a carried sequence's row is not + touched.""" + n = tl.program_id(0) + if tl.load(init_ptr + n) == 0: + h = tl.program_id(1) + i = tl.program_id(2) * BLOCK + tl.arange(0, BLOCK) + slot = tl.load(slot_ptr + n).to(tl.int64) + dst = ( + pool_ptr + slot * stride_ps + h * stride_ph + (i // K) * stride_pv + (i % K) * stride_pk + ) + tl.store(dst, tl.zeros([BLOCK], dtype=pool_ptr.dtype.element_ty), mask=i < PLANE) + + +def _chunk_inplace( + q, k, v, g, beta, scale, pool, indices, has_initial_state, cu_seqlens, output, linear_gate +) -> torch.Tensor: + """The short in-place indexed call: the fresh sequences' pool rows zeroed (one launch), + then FlashInfer's chunked kernel (`use_cp=False`) reading and writing the named rows of + the pool itself (`state_indices`, `output_state` the pool): no packed copy in or out.""" + import flashinfer + + num_seqs = indices.shape[0] + _, h, v_dim, k_dim = pool.shape + _zero_fresh_states_kernel[(num_seqs, h, triton.cdiv(v_dim * k_dim, _GATHER_BLOCK))]( + pool, indices, has_initial_state, *pool.stride(), + K=k_dim, PLANE=v_dim * k_dim, BLOCK=_GATHER_BLOCK, num_warps=4, + ) # fmt: skip + q3, k3, v3 = q.squeeze(0).contiguous(), k.squeeze(0).contiguous(), v.squeeze(0).contiguous() + out = ( + output.squeeze(0) + if output is not None + else q3.new_empty(q3.shape[0], v3.shape[1], v3.shape[2]) + ) + flashinfer.chunk_gated_delta_rule( + q=q3, + k=k3, + v=v3, + g=g.squeeze(0) if linear_gate else torch.exp(g.squeeze(0)), + beta=beta.squeeze(0).to(torch.float32), + scale=scale, + initial_state=pool, + output_final_state=True, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=False, + output=out, + output_state=pool, + use_cp=False, + state_indices=indices, + ) + return out.unsqueeze(0) + + def _chunk_carried( q, k, v, g, beta, scale, pool, indices, has_initial_state, cu_seqlens, output, linear_gate ) -> torch.Tensor: @@ -172,8 +296,10 @@ class FlashinferChunkGatedDeltaRule(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # spec: geometry (H, HV, K, V); seqlens; `slots`-row pool with `indices` - # (or _engine_slots); pool_dtype (fp32); output (None | "buffer"). Every cell is - # qwen3.8-flash-next-nvfp4's geometry, in place. + # (or _engine_slots); pool_dtype (fp32); output (None | "buffer" | "slice", a prefix + # view of a wider buffer); page_layers (the pool is the layer's view of the stock + # manager's state page, that many GDN states a slot); fi_path (the FlashInfer kernel + # reached). Every cell is in place: tp1's at 16/48, then one rank's of tp4/tep4/dep4. CELLS: tuple[Cell, ...] = tuple( Cell( why=f"qwen3.8 prefill into output=, 461-slot pool, seqlens {list(s)}", @@ -219,7 +345,48 @@ class FlashinferChunkGatedDeltaRule(OpWrapper): linear_gate=True, ), ), - ) + ) + ( + _rank_cell( + "tep4 4/12, fresh 8192-token prompt, CP path, 1874-slot bf16 page", + RANK4_GEOMETRY, (8192,), TEP4_SLOTS, (False,), "cp", + ), + _rank_cell( + "tep4 4/12, [200, 37, 91] carried/fresh/carried, CP path, 1874 slots", + RANK4_GEOMETRY, (200, 37, 91), TEP4_SLOTS, (True, False, True), "cp", + ), + _rank_cell( + "tp4 4/12, carried 8192-token chunk, CP path, 1874-slot bf16 page", + RANK4_GEOMETRY, (8192,), TP4_SLOTS, (True,), "cp", + ), + _rank_cell( + "dep4 16/48, carried 512-token prompt, 850-slot bf16 page", + QWEN38_GEOMETRY, (512,), DEP4_SLOTS, (True,), "chunked", + ), + _rank_cell( + "tep4 4/12 envelope, 8192 rows over 128 seqs of 63/65, chunked path", + RANK4_GEOMETRY, (63, 65) * 64, TEP4_SLOTS, (True, False) * 64, "chunked", + ), + _rank_cell( + "tep4 4/12, the carried 2-token snapshot chunk, in place, 1874-slot bf16 page", + RANK4_GEOMETRY, (2,), TEP4_SLOTS, (True,), "inplace", + ), + _rank_cell( + "tep4 4/12, fresh 5 + carried 2 tokens, in place, 1874 slots", + RANK4_GEOMETRY, (5, 2), TEP4_SLOTS, (False, True), "inplace", + ), + _rank_cell( + "tep4 4/12, one carried 64-token chunk (the in-place bound), 1874 slots", + RANK4_GEOMETRY, (64,), TEP4_SLOTS, (True,), "inplace", + ), + _rank_cell( + "tep4 4/12, one carried 65-token chunk, past the in-place bound: CP path", + RANK4_GEOMETRY, (65,), TEP4_SLOTS, (True,), "cp", + ), + _rank_cell( + "dep4 16/48, carried 2 + fresh 30 + carried 1 tokens, in place, 850 slots", + QWEN38_GEOMETRY, (2, 30, 1), DEP4_SLOTS, (True, False, True), "inplace", + ), + ) # fmt: skip note = """ Metadata consumed: none. @@ -230,11 +397,8 @@ class FlashinferChunkGatedDeltaRule(OpWrapper): si, True, False, qsl_long, False, False, None)` with `head_first=False` in the 12th slot, after `fused_gdn_post_conv` (q/k already normalized, contiguous `[1, T, H, 128]`; `g` fp32 log-space; `beta` fp32). Give it - context rows only: single-token rows are slow here. Measured per layer - at 8065 + 127 x 1 tokens: 1078.8 us with every row in this op, 399.4 us - with the 127 decode rows sent to FlashInfer's decode kernel - (`gdn_decode.gated_delta_rule_decode_pretranspose`) instead (both halves - plus the concatenation). + context rows only: single-token rows are slow here, and belong to + FlashInfer's decode kernel (`gdn_decode.gated_delta_rule_decode_pretranspose`). The target (`PrefillTarget._gdn_core`, `_gdn_core_verify`) passes the context rows by keyword after `gdn_prefill_conv_post(linear_gate=True)`, @@ -254,21 +418,61 @@ class FlashinferChunkGatedDeltaRule(OpWrapper): ragged context rows, and the target's call (300 + 1024 + 500 tokens, the middle one fresh, `has_initial_state`, `linear_gate`, `output=`). + Certified per rank, as `PrefillTarget._gdn_core` calls it (bf16 pool, + `has_initial_state`, `linear_gate`, `output=` a prefix view of a wider + buffer whose tail stays untouched). The pool is the layer's view of the + stock manager's state page (`PAGE_LAYERS` states a slot), at the + member's slot count: + tp4/tep4 (`RANK4_GEOMETRY`, 4 k / 12 v heads), 1874 slots: one fresh + 8192-token prompt, a carried 8192-token chunk, [200, 37, 91] + carried/fresh/carried, and 8192 rows over 128 sequences of 63/65 + tokens (the envelope: max_num_tokens over max_batch_size). + dep4 (16 / 48), 850 slots: a carried 512-token prompt. + The last slot is named at every one: offsets of 1.3e10 (tep4) and + 2.4e10 (dep4) elements, past int32; the gather and the scatter index the + slot in int64. + + Short calls run in place (`fi_path` `inplace`): a `has_initial_state` call of at most + `INPLACE_MAX_TOKENS` (64) context tokens zeroes its fresh sequences' pool rows (one + Triton launch, `_zero_fresh_states_kernel`; a carried row is not touched) and runs + FlashInfer's chunked kernel with `use_cp=False`, `state_indices` the slots and the + pool as both `initial_state` and `output_state`: the kernel reads and writes the named + rows of the (non-compact) pool itself, so neither the packed gather nor the scatter + runs. That is the 2-token snapshot step of chunked prefill + (`additional_snapshot_offsets_from_end: [2]`) on every GDN layer. Numerics: the + chunked kernel's sum, the state read in the pool's + dtype, kept fp32 inside and rounded once to the pool's dtype on the write, as the + packed path rounds it on the scatter. + + FlashInfer picks its kernel by `num_seqs x num_v_heads`: on sm_10x the + context-parallel pipeline (`cp_delta_rule_dsl_sm100`: t precompute, mn + precompute, fixup, prefill) when that times 4 is below the SM count + (`should_use_cp_host`), else the chunked kernel. GB300 has 152 SMs, so a + tp4/tep4 rank (12 v heads) runs CP for 1 to 3 context sequences and the + chunked kernel from 4; a dep4 rank (48) and tp1 always run the chunked + one. Both paths take the bf16 packed state. The cells hold the path + they reach (`fi_path`), and a CP fallback warning fails them. + + Per rank cell against the fp64 reference (bf16 pool), the CP pipeline and + the chunked kernel alike: o rel-L2 <= 3.77e-3, max abs 7.3e-4; named + slots' states rel-L2 <= 3.47e-3, max abs 3.9e-3. A value head reading the + wrong key head at 4/12 gives o rel-L2 1.41, which the band refuses. + Memory per call: a packed copy of the named slots in and out in the pool's dtype (2 x 3 MiB per sequence fp32, half that bf16), `exp(g)` (not with - `linear_gate`), and the - output. q/k/v/beta are not copied - when already contiguous / fp32. - - JIT and streams: the first call per process compiles the CuTe-DSL - kernel in-process (seconds), and the installed FlashInfer keeps one TMA - workspace per device, so two streams must not run this op concurrently. - Needs CUDA >= 13 torch (FlashInfer raises otherwise). - - Numerics (measured vs an fp64 recurrence, 1 x 8192 tokens, qwen3.8 - geometry): o rel-L2 4.06e-3, max abs 1.26e-3; state rel-L2 3.31e-3 -- - the Triton op's are - 4.00e-3 / 1.13e-3 / 3.26e-3. Not bitwise to it. A bf16 pool rounds each + `linear_gate`), and the output. q/k/v/beta are not copied when already + contiguous / fp32. + + JIT and streams: the first call per process and path compiles the + CuTe-DSL kernels in-process for the device it runs on (the cells ran + both paths on sm_103), and the installed + FlashInfer keeps one TMA workspace per device (the CP pipeline also a + cached initial-state workspace), so two streams must not run this op + concurrently. Needs CUDA >= 13 torch (FlashInfer raises otherwise). + + Numerics (vs an fp64 recurrence, 1 x 8192 tokens, qwen3.8 geometry): o + rel-L2 4.06e-3, max abs 1.26e-3; state rel-L2 3.31e-3 -- the Triton op's + are 4.00e-3 / 1.13e-3 / 3.26e-3. Not bitwise to it. A bf16 pool rounds each sequence's state to bf16 once per call, on the scatter. No cold-call defect: nothing in the call is autotuned. @@ -280,12 +484,20 @@ class FlashinferChunkGatedDeltaRule(OpWrapper): caller or the flag). Not certified: `state_workspace=`, an fp16 pool, `scale` other than the - default, fp16 operands, K or V other than 128, other head ratios, and what the - target never passes: `inplace_indexed_state_update=False` (final states - returned) and `use_qk_l2norm_in_kernel`. + default, fp16 operands, K or V other than 128, head layouts other than + 16/48 and 4/12, and what the target never passes: + `inplace_indexed_state_update=False` (final states returned) and + `use_qk_l2norm_in_kernel`. At 4/12: 2 context sequences on CP and 4 to + 127 on the chunked kernel (the cells hold 1 and 3, and 128); a dep4 rank + at more than 512 carried tokens; tp4's own slot count (tep4's stands in) + and any pool counted on the default target's own engines or under MTP, + kv_fp8/kv_nvfp4, chunked or reuse; + a call inside a CUDA graph (the target calls it eagerly). Hard limits - none probed. + T: max_num_tokens (8192) rows a call is certified at; chunked + prefill of a longer prompt is one call per chunk. + slot offsets: past int32 at both members' last slots (certified). """ def __call__( @@ -322,7 +534,7 @@ def __call__( and cu_seqlens is not None and g.dtype == torch.float32 ), "has_initial_state: the in-place indexed call only, q/k normalized by the caller" - o = _chunk_carried( + o = (_chunk_inplace if q.shape[1] <= INPLACE_MAX_TOKENS else _chunk_carried)( q, k, v, g, beta, scale, initial_state, initial_state_indices, has_initial_state, cu_seqlens, output, linear_gate, ) # fmt: skip diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py similarity index 76% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py index 7ed03ff9eb23..ce6852817a35 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py @@ -9,6 +9,9 @@ (tier-2 sourcing: the recurrence has no op boundary). One launch of the vendored Triton kernel ``_cached_replay_kernel`` -- plain ``triton.jit``, no autotuner; the Python body picks the launch configuration from the shapes. + +Called by the qwen4_exp target at every layout: the cells of tp1 (all 16 key / 48 value +heads), plus a tp4/tep4 rank's 4 key / 12 value heads over the tep4 engine's bf16 state page. """ from __future__ import annotations @@ -38,6 +41,13 @@ #: Layers per slot at which the last of `ENGINE_SLOTS` slots starts past 2^31 elements of #: the pool (`460 * L * 48 * 128 * 128 > 2**31` from L = 6). PAST_INT32_LAYERS = 8 +#: A tp4/tep4 rank's GDN heads: 16 key / 48 value split four ways. +RANK_K_HEADS, RANK_V_HEADS = 4, 12 +#: Recurrent-state slots of a tep4 engine rank (bf16 state, host-resident PLE table). +TEP4_SLOTS = 1874 +#: GDN layers whose states share one slot page in the stock manager (48 layers, every 4th +#: full attention): the slot stride is 36 layers' `[HV, V, K]` states. +PAGE_LAYERS = 36 def _cell_spec(path, n, t, **spec) -> dict: @@ -50,6 +60,25 @@ def _bf16_cell(why, path, n, t, **spec) -> Cell: ) +def _tep4_cell(why, path, n, t, **spec) -> Cell: + """A tp4/tep4 rank's heads over the stock tep4 engine's bf16 state page.""" + return Cell( + why=f"tep4 {why}", + spec=_cell_spec( + path, + n, + t, + state_dtype=torch.bfloat16, + H=RANK_K_HEADS, + HV=RANK_V_HEADS, + history_size=HISTORY_SIZE, + slots=TEP4_SLOTS, + page_layers=PAGE_LAYERS, + **spec, + ), + ) + + class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper): """One verify step of `N` sequences x `T` tokens through one GDN layer, for the cached-replay state layout. Per live row (`slot = state_indices[n] >= 0`), @@ -112,14 +141,22 @@ class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper): # zeroed history, as the manager leaves a slot after prefill), edge_slots # (rows 0 and n-1 on the pool's last and first slot), dup (rows given the # SAME slot, as CUDA-graph padding repeats one dummy slot; only the other - # rows and slots are held to the reference). Every cell passes `packed_qkv`, + # rows and slots are held to the reference), H / HV (key / value heads; tp1's + # 16 / 48 where absent), history_size (16, the only one certified), page_layers + # (the stock manager's coalesced page: the slot stride is that many layers' + # states, `slots` slots of it a reserved virtual range with memory mapped under + # the touched slots only, the cell's `layers` the top sub-pages; rows on + # both ends of the pool, then spread over it). Every cell passes `packed_qkv`, # separate raw `a`/`b` projections, no `output` and `launch_with_pdl=False`. # - # Four cells, one over the three-cell budget: the forward reaches two commit - # paths (in-kernel below 16 requests, all-layer from 16) at two checkpoint - # dtypes, and two edges no main-shape cell can carry -- CUDA-graph padding - # rows sharing one slot, and slot offsets past 2^31 elements at the engine's - # pool. + # tp1's four cells: the forward reaches two commit paths (in-kernel below 16 + # requests, all-layer from 16) at two checkpoint dtypes, and two edges no + # main-shape cell can carry -- CUDA-graph padding rows sharing one slot, and slot + # offsets past 2^31 elements at the engine's pool. Then a tp4/tep4 rank's six: + # the verify widths T = 2, 4, 8 (max_draft_len 1, 3, 7), both commit + # paths, and N = 128, where a bf16 state turns on the large-workload mapping of + # the replay and the commit; all at the tep4 engine's page, whose top slot sits + # 1.3e10 elements past the base. CELLS: tuple[Cell, ...] = ( Cell( why="all_layer n=32 T=8 full accept: commit every second step", @@ -150,6 +187,40 @@ class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper): steps=3, ), ), + _tep4_cell( + "small n=1 T=4: the in-kernel commit of the page's last slot (1873)", + "small", + 1, + 4, + steps=8, + ), + _tep4_cell( + "all_layer n=16 T=4: the smallest partitioned batch", "all_layer", 16, 4, steps=6 + ), + _tep4_cell( + "all_layer n=32 T=8: the widest verify below N=128", "all_layer", 32, 8, steps=6 + ), + _tep4_cell( + "all_layer n=128 T=4: bf16 at N=128, the large-workload mapping", + "all_layer", + 128, + 4, + steps=4, + ), + _tep4_cell( + "all_layer n=128 T=2: the BT=2 token tile (max_draft_len 1)", + "all_layer", + 128, + 2, + steps=6, + ), + _tep4_cell( + "all_layer n=128 T=8: the envelope's largest verify, 128 x 8 rows", + "all_layer", + 128, + 8, + steps=4, + ), ) note = """ @@ -176,6 +247,16 @@ class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper): when it does commit twice -- both silent at the op and far off the exact recurrence one step later. + The default target at tp4/tep4 (`DecodeTarget._gdn_verify -> _replay_verify`): a rank's + 4 key / 12 value heads, q/k/v the `[N, T, 4|12, 128]` views of the post-conv + `[N*T, 2560]` rows, `g`/`beta` `[N, T, 12]`, and `ssm_states` the layer's view of the + stock manager's bf16 state page, `[slots, 12, 128, 128]` at a slot stride of 36 + layers' states (7,077,888 elements): 1874 slots at the stock tep4 engine, whose top + slot sits 1.3e10 elements past the page's base. The history buffers are + `[slots, 2, 16, 12|4, 128]` (and `[slots, 2, 12, 16]` fp32); the manager's all-layer + commit stays on, so the call is partitioned from N = 16 exactly as above. At dep4 a + rank holds all 16 / 48 heads, tp1's geometry. + Metadata consumed: `state_indices` (one slot per row), and per slot `prev_num_accepted_tokens` and `cache_buf_idx`. `replay_work_items` / `n_writes` are only checked (shape, dtype); the kernel never reads them, their @@ -200,14 +281,17 @@ class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper): * the conv state is not this op's: under replay the conv still writes its per-step window and the manager promotes it. - The launch is one V tile, `num_warps=4`, `num_stages=1`, grid `(1, N * 48)`: the - tuned mappings need HV = 2H or 4H and never engage at HV/H = 3. The checkpoint's dtype - is the kernel's `STATE_FP32` specialization and nothing else below N = 128: an fp32 - checkpoint enters the history rebuild's tensor-core dots as a bf16 hi/lo pair, a bf16 - one as stored, and every commit is computed in fp32 and rounded once to the pool's - dtype. From N = 128 a bf16 state alone also turns on the L2-streaming mapping - (`main_num_stages=3`, a 5-stage commit), which no cell drives. `FLA_USE_FAST_OPS=1` at - import switches the kernel's `exp` to the approximate one; the cells run with it unset. + The kernels are Triton, JIT-compiled on first call for the running device; the cells + ran that JIT path on sm_103. The launch is one V tile, `num_warps=4`, grid + `(1, N * HV)`: the tuned mappings need HV = 2H or 4H and never engage at HV/H = 3 + (16 / 48 and 4 / 12 alike). The checkpoint's dtype is the kernel's `STATE_FP32` + specialization and nothing else below N = 128: an fp32 checkpoint enters the history + rebuild's tensor-core dots as a bf16 hi/lo pair, a bf16 one as stored, and every + commit is computed in fp32 and rounded once to the pool's dtype. From N = 128 a bf16 + state alone also turns on the large-workload mapping (`.cg` loads of the state and + the inputs, `main_num_stages=3`, and a 5-stage all-layer commit); the tep4 N = 128 + cells drive it. `FLA_USE_FAST_OPS=1` at import switches the kernel's `exp` to the + approximate one; the cells run with it unset. CUDA-graph replay on new data is bit for bit the eager call (the test captures N = 32 at T = 8 over an fp32 checkpoint, with the all-layer commit). The two commit @@ -215,16 +299,28 @@ class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper): certified against each other; each form the target calls is certified against the reference by the cells. - Certified: the all-layer commit over an fp32 checkpoint at N = 32, T = 8 and at - N = 16, T = 3 with 8 interleaved layers over the engine's 461-slot pool (the last - slot's offset past 2^31 elements); the in-kernel commit over a bf16 checkpoint at - N = 4, T = 8 and over an fp32 one at N = 8, T = 3 with three CUDA-graph padding rows - on one shared slot. T = 5..8 share one compiled token tile (`BT = next_pow2(T) = 8`); - the cells drive the BT = 4 and BT = 8 tiles, not BT = 2 (T = 2). At T = 8 and history + Certified at tp1's 16 / 48 heads: the all-layer commit over an fp32 checkpoint at + N = 32, T = 8 and at N = 16, T = 3 with 8 interleaved layers over the tp1 engine's + 461-slot pool (the last slot's offset past 2^31 elements); the in-kernel commit over a + bf16 checkpoint at N = 4, T = 8 and over an fp32 one at N = 8, T = 3 with three + CUDA-graph padding rows on one shared slot. T = 5..8 share one compiled token tile + (`BT = next_pow2(T) = 8`); these cells drive the BT = 4 and BT = 8 tiles, not + BT = 2 (T = 2). At T = 8 and history 16 a slot whose every draft is accepted appends rows 8..15 at `pnat + T == 16` without wrapping and wraps on the next step, so it commits every second step; the `accept="full"` cells drive that boundary (in-kernel and all-layer commit). + Certified at a tp4/tep4 rank's 4 / 12 heads over a bf16 checkpoint on the stock tep4 + engine's 1874-slot page, each cell several consecutive verify steps from warm + histories with random acceptance, rows on both ends of the pool: the in-kernel commit + at N = 1, T = 4 on the page's last slot; the all-layer commit at N = 16, T = 4, + N = 32, T = 8, and N = 128 at T = 2, 4 and 8 (the large-workload mapping, the BT = 2, + 4 and 8 token tiles, and the envelope's largest verify, 128 x 8 rows). The gates: + output elements within 4 bf16 ulps of their token block and relative RMS within 1.5 + ulps; the bf16 checkpoint within 2 ulps of its head's absmax; `old_u` within 4 ulps; + every cell lands at most 0.84 of its gate. A reference with the key heads reversed + (value heads 0..2 on key head 3) is refused by the output gate. + Loud refusals: work items without `use_all_layer_commit` (ValueError) and the reverse (ValueError); one of `replay_work_items` / `n_writes` alone (ValueError); a strided `output`; a strided head axis of `g`; `history_size < @@ -243,12 +339,20 @@ class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper): harmless: the entry's input guard makes it contiguous. Not certified: pre-gated `g`/`beta` (`A_log=None`); `use_qk_l2norm_in_kernel= - False`; an `ssm_states` other than fp32 or bf16, or a bf16 one at N >= 128; HV/H other - than 3, K or V other than 128; T = 2, or T outside 2..8 (T = 9 raises: the kernel's - bare `assert T <= 8`), at either `ssm_states` dtype; `history_size` other than 16; `old_u` / - `old_k` other than bf16, `old_G` other than fp32; a `scale`; any launch-knob - override (`block_v`, `num_warps`, `use_l2_*`, `main_num_stages`); an output - dtype other than bf16; int64 `state_indices`; N above 32. + False`; an `ssm_states` other than fp32 or bf16; at 16 / 48 heads a bf16 one at + N >= 128, N above 32, T = 2, and any slot count but the cells' (dep4's own pool is not + driven); HV/H other than 3, K or V other than 128; T = 3, 5, 6 and 7 at 4 / 12 (they + share the BT = 4 and 8 tiles), T outside 2..8 (T = 9 raises: the kernel's bare + `assert T <= 8`); N above 128; `history_size` other than 16; `old_u` / `old_k` + other than bf16, `old_G` other than fp32; a `scale`; any launch-knob override + (`block_v`, `num_warps`, `use_l2_*`, `main_num_stages`); an output dtype other than + bf16; int64 `state_indices`; `g`/`beta` as strided column views of the input + projection (the cells pass them contiguous). At 4 / 12: the commit over the engine's + 36 layers (the cells commit 2 sub-pages of the 36-layer page; a 36-layer history puts + `old_u`'s layer offsets past 2^31 elements, addressed in int64 by the commit's source + but not driven); a tp4 engine's slot count and the MTP engine's at tep4 (neither + measured); CUDA-graph capture and padding rows sharing a slot (both certified at + 16 / 48 only). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_gated_norm_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_gated_norm_pdl.py similarity index 83% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_gated_norm_pdl.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_gated_norm_pdl.py index e9179d141741..2112962a1b7a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_gated_norm_pdl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_gated_norm_pdl.py @@ -3,10 +3,8 @@ """One Gated-DeltaNet decode step per row, and the gated RMS norm on its output, in one PDL launch. -Tier 5 (CuTe DSL written for this target; the recurrence is the gated delta rule of -FlashInfer's `gdn_decode.gated_delta_rule_decode_pretranspose`, reimplemented here). At one -decode row the step reads and writes each v-head's 64 KB fp32 state once and -otherwise does a few hundred flops per state row; the rest of its time is latency. +A CuTe DSL kernel; the recurrence is the gated delta rule of FlashInfer's +`gdn_decode.gated_delta_rule_decode_pretranspose`, reimplemented here. The state is not produced by anything in the step ahead of this kernel -- the pool row is this layer's own and was last written by this kernel one step ago, and the slot index is filled before the step's graph is launched -- so it is read @@ -35,10 +33,11 @@ from cutlass.cutlass_dsl import dsl_user_op from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.norm.rms_norm_gated_token_major import ( + rms_norm_gated_token_major as _gated_norm, +) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl -from ..norm.rms_norm_gated_token_major import rms_norm_gated_token_major as _gated_norm - #: Decode rows per launch this entry is certified for. MAX_ROWS = 4 #: CTAs per v-head (one cluster) and warps per CTA; their product must be <= 32. @@ -386,9 +385,11 @@ class GdnDecodeGatedNormPdl(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: rows; pool (default a contiguous 461-slot pool | "coalesced": the - # cache manager's pool, slot stride 36 layers' states, slots 100 and 101); - # pad (row 1 given slot -1). + # spec: rows; k_heads / v_heads (default 16 / 48); pool (default a contiguous 461-slot + # pool | "coalesced": the cache manager's page, slot stride 36 layers' states -- 102 slots + # (slots 100 and 101), or with pool_slots the engine's measured slot count, the view the + # last layer's sub-page and the top slots named); pad (row 1 given slot -1; at one row, a + # second call over slot -1). CELLS: tuple[Cell, ...] = ( Cell(why="qwen3.8 decode, 1 row (con 1), 461-slot pool", spec=dict(rows=1)), Cell(why="qwen3.8 decode, 4 rows (MAX_ROWS), scattered slots", spec=dict(rows=4)), @@ -396,6 +397,25 @@ class GdnDecodeGatedNormPdl(OpWrapper): why="2 rows, coalesced pool (slot stride 36 layers, slot 100), row 1 a pad slot (-1)", spec=dict(rows=2, pool="coalesced", pad=True), ), + Cell( + why="tp4/tep4 decode, 1 row at 4 k / 12 v heads, slot 1873 of tep4's 1874-slot " + "engine page, then a pad slot", + spec=dict(rows=1, k_heads=4, v_heads=12, pool="coalesced", pad=True, pool_slots=1874), + ), + Cell( + why="tp4/tep4, 2 rows at 4 / 12 heads, the engine page's top slots, row 1 a pad " + "slot (-1)", + spec=dict(rows=2, k_heads=4, v_heads=12, pool="coalesced", pad=True, pool_slots=1874), + ), + Cell( + why="tp4/tep4, 4 rows (MAX_ROWS) at 4 / 12 heads, scattered over the 1874-slot page", + spec=dict(rows=4, k_heads=4, v_heads=12, pool="coalesced", pool_slots=1874), + ), + Cell( + why="dep4 decode, 1 row at 16 / 48 heads, slot 849 of dep4's 850-slot engine page, " + "then a pad slot", + spec=dict(rows=1, k_heads=16, v_heads=48, pool="coalesced", pad=True, pool_slots=850), + ), ) note = """ @@ -406,7 +426,7 @@ class GdnDecodeGatedNormPdl(OpWrapper): `dt_bias` fp32 `[HV]`; `norm_weight` bf16 `[V]`; `V % 32 == 0` (rows a warp: `V / 32`), `K % 32 == 0`; int32 `slots` `[rows]`. **The pool rows and `slots` must not be written by any kernel of the same step queued ahead of this one**: both are read - before `griddepcontrol.wait`. That holds for this target's decode steps (each + before `griddepcontrol.wait`. That holds for the default target's decode steps (each layer's pool is written only by its own recurrence, `slots` is filled before the step's graph launches) and is why the entry is not for verify or prefill steps. @@ -415,10 +435,40 @@ class GdnDecodeGatedNormPdl(OpWrapper): rounded once. Inside torch's default band for bf16 / fp32 of the fp64 composition. Deterministic. Graph-capturable after the first call per geometry (compiled then). - Certified: 1 and 4 rows over a contiguous 461-slot pool, and 2 rows over the cache - manager's coalesced pool (slot stride 36 layers' states, slot 100 ~2.8G elements past - the base) with the second row a pad slot (-1). Under graph replay the op reads the - projection after `griddepcontrol.wait` (1 row). Not certified: 3 rows, `pdl=False`. + Source: the CuTe DSL kernel in this file, JIT-compiled (`cute_ext.compile`) for the + running device at the first call per geometry; the cells ran that path on sm_103. + + The default target: its routing admits a bf16 recurrent-state cache only, and this kernel takes + an fp32 pool only -- a bf16 pool fails its compile loudly ("`copy` currently only + supports equal source and destination element type bit width") -- so the default target's + decode step does not launch it. Its bf16 decode recurrence is + `gdn_decode_rows_gated_norm`, which reads this entry's `reference` and + `_recurrence_reference` (both round the state to the pool's own dtype) and not its + kernel. + + Certified, sm_103: + - tp1's 16 / 48 heads: 1 and 4 rows over a contiguous 461-slot pool, and 2 rows over a + 102-slot coalesced page (slot stride 36 layers' states, slot 100 ~2.8G elements past + the base) with the second row a pad slot (-1). + - A tp4/tep4 rank's 4 / 12 heads (proj width 4120): 1 row, 2 rows (row 1 a pad) and 4 + rows over tep4's engine page -- 1874 slots, slot stride 36 layers' states (7077888 + elements), the view the last GDN layer's sub-page; slot 1873 sits ~1.3e10 elements + past the base. + - dep4's 16 / 48 heads: 1 row at slot 849 of its 850-slot engine page (~2.4e10 elements + past the base). + - A one-row call over a pad slot (-1) writes zeros and no state; under graph replay the + op reads the projection after `griddepcontrol.wait`, at 16 / 48 and at 4 / 12 heads. + Against the fp64 reference: at 4 / 12 heads every output element is bitwise its + bf16-rounded reference and the state is within 1.2e-7 (relative L2 <= 7.8e-8); at + 16 / 48 at most 5 of 24576 output elements differ (max 9.8e-4, relative L2 1.2e-5). On + the engine pages every other byte of a touched slot's page stays unchanged. + + The engine slot counts are the stock engines' with a bf16 state cache at + `max_batch_size` 128 (the PLE / GDN slot space the manager sizes from memory); an fp32 + pool in that memory holds about half as many, so the certified slots bound it. + + Not certified: 3 rows; `pdl=False`; tp4's own engine slot count (tep4's 1874 stands in: + a tp4 rank's GDN pools are the same shape); a bf16 pool (refused, above). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_rows_gated_norm.py similarity index 81% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_rows_gated_norm.py index 7fb0f8144a6b..c3c9d94b7133 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_decode_rows_gated_norm.py @@ -4,10 +4,9 @@ over the fp32 or bf16 state pool: the decode step at batch sizes past the cluster kernel's over an fp32 pool, and at every batch size over a bf16 pool. -Tier 6 (CuTe DSL written for this target; the recurrence is the gated delta rule of -FlashInfer's `gdn_decode.gated_delta_rule_decode_pretranspose`, reimplemented here). At a -decode batch the step is the pool's traffic -- each (row, v-head) state, 64 KB fp32 or -32 KB bf16, read and written once -- so a CTA owns one whole state and streams it: warp w +A CuTe DSL kernel; the recurrence is the gated delta rule of FlashInfer's +`gdn_decode.gated_delta_rule_decode_pretranspose`, reimplemented here. A CTA owns one +(row, v-head) state, 64 KB fp32 or 32 KB bf16, read and written once, and streams it: warp w the value rows [w * RPW, (w + 1) * RPW), lane l the key columns [4l, 4l + 4) of each, so a warp's access is one contiguous state row. A bf16 state is widened to fp32 on the load and rounded back (nearest even) on the store. The state is loaded before @@ -22,8 +21,7 @@ output and state): the state rows loaded into registers (128 registers a thread, two CTAs an SM), or -- from the batch where the register variant needs a second wave -- each warp's rows brought into shared memory by one bulk copy (TMA) issued as soon as the slot is known, -which holds no registers while in flight (64 a thread, three to four CTAs an SM): the pool -streams at 6.8 TB/s read + write instead of 5.3-5.7. +which holds no registers while in flight (64 a thread, three to four CTAs an SM). """ from __future__ import annotations @@ -38,19 +36,50 @@ from cutlass.cutlass_dsl import T, dsl_user_op from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_decode_gated_norm_pdl import ( + _recurrence_reference, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_decode_gated_norm_pdl import ( + gdn_decode_gated_norm_pdl as _one_launch, +) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr -from .gdn_decode_gated_norm_pdl import _recurrence_reference -from .gdn_decode_gated_norm_pdl import gdn_decode_gated_norm_pdl as _one_launch - #: Key columns a lane (one 16-byte fp32 or 8-byte bf16 vector of a state row). _PL = 4 #: Warps a CTA; a warp owns `V // _WARPS` value rows of the head's state. _WARPS = 8 -#: Largest recurrent-pool slot count the cache manager was seen to allocate for this -#: checkpoint; the kernel addresses a slot by its index alone, any pool size alike. +#: Largest recurrent-pool slot count the cache manager allocates for this checkpoint; the +#: kernel addresses a slot by its index alone, any pool size alike. ENGINE_SLOTS = 461 +#: Recurrent-state slots one rank's engine allocates for this checkpoint (bf16 state cache, +#: `max_batch_size` 128), counted at tep4 and dep4. A tp4 rank's GDN state has tep4's +#: per-slot shape (4 k / 12 v heads), so tp4 takes tep4's count. +TEP4_SLOTS = 1874 +DEP4_SLOTS = 850 +TP4_SLOTS = TEP4_SLOTS + +#: The rank geometries the engine cells take: a tp4/tep4 rank's 4 k / 12 v heads on its +#: 4120-wide projection, a dep4 rank's whole 16 / 48 on 16480, each over its member's bf16 +#: engine page. +_TEP4 = dict( + k_heads=4, + v_heads=12, + proj_width=4120, + state="bf16", + pool="coalesced", + pool_slots=TEP4_SLOTS, + member="tep4", +) +_DEP4 = dict( + k_heads=16, + v_heads=48, + proj_width=16480, + state="bf16", + pool="coalesced", + pool_slots=DEP4_SLOTS, + member="dep4", +) def _iv(x, loc, ip): @@ -655,8 +684,7 @@ def _use_bulk(rows, v_heads, device, state_bf16=True) -> bool: Over a bf16 pool, from the batch the register variant cannot hold in one wave (two of its CTAs an SM): below that the register variant's lead (its rows are in registers when the wait returns) is worth more than the occupancy. An fp32 pool keeps the register variant: - its 64 state registers a thread already hold the CTA's 64 KB in flight (6.9 TB/s; the - bulk copy 6.85). + its 64 state registers a thread already hold the CTA's 64 KB in flight. """ sms = torch.cuda.get_device_properties(device).multi_processor_count return state_bf16 and rows * v_heads > 2 * sms @@ -779,9 +807,11 @@ class GdnDecodeRowsGatedNorm(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: rows; pool (default a contiguous 461-slot pool | "coalesced": the - # cache manager's pool, slot stride 36 layers' states, slots 100 and 101); - # pad (row 1 given slot -1). + # spec: rows; k_heads / v_heads (default 16 / 48) and proj_width (their projection's); + # state (default fp32 | "bf16"); pool (default a contiguous 461-slot pool | + # "coalesced": the cache manager's page, slot stride 36 layers' states -- 102 slots + # (slots 100 and 101), or with pool_slots an engine's slot count, the view the last + # layer's sub-page); pad (row 1 given slot -1); member (the layout the counts are from). CELLS: tuple[Cell, ...] = ( Cell(why="qwen3.8 decode, 128 rows (con 128), 461-slot pool", spec=dict(rows=128)), Cell(why="qwen3.8 decode, 3 rows (past the conv-fused path)", spec=dict(rows=3)), @@ -798,6 +828,38 @@ class GdnDecodeRowsGatedNorm(OpWrapper): why="bf16 pool, 2 rows, coalesced: slot stride 36 layers, slot 100", spec=dict(rows=2, pool="coalesced", state="bf16"), ), + Cell( + why="tep4 decode, 1 row (con 1) at 4 k / 12 v heads, slot 1873 of the 1874-slot " + "bf16 engine page", + spec=dict(_TEP4, rows=1), + ), + Cell( + why="tep4, 2 rows at 4 / 12 heads: slots 1873 and 0 of the 1874-slot bf16 page", + spec=dict(_TEP4, rows=2), + ), + Cell( + why="tep4, 3 rows at 4 / 12 heads over the 1874-slot bf16 engine page", + spec=dict(_TEP4, rows=3), + ), + Cell( + why="tep4, 128 rows (con 128; the bulk-copy variant) at 4 / 12 heads over the " + "1874-slot bf16 page", + spec=dict(_TEP4, rows=128), + ), + Cell( + why="tp4, 128 rows at 4 / 12 heads over its 1874-slot bf16 page (tep4's count), " + "row 1 a pad slot (-1)", + spec=dict(_TEP4, rows=128, pad=True, pool_slots=TP4_SLOTS, member="tp4"), + ), + Cell( + why="dep4 decode, 32 rows at 16 / 48 heads (proj 16480) over the 850-slot bf16 " + "engine page", + spec=dict(_DEP4, rows=32), + ), + Cell( + why="dep4, 128 rows (max_batch_size) at 16 / 48 heads over the 850-slot bf16 page", + spec=dict(_DEP4, rows=128), + ), ) note = """ @@ -809,19 +871,16 @@ class GdnDecodeRowsGatedNorm(OpWrapper): `V / 8` in {4, 8, 16, 32}; `A_log`, `dt_bias` fp32 `[HV]`; `norm_weight` bf16 `[V]`; int32 `slots` `[rows]`. **The pool rows and `slots` must not be written by any kernel of the same step queued ahead of this one**: both are read before - `griddepcontrol.wait` -- as `gdn_decode_gated_norm_pdl` states, this target's decode + `griddepcontrol.wait` -- as `gdn_decode_gated_norm_pdl` states, the default target's decode steps keep that (a layer's pool is written only by its own recurrence, `slots` is filled before the step's graph launches), and verify and prefill steps do not use the entry. Launch: one CTA of 8 warps a (v-head, row), `HV * rows` CTAs, PDL (off with TRTLLM_ENABLE_PDL=0); a bf16 pool's state by bulk copy into shared memory from - `rows * HV > 2 * SMs` (7 rows here; an fp32 pool, or a layout the copy cannot take, - stays on registers), both variants compiled on the first call. GB300, bf16 pool, a CUDA - graph of 36 layers each streaming its own pool from HBM, PDL on / off, us a layer - (register-only kernel -> this): 256 rows 152.4 / 141.6 -> 118.8 / 119.4 (6.8 TB/s read + - write), 128 rows 77.7 / 72.6 -> 61.1 / 61.7, 32 rows 21.7 / 20.7 -> 17.9 / 18.5, 8 rows - 7.2 / 7.5 -> 6.5 / 6.9, 1 row 2.6 / 3.5 (the register variant). + `rows * HV > 2 * SMs` (GB300's 152 SMs: 7 rows at 48 v heads, 26 at a tp4/tep4 rank's + 12; an fp32 pool, or a layout the copy cannot take, stays on registers), both variants + compiled on the first call. Numerics: fp32 inside; each dot product over K is a lane's 4-term FMA chain and then the transposed reduction's tree; `o` rounded to bf16 before the norm, the output @@ -830,14 +889,52 @@ class GdnDecodeRowsGatedNorm(OpWrapper): the pool's dtype of the composition's float64 reference, not bitwise FlashInfer's kernel. Batch-invariant: a row's output and state do not depend on the other rows. - Certified: 128 and 3 rows (the fewest a decode step sends here once the conv-fused - path stops) over a contiguous 461-slot pool, and 2 rows over the cache manager's - coalesced pool (slot stride 36 layers' states, slot 100 ~2.8G elements past the - base) with the second row a pad slot (-1); the same over a bf16 pool. Under graph - replay the op reads the projection after `griddepcontrol.wait` (16 rows). The bulk-copy - and register variants bitwise equal (outputs and state) at 1-33 rows over contiguous bf16 - and fp32 pools and the coalesced pool (slot 100 by the bulk copy), pad rows included. Not - certified: `pdl=False`. + Source: the CuTe DSL kernel in this file, JIT-compiled (`cute.compile`) for the running + device at the first call per geometry; the cells ran that path on sm_103 (GB300). + + The default target: its routing admits a bf16 recurrent-state cache only, so its decode step + runs this entry over the bf16 pool at every batch size (the fp32-only one-launch kernel, + `gdn_decode_gated_norm_pdl`, is never launched here; its references are this entry's). + A tp4/tep4 rank holds 4 k / 12 v heads, its projection `[rows, 4120]`; a dep4 rank all + 16 / 48, `[rows, 16480]`. The pool is the stock manager's coalesced page: every GDN + layer's state of a slot in one page, slot stride 36 layers' states (7077888 elements at + 12 v heads, 28311552 at 48), a layer's view its sub-page. + + Certified, sm_103: + - tp1's 16 / 48 heads. fp32 pool: 128 and 3 rows (the fewest a decode step sends here + once the conv-fused path stops) over a contiguous 461-slot pool, and 2 rows over a + 102-slot coalesced page (slot 100 ~2.8G elements past the base) with the second row a + pad slot (-1). bf16 pool: 128, 1 and 33 rows (one pad slot) over the contiguous pool, + and 2 rows over the coalesced page. + - A tp4/tep4 rank's 4 / 12 heads over tep4's bf16 engine page (1874 slots, the view the + last GDN layer's sub-page; slot 1873 sits ~1.3e10 elements past the base): 1, 2, 3 + rows (register variant) and 128 rows (bulk copy); tp4 128 rows with a pad slot over + the same 1874 slots. + - dep4's 16 / 48 heads over its bf16 engine page (850 slots; slot 849 ~2.4e10 elements + past the base): 32 rows and 128 rows (`max_batch_size`), both by the bulk copy. + - Every other byte of a touched slot's page stays unchanged; a pad row's output is zero. + - Under graph replay the op reads the projection after `griddepcontrol.wait`: 16 rows at + 16 / 48, 3 rows (register) and 128 rows (bulk copy) at 4 / 12. + - Batch-invariant over a bf16 pool, bitwise, at 1-256 rows across the variant switch, at + 16 / 48 and 4 / 12. + - The bulk-copy and register variants bitwise equal (outputs and state), with and without + the norm, at 1-33 rows over contiguous bf16 and fp32 pools and the tp1 coalesced page, + and at 4 / 12 heads over a contiguous bf16 pool (3 and 128 rows) and tep4's engine page + (3 rows at slots 1873, 0, 5). + Against the fp64 reference, torch's default bands: at 4 / 12 heads at most 27 of 196608 + output elements differ in value from the bf16-rounded reference (max 3.9e-3; relative + L2 <= 2.7e-5 at 128 rows, 5.2e-5 at 3 rows, where one element of 4608 differs by + 2.0e-3), state relative L2 <= 2.4e-5 (max 2.0e-3); dep4 at most + 93 of 786432 (max 7.8e-3, relative L2 3.1e-5), state relative L2 <= 1.7e-5. + + The engine slot counts are at a bf16 state cache, `max_batch_size` 128 and + `kv_cache_free_gpu_memory_fraction` 0.6; a tp4 rank's GDN pools have tep4's per-slot + shape, so its cell takes tep4's count. The kernel addresses a slot as + `Int64(slot) * s_slot`: any slot count an int32 index can name. + + Not certified: `pdl=False`; the slot counts of the members' variants (kv_fp8, kv_nvfp4, + chunked, reuse at dep4; MTP at tep4, `max_seq_len` 16384) and of tp4's own engine; this + target's own call, which its modeling does not make yet. """ def __call__( @@ -941,6 +1038,20 @@ class GdnDecodeRows(OpWrapper): why="bf16 pool: 3 rows, one pad slot (-1)", spec=dict(rows=3, pad=True, out="row_range", state="bf16"), ), + Cell( + why="tp4/tep4 mixed step: 3 generation rows at 4 k / 12 v heads (proj 4120), bf16 " + "pool, into a row range", + spec=dict( + rows=3, k_heads=4, v_heads=12, proj_width=4120, state="bf16", out="row_range" + ), + ), + Cell( + why="tp4/tep4 mixed step: 127 generation rows (max_batch_size with one context " + "row) at 4 / 12 heads", + spec=dict( + rows=127, k_heads=4, v_heads=12, proj_width=4120, state="bf16", out="row_range" + ), + ), ) note = """ @@ -950,8 +1061,17 @@ class GdnDecodeRows(OpWrapper): writes the generation rows' slots (the context rows' reset and chunked recurrence touch only theirs). Numerics: the recurrence of `GdnDecodeRowsGatedNorm` (the same kernel). - Certified: a mixed step's 121 generation rows, 3 rows with a pad slot (-1), and 2 rows - over the cache manager's coalesced pool, each into a row range of the step output. + Certified, sm_103: at tp1's 16 / 48 heads a mixed step's 121 generation rows, 3 rows with + a pad slot (-1), and 2 rows over the cache manager's coalesced pool, each into a row range + of the step output (fp32 and bf16 pools). At a tp4/tep4 rank's 4 / 12 heads (proj 4120) + over a contiguous bf16 pool, 3 rows and 127 (`max_batch_size` with one context row), into + a row range: output at most 12 of 195072 elements off the bf16-rounded fp64 reference + (max 3.1e-5), state relative L2 <= 9.8e-6. The bitwise bulk / register agreement with + `norm` off covers this entry's launch at 4 / 12 over tep4's engine page. + + Not certified: this entry's own call over an engine's page (a mixed step's eager path at + tp4/tep4 or dep4), `pdl=False`; at dep4 a mixed step runs the 16 / 48 cells above, over + tp1's pools. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_prefill_conv_post.py similarity index 79% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_prefill_conv_post.py index 99ea41dfd83d..acb71d5577f5 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/ssm/gdn_prefill_conv_post.py @@ -5,14 +5,9 @@ the SiLU, the q/k split and L2 norm, the v copy and the decay/beta gates fused in, and each sequence's conv window published to the paged pool. -Tier 6 (CuTe DSL written for this target) for the conv pass, Tier 5 (Triton) -for the gates and the window publish. The three-call -composition it replaces -- `extract_transpose_prefill_slice` (token-major -> -channel-major copy), `causal_conv1d_fwd` (channel-major conv, in place) and -`fused_gdn_post_conv` (split + norm + gates) -- moves the `[T, 10240]` -activation through HBM three times; no ``torch.ops.trtllm.*`` op, in-tree -callable or installed FlashInfer API fuses the conv with the norm, and the -channel-last `causal_conv1d_fwd` kernel alone is slower than this whole pass. +The conv pass is CuTe DSL, the gates and the window publish Triton. The conv pass +reads a pool slot through a 64-bit base address (on a dep4 rank's stock conv page a +slot's offset passes 2^31 elements from slot 1942). """ from __future__ import annotations @@ -37,17 +32,28 @@ #: The target's one GDN input-projection output: [q|k|v | z | a | b]. IN_PROJ_WIDTH = QKV_DIM + NUM_V_HEADS * HEAD_V_DIM + 2 * NUM_V_HEADS # 16480 CONV_WIDTH = 4 -#: Largest conv-pool slot count the cache manager was seen to allocate for this -#: checkpoint (the `causal_conv1d_fwd` entry's ENGINE_SLOTS). +#: Largest conv-pool slot count the cache manager allocates for this checkpoint +#: (461 slots on one GB300). ENGINE_SLOTS = 461 +#: One tp4/tep4 rank's share: the heads split four ways, its own `[q|k|v|z|a|b]` +#: input projection 2560 + 1536 + 2 * 12 = 4120 wide. +RANK4_K_HEADS, RANK4_V_HEADS = NUM_K_HEADS // 4, NUM_V_HEADS // 4 +RANK4_IN_PROJ_WIDTH = (2 * RANK4_K_HEADS + 2 * RANK4_V_HEADS) * HEAD_K_DIM + 2 * RANK4_V_HEADS +#: Conv-pool slots the engine binds per rank (n-gram table host-resident): 1874 at +#: tep4, 850 at dep4. A tp4 rank holds tep4's GDN state, so tep4's count bounds it. +TEP4_SLOTS, DEP4_SLOTS = 1874, 850 +TP4_SLOTS = TEP4_SLOTS +#: The stock manager's conv page: one slot holds all 36 GDN layers' `[C, 3]` states, +#: so a layer's view has slot stride `36 * C * 3` and the last layer sits `35 * C * 3` +#: elements into the page. +PAGE_GDN_LAYERS = 36 A_LOG_RANGE = (-5.5625, 3.5625) DT_BIAS_RANGE = (-8.9375, 19.25) SOFTPLUS_THRESHOLD = 20.0 #: Channels per program of the window publish. _PUBLISH_BLOCK_C = 512 -#: Tokens per program of the gates: at 16384 tokens on GB300 4.4 us (6.1 at 64, 5.3 at -#: 16 or 32). +#: Tokens per program of the gates. _GATES_BLOCK_T = 8 @@ -287,6 +293,22 @@ def _cp_async16(smem_addr, gmem_addr, src_bytes, *, loc=None, ip=None): ) +@dsl_user_op +def _ld_bf16(addr, *, loc=None, ip=None): + """One bf16 from global memory at a 64-bit byte address, as fp32.""" + return cutlass.Float32( + _asm1( + T.f32(), + [cutlass.Int64(addr)], + "{ .reg .b32 r; ld.global.u16 r, [$1]; shl.b32 r, r, 16; mov.b32 $0, r; }", + "=f,l", + loc, + ip, + side_effects=True, + ) + ) + + @dsl_user_op def _cp_async_wait_all(*, loc=None, ip=None): _asm1(T.i32(), [], "{ cp.async.wait_all; mov.u32 $0, 0; }", "=r", loc, ip, side_effects=True) @@ -341,10 +363,9 @@ def __call__( stream: cuda.CUstream, ): weight = cute.make_tensor(w_ptr, cute.make_layout((self.qkv, 4), stride=(w_s0, w_s1))) - # The pool's slot extent is never read: only `cache_indices` slots are addressed. - pool = cute.make_tensor( - st_ptr, cute.make_layout((1 << 20, self.qkv, 3), stride=(st_s0, st_s1, st_s2)) - ) + # The pool goes in as its base address and strides: a carried column is read at + # `base + slot * st_s0` in 64 bits (on the stock page that passes 2^31 elements + # from slot 1942 at 16 / 48 heads, 7768 at 4 / 12) plus its in-slot offset in 32. qsl = cute.make_tensor(qsl_ptr, cute.make_layout((num_seqs + 1,), stride=(1,))) slots = cute.make_tensor(slot_ptr, cute.make_layout((num_seqs,), stride=(1,))) init = cute.make_tensor(init_ptr, cute.make_layout((num_seqs,), stride=(1,))) @@ -353,7 +374,10 @@ def __call__( self.kernel( x_ptr.toint(), weight, - pool, + st_ptr.toint(), + st_s0, + st_s1, + st_s2, qsl, slots, init, @@ -377,7 +401,10 @@ def kernel( self, x_base: cutlass.Int64, mW: cute.Tensor, - mState: cute.Tensor, + st_base: cutlass.Int64, + st_s0: cutlass.Int32, + st_s1: cutlass.Int32, + st_s2: cutlass.Int32, mQsl: cute.Tensor, mSlot: cute.Tensor, mInit: cute.Tensor, @@ -507,6 +534,7 @@ def kernel( s0 = mQsl[lo2] slot = mSlot[sq] init = cutlass.Int32(mInit[sq]) + slot_base = st_base + cutlass.Int64(slot) * cutlass.Int64(st_s0) * 2 for v in cutlass.range_constexpr(_VEC): a = cutlass.Float32(0.0) for kk in cutlass.range_constexpr(4): @@ -516,7 +544,11 @@ def kernel( xv = cutlass.Float32(0.0) if init != 0: if slot != pad_slot_id: - xv = mState[slot, c0 + v, 3 + pos - s0].to(cutlass.Float32) + xv = _ld_bf16( + slot_base + + cutlass.Int64((c0 + v) * st_s1 + (3 + pos - s0) * st_s2) + * 2 + ) a = _fma_rn(xv, w[v, kk], a) acc[v] = a y0, y1 = _silu2(acc[0], acc[1]) @@ -587,13 +619,15 @@ def _gdn_gates_kernel( def _cute_conv_pass_ok(x, weight, conv_states, num_k_heads, head_k_dim, num_v_heads) -> bool: """The CuTe conv pass's preconditions: bf16, 128-channel heads, 16-byte aligned rows - of `x` (its `cp.async` copies), int32 element strides.""" + of `x` (its `cp.async` copies), int32 element strides and an int32 offset inside a + pool slot (a slot's base is 64-bit).""" + in_slot = (conv_states.shape[1] - 1) * conv_states.stride(1) + 2 * conv_states.stride(2) return ( x.dtype == torch.bfloat16 and head_k_dim == _VEC * 32 and x.data_ptr() % 16 == 0 and (x.stride(0) * 2) % 16 == 0 - and max(*weight.stride(), *conv_states.stride()) < 2**31 + and max(*weight.stride(), *conv_states.stride(), in_slot) < 2**31 ) @@ -749,6 +783,24 @@ def _engine_slots(num_seqs: int) -> tuple[int, ...]: return slots +def _rank_cell(rank4: bool, lengths, initial, pool_slots, slots, linear_gate) -> dict: + """A rank's call: tp4/tep4 (`rank4`, 4 k / 12 v heads over its 4120-wide projection) or + dep4 (all 16 / 48 heads, 16480), over the stock conv page at `pool_slots`.""" + heads = (RANK4_K_HEADS, RANK4_V_HEADS) if rank4 else (NUM_K_HEADS, NUM_V_HEADS) + return dict( + num_k_heads=heads[0], + num_v_heads=heads[1], + head_dim=HEAD_K_DIM, + x_row_stride=RANK4_IN_PROJ_WIDTH if rank4 else IN_PROJ_WIDTH, + lengths=lengths, + initial=initial, + pool_slots=pool_slots, + page_layers=PAGE_GDN_LAYERS, + slots=slots, + linear_gate=linear_gate, + ) + + class GdnPrefillConvPost(OpWrapper): """Causal conv + SiLU of the context rows, then the GDN q/k/v/g/beta. @@ -786,7 +838,10 @@ class GdnPrefillConvPost(OpWrapper): # spec: lengths; slots (pool rows used, None -> _engine_slots); pool_slots; # initial; layout "in_proj" (x and a/b column views of one [T, 16480] # buffer, the target's call) | "rows" (x/a/b the first rows of a longer - # projection, as a mixed step's context half), extra_rows. + # projection, as a mixed step's context half), extra_rows. A rank's cell + # (`_rank_cell`) also names num_k_heads, num_v_heads, head_dim, x_row_stride + # (its own projection's width) and page_layers (the pool is the last layer's + # view of the stock page); without them, the tp1 geometry over a two-layer buffer. CELLS: tuple[Cell, ...] = ( Cell( why="target call, engine 461-slot pool: 37+1+128 tokens at slots 460, 0, 230", @@ -824,6 +879,58 @@ class GdnPrefillConvPost(OpWrapper): initial=tuple(i % 3 != 0 for i in range(128)), ), ), + Cell( + why="tep4 rank, 4/12 heads at stride 4120: fresh 8192-token context, last of 1874 slots", + spec=_rank_cell(True, (8192,), (False,), TEP4_SLOTS, (TEP4_SLOTS - 1,), True), + ), + Cell( + why="tep4 rank: 37+1+128 tokens, carry/fresh/carry, slots 1873, 0, 937 of 1874", + spec=_rank_cell( + True, + (37, 1, 128), + (True, False, True), + TEP4_SLOTS, + (TEP4_SLOTS - 1, 0, TEP4_SLOTS // 2), + False, + ), + ), + Cell( + why="tp4 rank (tep4's 4/12 heads and 1874-slot bound): fresh 8192 tokens at slot 0", + spec=_rank_cell(True, (8192,), (False,), TP4_SLOTS, (0,), True), + ), + Cell( + why="dep4 rank, 16/48 heads: two 512-token chunks, the first carrying, 850 slots", + spec=_rank_cell( + False, + (512, 512), + (True, False), + DEP4_SLOTS, + (DEP4_SLOTS - 1, DEP4_SLOTS // 2), + True, + ), + ), + Cell( + why="dep4 rank at max_num_tokens: one carrying 8192-token chunk, last of 850 slots", + spec=_rank_cell(False, (8192,), (True,), DEP4_SLOTS, (DEP4_SLOTS - 1,), True), + ), + # Pools past the int32 slot offset (slot 1942 at 16 / 48 heads, 7768 at 4 / 12), + # as larger dep4 pools have (kv_nvfp4, reuse, a 0.85 KV fraction). + Cell( + why="dep4 rank, 2048-slot pool: 37+1+128 tokens at slots 2047, 1942, 0, past int32", + spec=_rank_cell(False, (37, 1, 128), (True, False, True), 2048, (2047, 1942, 0), True), + ), + Cell( + why="dep4 rank, 4096-slot pool: a carrying 8192-token chunk at slot 4095, past 2^32", + spec=_rank_cell(False, (8192,), (True,), 4096, (4095,), True), + ), + Cell( + why="tp4/tep4 rank, 8192-slot pool: a fresh 8192-token context at slot 8191", + spec=_rank_cell(True, (8192,), (False,), 8192, (8191,), True), + ), + Cell( + why="tp4/tep4 rank, 8192-slot pool: carrying 8192 tokens at slot 8191, past int32", + spec=_rank_cell(True, (8192,), (True,), 8192, (8191,), True), + ), ) note = """ @@ -844,22 +951,33 @@ class GdnPrefillConvPost(OpWrapper): the factor FlashInfer's kernel takes, bitwise the `torch.exp` the chunk op would otherwise launch. + The default target (tp4, tep4, dep4) makes the same call per rank, from its + prefill GDN core: a tp4 or tep4 rank over its own `[T, 4120]` projection + (4 key / 12 value heads, conv dim 2560: its conv weight rows, its + `A_log`/`dt_bias` entries and its conv pool are the rank's shares), a dep4 + rank over the whole `[T, 16480]` projection. The pool is the layer's view + of the stock manager's conv page (`get_conv_states` at tp4, tep4 and + dep4): one slot holds all 36 GDN layers' `[C, 3]` states, so the view + has strides `(36 * C * 3, 3, 1)` -- + (276480, 3, 1) at C 2560, (1105920, 3, 1) at C 10240 -- and layer `i` + starts `i * C * 3` elements into the page. The target's forward does + not call this variant yet; the rank cells mirror tp1's call at the + rank's widths. + Three launches. The conv pass (CuTe DSL, `_ConvPassKernel`): grid - `80 heads x cdiv(T, 64)`, 128 threads; a CTA copies its 67-row tile - (3 history rows) of the head's 128 channels into shared memory with + `heads x cdiv(T, 64)` (80 heads at tp1 and dep4, 20 on a tp4/tep4 + rank), 128 threads; a CTA copies its 67-row tile (3 history rows) of + the head's 128 channels into shared memory with 16-byte `cp.async`, then a thread slides the width-4 window over 16 tokens of 4 channels, each element converted once; each thread finds its first token's sequence by a binary search of `query_start_loc` (sync-free, capture-safe), and a thread whose 16 tokens and windows lie inside one sequence takes a straight-line branch, the rest read the pool - token by token. The gates (Triton): one row-coalesced `[8, 48]` tile a - program (`exp(g)` with `linear_gate`, libdevice `expf`). Then a - per-sequence window publish (Triton), after both, so no - program reads a window another has already advanced. - - Speed (one layer, 461-slot pool, in_proj views, us, graph replay over 4 - layers' inputs): 16384 tokens / 3 contexts 136 (the three-call chain it - replaced 528). + token by token (a carried column at the slot's 64-bit base address plus + its 32-bit in-slot offset, one `ld.global.u16`). The gates (Triton): one + row-coalesced `[8, 48]` tile a program (`exp(g)` with `linear_gate`, + libdevice `expf`). Then a per-sequence window publish (Triton), after + both, so no program reads a window another has already advanced. Numerics: the conv result is rounded to bf16 before the norm reads it, as the unfused path stores it. Conv taps: fp32 FMAs from zero in tap order @@ -869,7 +987,7 @@ class GdnPrefillConvPost(OpWrapper): `1 + e <= 2^126`; below acc = -87.3 it returns -0 where those return a value under 1e-36. q/k norm: the 128-channel fp32 sum of squares as each thread's pairs plus a 32-lane butterfly, `sqrt.approx.ftz`, `div.full`. - Against the unfused path (measured at the target-call cell): q/k/v within + Against the unfused path (at the target-call cell): q/k/v within one bf16 ulp on 1e-6..1e-4 of elements, g/beta within the fp32 band, conv pool bitwise. @@ -879,6 +997,36 @@ class GdnPrefillConvPost(OpWrapper): 64 tokens, some shorter than the window); and the target's call with `linear_gate` (3000 + 8192 + 5000 tokens, the middle one fresh). + Certified per rank, over the stock conv page's last-layer view at the + engine's per-rank slot counts (`TEP4_SLOTS` 1874, `DEP4_SLOTS` 850; + `TP4_SLOTS` is tep4's), each also checked to leave the page's other 35 + layers untouched: + - 4 / 12 heads, stride 4120 (a tp4/tep4 rank, its own compile): a + fresh 8192-token context at the last slot and at slot 0 (with + `linear_gate`); 37 + 1 + 128 tokens carrying / fresh / carrying at + slots 1873, 0 and 937 (without). + - 16 / 48 heads, stride 16480 (a dep4 rank): two 512-token chunks, the + first carrying, at slots 849 and 425; one carrying 8192-token chunk + (max_num_tokens) at slot 849; both with `linear_gate`. + Against the reference at these cells: q/k/v differ by one bf16 ulp on + 1.2e-5 to 1.3e-4 of elements (largest 1.6e-2, on v near 4), within + torch's default bf16 band; g and beta inside the fp32 band; conv pool + bitwise. + + Certified past the int32 slot offset, on the same page (larger dep4 + pools: kv_fp8 binds 1217 slots, kv_nvfp4 1498; reuse and a 0.85-fraction + fp8 engine are uncounted), all with `linear_gate`: + - 16 / 48 heads, a 2048-slot pool: 37 + 1 + 128 tokens carrying / + fresh / carrying at slots 2047, 1942 and 0; a 4096-slot pool: one + carrying 8192-token chunk at slot 4095 (offset past 2^32). + - 4 / 12 heads, an 8192-slot pool: an 8192-token context at slot 8191, + fresh and carrying. + Same bands as above (q/k/v one-ulp flips on 2.9e-6 to 2.5e-5 of + elements). The carrying cells need the 64-bit slot base: an int32 slot + offset is an illegal memory access there, or at slot 4095 a wrapped + offset inside the page and a silently wrong q/k/v. The fresh 8192-slot + cell reads no carry. + A `pad_slot_id` context reads no carry and publishes nothing; its outputs are computed from a zero window (the unfused conv leaves those columns unconvolved instead -- nothing consumes either). The engine gives every @@ -892,19 +1040,29 @@ class GdnPrefillConvPost(OpWrapper): `has_initial_state`. Loud (`__call__`): `head_k_dim != head_v_dim`, conv width other than 4, and what the conv pass cannot take: an `x` not bf16, heads other than 128 channels, a row of `x` not 16-byte aligned, a weight - or pool stride from 2^31 elements. + or pool stride, or an offset inside one pool slot, from 2^31 elements. Not checked: `cache_indices` within the pool and distinct; `query_start_loc` non-decreasing and ending at `T`. - Not certified: fp16/fp32 activations, other head geometries, bias, a row - stride of `x` other than the projection's 16480 (the conv pass compiles one - kernel per row stride), zero-token contexts, `a + dt_bias` past the softplus - threshold on every head. + Compiled at first call, per `(num_k_heads, num_v_heads, head_k_dim, row + stride of x)` (`_CONV_PASS`): the conv pass by CuTe DSL (TVM-FFI entry, + `--opt-level 2`), the gates and publish by Triton. The cells ran both JIT + paths on sm_103 (GB300). + + Not certified: fp16/fp32 activations; head geometries other than 16 / 48 + and 4 / 12; bias; a row stride of `x` other than the projections' 16480 + and 4120 (the conv pass compiles one kernel per row stride); zero-token + contexts; `a + dt_bias` past the softplus threshold on every head; pools + past 4096 slots at 16 / 48 heads or 8192 at 4 / 12; tp4's own pool. Hard limits - element offsets 64-bit for the rows of x and the outputs; weight and - pool strides int32 (refused from 2^31) not probed beyond 8192 tokens. + element offsets 64-bit for the rows of x, the outputs and a pool + slot's base (`slot * stride(0)`, in the conv pass and the publish); + weight and pool strides and the offset inside a slot int32 (refused + from 2^31) not probed beyond 8192 tokens. + pool slot `cache_indices` is int32, so a slot index below 2^31; driven + to offsets past 2^32 elements (slot 4095 at 16 / 48 heads, 4.5e9). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/vision/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/vision/__init__.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/vision/patch_embed_conv3d.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/vision/patch_embed_conv3d.py new file mode 100644 index 000000000000..d0f2eabf34aa --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/ops/vision/patch_embed_conv3d.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Vision patch embedding: one 3-D convolution per pixel patch, made as the GEMM it is. + +With one whole patch per row and the stride equal to the kernel, the convolution +has one output position per row: `out = rows @ filters.T + bias`, where `rows` and +`filters` are views of `pixel_values` and `weight`. `__call__` is one +`torch.ops.trtllm.cublas_mm` call (the C++ thop in `thop/cublasScaledMM.cpp`, +cuBLASLt, bias fused into the epilogue) between views that launch nothing. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +#: This checkpoint: in_channels 3, kernel (temporal_patch_size 2, patch_size +#: 16, 16), out_channels 1152 -- weight [1152, 3, 2, 16, 16], a 1536-pixel row. +IN_CHANNELS = 3 +KERNEL = (2, 16, 16) +EMBED_DIM = 1152 + + +class PatchEmbedConv3d(OpWrapper): + """`out[n, k] = bias[k] + sum_{c,t,i,j} pixel[n, c, t, i, j] * weight[k, c, t, i, j]`. + + One whole `[in_channels, kt, kh, kw]` patch per row of `pixel_values`, flat + `[N, C*kt*kh*kw]` (what the image processor emits, channel-slowest) or + explicit `[N, C, kt, kh, kw]`; the kernel extent is read from `weight`. + Cross-correlation (no kernel flip), like `nn.Conv3d`. Returns a fresh + contiguous `[N, out_channels]` in the input dtype. + + Fusion boundary: the patch projection and its bias. Patchification, the + position embedding added after it, and the row order (rows pass through in + the order they arrive -- the merge-block order the merger needs) are the + caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: num_patches, dtype. The checkpoint's geometry, flat contiguous + # rows and a bias: what the vision tower passes. The tower is replicated, + # so these are every member's shapes (rank-invariant). + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"checkpoint geometry, {n} patches bf16", + spec=dict(num_patches=n, dtype=torch.bfloat16), + ) + for n in (64, 4096) + ) + ( + Cell( + why="tower call ceiling: 8192 tokens x merge 2x2 = 32768 patches bf16", + spec=dict(num_patches=32768, dtype=torch.bfloat16), + ), + ) + + note = """ + Source: `torch.ops.trtllm.cublas_mm`, compiled into the engine's torch-op + library (no JIT path). One kernel launch per call; cuBLASLt picks the + kernel by row count, each with the bias in its epilogue. + + Metadata consumed: none. `pixel_values` is not mutated. Several images, + and a video's frame groups, are just more rows; there is no batch axis. + + Measured on GB300 against `reference` at the three cells: 99.91-99.92% of + outputs bitwise equal, the rest one-ulp rounding flips of the element + (at most 4.4 bf16 ulps at the output RMS); relative L2 <= 9.4e-5; the worst + element at 0.49 of torch's default bf16 band. Each cell also holds + `reference`, in fp32, to an fp32 `F.conv3d` of the 5-D patches (TF32 off, + max |d| 2.0e-6) -- the flatten order `__call__` and `reference` share is + `nn.Conv3d`'s. + + `torch.nn.functional.conv3d` (cuDNN) on the same patches adds the bias + separately in bf16, rounding a second time (relative L2 2.85e-3 against + `reference`). + + Row-count invariance is not claimed: the first 64 rows of a 4096-row call + measured bitwise equal to a 64-row call, but the kernel changes with the + row count. Check rows against `reference`, never against a call of a + different size. CUDA-graph capturable (replayed at 64 and 32768 patches); + the output is a fresh allocation. + + Silent, refused by `is_valid` (each measured unguarded): a row that is not + exactly one patch; row-strided pixel rows (2-D, or 5-D with a batch + stride) and a weight sliced along its output channels, read as if dense + (10^5 times outside the band); a bias in another dtype (fp32 or fp16 with + bf16 input) or strided (5000 times outside); a bias shorter than + `out_channels`, read past its end without an error; fp32 input with a + bias, which drops the bias. Loud, not guarded: a channels-last or + transposed operand fails the view, and a weight dtype unlike the input's + fails in cuBLASLt (CUBLAS_STATUS_NOT_SUPPORTED). + + Not certified: dtypes other than bf16 (fp16 with an fp16 bias measured + inside the band, not a cell); other geometries; `bias=None`; patch counts + below 64, between the cells, and past 32768. + """ + + def __call__( + self, + pixel_values: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + out_channels, in_channels, kt, kh, kw = weight.shape + rows = pixel_values.view(-1, in_channels * kt * kh * kw) + filters = weight.view(out_channels, -1) + return torch.ops.trtllm.cublas_mm(rows, filters.t(), bias, None, 0, None) + + def reference( + self, + pixel_values: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """True-fp32 matmul of the flattened rows against the flattened filter.""" + rows = pixel_values.reshape(pixel_values.shape[0], -1).float() + filters = weight.reshape(weight.shape[0], -1).float() + with true_fp32_matmul(): + out = rows @ filters.t() + if bias is not None: + out = out + bias.float() + return out.to(pixel_values.dtype) + + def is_valid( + self, + pixel_values: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> None: + # The view reinterprets anything whose element count divides: a whole + # [1, 3, 2, 64, 64] volume becomes 16 "patches" -- a [16, 1152] result + # of the right shape from the wrong pixels -- and [N, 3072] rows become + # 2N, without an error. + _, in_channels, kt, kh, kw = weight.shape + patch_numel = in_channels * kt * kh * kw + if pixel_values.dim() == 2: + one_patch_per_row = pixel_values.shape[1] == patch_numel + elif pixel_values.dim() == 5: + one_patch_per_row = tuple(pixel_values.shape[1:]) == (in_channels, kt, kh, kw) + else: + one_patch_per_row = False + assert one_patch_per_row, ( + f"each row must be exactly one patch: pixel_values must be " + f"[num_patches, {patch_numel}] or " + f"[num_patches, {in_channels}, {kt}, {kh}, {kw}] for this weight; got " + f"{tuple(pixel_values.shape)}" + ) + # cuBLASLt reads the rows as dense row-major and the filter as dense + # column-major: a row-strided view, or a weight sliced along its output + # channels, still views to 2-D and is read as if dense. + assert pixel_values.is_contiguous(), ( + "pixel_values must be contiguous; a strided view is read as dense rows" + ) + assert weight.is_contiguous(), "weight must be contiguous; a strided view is read as dense" + if bias is not None: + # The epilogue reads the bias as `out_channels` dense elements of + # the output dtype, and drops it altogether for fp32 input. + assert bias.dtype == pixel_values.dtype, ( + f"bias dtype {bias.dtype} must equal the input dtype {pixel_values.dtype}" + ) + assert bias.shape == (weight.shape[0],) and bias.is_contiguous(), ( + f"bias must be a contiguous [{weight.shape[0]}] tensor; got " + f"{tuple(bias.shape)} with stride {bias.stride()}" + ) + assert pixel_values.dtype != torch.float32, ( + "the bias is silently dropped for fp32 input" + ) + + +patch_embed_conv3d = PatchEmbedConv3d() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/__init__.py new file mode 100644 index 000000000000..eba053a6d445 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Qwen3.8-Flash-Next-NVFP4 / sm_103 / default: tp4, tep4, dep4 and dep2.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py new file mode 100644 index 000000000000..bb040f05c87b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py @@ -0,0 +1,6715 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ModelingV2 target: Qwen3.8-Flash-Next-NVFP4 (hf-fc694b5_orig) / sm_103 / default. + +One target for the layouts `routing.py` admits in one node: one rank (tp1, every head and +expert); at four ranks tp4 (attention TP4, routed experts split by intermediate), tep4 +(attention TP4, experts EP4) and dep4 (attention DP, experts EP4); at two ranks dep2 (attention +DP, experts EP2), the layout of both disaggregated-serving engines, which is dep4 at half the +group. The core holds the static +structure (this rank's share of the weights, derived operands, per-manager pool views); +`PrefillTarget` (any batch with context rows) and `DecodeTarget` (generation only) run a step +over it, passing every argument to every op at the call site. + +The forward is cut per rank; what the layout settles is settled once, at construction: + +* **Attention and Gated-DeltaNet heads** are this rank's share (`weights.rank_dims`, + `pool.pool_geometry`): at attention TP4 six q heads, one of the two KV heads, four GDN key + and twelve value heads. The fused attention input projection row is + `[q | k | v | gate | index q, k | v slot]`, 4352 wide; the QSA indexer is replicated. +* **The mixer output is a partial sum** under attention TP: o_proj and out_proj read this + rank's heads. One plain sum (`comm/allreduce`, op `NONE`) completes it before the + hyper-connection inject, which every rank runs on the same reduced rows. Under attention DP + a rank holds every head and serves its own requests, so there is nothing to sum. +* **Routed experts** run on the trtllm-gen NVFP4 runner entered with the router logits + (`fp4_block_scale_moe_renorm_routed`: its own Renormalize routing over the whole routing + space). A rank holds an EP window of + `num_experts / moe_ep_size` whole experts (tep4, dep4), or every expert at its quarter of the + intermediate (tp4: 160 of 640 rows, stored zero-padded to 256, `weights.EXPERT_INTER_ALIGN`, + with zero codes and scales, so the padding adds nothing). Under attention TP every rank holds + every row and the partial output -- the window's, or the intermediate slice's -- is summed the + same way. Under attention DP each rank holds its own rows: they are padded to the group's + largest row count, gathered (`comm/allgather`), routed and run through the window, and + reduce-scattered back (`comm/reducescatter`), each row's window partials summed once + (`_expert_sum_op`). The router and the bf16 shared expert are replicated and run on this + rank's rows, so the block output `routed + sigmoid(gate) * shared` is formed by the next + hyper-connection site on the reduced routed rows. +* **Hyper-connections, PLE, the embedding** run replicated on every rank's rows; + `lm_head` is the shell's (vocab-parallel under attention TP, gathered by the shell's head; + whole under attention DP, over this rank's own rows). + +Each hyper-connection site is `qwen4_exp_hc_inject_norm` (the previous block's injection +fused with this site's grouped norm; after a MoE block it also forms that block's output) +and then the mixer: the first projection with its activations (one tcgen05 launch, +`qwen4_exp_hc_decode_tc_first`, up to `_HC_TC_FIRST_MAX_ROWS` rows; past them `cublas_mm`, +fp32 out, then `qwen4_exp_hc_lowrank_act`) and `qwen4_exp_hc_up_gate_mean`. The MoE site's +routed input is that mixer output quantized to NVFP4 with linear block scales at the layer's +FC1 input global scale: in the mixer's own epilogue under attention TP +(`qwen4_exp_hc_up_gate_mean_fp4`), on the gathered rows under attention DP (`fp4_quantize`). +There is **no** `model.language_model.norm.weight`: the stack-final mixer is the final norm. + +**One rank (tp1).** With nothing to sum or exchange, a rank holding the whole model takes forms +of its own, settled at construction (`one_rank`): the decode projections' one-rank dispatch +(`_one_rank_projection`); the hyper-connection sites' row-count forms (`hc_fast_sites`: the +decode GEMV and tcgen05 rows pairs, the prefill rstd sites), which the other layouts leave +out so that a request's tokens do not depend on what its steps carry beside it; a trtllm-gen +decode MoE whose finalize the next site folds in (`moe_fold`); the CuteDSL copy of the +routed experts for prefill steps past `_CUTE_DSL_MOE_MIN_ROWS` (`moe_prefill_copy`), or the +CuteDSL pipeline for every step (`moe_config.backend: CUTEDSL`, `_cute_dsl_routed`); an fp32 +recurrent state, whose captured single-token decode runs `gdn_decode_gated_norm_pdl`; the +decode draft layer and its FP8 head; the NVFP4 pool's context rows attended a group of +requests at a time (`_nvfp4_context_plan`); and the opt-in FP8 / NVFP4 copies of the dense +projections and heads (`_DENSE_FP8_ENV`, `_GDN_INPROJ_NVFP4_ENV`, `_TARGET_HEAD_NVFP4_ENV`, +`_MTP_HEAD_NVFP4_ENV`, `_MTP_MOE_NVFP4_ENV`). + +**QSA sparse attention.** The index plane is maintained on every step so a request can +cross `dense_seq_len_threshold` later without rebuilding its prefix; above it the sparse +core (index scores, block top-k and expansion, sparse GQA) replaces the dense op, and the +main K/V append becomes this layer's own call. The raw index key is stored before anything +normalizes it: the compression applies the norm and the rotation itself. + +**KV storage.** The paged K/V pool holds the checkpoint's bf16, fp8-e4m3 under +`kv_cache_config.dtype: fp8` (`QM_FP8_KV_CACHE`) or NVFP4 under `nvfp4` (`QM_NVFP4_KV_CACHE`, +e2m1 pairs with e4m3 block scales in a second pool, second-level scales 1.0). The NVFP4 +attention kernels exist for the paged-context path only. The QSA index and position planes +keep their own types under all three. + +**Breakable prefill CUDA graph.** With `prefill_cuda_graph_backend: breakable` the engine +captures the core's `forward` (the body, not the shell) for a step with context rows, pads +the step up to the nearest bucket and replays it without calling the body. The shell +therefore calls `begin_step` before the body on every step: the asserts, the cache-manager +bind, the target choice and the prefill target's host bookkeeping all live there. The +prefill forward is row-wise over the `T` rows it is handed, with four breaks +(`catalog/graph/eager_on_graph`) that read the step: the vision splice, the PLE block, and +each layer's attention core and Gated-DeltaNet recurrence. Each writes rows `[:n]` of a fixed +output and zeroes rows `[n:]`. + +**Host offload.** By default the PLE n-gram table lives in host memory the device addresses +(`weights._host_resident_empty`), one copy per rank; `TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD=0` +makes it device-resident. Output is identical. + +**One-model MTP-Eagle.** With a `speculative_config` the core also holds the checkpoint's MTP +module (`MTPLayer`), replayed `max_draft_len` times by the engine's spec worker, and the trunk +runs every generation request's `max_draft_len + 1` verify rows: + +* The draft layer is one more full-attention QSA layer, cut like a trunk attention layer (this + rank's q heads, its KV head, the replicated indexer; every head under attention DP), over the + manager the draft step is handed: under attention TP that is a separate draft KV cache, whose + K/V, index and position pages hold the draft layer alone; under attention DP the engine + builds none and the trunk's manager maps the draft layer too (`_bind_draft_cache_manager`). + Its front end, router, shared expert and hyper-connections are replicated; its FP8 + block-scale experts run as this rank's window of `num_experts / tp_size` experts + (`fp8_block_scale_moe_runner`'s window). Under attention TP the o_proj and expert partials + are completed by the same plain sum as the trunk's. That window is the same at tp4 as at + tep4: the FP8 experts' 128x128 weight blocks do not take the trunk's 160-row intermediate + split, and a sum of whole-expert windows is the layer's output whichever way the trunk's + experts are cut. +* Under attention DP each rank drafts for its own requests, and the draft MoE runs the trunk's + round trip (`_expert_sum_op`): the group's rows gathered, routed and run through this rank's + window, reduce-scattered back. The spec worker hands each draft step the group's row counts + (the trunk step's tokens on the first, one row a request on the rest, attention-DP dummy + requests included) without writing them into the metadata; `_dp_pad_rows` pads this rank's + rows by them, so every rank issues the same collectives at every draft step. +* A verify step's Gated-DeltaNet rows take the multi-token conv with per-step window snapshots + and the stock Triton cached replay over the bf16 pool, which commits a full history in the + kernel below `CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE` requests and leaves it to the + manager's all-layer commit from there on, exactly as the stock layer hands it work items. +* What a verify row writes that the runtime does not roll back -- the PLE n-gram history and + conv state, every attention layer's QSA index-K plane and the shared position plane -- is + staged during the forward and rewritten to the accepted prefix by `SpeculativeRollback`, + which the spec worker calls after its draft loop. +* IndexShare (`index_share_for_mtp_iteration`) runs the draft indexer once per draft loop: + the draft-extend step captures each request's selection and the later draft steps rebuild + theirs from it (`qsa_mtp_reuse_topk`). + +Not assembled here: the legacy per-step recurrent verify (`TRTLLM_USE_GDN_REPLAY=0`), which +routing declines and the bind refuses. + +Integration facts this file depends on: + +1. The engine rebuilds the composite config and normalizes its text tower + (`config_utils.load_pretrained_config`). The shell adopts that text config as + `model_config.pretrained_config` before building the core, because the recurrent-state and + PLE pool geometry are read from it. +2. `position_ids` arrives in the mRoPE layout (`[3, 1, T]`) even for text-only batches. The + main rope reads all three rows; the QSA index plane takes coordinate 0 only. +""" + +from __future__ import annotations + +import math +import os +import re +from typing import NamedTuple + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tensorrt_llm._torch._experimental.modeling_v2._router_index import step_contract_enabled +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope as mtp_fused_qk_norm_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import ( + QM_FP8_KV_CACHE, + QM_NVFP4_KV_CACHE, + thop_attention, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.allgather import allgather +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.allreduce import ( + allreduce, + allreduce_group_args, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.reducescatter import reducescatter +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import ( + cublas_mm as mtp_cublas_mm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.graph.eager_on_graph import ( + eager_on_graph, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( + MAX_CERTIFIED_TOKENS as _MOE_MAX_ROWS, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.fused_sigmoid_mul_inplace import ( + fused_sigmoid_mul_inplace, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.fused_sigmoid_mul_inplace import ( + fused_sigmoid_mul_inplace as mtp_fused_sigmoid_mul_inplace, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.silu_and_mul_rows import ( + silu_and_mul_rows, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_canonical_selection import ( + qsa_canonical_selection, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_expand_block_indices import ( + qsa_expand_block_indices, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_mtp_reuse_topk import ( + qsa_mtp_reuse_topk, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_nvfp4_pages_to_bf16 import ( + qsa_nvfp4_pages_to_bf16, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_index_scores import ( + qsa_paged_index_scores, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_index_store import ( + qsa_paged_index_store, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_kv_store import ( + qsa_paged_kv_store, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_kv_store_nvfp4 import ( + qsa_paged_kv_store_nvfp4, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa import ( + qsa_paged_sparse_gqa, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_decode import ( + qsa_paged_sparse_gqa_decode, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_decode import ( + supports as _sparse_decode_supports, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4 import ( + qsa_paged_sparse_gqa_nvfp4, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4_decode import ( # noqa: E501 + qsa_paged_sparse_gqa_nvfp4_decode, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4_verify import ( # noqa: E501 + qsa_paged_sparse_gqa_nvfp4_verify, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4_verify import ( # noqa: E501 + takes_grouped as _verify_takes_grouped, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_prefill_compress import ( + qsa_prefill_compress, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rollback_planes import ( + plane_table as _plane_table, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rollback_planes import ( + qsa_rollback_gather, + qsa_rollback_restore, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_sparse_gqa_rowtile import ( + qsa_sparse_gqa_rowtile as one_rank_sparse_gqa_rowtile, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_topk_expand_decode import ( + qsa_topk_expand_decode, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_topk_expand_decode import ( + supports as _topk_expand_supports, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_topk_select import ( + qsa_topk_select, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qwen4_exp_attn_decode_prologue import ( # noqa: E501 + SCORE_ROWS as _PROLOGUE_SCORE_ROWS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qwen4_exp_attn_decode_prologue import ( # noqa: E501 + qwen4_exp_attn_decode_prologue as one_rank_attn_decode_prologue, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qwen4_exp_qk_norm_rope import ( + qwen4_exp_qk_norm_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.spec_rollback_commit import ( + spec_rollback_commit, + spec_rollback_stage, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.verify_step_setup import ( + _BLOCK_S as _VERIFY_SETUP_MAX_SEQS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.verify_step_setup import ( + verify_step_setup, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.bf16_gemv_pdl import ( + MAX_ROWS as _GEMV_MAX_ROWS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.bf16_gemv_pdl import ( + TILING as _GEMV_TILING, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.bf16_gemv_pdl import ( + bf16_gemv_pdl, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.fp8_gemv import ( + fp8_gemv_conv1d_op, + fp8_gemv_conv1d_verify_op, + fp8_gemv_op, + pack_from_bf16, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.fp8_per_tensor_mm import ( + fp8_per_tensor_mm, + fp8_per_tensor_weight, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.low_m_bf16_mm import ( + low_m_bf16_mm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.mid_rows_mm import ( + MidWeight, + mid_rows_mm_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.mid_rows_mm import ( + precompile as mid_rows_precompile, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.nvfp4_gemv import ( + nvfp4_gemv_op, + nvfp4_split_gemv_conv1d_op, + nvfp4_split_gemv_conv1d_verify_op, + nvfp4_split_gemv_op, + precompile_split, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.nvfp4_gemv import ( + pack_from_bf16 as nvfp4_pack_from_bf16, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.nvfp4_gemv import ( + pack_split_from_bf16 as nvfp4_pack_split_from_bf16, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.nvfp4_gemv import ( + precompile as nvfp4_gemv_precompile, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_decode_first_cute import ( + qwen4_exp_hc_decode_first_cute, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_decode_mix_fp4_cute import ( + qwen4_exp_hc_decode_mix_cute, + qwen4_exp_hc_decode_mix_fp4_cute, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_decode_rows_tc import ( + MAX_ROWS as _HC_DECODE_ROWS_MAX, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_decode_rows_tc import ( + qwen4_exp_hc_decode_rows_first, + qwen4_exp_hc_decode_rows_mix, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_first_rescale_tc import ( + qwen4_exp_hc_first_rescale_tc, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_inject_norm import ( + qwen4_exp_hc_inject_norm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_inject_rstd import ( + qwen4_exp_hc_inject_rstd, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_lowrank_act import ( + qwen4_exp_hc_lowrank_act, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_up_gate_mean import ( + qwen4_exp_hc_up_gate_mean, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_up_gate_mean_rstd import ( + qwen4_exp_hc_up_gate_mean_rstd, + qwen4_exp_hc_up_gate_mean_rstd_fp4, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_post_inject import ( + qwen4_exp_post_inject, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import ( # noqa: E501 + cute_dsl_nvfp4_gather_grouped_gemm_act_fusion, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.cute_dsl_nvfp4_grouped_gemm import ( + cute_dsl_nvfp4_grouped_gemm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_block_scale_moe_renorm_permuted import ( + fp4_block_scale_moe_renorm_permuted, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_block_scale_moe_renorm_routed import ( + fp4_block_scale_moe_renorm_routed, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_moe_renorm_cute_dsl import ( + fp4_moe_renorm_cute_dsl, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp8_block_scale_moe_runner import ( + fp8_block_scale_moe_runner, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fused_sigmoid_gate_mul_add import ( + fused_sigmoid_gate_mul_add, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.moe_sort import moe_sort +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.moe_unpermute_inplace import ( + moe_unpermute_inplace, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.renorm_moe_routing_op import ( + renorm_moe_routing_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.norm.flashinfer_gemma_rmsnorm import ( + flashinfer_gemma_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.norm.rms_norm_gated_token_major import ( + rms_norm_gated_token_major, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ple_gate_conv_inject_norm import ( + ple_gate_conv_inject_norm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ple_ngram_embed import ( + ple_ngram_embed, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.quantization.bf16_to_fp32 import ( + bf16_to_fp32, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.quantization.fp8_quantize_1x128 import ( + fp8_quantize_1x128, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.causal_conv1d_triton_update import ( + causal_conv1d_triton_update, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.causal_conv1d_update import ( + causal_conv1d_update, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.cu_seqlens_from_seq_lens import ( + cu_seqlens_from_seq_lens, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.flashinfer_chunk_gated_delta_rule import ( + INPLACE_MAX_TOKENS as _GDN_INPLACE_MAX_TOKENS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.flashinfer_chunk_gated_delta_rule import ( + flashinfer_chunk_gated_delta_rule, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.fused_recurrent_gated_delta_rule_cached_replay_update import ( # noqa: E501 + fused_recurrent_gated_delta_rule_cached_replay_update, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_decode_gated_norm_pdl import ( + MAX_ROWS as _GDN_DECODE_FUSED_MAX_ROWS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_decode_gated_norm_pdl import ( + gdn_decode_gated_norm_pdl, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_decode_rows_gated_norm import ( + gdn_decode_rows, + gdn_decode_rows_gated_norm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_prefill_conv_post import ( + gdn_prefill_conv_post, +) +from tensorrt_llm._torch.attention.backends.fmha.phased import get_spec_decoding_position_offsets +from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM +from tensorrt_llm._torch.modules.fla.cached_replay import CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE +from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor +from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks +from tensorrt_llm._torch.speculative import get_spec_worker, should_use_separate_draft_kv_cache +from tensorrt_llm._torch.utils import ActivationType +from tensorrt_llm._utils import mpi_disabled +from tensorrt_llm.functional import AllReduceFusionOp +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.mode import QuantAlgo + +from . import pool as _pool +from . import vision as _vision +from . import weights as _weights +from .ops.attention.qsa_fp8_pages_to_bf16__prefill import qsa_fp8_pages_to_bf16 +from .ops.attention.qsa_sparse_gqa_rowtile__prefill import qsa_sparse_gqa_rowtile +from .ops.attention.qwen4_exp_attn_decode_prologue__decode import qwen4_exp_attn_decode_prologue +from .ops.hc.qwen4_exp_hc_decode_tc_first__decode import MAX_ROWS as _HC_TC_FIRST_ROWS +from .ops.hc.qwen4_exp_hc_decode_tc_first__decode import precompile as hc_tc_first_precompile +from .ops.hc.qwen4_exp_hc_decode_tc_first__decode import qwen4_exp_hc_decode_tc_first +from .ops.hc.qwen4_exp_hc_up_gate_mean_fp4__moe import qwen4_exp_hc_up_gate_mean_fp4 +from .ops.ssm.gdn_cached_replay_cute__verify import gdn_cached_replay_cute + +# The GPU architecture this target is certified for; a direct instantiation could +# bypass routing. +_SM = (10, 3) + +# `moe_config.backend` request -> the routed-expert backend the trunk runs. +_TRUNK_MOE_BACKENDS = {"AUTO": "TRTLLM", "TRTLLM": "TRTLLM", "CUTEDSL": "CUTEDSL"} +# The CuteDSL grouped GEMMs' MMA tile M, which is also the tile `moe_sort` pads each expert's +# rows to: the GEMMs walk the sort's tile tables. sm_103 admits 128 and 256. +_CUTEDSL_TILE_SIZE = 128 +# The longest MTP draft the CuteDSL pipeline is certified at: its cells capture up to 128 rows +# (32 requests of 4 verify rows). +_CUTEDSL_MAX_DRAFT_LEN = 3 + +# Rows per request of the widest verify step (`max_draft_len + 1`), and of the widest +# one without the GDN cached replay (the per-step recurrent states). +_MAX_VERIFY_WIDTH = _pool.MAX_VERIFY_WIDTH + +# Recurrent-state cache dtypes the target serves: at one rank the checkpoint's fp32 too. +_SSM_STATE_DTYPES = (torch.bfloat16,) +_ONE_RANK_SSM_STATE_DTYPES = (torch.float32, torch.bfloat16) + +# QSA backend metadata buffers rebuilt every step. Their absence means the engine did not +# select the QSA metadata and the index plane would never be written. +_QSA_FIELDS = ( + "qsa_block_table", + "qsa_req_idx_per_token", + "qsa_logical_positions", + "qsa_sequence_lengths", + "qsa_visible_blocks", + "qsa_topk_indices", + "qsa_topk_row_starts", +) + +# Fields read off `attn_metadata.mamba_metadata`. Absence is silent: a smoke run passes +# with no recurrent state at all. +_MAMBA_FIELDS = ( + "state_indices", + "has_initial_states", + "query_start_loc", + "query_start_loc_long", +) + +# thop_attention's constant arguments: in-kernel RoPE off (the rope runs before the op), +# causal mask, no MLA, sparse or cross surface. +_CALL_CONSTANTS = dict( + output_sf=None, + k=None, + v=None, + is_fused_qkv=True, + update_kv_cache=True, + attention_input_type=0, + is_mla_enable=False, + mask_type=1, + q_scaling=1.0, + attention_sinks=None, + position_embedding_type=0, + rotary_inv_freq=None, + rotary_cos_sin=None, + rope_dim=0, + rope_base=10000.0, + rope_scale_type=0, + rope_scale=1.0, + rope_short_m_scale=1.0, + rope_long_m_scale=1.0, + rope_max_positions=1024, + rope_original_max_positions=1024, + out_scale=None, + latent_cache=None, + q_pe=None, + q_lora_rank=None, + kv_lora_rank=None, + qk_nope_head_dim=None, + qk_rope_head_dim=None, + v_head_dim=None, + rope_append=None, + chunked_prefill_buffer_batch_size=1, + attention_chunk_size=None, + softmax_stats_tensor=None, + sparse_kv_indices=None, + sparse_kv_offsets=None, + sparse_attn_indices=None, + sparse_attn_offsets=None, + sparse_attn_indices_block_size=0, + mrope_rotary_cos_sin=None, + mrope_position_deltas=None, + helix_position_offsets=None, + helix_is_inactive_rank=None, + cross_kv=None, + relative_attention_bias=None, + relative_attention_max_distance=0, + kv_norm_weight=None, + kv_norm_eps=1e-6, + skip_correction_threshold=0.0, +) + +_SF_VEC = _weights.SF_VEC +#: Rows up to which the MoE router runs `low_m_bf16_mm` (FlashInfer's low-M direct +#: and split-K kernels, PDL); more rows run cuBLAS. 32 is the top of both FlashInfer kernels' +#: policy. +_ROUTER_LOW_M_MAX_ROWS = 32 # capacity: fixed: a dispatch threshold, not a buffer size +_ROUTING_METHOD_RENORM = 1 +_ACT_TYPE_SWIGLU = 0 +_PAD_SLOT_ID = -1 +#: Whether a captured decode step's projections on `low_m_bf16_mm` launch with programmatic +#: dependent launch (their FlashInfer kernels stream the weight before waiting on the +#: activation's producer; the cuBLAS branch ignores it). +_DECODE_PROJ_PDL = True +#: The K at which `_decode_projection` runs `bf16_gemv_pdl` up to `_GEMV_MAX_ROWS` rows: a +#: tp4/tep4 rank's o and GDN out projections. +_GEMV_TP4_K = 1536 +#: Widest decode projection that `_one_rank_projection` sends to `bf16_gemv_pdl` at 2.. +#: `_GEMV_MAX_ROWS` rows: the shared expert's one-column gate logit (~3 us a layer at 4 rows +#: under `low_m_bf16_mm`). +_GEMV_NARROW_N = 8 # capacity: fixed: a dispatch threshold, not a buffer size +#: Whether a verify step's conv signals its dependents at its start and the CuTe replay after it +#: launches with programmatic dependent launch: the replay then streams the checkpoint and the +#: history (which nothing in the step writes ahead of it) while the conv runs, and reads the +#: conv's output only after `griddepcontrol.wait`. +_GDN_REPLAY_PDL = True +_SOFTPLUS_THRESHOLD = 20.0 +_L2_NORM_EPS = 1e-6 +_PLE_CLAMP_MIN = 1e-6 +#: Row counts up to which a hyper-connection mixer runs its first half (projection and +#: activations) as the one tcgen05 launch `qwen4_exp_hc_decode_tc_first` instead of cuBLAS +#: split-K (+ its reduce) + `qwen4_exp_hc_lowrank_act`: the entry's MAX_ROWS (two 128-row +#: MMA tiles). Every decode batch and MTP verify step a member captures is inside it +#: (the HC is replicated, so the rows are this rank's tokens); larger steps keep cuBLAS. +_HC_TC_FIRST_MAX_ROWS = _HC_TC_FIRST_ROWS +#: At one rank (`hc_fast_sites`), decode rows up to which a hyper-connection site runs as +#: `qwen4_exp_hc_decode_first_cute` + `qwen4_exp_hc_decode_mix_cute` (or its NVFP4 variant at an +#: MoE site): two GEMV launches, which win at one row and lose from two (GB300 micro-bench). From +#: there up to `_HC_DECODE_ROWS_MAX` the tcgen05 pair `qwen4_exp_hc_decode_rows_first` + +#: `qwen4_exp_hc_decode_rows_mix`; larger steps, and every step at the other layouts, +#: `qwen4_exp_hc_inject_norm` and the mixer. The two fast pairs form the mixer's low-rank +#: activations within 1-2 bf16 ulp of that path, not bitwise, so a request's tokens would depend +#: on how many rows its step carries; the other layouts keep one path at every batch size. +_HC_DECODE_GEMV_MAX_ROWS = 1 # capacity: fixed: dispatch threshold, not a buffer bound +#: Decode rows from which a folded MoE block (`DecodeTarget._moe`) launches its router before +#: the shared expert: launched after the side stream's gate|up GEMM its CTAs waited for SMs +#: behind that GEMM's. GB300 micro-bench (24 layers of cold weights in one graph), us a block, +#: shared expert first vs router first: 2 rows 21.7 vs 22.4, 8 rows 46.5 vs 45.7, 32 rows +#: 112.9 vs 112.1, 128 rows 206.7 vs 204.4, 256 rows 275.6 vs 271.7. +_ROUTER_FIRST_MIN_ROWS = 8 # capacity: fixed: a dispatch threshold, not a buffer size +#: At one rank (`hc_fast_sites`), prefill steps of at least this many rows run each trunk +#: hyper-connection site on the un-normalized residual: `qwen4_exp_hc_inject_rstd` stores +#: `hc_out` and a `[T, C]` rstd instead of the `[T, C*H]` normed stream, and the rstd mixer +#: (`_hc_mixer_rstd`, `_hc_moe_mixer_rstd`) forms `normed` on chip. Smaller steps, and every step +#: at the other layouts, keep `qwen4_exp_hc_inject_norm` and the stored-normed mixer, which the +#: rstd mixer does not match bitwise (a generation row in a mixed step would round differently +#: from the same row in a decode step). GB300 micro-bench, us a site +#: (stored normed vs rstd): 2048 rows 57.0 vs 56.7, 3072 rows 82.4 vs 76.1, 4096 rows 104.0 +#: vs 90.8, 8192 rows 186.2 vs 164.9, 16384 rows 360.1 vs 293.7. +_HC_RSTD_MIN_ROWS = 3072 # capacity: fixed: a dispatch threshold, not a buffer size +#: Decode rows up to which a bf16 pool's gated sparse GQA runs the one-launch decode kernel +#: (`qsa_paged_sparse_gqa_decode`) at one rank, and the index scores prefetch their block table. +_SPARSE_DECODE_MAX_ROWS = 64 # capacity: fixed: a dispatch threshold +_QSA_POSITION_AXES = _pool.QSA_POSITION_AXES +# The compressed index keys rotate at coordinate 0 alone (no mRoPE section), as the index +# query does. +_QSA_MROPE_SECTION = None +# The draft experts' activation scales are fp32, not the UE8M0 power-of-two form. +_MTP_MOE_USE_UE8M0 = False +# Rows per draft FP8 runner call: its 128-expert windows are certified up to one 8192-row +# chunk (`fp8_block_scale_moe_runner`). The rows an attention-DP group gathers (up to +# its size times a rank's) run in chunks of it. +_MTP_MOE_MAX_ROWS = 8192 + +# A context launch of at least this many rows over a bf16 or fp8 pool takes +# `qsa_sparse_gqa_rowtile` (tensor cores, K/V tiles shared by a group of rows); below it the +# per-row kernel is faster. +_ROWTILE_MIN_ROWS = 2048 # capacity: fixed: a dispatch threshold, not a buffer size +# Over an fp8 pool the launch reads its requests' pages widened into a bf16 scratch; a launch +# whose scratch would pass this many bytes stays on the per-row kernel. +_FP8_ROWTILE_MAX_SCRATCH = 1 << 30 # capacity: fixed: bounds a per-call scratch allocation +# At one rank a step with more context rows than this over an NVFP4 pool attends them on the +# row-tile pass over a bf16 scratch of their requests' pages (`_sparse_gqa_nvfp4_context`): +# 16384 context rows 1.36 ms a layer there against 77-948 ms for the NVFP4 per-row kernel. +_NVFP4_ROWTILE_MIN_ROWS = 256 # capacity: fixed: a dispatch threshold, not a buffer size +# The bf16 K/V scratch those context rows read holds every page of each context request's KV +# history. Consecutive context requests are dequantized and attended in groups whose pages fit +# this many bytes; a request with more pages forms a group of its own. +_NVFP4_CONTEXT_SCRATCH_BYTES = 128 << 20 # capacity: fixed: a scratch memory budget + +#: Set to 1 to run the GDN input / output and attention input / output projections in FP8 +#: (`fp8_per_tensor_mm`: e4m3 weights with a per-tensor scale, the activation quantized per +#: call with a dynamic per-tensor scale) wherever a call reaches `_DENSE_FP8_MIN_ROWS` rows, and +#: decode rows on `fp8_gemv` over packed FP8 copies with 128x128 block scales (`_FP8_GEMV_ROWS`) +#: or `mid_rows_mm` (`_MID_MM_ROWS`). A lossy mode beside the checkpoint's bf16 projections +#: (gsm8k-gated), certified at tp1; off by default and at every other layout. +_DENSE_FP8_ENV = "TRTLLM_MODELING_V2_DENSE_FP8" + +#: Set to 1 to give each GDN layer an NVFP4 copy of its input projection (e2m1 values, one e4m3 +#: scale per 16 hidden elements of a row, an fp32 global scale; `nvfp4_gemv.pack_split`, built +#: at load from the bf16 weight) and run its decode rows -- up to `_INPROJ_NVFP4_MAX_ROWS` -- on +#: `fp8_gemv`'s kernel over it, the conv step as the GEMV's epilogue in a captured step +#: (single-token and verify). Lossy (gsm8k-gated), certified at tp1; off by default. +_GDN_INPROJ_NVFP4_ENV = "TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4" + +#: Rows from which a projection takes the per-tensor FP8 GEMM, per projection kind: the input +#: projections from 2 rows (one row stays on the GEMV entries, whose PDL overlap the two +#: quantization launches would break), the K=6144 output projections only where the GEMM saving +#: outgrows the activation quantization's two launches (GB300 micro-bench: cuBLASLt bf16 8-10 us +#: vs FP8 6-7 us + ~4 us below ~512 rows). +_DENSE_FP8_MIN_ROWS = {"in": 2, "out": 512} # capacity: fixed: dispatch thresholds + +#: Row ranges `(lo, hi)` in which a projection with a packed FP8 decode copy runs `fp8_gemv` +#: (FP8 weights, bf16 activations, one PDL launch), per kind; each upper end is that kind's +#: crossover against the next entry (GB300 micro, us): in 8.7 vs ~9 `mid_rows_mm` at 16 rows; +#: out 6.4 vs 8.4 bf16 at 16; gate|up 2.6 vs 4.0 at 16; down 1.7 vs 2.9 at 8. +_FP8_GEMV_ROWS = {"in": (1, 16), "out": (1, 16), "sh_gu": (1, 16), "sh_dn": (1, 8)} # dispatch + +#: Row ranges in which a projection runs `mid_rows_mm` (tcgen05, the per-tensor FP8 codes or the +#: bf16 weight) between `fp8_gemv` and the GEMM entries. +_MID_MM_ROWS = { + "in": (17, 256), + "out": (17, 128), + "sh_gu": (17, 64), + "sh_dn": (0, -1), + "router": (17, 64), +} # capacity: fixed: dispatch thresholds + +#: `_mid` weight-name patterns -> the `_MID_MM_ROWS` kind their projections dispatch as. +_MID_KINDS = ( + (r"l\d+_(attn_in_w|in_proj_w)", "in"), + (r"l\d+_(o_w|out_w)", "out"), + (r"l\d+_sh_gu", "sh_gu"), + (r"l\d+_sh_dn", "sh_dn"), + (r"l\d+_router", "router"), +) + +#: Rows up to which a GDN layer's decode in_proj reads the NVFP4 copy (`_GDN_INPROJ_NVFP4_ENV`). +#: GB300 micro (us, conv included, behind a PDL predecessor): 6.4 / 6.1 at 1 / 4 rows against +#: fp8_gemv_conv1d's 7.9 / 7.7; alone 7.1 against 7.9 at 8 rows, 9.1 against 8.6 at 16. +_INPROJ_NVFP4_MAX_ROWS = 8 # capacity: fixed: a dispatch threshold, at most fp8_gemv's MAX_ROWS + +#: Rows up to which a captured verify step's conv runs in the in_proj GEMV's epilogue +#: (`fp8_gemv_conv1d_verify`, certified at 2..16 rows); more rows take the Triton conv. +_VERIFY_CONV_FUSED_MAX_ROWS = 16 # capacity: fixed: the entry's certified rows + +#: Rows up to which a captured single-token step runs the conv step as `fp8_gemv_conv1d`'s +#: epilogue over the packed FP8 in_proj copy (dense FP8 mode). +_FP8_CONV_FUSED_MAX_ROWS = 8 # capacity: fixed: a dispatch threshold + + +def _switch(env: str) -> bool: + mode = os.environ.get(env, "0") + assert mode in ("0", "1"), f"{env}={mode!r}: expected 0 or 1" + return mode == "1" + + +def _gemv_ready(x: torch.Tensor) -> bool: + """`fp8_gemv`'s activation layout: unit column stride, rows 16-byte aligned.""" + return x.stride(1) == 1 and x.stride(0) % 8 == 0 and x.data_ptr() % 16 == 0 + + +def _mid_rows(kind: str, rows: int) -> bool: + lo, hi = _MID_MM_ROWS[kind] + return lo <= rows <= hi + + +def _dense_mm(x, weight_t, fp8, gemm, kind: str, packed=None, mid=None) -> torch.Tensor: + """A dense projection: `fp8_gemv` over the packed decode copy (`packed`) within + `_FP8_GEMV_ROWS[kind]` rows, `mid_rows_mm` over `mid` (a `MidWeight`) within + `_MID_MM_ROWS[kind]` rows, the per-tensor FP8 GEMM (`fp8` = `(codes_t, scale)`) from + `_DENSE_FP8_MIN_ROWS[kind]` rows, else the step's bf16 entry `gemm` -- every call when the + dense FP8 mode is off (all copies None).""" + rows = x.shape[0] + lo, hi = _FP8_GEMV_ROWS.get(kind, (1, 0)) + if packed is not None and lo <= rows <= hi and _gemv_ready(x): + return fp8_gemv_op(x, packed) + if mid is not None and _mid_rows(kind, rows) and (not mid.fp8 or _gemv_ready(x)): + return mid_rows_mm_op(x, mid) + if fp8 is not None and rows >= _DENSE_FP8_MIN_ROWS[kind]: + return fp8_per_tensor_mm(x, *fp8) + return gemm(x, weight_t) + + +def _draft_mm(x: torch.Tensor, weight_t: torch.Tensor, packed, gemm) -> torch.Tensor: + """A draft-layer projection: `fp8_gemv` over its packed FP8 copy (`_MTP_DENSE_FP8_ENV`) at up + to `_MTP_DENSE_FP8_MAX_ROWS` rows, else `gemm` over the bf16 `weight_t`.""" + if packed is not None and x.shape[0] <= _MTP_DENSE_FP8_MAX_ROWS and _gemv_ready(x): + return fp8_gemv_op(x, packed) + return gemm(x, weight_t) + + +def _precompile_mid(mid: dict) -> None: + """Compile every `mid_rows_mm` variant the decode dispatch can reach, once per weight + geometry: a variant first met by a serving step would stall it for the compile.""" + seen = {} + for name, mw in mid.items(): + kind = next(k for rx, k in _MID_KINDS if re.fullmatch(rx, name)) + lo, hi = _MID_MM_ROWS[kind] + key = (mw.n, mw.k, mw.fp8, None if mw.scale is None else mw.scale.numel()) + if lo <= hi and key not in seen: + seen[key] = mid_rows_precompile(mw, max(lo, 1), hi) + + +#: Set to 1 to give the FP8 draft head an NVFP4 copy of its lm_head (e2m1 values, one e4m3 +#: scale per 16 hidden elements of a vocabulary row, an fp32 global scale; built at load) and run +#: its decode rows -- up to `_FP8_GEMV_HEAD_ROWS` -- on `nvfp4_gemv` (GB300, 248320 x 2560: 54-57 +#: us at 1-8 rows, 66-67 at 9-16, against fp8_gemv's 87-90 and 96-98). The verify head is +#: untouched: outputs are unchanged and the acceptance length is the cost to watch. Certified at +#: tp1; 0 (default): `fp8_gemv`. +_MTP_HEAD_NVFP4_ENV = "TRTLLM_MODELING_V2_MTP_HEAD_NVFP4" + +#: Set to 1, under the dense FP8 mode, to run the target head's decode rows (up to +#: `_FP8_GEMV_HEAD_ROWS`) on `nvfp4_gemv` over an NVFP4 copy of lm_head -- the draft head's under +#: `_MTP_HEAD_NVFP4_ENV`, else one built at load -- instead of `fp8_gemv`. Lossy: the verify +#: logits change. Certified at tp1; 0 (default): `fp8_gemv`. +_TARGET_HEAD_NVFP4_ENV = "TRTLLM_MODELING_V2_TARGET_HEAD_NVFP4" + +#: Rows up to which the vocabulary projection (the draft head; under the dense FP8 mode the +#: verify head too) runs `fp8_gemv` on the packed copy; more rows take cuBLASLt over the +#: per-row-scaled copy. GB300, 248320 x 2560: fp8_gemv 86-98 us at 1-16 rows against 107-134, +#: then 142 at 20 rows and 172 at 32 against a flat ~108. +_FP8_GEMV_HEAD_ROWS = 16 # capacity: fixed: a dispatch threshold, at most fp8_gemv's MAX_ROWS + +#: Rows up to which the FP8 vocabulary projection's rows past `_FP8_GEMV_HEAD_ROWS` run +#: `mid_rows_mm` over the per-row-scaled copy (activation quantized per row); more rows keep +#: cuBLASLt. GB300, 248320 x 2560, us at 17 / 32 / 64 / 128 rows: 99.7 / 101.5 / 106.0 / 125.8 +#: against 110.0 / 111.9 / 107.7 / 137.5. +_MID_HEAD_ROWS = 128 # capacity: fixed: a dispatch threshold, at most mid_rows_mm's MAX_ROWS + +#: Set to 1 to run the draft layer's dense projections (`fc_e`, `fc_h`, the fused attention +#: input, `o`, the shared expert's gate|up and down) on `fp8_gemv` over packed FP8 copies at decode +#: rows (up to `_MTP_DENSE_FP8_MAX_ROWS`). The draft only proposes tokens, so outputs are unchanged; +#: only the acceptance length can move. Certified at tp1; 0 (default): bf16. +_MTP_DENSE_FP8_ENV = "TRTLLM_MODELING_V2_MTP_DENSE_FP8" +_MTP_DENSE_FP8_MAX_ROWS = 16 # capacity: fixed: a dispatch threshold, at most fp8_gemv's MAX_ROWS + +#: Set to 1 to run the draft layer's routed experts in NVFP4 on the trunk's trtllm-gen runner: at +#: load the checkpoint's FP8 128x128 block-scaled experts are requantized (per-expert weight global +#: scales, 16-element e4m3 block scales) into the trunk's layout and the FP8 copies are freed. The +#: draft only proposes tokens, so this can change the acceptance length, not outputs. Certified at +#: tp1; 0 (default): the FP8 block-scale runner. +_MTP_MOE_NVFP4_ENV = "TRTLLM_MODELING_V2_MTP_MOE_NVFP4" + +#: The NVFP4 draft MoE's two activation global scales are static, `448 * 6 / A` with `A` this factor +#: times the largest calibrated amax of the trunk's routed experts (its `input_scale`s, FC1 input and +#: FC2 input; the draft layer carries none). A global scale only has to keep the e4m3 block scales +#: inside their range, so the margin costs nothing until blocks are ~1e-3 of `A`. +_MTP_MOE_NVFP4_FC1_MARGIN = 8.0 +_MTP_MOE_NVFP4_FC2_MARGIN = 4.0 + +_PLE_HOST_OFFLOAD_ENV = "TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD" +_PLE_HOST_OFFLOAD_OFF = frozenset({"0", "false", "no", "off"}) + +#: At one rank, context-step MoE launches of at least this many rows (and at most +#: `_CUTE_DSL_MOE_MAX_ROWS`, the op's certified top) run on the CuTe DSL MoE +#: (`fp4_moe_renorm_cute_dsl`, over the second expert copy `weights.py` keeps in its layout); +#: smaller ones, every decode one among them, keep the trtllm-gen runner. GB300 micro-benches, +#: us a layer, CuTe DSL vs trtllm-gen: 2048 rows 275 vs 273, 8192 rows 451 vs 555, 16384 rows +#: 656 vs 882. +_CUTE_DSL_MOE_MIN_ROWS = 4096 # capacity: fixed: a dispatch threshold, not a buffer size +_CUTE_DSL_MOE_MAX_ROWS = 16384 # capacity: fixed: the op's certified rows +# Whether one rank keeps that second copy (about 1.3 GiB a layer). No engine arg owns it, hence +# an env var read at construction; "off" (the spellings above) drops the copy and every launch +# runs on trtllm-gen. +_MOE_PREFILL_COPY_ENV = "TRTLLM_MODELING_V2_MOE_PREFILL_COPY" + + +def _moe_prefill_copy(one_rank: bool, trunk_moe_backend: str, max_num_tokens: int) -> bool: + """Whether the trunk keeps the CuteDSL-layout expert copy: at one rank only (every expert + is local, the layout the CuteDSL MoE is certified at), only under trtllm-gen, only when a + step's rows (at most `max_num_tokens`) can reach the dispatch threshold, and not when the + env var turns it off.""" + return ( + one_rank + and trunk_moe_backend == "TRTLLM" + and max_num_tokens >= _CUTE_DSL_MOE_MIN_ROWS + and os.environ.get(_MOE_PREFILL_COPY_ENV, "1").strip().lower() not in _PLE_HOST_OFFLOAD_OFF + ) + + +def _ple_host_offload_enabled() -> bool: + """Whether the n-gram table is declared in host memory (the default).""" + value = os.environ.get(_PLE_HOST_OFFLOAD_ENV, "1").strip().lower() + return value not in _PLE_HOST_OFFLOAD_OFF + + +def _moe_layer_scalars(w, i: int, start: int, count: int) -> dict[str, torch.Tensor]: + """Layer `i`'s routed-expert scalars for the trtllm-gen runner over this rank's window + `[start, start + count)`, built from the checkpoint's `input_scale` / `weight_scale_2` + rows (declared over the whole routing space on every rank). + + The checkpoint stores modelopt reciprocals (`input_scale = 1 / g_act`, `weight_scale_2 = + 1 / g_w`), so a quantizer's global scale is `1 / input_scale` and the runner's scalars + are products of the two.""" + isc1, isc1_up = w[f"l{i}_e_isc1"], w[f"l{i}_e_isc1_up"] + assert torch.equal(isc1, isc1_up), ( + f"layer {i}: expert gate/up input_scale differ; the fused FC1 reads one activation " + "quantization for both halves" + ) + assert torch.equal(w[f"l{i}_e_ws2_1"], w[f"l{i}_e_ws2_1_up"]), ( + f"layer {i}: expert gate/up weight_scale_2 differ; the fused FC1 carries one alpha " + "per expert for both halves" + ) + # One activation quantization feeds all experts, on every rank. + assert torch.equal(isc1.min(), isc1.max()), ( + f"layer {i}: the routed experts do not share one gate_proj.input_scale, so a single " + "activation quantization cannot feed them all" + ) + e_isc2, ws2_2 = w[f"l{i}_e_isc2"], w[f"l{i}_e_ws2_2"] + gate1 = isc1 * w[f"l{i}_e_ws2_1"] + window = slice(start, start + count) + return dict( + # The routed FC1 activation global scale (one per layer). + global_scale=(1.0 / isc1[:1]).contiguous(), + # [count] each. Swapping the first two is finite and silent. + output1_scale_scalar=(gate1 / e_isc2)[window].contiguous(), + output1_scale_gate_scalar=gate1[window].contiguous(), + output2_scale_scalar=(e_isc2 * ws2_2)[window].contiguous(), + ) + + +def _moe_layer_scalars_cd(w, i: int) -> dict[str, torch.Tensor]: + """Layer `i`'s routed-expert scalars for the CuteDSL MoE over the second expert copy (one + rank, every expert; `_moe_layer_scalars` checks the checkpoint facts they rest on). FC1's + per-expert alpha is `input_scale * weight_scale_2`, its NVFP4 input the same activation + quantization as the trtllm-gen runner's; FC2's input is quantized at one global scale, + `1 / max(input_scale)` over the layer's experts, and its alpha carries that maximum.""" + isc2_max = w[f"l{i}_e_isc2"].max() + return dict( + alpha1=(w[f"l{i}_e_isc1"] * w[f"l{i}_e_ws2_1"]).contiguous(), + alpha2=(isc2_max * w[f"l{i}_e_ws2_2"]).contiguous(), + global_sf=(1.0 / isc2_max).reshape(1).contiguous(), + ) + + +class _Replay(NamedTuple): + """The GDN cached replay's history buffers and bookkeeping of one manager.""" + + old: dict # GDN layer -> (old_x, old_B, old_dt, old_dA_cumsum) + cache_buf_idx: torch.Tensor # [slots] int32: the history half each slot writes next + prev_num_accepted: torch.Tensor # [slots] int32: tokens pending since the checkpoint + history: int # rows a history half holds + width: int # the one verify width a step may carry per generation request + + +class _Bound(NamedTuple): + """What the core reads off one KV cache manager: the manager, this rank's pool views + (`pool.RankPools`) and, under MTP, the verify step's runtime-owned buffers: the per-step + conv windows (`int_conv`, by GDN layer), the rows `update_mamba_states` promotes them + from (`int_idx`) and the cached replay. Valid only for that manager.""" + + cm: object + pools: _pool.RankPools + int_idx: torch.Tensor | None = None + int_conv: dict | None = None + replay: _Replay | None = None + + +class _DraftBound(NamedTuple): + """The draft layer's planes in the manager its draft step is handed: its QSA index plane, + that manager's position plane, its paged K/V and, under IndexShare, its row of the + metadata's shared selection buffers.""" + + cm: object + index_cache: torch.Tensor + position_cache: torch.Tensor + kv: _pool.KvPlanes + share_slot: int | None + + +class Qwen4ExpDefaultModelingV2(DecoderModel): + def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig): + """`composite_config` is passed in because the text sub-config adopted as + `model_config.pretrained_config` carries neither the vision half nor the media + token ids.""" + super().__init__(model_config) + cfg = model_config.pretrained_config + mapping = model_config.mapping + assert torch.cuda.get_device_capability() == _SM, ( + f"target certified on sm_{_SM[0]}{_SM[1]}, running on " + f"sm_{''.join(map(str, torch.cuda.get_device_capability()))}" + ) + assert getattr(cfg, "model_type", None) == "qwen4_exp_text", ( + f"the core must be built on the normalized text config; got model_type " + f"{getattr(cfg, 'model_type', None)!r}" + ) + self.mapping = mapping + # The ranks attention heads and the recurrent heads are split over: the TP group, + # or one rank under attention DP. + self.attn_tp_size = _pool.attention_tp_size(mapping) + # Under attention DP each rank serves its own requests and the MoE blocks run over the + # rows the whole group holds (`_expert_sum_op`). The group is the TP group, which spans + # the DP ranks; its size is the world's, never a constant. + self.dp_size = mapping.tp_size if mapping.enable_attention_dp else 1 + self.dp_rank = mapping.tp_rank + self.dp_group = list(mapping.tp_group) + # One-model MTP-Eagle (routing admits no other speculation): the checkpoint's one MTP + # module, built only when the engine asked for speculation. + spec_config = getattr(model_config, "spec_config", None) + self.mtp_enabled = spec_config is not None + # The draft layer's FP8 block-scale experts are whole-expert windows over the TP group + # at every layout, tp4 included (`weights.rank_dims`, `mtp_local_experts`). + # Rows per generation request of the widest verify step (`max_draft_len + 1`); 1 + # without MTP. + self.max_verify_rows = spec_config.max_draft_len + 1 if self.mtp_enabled else 1 + # The most rows one forward carries (a draft step's included): sizes the draft + # layer's unit injection weight (`_MtpWeights.unit_inject`). + self.max_num_tokens = model_config.max_num_tokens + + self.num_layers = cfg.num_hidden_layers + self.hidden = cfg.hidden_size + self.vocab = cfg.vocab_size + self.eps = cfg.rms_norm_eps + layer_types = list(cfg.layer_types) + self.is_attn = [t == "full_attention" for t in layer_types] + self.attn_layers = [i for i, a in enumerate(self.is_attn) if a] + self.gdn_layers = [i for i, a in enumerate(self.is_attn) if not a] + + # ---- hyper-connections ------------------------------------------ + self.hc_count = cfg.hc_count + self.hc_lowrank = cfg.hc_lowrank + self.stream = self.hc_count * self.hidden + # Columns of the padded first-projection weight (`_fuse_hc_first_projections`) the + # first projection multiplies: the real `[down; block_inject]` rows rounded up to 32. + self.hc_first_cols = -(-(self.hc_lowrank + self.hc_count) // 32) * 32 + + # ---- RoPE (partial, neox pairs, interleaved mRoPE) ----------------- + self.head_dim = cfg.head_dim + rope = getattr(cfg, "rope_parameters", None) or getattr(cfg, "rope_scaling", None) + self.theta = float(rope["rope_theta"]) + partial = float(rope.get("partial_rotary_factor", cfg.partial_rotary_factor)) + self.rotary_dim = int(self.head_dim * partial) + self.mrope_declared = str(rope.get("type", rope.get("rope_type"))) == "mrope" + section = list(rope["mrope_section"]) + self.mrope_section1 = int(section[1]) + self.mrope_section2 = int(section[2]) + self.max_positions = cfg.max_position_embeddings + + # ---- QSA sparse index (replicated on every rank) ------------------- + self.index_heads = cfg.indexer_n_heads + self.index_kv_heads = cfg.indexer_kv_heads + self.index_head_dim = cfg.indexer_head_dim + self.token_topk = cfg.indexer_budget + self.compress_ratio = cfg.indexer_compress_ratio + self.block_topk = self.token_topk // self.compress_ratio + # Dense-attention cutoff: QSA resolves max(token_topk, seq_len_threshold); checked + # against the engine's value because the CUDA-graph runner splits its short and long + # decode families on it. + self.dense_threshold = self.token_topk + sparse_cfg = getattr(model_config, "sparse_attention_config", None) + assert getattr(sparse_cfg, "algorithm", None) == "qsa", ( + f"this target is a QSA sparse-attention assembly and the engine resolved " + f"sparse_attention_config={sparse_cfg!r}; QSA is this architecture's own default" + ) + qsa_params = sparse_cfg.to_sparse_params(pretrained_config=cfg) + assert ( + qsa_params.index_n_heads, + qsa_params.index_kv_heads, + qsa_params.index_head_dim, + qsa_params.token_topk, + qsa_params.compress_ratio, + ) == ( + self.index_heads, + self.index_kv_heads, + self.index_head_dim, + self.token_topk, + self.compress_ratio, + ), ( + f"the engine lowered the QSA index geometry to {qsa_params!r}; the side caches " + "are sized from the engine's copy and the kernels are called with this one" + ) + assert qsa_params.dense_seq_len_threshold == self.dense_threshold, ( + f"the engine's resolved QSA cutoff is {qsa_params.dense_seq_len_threshold} and " + f"this forward branches at {self.dense_threshold}: the runner captures its short " + "decode graph family at the engine's number" + ) + # IndexShare for the MTP draft loop (`index_share_for_mtp_iteration`), as the engine + # resolved it: the same lowering sizes the metadata's shared selection buffers. + self.mtp_index_share = bool(qsa_params.mtp_index_share) + self.index_width = (self.index_heads + self.index_kv_heads) * self.index_head_dim + # The index query's `[q | k | v]` scratch for its norm + rope; the spare head is the + # v slot the op needs but never reads. + self.index_scratch_width = self.index_width + self.index_head_dim + self.softmax_scale = 1.0 / math.sqrt(self.head_dim) + + # ---- the MTP module ------------------------------------------------ + # Global layer index of the draft layer (the KV managers map it there, and the + # weight table keys it there); its rope base is its own block's, else the trunk's. + self.mtp_layer_idx = self.num_layers + mtp_block = getattr(cfg, "mtp", None) + self.mtp_theta = ( + float(mtp_block.get("rope_theta", self.theta)) + if isinstance(mtp_block, dict) + else self.theta + ) + + # ---- gated deltanet ---------------------------------------------- + self.head_k_dim = cfg.linear_key_head_dim + self.head_v_dim = cfg.linear_value_head_dim + self.conv_kernel = cfg.linear_conv_kernel_dim + + # ---- MoE ---------------------------------------------------------- + self.num_experts = cfg.num_experts + self.topk = cfg.num_experts_per_tok + self.moe_inter = cfg.moe_intermediate_size + self.shared_inter = cfg.shared_expert_intermediate_size + requested = model_config.moe_backend_requested.upper() + assert requested in _TRUNK_MOE_BACKENDS, ( + f"moe_config.backend={model_config.moe_backend_requested!r} is not implemented by " + f"this target; the supported requests are {sorted(_TRUNK_MOE_BACKENDS)}" + ) + # Fixes the load-time layout of the expert stacks (`weights.py`). + self.trunk_moe_backend = _TRUNK_MOE_BACKENDS[requested] + # The CuteDSL pipeline runs every expert at once: one rank only. + assert self.trunk_moe_backend != "CUTEDSL" or mapping.world_size == 1, ( + "moe_config.backend=CUTEDSL is implemented at one rank (tp1) only; this layout " + f"has {mapping.world_size} ranks" + ) + assert self.trunk_moe_backend != "CUTEDSL" or ( + self.max_verify_rows - 1 <= _CUTEDSL_MAX_DRAFT_LEN + ), ( + f"max_draft_len={self.max_verify_rows - 1} with moe_config.backend=CUTEDSL: the " + f"CuteDSL expert pipeline is certified at max_draft_len up to {_CUTEDSL_MAX_DRAFT_LEN}" + ) + # At one rank the CuteDSL-layout copy of the experts for prefill-sized launches + # (`_moe_prefill_copy`); `weights.py` declares it from this. + self.moe_prefill_copy = _moe_prefill_copy( + mapping.world_size == 1, self.trunk_moe_backend, model_config.max_num_tokens + ) + # The draft experts' NVFP4 copy at one rank (`_mtp_moe_nvfp4`), built once: it frees + # the FP8 experts, and the shell's post-load walk reaches `post_load_weights` again. + self._mtp_moe_nvfp4 = None + + # ---- PLE ----------------------------------------------------------- + # One-based in the config. + self.ple_layer = list(cfg.ple_layer_ids)[0] - 1 + self.ngram_size = cfg.ngram_size + self.ple_dilation = self.ngram_size + self.ngram_vocab_divisor = cfg.make_ngram_vocab_size_divisible_by + eos = cfg.eos_token_id + self.eos_token_id = int(eos[0] if isinstance(eos, (list, tuple)) else eos) + self.ple_table_dtype = _weights.PLE_TABLE_DTYPE + self.ple_host_offload = _ple_host_offload_enabled() + # CPU allocation behind the offloaded table; the `ple_table` parameter is a + # device-address view that owns none of those pages. + self._ple_table_host = None + + # ---- the vision half (the weight table declares the tower) --------- + vcfg = getattr(composite_config, "vision_config", None) + assert vcfg is not None, "the composite config carries no vision_config" + self.vision_config = vcfg + self.image_token_id = int(composite_config.image_token_id) + self.video_token_id = int(composite_config.video_token_id) + + # ---- KV pool storage ------------------------------------------------ + self.kv_cache_quant_algo = model_config.quant_config.kv_cache_quant_algo + assert self.kv_cache_quant_algo in (None, QuantAlgo.FP8, QuantAlgo.NVFP4), ( + f"this target stores K/V as the checkpoint's bf16, as fp8-e4m3 or as NVFP4; the " + f"engine resolved kv_cache_quant_algo {self.kv_cache_quant_algo!r}" + ) + self.quant_mode = { + None: 0, + QuantAlgo.FP8: QM_FP8_KV_CACHE, + QuantAlgo.NVFP4: QM_NVFP4_KV_CACHE, + }[self.kv_cache_quant_algo] + + # ---- this rank's share ---------------------------------------------- + d = _weights.MODEL_WEIGHTS.dims(self) + assert d.expert_split in _weights.LOADED_EXPERT_SPLITS, d.expert_split + self.heads_q = d.heads_q + self.heads_kv = d.heads_kv + self.kv_head_start = d.kv_head_start + self.q_width = d.q_width + self.kv_width = d.kv_width + self.qkv_width = self.q_width + 2 * self.kv_width + # The fused input projection's output row, `[q | k | v | gate | index q, k | v slot]` + # (`weights._attn_in_w`). Its leading columns are the packed QKV the main norm + rope + # rewrites in place; the rope addresses a row as `(Hq + Hk + Hv) * head_dim` and + # never touches the v heads, so everything after k is "v" to it. + self.attn_in_width = d.attn_in_rows + assert self.attn_in_width == self.qkv_width + self.q_width + self.index_scratch_width + assert self.attn_in_width % self.head_dim == 0, self.attn_in_width + self.attn_in_tail_heads = self.attn_in_width // self.head_dim - self.heads_q - self.heads_kv + self.index_col = self.qkv_width + self.q_width + self.num_k_heads = d.num_k_heads + self.num_v_heads = d.num_v_heads + self.key_dim = d.key_dim + self.value_dim = d.value_dim + self.conv_dim = d.conv_dim + # Column widths of the one GDN input projection's output: [q|k|v | z | a | b]. + self.in_proj_split = d.in_proj_split + self.local_experts = d.local_experts + self.expert_offset = d.expert_offset + # The runner's `intermediate_size`: each held expert's intermediate as stored, this + # rank's slice zero-padded at tp4 (`weights.EXPERT_INTER_ALIGN`), whole otherwise. + self.expert_inter = d.expert_inter + # The draft layer's FP8 experts: this rank's window over the whole TP group. + self.mtp_local_experts = d.mtp_local_experts + self.mtp_expert_offset = d.mtp_expert_offset + + # Every weight but the n-gram table; meta-init intercepts torch.empty. + self.w = _weights.MODEL_WEIGHTS.declare(d) + + # Derived once after load (`post_load_weights`). + self._moe_scalars: list | None = None + self._gdn_gating: dict | None = None + self._cos_sin: torch.Tensor | None = None + self._hc_first: dict | None = None + # The attention op's KV scales: fp32 [1] ones over an fp8 pool, fp32 [3] ones (q, K, + # V) over an NVFP4 one, None over bf16. + self.kv_scale_orig_quant: torch.Tensor | None = None + self.kv_scale_quant_orig: torch.Tensor | None = None + # The shared expert's stream (see `_moe_block`); made after load, because a stream + # created inside a graph capture is not capturable. + self._side_stream: torch.cuda.Stream | None = None + # The tensor-parallel sum of a block's partial output (`_tp_sum_op`), and the routed + # experts' completion over this rank's rows (`_expert_sum_op`). + self.tp_sum = None + self.expert_sum = None + self._prefill: "PrefillTarget | None" = None + self._decode: "DecodeTarget | None" = None + # Everything derived from the KV cache manager (`_bind_cache_manager`). + self._bound: _Bound | None = None + self._binds = 0 + # Set by `begin_step`, cleared by `forward`. + self._step_begun = False + # MTP module operands (`_mtp_views`) and its index rope table; None with speculation + # off. + self._mtp: "_MtpWeights | None" = None + self._mtp_cos_sin: torch.Tensor | None = None + # The verify step's side-cache rollback, registered with the spec worker. + self.spec_rollback: SpeculativeRollback | None = ( + SpeculativeRollback() if self.mtp_enabled else None + ) + # The draft layer's own bind, keyed on the manager its draft step sees. + self._draft_bound: _DraftBound | None = None + + # ---------------------------------------------------------------- load --- + + def _fuse_hc_first_projections(self) -> None: + """Concatenate each hyper-connection's `down` and `block_inject` weights into one + `[rows, stream]` tensor, zero-padded to a multiple of 128 rows, so a site's first + projection is one GEMM: the trunk's sites, keyed `(layer, tag)` and `"fhc"`, and with + MTP the draft layer's, keyed `"mtp_ahc"`, `"mtp_mhc"` and `"mtp_fhc"`. The + checkpoint's parameters are rebound as row views of the padded tensors.""" + w = self.w + + def fuse(*params): + rows = sum(t.shape[0] for t in params) + padded = -(-rows // 128) * 128 + pad = torch.zeros( + [padded - rows, self.stream], dtype=params[0].dtype, device=params[0].device + ) + both = torch.cat([t.data for t in params] + [pad], 0) + start = 0 + for t in params: + n = t.shape[0] + t.data = both[start : start + n] + start += n + return both + + self._hc_first = { + (i, tag): fuse(w[f"l{i}_{tag}_down"], w[f"l{i}_{tag}_inject"]) + for i in range(self.num_layers) + for tag in ("ahc", "mhc") + } + self._hc_first["fhc"] = fuse(w["fhc_down"]) + if self.mtp_enabled: + p = f"l{self.mtp_layer_idx}_mtp_" + self._hc_first["mtp_ahc"] = fuse(w[f"{p}ahc_down"], w[f"{p}ahc_inject"]) + self._hc_first["mtp_mhc"] = fuse(w[f"{p}mhc_down"], w[f"{p}mhc_inject"]) + self._hc_first["mtp_fhc"] = fuse(w[f"{p}fhc_down"]) + + def post_load_weights(self) -> None: + """Derive what the targets read, then build them, now that the weights are real. + + Not an override (DecoderModel has no such hook); named for the shell hook.""" + w = self.w + device = w["embed"].device + self._moe_scalars = [ + _moe_layer_scalars(w, i, self.expert_offset, self.local_experts) + for i in range(self.num_layers) + ] + # The CuteDSL scalars: of the second expert copy (`moe_prefill_copy`), or of the + # experts themselves on the CuteDSL backend. + self._moe_scalars_cd = ( + [_moe_layer_scalars_cd(w, i) for i in range(self.num_layers)] + if self.moe_prefill_copy or self.trunk_moe_backend == "CUTEDSL" + else None + ) + # The delta-rule entries index A_log and dt_bias as `ptr + head` and certify fp32; + # the checkpoint is bf16. + self._gdn_gating = { + i: ( + w[f"l{i}_A_log"].float().contiguous(), + w[f"l{i}_dt_bias"].float().contiguous(), + ) + for i in self.gdn_layers + } + # Reconcile the table height against the checkpoint's id-space tensors: an index past + # the end is a device-side assert that poisons the process. + table_rows = w["ple_table"].shape[0] + used_rows = int(w["ple_offs"][-1]) + int(w["ple_vocab"][-1]) + assert used_rows <= table_rows, ( + f"the n-gram table has {table_rows} rows but its per-head offsets and vocabulary " + f"sizes address {used_rows}" + ) + assert table_rows % self.ngram_vocab_divisor == 0, (table_rows, self.ngram_vocab_divisor) + self._cos_sin = self._index_rope_table(self.theta, device) + # Before the draft layer's views, which read its fused first projections. + self._fuse_hc_first_projections() + if self.mtp_enabled: + # The draft layer's index plane rotates at the module's own rope base, like its + # attention rope. + self._mtp_cos_sin = ( + self._cos_sin + if self.mtp_theta == self.theta + else self._index_rope_table(self.mtp_theta, device) + ) + self._mtp = _mtp_views(self) + self._build_draft_one_rank(device) + if self.quant_mode in (QM_FP8_KV_CACHE, QM_NVFP4_KV_CACHE): + # Built here, not at construction: that runs under meta init. + scales = 3 if self.quant_mode == QM_NVFP4_KV_CACHE else 1 + self.kv_scale_quant_orig = torch.ones(scales, dtype=torch.float32, device=device) + self.kv_scale_orig_quant = 1.0 / self.kv_scale_quant_orig + self._build_dense_copies(device) + # The CPU-only load checks build the targets over host weights; they launch nothing. + if device.type == "cuda": + # Every first-half variant `_hc_mixer` launches up to `_HC_TC_FIRST_MAX_ROWS` rows, + # compiled here so that no serving step (or graph capture) meets a compile. + hc_tc_first_precompile(self.stream, self.hc_lowrank, self.hc_count, device) + self._report_ple_residency() + self._side_stream = torch.cuda.Stream(device=device) + # The prefill CuteDSL MoE's output memset, beside its FC1 (`fp4_moe_renorm_cute_dsl`): + # its own stream, so the join before FC2 does not also wait for the shared expert. + self._moe_memset_stream = torch.cuda.Stream(device=device) + self.tp_sum = self._tp_sum_op() + self.expert_sum = self._expert_sum_op() + # A decode step's bf16 projection entry, and its GDN in_proj's: at tp1 the one-rank + # dispatch (`_one_rank_projection`), at the other layouts the shapes' own. + one_rank = self.mapping.world_size == 1 + self.decode_mm = _one_rank_projection if one_rank else _decode_projection + # At one rank a trtllm-gen decode MoE leaves its finalize to the next site + # (`DecodeTarget._moe`). + self.moe_fold = one_rank and self.trunk_moe_backend == "TRTLLM" + # At one rank the hyper-connection sites' fast forms by row count (the decode GEMV and + # tcgen05 rows pairs, the prefill rstd sites); at the other layouts one form serves every + # batch size, so a request's tokens do not depend on what its steps carry beside it. + self.hc_fast_sites = one_rank + self.decode_in_mm = ( + _one_rank_projection + if one_rank + else (lambda mat_a, mat_b: low_m_bf16_mm(mat_a, mat_b, _DECODE_PROJ_PDL)) + ) + self._prefill = PrefillTarget(self) + self._decode = DecodeTarget(self) + # The FP8 draft head at tp1, set by the shell's `post_load_weights`. + self.draft_head = None + + def _build_draft_one_rank(self, device) -> None: + """At one rank the draft layer's operands for `DecodeTarget._mtp_layer`: the fused + `[q | k | v | gate | index | pad]` attention input (`weights._attn_in_w`), with + `_MTP_MOE_NVFP4_ENV` its routed experts in NVFP4, and with `_MTP_DENSE_FP8_ENV` the + packed FP8 copies of its dense projections; None / empty at every other layout.""" + self._mtp_attn_in = None + self._mtp_gemv8 = {} + if self.mapping.world_size != 1: + return + w = self.w + p = f"l{self.mtp_layer_idx}_mtp_" + attn_in = _weights._attn_in_w( + (w[f"{p}q_w"], w[f"{p}k_w"], w[f"{p}v_w"], w[f"{p}idx_w"]), self + ) + self._mtp_attn_in = attn_in.t() + if device.type != "cuda": + return + if _switch(_MTP_MOE_NVFP4_ENV): + if self._mtp_moe_nvfp4 is None: + self._mtp_moe_nvfp4 = _mtp_moe_nvfp4(self) + self._mtp = self._mtp._replace(**self._mtp_moe_nvfp4) + if _switch(_MTP_DENSE_FP8_ENV): + self._mtp_gemv8 = { + "fc_e": pack_from_bf16(w[f"{p}fc_e"]), + "fc_h": pack_from_bf16(w[f"{p}fc_h"]), + "attn_in": pack_from_bf16(attn_in), + "o": pack_from_bf16(w[f"{p}o_w"]), + "sh_gu": pack_from_bf16(w[f"{p}sh_gu"][: 2 * self.shared_inter]), + "sh_dn": pack_from_bf16(w[f"{p}sh_dn"]), + } + + def _build_dense_copies(self, device) -> None: + """The trunk's dense FP8 copies (`_DENSE_FP8_ENV`) beside the bf16 weights the one-row + GEMV and fused paths keep reading, and the GDN in_proj's NVFP4 copies + (`_GDN_INPROJ_NVFP4_ENV`), with every decode variant compiled now; both only at tp1, + the layout they are certified at, and empty otherwise.""" + w = self.w + one_rank = self.mapping.world_size == 1 + self._dense8, self._gemv8, self._mid, self._gemv_nv = {}, {}, {}, {} + if one_rank and _switch(_DENSE_FP8_ENV): + for name in w: + if re.fullmatch(r"l\d+_(attn_in_w|o_w|in_proj_w|out_w)", name): + codes, scale = fp8_per_tensor_weight(w[name]) + self._dense8[name] = (codes.t(), scale) + self._gemv8[name] = pack_from_bf16(w[name]) + self._mid[name] = MidWeight(codes, scale) + elif re.fullmatch(r"l\d+_sh_gu", name): + # The shared expert's gate|up without its gate-logit row, which stays bf16. + self._gemv8[name] = pack_from_bf16(w[name][: 2 * self.shared_inter]) + self._mid[name] = MidWeight(w[name][: 2 * self.shared_inter]) + elif re.fullmatch(r"l\d+_sh_dn", name): + self._gemv8[name] = pack_from_bf16(w[name]) + self._mid[name] = MidWeight(w[name]) + elif re.fullmatch(r"l\d+_router", name): + self._mid[name] = MidWeight(w[name]) + if device.type == "cuda": + _precompile_mid(self._mid) + if one_rank and _switch(_GDN_INPROJ_NVFP4_ENV) and device.type == "cuda": + for name in w: + if re.fullmatch(r"l\d+_in_proj_w", name): + self._gemv_nv[name] = nvfp4_pack_split_from_bf16(w[name]) + if self._gemv_nv: + precompile_split( + next(iter(self._gemv_nv.values())), + _INPROJ_NVFP4_MAX_ROWS, + self.conv_dim, + _PAD_SLOT_ID, + tuple(range(2, self.max_verify_rows + 1)), + ) + + def _tp_sum_op(self): + """The sum that completes a block's per-rank partial output, settled once: under + attention TP the plain all-reduce over the TP group (`comm/allreduce`, op `NONE`), + which every rank gets the same bits back from; otherwise nothing to sum.""" + if self.attn_tp_size == 1: + return lambda x: x + group = allreduce_group_args(self.mapping) + + def tp_sum(x: torch.Tensor) -> torch.Tensor: + return allreduce( + x, + None, + None, + scale=None, + bias=None, + op=int(AllReduceFusionOp.NONE), + eps=self.eps, + trigger_completion_at_end=False, + **group, + ) + + return tp_sum + + def _expert_sum_op(self): + """How a MoE block's routed experts become the rows this rank owns, settled once. + The op is `expert_sum(run, x, pad_rows)`, where `run(rows)` is this rank's expert + window over `rows` and `pad_rows` is `_dp_pad_rows`'s count. + + Under attention TP (or at one rank) every rank holds every row: the window's partial + over `x`, completed by `tp_sum`. Under attention DP each rank holds its own rows, as + gpt-oss's and deepseek's default targets run it: pad them to the group's largest row + count, gather the group's rows (`comm/allgather`, in rank order), run the window over + all of them, reduce-scatter (`comm/reducescatter`), which sums each row's window + partials once and hands this rank its slice, and drop the padding. Routing on the + gathered rows makes every rank's routing of a row the same, so each of its top-k + experts falls in exactly one window.""" + if self.dp_size == 1: + tp_sum = self.tp_sum + + def window_sum(run, x: torch.Tensor, pad_rows: int) -> torch.Tensor: + return tp_sum(run(x)) + + return window_sum + group = self.dp_group + # Under MPI the plain symbols; with MPI disabled the `_pg` ones over the TP process + # group, as `allreduce_group_args` routes them. + process_group = self.mapping.tp_group_pg.boxed() if mpi_disabled() else None + + def round_trip(run, x: torch.Tensor, pad_rows: int) -> torch.Tensor: + rows = x.shape[0] + padded = x if not pad_rows else nn.functional.pad(x, [0, 0, 0, pad_rows]) + gathered = allgather(padded, None, group, process_group) + mine = reducescatter(run(gathered), None, group, process_group) + return mine if not pad_rows else mine[:rows] + + return round_trip + + def _dp_pad_rows(self, md, num_tokens: int) -> int: + """Rows this rank pads its MoE input with before the attention-DP round trip: up to + the group's largest per-rank token count, so both collectives run in their uniform + form. Zero without attention DP. + + A decode CUDA graph pads every rank to the captured batch, and a prefill CUDA graph + every rank to one bucket (the engine then reports the bucket for every rank), so a + captured step's counts are equal and its padding is zero rows.""" + if self.dp_size == 1: + return 0 + counts = [int(n) for n in md.all_rank_num_tokens] + assert len(counts) == self.dp_size, (counts, self.dp_size) + assert counts[self.dp_rank] == num_tokens, ( + f"rank {self.dp_rank} holds {num_tokens} rows but all_rank_num_tokens says " + f"{counts[self.dp_rank]} ({counts}); the collectives would gather the wrong split" + ) + return max(counts) - num_tokens + + def _index_rope_table(self, theta: float, device) -> torch.Tensor: + """The index compression's rope table at `theta` in the QSA kernels' layout: + `[max_positions, 2 * rotary_pairs]` fp32, cos block then sin block.""" + exponent = torch.arange(0, self.rotary_dim, 2, dtype=torch.float64, device=device) + inv_freq = (1.0 / (theta ** (exponent / self.rotary_dim))).to(torch.float32) + positions = torch.arange(self.max_positions, dtype=torch.float32, device=device) + angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0) + return torch.cat([torch.cos(angles), torch.sin(angles)], dim=1).contiguous() + + def _report_ple_residency(self) -> None: + """Check where the n-gram table ended up, and report it once: offload is + output-identical, so the pointer's memory kind read back from the driver is the + signal that cannot be faked.""" + table = self.w["ple_table"] + kind = _weights.memory_kind(table.data_ptr()) + expected = ( + _weights.CUDA_MEMORY_TYPE_HOST + if self.ple_host_offload + else _weights.CUDA_MEMORY_TYPE_DEVICE + ) + assert kind == expected, ( + f"the n-gram table was declared with {_PLE_HOST_OFFLOAD_ENV}=" + f"{'1' if self.ple_host_offload else '0'}, but the driver reports cudaMemoryType " + f"{kind} for the pointer the forward gathers from" + ) + mib = 1024.0 * 1024.0 + free, total = torch.cuda.mem_get_info() + print( + f"ple n-gram table: {tuple(table.shape)} {table.dtype}, " + f"{table.numel() * table.element_size() / mib:.1f} MiB, " + f"{'HOST-resident (device-addressable)' if self.ple_host_offload else 'DEVICE-resident'}" + f"; after load: allocated {torch.cuda.memory_allocated() / mib:.1f} MiB, " + f"driver-used {(total - free) / mib:.1f} MiB of {total / mib:.1f} MiB", + flush=True, + ) + + # ---------------------------------------------------------------- bind --- + + def _bind_cache_manager(self, md, position_ids) -> None: + """Bind this rank's pool views of the step's KV cache manager, checked against the + geometry the weights were cut to. + + Keyed on the manager, not a first-forward flag: capacity estimation builds a + throwaway manager and then the real one on the same module, and a view held across + the swap dangles silently. The contract checks re-run on every bind.""" + cm = md.kv_cache_manager + assert cm is not None, "a paged KV cache is required" + assert not (self._binds and torch.cuda.is_current_stream_capturing()), ( + f"the KV cache manager changed to {type(cm).__name__} during CUDA graph capture" + ) + self._binds += 1 + missing = [name for name in _QSA_FIELDS if not hasattr(md, name)] + assert not missing, ( + f"metadata fields missing: {missing}; the engine did not select QSA metadata" + ) + assert position_ids.dtype == torch.int32, position_ids.dtype + assert self.mrope_declared, "this target takes the mRoPE position layout" + assert position_ids.dim() == 3 and position_ids.shape[0] == _QSA_POSITION_AXES, ( + f"expected the [{_QSA_POSITION_AXES}, 1, T] mRoPE position buffer, got " + f"{tuple(position_ids.shape)}" + ) + pools = {row[0] for row in md.host_kv_cache_pool_mapping.tolist()} + assert pools <= {0, 1}, ( + f"KV addressing over {len(pools)} pools is not certified; layer->pool ids " + f"{sorted(pools)}" + ) + assert self.quant_mode != QM_NVFP4_KV_CACHE or md.use_paged_context_fmha, ( + "an NVFP4 KV cache is read by the paged-context attention kernels only, which the " + "engine enables for chunked prefill, KV cache block reuse or speculative decoding" + ) + g = _pool.pool_geometry( + self.model_config.pretrained_config, + self.mapping, + tokens_per_block=md.tokens_per_block, + kv_cache_quant_algo=self.kv_cache_quant_algo, + dtype=self.w["embed"].dtype, + ) + # The manager's per-rank share is the one the weights were cut to. + assert (g.heads_kv, g.kv_head_start, g.num_k_heads, g.num_v_heads, g.conv_dim) == ( + self.heads_kv, + self.kv_head_start, + self.num_k_heads, + self.num_v_heads, + self.conv_dim, + ), ( + f"the KV cache manager's share {g} disagrees with the weights' heads " + f"(kv {self.heads_kv} from {self.kv_head_start}, GDN {self.num_k_heads} / " + f"{self.num_v_heads}, conv {self.conv_dim})" + ) + # `qsa_topk_indices` rows are addressed as `block_topk`-wide rows. + assert md.qsa_topk_indices.shape[1] == self.block_topk, ( + f"the engine sized the QSA selection buffer {md.qsa_topk_indices.shape[1]} wide " + f"and this target selects {self.block_topk} blocks per row" + ) + mm = md.mamba_metadata + assert mm is not None and mm is not False, ( + "attn_metadata.mamba_metadata is absent: the runtime allocated no recurrent-state pool" + ) + mm_missing = [name for name in _MAMBA_FIELDS if not hasattr(mm, name)] + assert not mm_missing, f"mamba metadata fields missing: {mm_missing}" + rank_pools = _pool.take_pools(cm, g, self.attn_layers, self.gdn_layers, self.ple_layer) + # No sliding window: the targets pass `max_seq_len` as the window. + windows = cm.max_attention_window_vec + attn_windows = {windows[off] for off in rank_pools.pool_idx.values()} + assert attn_windows <= {None, md.max_seq_len}, ( + f"the attention layers' pool windows are {sorted(map(repr, attn_windows))}; a " + "sliding window is not implemented here" + ) + assert md.effective_beam_width == 1, "beam search is not implemented" + assert md.cache_indirection is None, "beam search is not implemented" + assert not md.is_cross, "cross attention is not implemented" + # The pool sizes the entries' page- and slot-addressed cells have to cover, reported + # on every rank (the executor's workers keep the info log level off). + first_attn, first_gdn = self.attn_layers[0], self.gdn_layers[0] + print( + f"modeling_v2 qwen4_exp default target: rank {self.mapping.rank} binds " + f"{rank_pools.kv[first_attn].k.shape[0]} K/V pages, " + f"{rank_pools.index_cache[first_attn].shape[0]} QSA index pages, " + f"{rank_pools.conv[first_gdn].shape[0]} recurrent-state slots and " + f"{rank_pools.ple_state[0].shape[0]} PLE slots (bind {self._binds})", + flush=True, + ) + bound = _Bound(cm, rank_pools) + if self.spec_rollback is not None: + self.spec_rollback.set_index_planes( + [rank_pools.index_cache[i] for i in self.attn_layers] + ) + if self.mtp_enabled: + bound = bound._replace(**self._bind_step_buffers(cm, g, mm, rank_pools)) + self._bound = bound + + def _bind_step_buffers(self, cm, g: _pool.PoolGeometry, mm, pools: _pool.RankPools) -> dict: + """The verify step's runtime-owned buffers of one manager: every verify row's conv + window goes to a per-step buffer that `update_mamba_states` promotes from, and the + recurrent state is the cached replay (a checkpoint per slot plus a compact history of + the tokens accepted since it). Returns the `_Bound` fields it adds.""" + _pool.require( + cm, + ("is_speculative", "mamba_layer_cache", "get_replay_state_update_metadata"), + "binds the verify step's per-step state buffers and replay metadata", + ) + assert cm.is_speculative(), ( + f"the KV cache manager {type(cm).__name__} reports is_speculative() False, so it " + "allocated none of the per-step state buffers the verify path writes" + ) + layer_caches = {i: cm.mamba_layer_cache(i) for i in self.gdn_layers} + int_conv = {i: cache.intermediate_conv_window for i, cache in layer_caches.items()} + int_idx = getattr(cm, "intermediate_state_indices", None) + g0 = self.gdn_layers[0] + _pool.check_step_buffers(g, int_conv[g0], pools.conv[g0], int_idx) + meta = cm.get_replay_state_update_metadata() + if meta is None: + raise NotImplementedError( + "the KV cache manager runs the legacy per-step recurrent verify path " + "(TRTLLM_USE_GDN_REPLAY=0), which this target does not assemble; leave " + "TRTLLM_USE_GDN_REPLAY unset for the cached replay" + ) + assert not getattr(cm, "use_gdn_flashinfer_ring_replay", False), ( + "the manager laid its replay history out for the FlashInfer ring kernel; this " + "target runs the double-buffered Triton cached replay" + ) + assert layer_caches[g0].intermediate_ssm is None, ( + "the cache manager runs the cached-replay verify path and also allocated a " + "per-step recurrent buffer; the runtime allocates one or the other" + ) + ssm0 = pools.ssm[g0] + history, width = meta.replay_history_size, meta.replay_step_width + _pool.check_replay_state( + history, width, meta.prev_num_accepted_tokens, meta.cache_buf_idx, ssm0 + ) + old = {} + for i, cache in layer_caches.items(): + old[i] = (cache.old_x, cache.old_B, cache.old_dt, cache.old_dA_cumsum) + _pool.check_replay_history(g, i, history, ssm0.shape[0], *old[i][:3]) + replay = _Replay(old, meta.cache_buf_idx, meta.prev_num_accepted_tokens, history, width) + return dict(int_idx=int_idx, int_conv=int_conv, replay=replay) + + def _bind_draft_cache_manager(self, md) -> None: + """Bind the draft layer to the manager its draft step sees: its own index plane, that + manager's position plane and its paged K/V. + + Under attention TP that is the separate draft KV cache the worker swaps in for the + draft loop, whose pools hold the draft layer's pages and none of the trunk's; + otherwise the trunk's manager, which maps the draft layer as well. Keyed on the + manager like `_bind_cache_manager`, and never first bound under capture: the runner + runs every captured batch shape eagerly first.""" + cm = md.kv_cache_manager + assert not torch.cuda.is_current_stream_capturing(), ( + f"the draft layer is not bound to KV cache manager {type(cm).__name__} (first use, " + "or the manager changed since its last bind) and CUDA graph capture is active" + ) + _pool.require( + cm, + ("get_index_k_buffer", "get_qsa_position_buffer", "get_buffers", "layer_offsets"), + "reads the draft layer's QSA index plane, position plane and paged K/V", + ) + g = _pool.pool_geometry( + self.model_config.pretrained_config, + self.mapping, + tokens_per_block=md.tokens_per_block, + kv_cache_quant_algo=self.kv_cache_quant_algo, + dtype=self.w["embed"].dtype, + ) + layer = self.mtp_layer_idx + index_cache = cm.get_index_k_buffer(layer) + position_cache = cm.get_qsa_position_buffer() + _pool.check_position_plane(g, position_cache) + kv = _pool.kv_planes(cm, g, layer, index_cache) + share_slot = None + if self.mtp_index_share: + # The metadata sizes its shared selection buffers from the same resolved knob, + # one row per sparse layer of whichever manager is bound. + assert md.sparse_metadata_params.mtp_index_share, ( + "this target resolved IndexShare on but the engine's QSA metadata did not" + ) + slots = cm.qsa_shared_topk_slots + _pool.check_shared_selection(md.qsa_shared_topk_indices, slots, layer, self.block_topk) + share_slot = slots[layer] + print( + f"modeling_v2 qwen4_exp default target: rank {self.mapping.rank} binds the draft layer to " + f"{type(cm).__name__} ({'the trunk manager' if cm is self._bound.cm else 'separate'}" + f"): {kv.k.shape[0]} K/V pages, {index_cache.shape[0]} QSA index pages" + f"{'' if share_slot is None else f', IndexShare slot {share_slot}'}", + flush=True, + ) + self._draft_bound = _DraftBound(cm, index_cache, position_cache, kv, share_slot) + + def _draft_step_bind(self, md) -> tuple[_DraftBound, int]: + """The draft layer's bind for this draft step's manager, and the layer's pool offset + in it (what `thop_attention` addresses), checked on every call: the worker swaps the + manager and the page tables in for the draft loop, and a page table still decoding + the trunk's offsets would read and write another pool's pages.""" + cm = md.kv_cache_manager + assert cm is not None, "a paged KV cache is required" + offsets = cm.layer_offsets + assert self.mtp_layer_idx in offsets, ( + f"the MTP module's attention is global layer index {self.mtp_layer_idx}, which the " + f"cache manager {type(cm).__name__} does not map (it maps {len(offsets)} layers); " + "the engine allocated no KV layer for the draft module" + ) + layer_idx = offsets[self.mtp_layer_idx] + window = cm.max_attention_window_vec[layer_idx] + assert window is None or window == md.max_seq_len, ( + f"the MTP layer's pool window is {window!r} but this call passes max_seq_len " + f"{md.max_seq_len} as the attention window; a sliding window is not implemented" + ) + if self._draft_bound is None or cm is not self._draft_bound.cm: + self._bind_draft_cache_manager(md) + assert cm is self._bound.cm or md.qsa_block_table is getattr( + md, "qsa_draft_block_table", None + ), ( + "the draft step is bound to a separate draft KV cache manager but qsa_block_table " + "is not that manager's table: the engine's QSA metadata did not swap its page " + "table for the draft forward" + ) + return self._draft_bound, layer_idx + + # -------------------------------------------------------------- forward --- + + def _select_target(self, attn_metadata) -> "PrefillTarget | DecodeTarget": + """Which target runs this step: decode if `num_contexts == 0`, else prefill.""" + return self._decode if attn_metadata.num_contexts == 0 else self._prefill + + def begin_step( + self, + attn_metadata: AttentionMetadata, + input_ids: torch.IntTensor | None, + position_ids: torch.IntTensor | None, + inputs_embeds: torch.FloatTensor | None = None, + lora_params: dict | None = None, + spec_metadata=None, + vision: tuple | None = None, + ) -> None: + """Check what holds for every step, bind the cache manager, choose the target and + let the prefill target prepare the step. + + The shell calls it before `forward`, on every step: a breakable prefill CUDA graph + captures `forward` and replays it without calling it, so whatever a step reads from + the host is read here. `vision` is `(rows, mm_token_indices)` when the step carries + media.""" + assert self._prefill is not None and self._decode is not None, ( + "load_weights must run before forward" + ) + assert position_ids is not None + assert isinstance(attn_metadata, TrtllmAttentionMetadata) + assert lora_params is None, "LoRA is not implemented by this target" + # The PLE layer hashes token ids, so they are required on every step. + assert input_ids is not None, "this target's PLE layer hashes token ids" + assert inputs_embeds is None, ( + "this target computes its token embedding itself; an engine-supplied " + "inputs_embeds is not implemented" + ) + # `spec_metadata` is present exactly when drafting. One-model MTP-Eagle captures no + # trunk layer, so the trunk owes no capture hook. + assert (spec_metadata is not None) == self.mtp_enabled, ( + f"spec_metadata {'arrived' if spec_metadata is not None else 'is absent'} but the " + f"model was built with mtp_enabled={self.mtp_enabled}; the draft path and the " + "weights it loads are decided at construction" + ) + if self._bound is None or attn_metadata.kv_cache_manager is not self._bound.cm: + self._bind_cache_manager(attn_metadata, position_ids) + target = self._select_target(attn_metadata) + if target is self._prefill: + target.prepare(attn_metadata, input_ids, position_ids, vision, spec_metadata) + else: + assert vision is None, "a generation-only step carries vision rows" + self._step_begun = True + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: torch.IntTensor | None = None, + position_ids: torch.IntTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + lora_params: dict | None = None, + **kwargs, + ): + """Run the step `begin_step` prepared on the target it chose. Under a breakable + prefill CUDA graph this is the body the engine captures; nothing here may read the + step from the host. Returns the logits input, or under speculation the decode + target's `(final, hc)` pair and the prefill target's one packed tensor + (`split_spec_output`).""" + assert self._step_begun, ( + "the shell calls begin_step before every forward; this forward has no prepared step" + ) + self._step_begun = False + return self._select_target(attn_metadata).forward( + attn_metadata, input_ids=input_ids, position_ids=position_ids, **kwargs + ) + + def split_spec_output(self, output) -> tuple[torch.Tensor, torch.Tensor]: + """`(final, hc)` out of what `forward` returned under speculation: `final` is the + `[T, hidden]` mix `lm_head` reads and `hc` the `[T, hc_count * hidden]` stream the + MTP module reads. The decode target returns the pair; the prefill target returns one + `[T, hidden + stream]` tensor (a breakable prefill graph's body returns one tensor), + of which the two are column views.""" + if isinstance(output, tuple): + return output + final, hc = torch.split(output, [self.hidden, self.stream], -1) + return final, hc + + +# --------------------------------------------------------------- weight views --- + + +class _LayerWeights(NamedTuple): + """One trunk layer's weights as the views the ops read (GEMM operands are `.t()` views). + Attention fields are None on a Gated-DeltaNet layer and vice versa.""" + + ahc_norm: torch.Tensor + ahc_first: torch.Tensor + ahc_up: torch.Tensor + mhc_norm: torch.Tensor + mhc_first: torch.Tensor + mhc_up: torch.Tensor + attn_in: torch.Tensor | None + o: torch.Tensor | None + q_norm: torch.Tensor | None + k_norm: torch.Tensor | None + idx_qn: torch.Tensor | None + idx_kn: torch.Tensor | None + in_proj: torch.Tensor | None + out: torch.Tensor | None + conv_w: torch.Tensor | None + a_log: torch.Tensor | None + dt_bias: torch.Tensor | None + gnorm: torch.Tensor | None + router: torch.Tensor + moe_gscale: torch.Tensor + fc1_w: torch.Tensor + fc1_s: torch.Tensor + fc2_w: torch.Tensor + fc2_s: torch.Tensor + out1_scale: torch.Tensor + out1_gate_scale: torch.Tensor + out2_scale: torch.Tensor + # `[gate | up | gate logit | pad]` (`weights._cat_rows`), and the down projection. + sh_gug: torch.Tensor + sh_dn: torch.Tensor + # Its gate|up rows and its gate-logit row, as row views of `sh_gug`. + sh_gu: torch.Tensor + sh_gate: torch.Tensor + # The dense FP8 mode's copies (`_DENSE_FP8_ENV`), None without it: per-tensor FP8 + # `(codes_t, scale)`, packed `fp8_gemv` copies and `mid_rows_mm` views; and the GDN + # in_proj's NVFP4 copy (`_GDN_INPROJ_NVFP4_ENV`). + attn_in8: tuple | None = None + o8: tuple | None = None + in_proj8: tuple | None = None + out8: tuple | None = None + attn_in_g: object | None = None + o_g: object | None = None + in_proj_g: object | None = None + out_g: object | None = None + sh_gu_g: object | None = None + sh_dn_g: object | None = None + in_proj_nv: object | None = None + attn_in_m: object | None = None + o_m: object | None = None + in_proj_m: object | None = None + out_m: object | None = None + sh_gu_m: object | None = None + sh_dn_m: object | None = None + router_m: object | None = None + # The routed experts' CuteDSL-layout copy and its scalars (`_moe_layer_scalars_cd`), which + # prefill-sized launches run on at one rank (`moe_prefill_copy`); None otherwise. + fc1_w_cd: torch.Tensor | None = None + fc1_s_cd: torch.Tensor | None = None + fc2_w_cd: torch.Tensor | None = None + fc2_s_cd: torch.Tensor | None = None + alpha1_cd: torch.Tensor | None = None + alpha2_cd: torch.Tensor | None = None + global_sf_cd: torch.Tensor | None = None + # On the CuteDSL backend the experts' own scalars (`_moe_layer_scalars_cd`); None on + # trtllm-gen. + alpha1: torch.Tensor | None = None + alpha2: torch.Tensor | None = None + global_sf: torch.Tensor | None = None + # `1 + ahc_norm` / `1 + mhc_norm` in fp32, contiguous (`_norm_w1`): the rstd mixers' + # operands (`_HC_RSTD_MIN_ROWS`); set by `PrefillTarget` only. + ahc_norm_w1: torch.Tensor | None = None + mhc_norm_w1: torch.Tensor | None = None + + +def _layer_views(core: Qwen4ExpDefaultModelingV2) -> list[_LayerWeights]: + """Every trunk layer's `_LayerWeights`, read off the realized weights.""" + w = core.w + if getattr(core, "trunk_moe_backend", "TRTLLM") == "CUTEDSL": + w_dtype, s_dtype = torch.float4_e2m1fn_x2, torch.uint8 + else: + w_dtype, s_dtype = torch.uint8, torch.float8_e4m3fn + layers = [] + for i in range(core.num_layers): + attn: dict = dict(attn_in=None, o=None, q_norm=None, k_norm=None, idx_qn=None, idx_kn=None) + gdn: dict = dict(in_proj=None, out=None, conv_w=None, a_log=None, dt_bias=None, gnorm=None) + if core.is_attn[i]: + attn = dict( + attn_in=w[f"l{i}_attn_in_w"].t(), + o=w[f"l{i}_o_w"].t(), + q_norm=w[f"l{i}_q_norm"], + k_norm=w[f"l{i}_k_norm"], + idx_qn=w[f"l{i}_idx_qn"], + idx_kn=w[f"l{i}_idx_kn"], + ) + else: + a_log, dt_bias = core._gdn_gating[i] + gdn = dict( + in_proj=w[f"l{i}_in_proj_w"].t(), + out=w[f"l{i}_out_w"].t(), + # [conv_dim, 1, width] -> [conv_dim, width]: a view (unit width stride). + conv_w=w[f"l{i}_conv_w"].reshape(core.conv_dim, core.conv_kernel), + a_log=a_log, + dt_bias=dt_bias, + gnorm=w[f"l{i}_gnorm"], + ) + scalars = core._moe_scalars[i] + cute_dsl: dict = {} + if getattr(core, "trunk_moe_backend", "TRTLLM") == "CUTEDSL": + scalars_cd = core._moe_scalars_cd[i] + cute_dsl = dict( + alpha1=scalars_cd["alpha1"], + alpha2=scalars_cd["alpha2"], + global_sf=scalars_cd["global_sf"], + ) + elif getattr(core, "moe_prefill_copy", False): + scalars_cd = core._moe_scalars_cd[i] + cute_dsl = dict( + fc1_w_cd=w[f"l{i}_fc1_w_cd"], + fc1_s_cd=w[f"l{i}_fc1_s_cd"], + fc2_w_cd=w[f"l{i}_fc2_w_cd"], + fc2_s_cd=w[f"l{i}_fc2_s_cd"], + alpha1_cd=scalars_cd["alpha1"], + alpha2_cd=scalars_cd["alpha2"], + global_sf_cd=scalars_cd["global_sf"], + ) + layers.append( + _LayerWeights( + ahc_norm=w[f"l{i}_ahc_norm"], + ahc_first=core._hc_first[(i, "ahc")].t(), + ahc_up=w[f"l{i}_ahc_up"].t(), + mhc_norm=w[f"l{i}_mhc_norm"], + mhc_first=core._hc_first[(i, "mhc")].t(), + mhc_up=w[f"l{i}_mhc_up"].t(), + **attn, + **gdn, + router=w[f"l{i}_router"].t(), + moe_gscale=scalars["global_scale"], + # The expert stacks are uint8 storage; each backend reads the packed nibbles + # and the block scales as its own dtypes. + fc1_w=w[f"l{i}_fc1_w"].view(w_dtype), + fc1_s=w[f"l{i}_fc1_s"].view(s_dtype), + fc2_w=w[f"l{i}_fc2_w"].view(w_dtype), + fc2_s=w[f"l{i}_fc2_s"].view(s_dtype), + out1_scale=scalars["output1_scale_scalar"], + out1_gate_scale=scalars["output1_scale_gate_scalar"], + out2_scale=scalars["output2_scale_scalar"], + sh_gug=w[f"l{i}_sh_gu"].t(), + sh_dn=w[f"l{i}_sh_dn"].t(), + sh_gu=w[f"l{i}_sh_gu"][: 2 * core.shared_inter].t(), + sh_gate=w[f"l{i}_sh_gu"][2 * core.shared_inter : 2 * core.shared_inter + 1].t(), + **{ + field: getattr(core, copies, {}).get(f"l{i}_{name}") + for field, copies, name in ( + ("attn_in8", "_dense8", "attn_in_w"), + ("o8", "_dense8", "o_w"), + ("in_proj8", "_dense8", "in_proj_w"), + ("out8", "_dense8", "out_w"), + ("attn_in_g", "_gemv8", "attn_in_w"), + ("o_g", "_gemv8", "o_w"), + ("in_proj_g", "_gemv8", "in_proj_w"), + ("out_g", "_gemv8", "out_w"), + ("sh_gu_g", "_gemv8", "sh_gu"), + ("sh_dn_g", "_gemv8", "sh_dn"), + ("in_proj_nv", "_gemv_nv", "in_proj_w"), + ("attn_in_m", "_mid", "attn_in_w"), + ("o_m", "_mid", "o_w"), + ("in_proj_m", "_mid", "in_proj_w"), + ("out_m", "_mid", "out_w"), + ("sh_gu_m", "_mid", "sh_gu"), + ("sh_dn_m", "_mid", "sh_dn"), + ("router_m", "_mid", "router"), + ) + }, + **cute_dsl, + ) + ) + return layers + + +class _StackWeights(NamedTuple): + """The weights outside the layer loop: the stack-final mixer (this checkpoint's final + norm) and the PLE side path.""" + + fhc_norm: torch.Tensor + fhc_first: torch.Tensor + fhc_up: torch.Tensor + ple_key: torch.Tensor + ple_val: torch.Tensor + ple_nk_w: torch.Tensor + ple_nq_w: torch.Tensor + ple_nc_w: torch.Tensor + ple_conv_w: torch.Tensor + ple_mult: torch.Tensor + ple_vocab: torch.Tensor + ple_offs: torch.Tensor + ple_tscale: torch.Tensor + # `1 + fhc_norm` in fp32 (see `_LayerWeights.ahc_norm_w1`); set by `PrefillTarget` only. + fhc_norm_w1: torch.Tensor | None = None + + +def _norm_w1(norm: torch.Tensor) -> torch.Tensor: + """A grouped norm's `(1 + w)` in fp32, contiguous: what `qwen4_exp_hc_up_gate_mean_rstd` + multiplies the scaled residual by (the `1 + w` `qwen4_exp_hc_inject_norm` forms per + element, bitwise).""" + return (1.0 + norm.float()).contiguous() + + +def _stack_views(core: Qwen4ExpDefaultModelingV2) -> _StackWeights: + w = core.w + return _StackWeights( + fhc_norm=w["fhc_norm"], + fhc_first=core._hc_first["fhc"].t(), + fhc_up=w["fhc_up"].t(), + ple_key=w["ple_key_w"].t(), + ple_val=w["ple_val_w"].t(), + # The three PLE norm weights whole: the fused tail normalizes each stream inside one + # program. + ple_nk_w=w["ple_nk"], + ple_nq_w=w["ple_nq"], + ple_nc_w=w["ple_nc"], + ple_conv_w=w["ple_conv_w"], + ple_mult=w["ple_mult"], + ple_vocab=w["ple_vocab"], + ple_offs=w["ple_offs"], + # The table's dequantize scale at the gathered rows' rank (3), all 1s. + ple_tscale=w["ple_tscale"].reshape(1, 1, 1), + ) + + +class _MtpWeights(NamedTuple): + """The MTP module's weights at this rank's share (keyed at its layer index, + `num_hidden_layers`), as the views the ops read: GEMM operands are `.t()` views. Each + hyper-connection site is the trunk's: its whole `[stream]` grouped-norm weight, its + `[down; block_inject; 0]` first projection (`_fuse_hc_first_projections`) and its up + projection; `unit_inject` is the front end's per-stream `x_h + x_e` as the first site's + injection at unit weight.""" + + pre_e_norm: torch.Tensor + pre_h_norm: torch.Tensor + fc_e: torch.Tensor + fc_h: torch.Tensor + ahc_norm: torch.Tensor + ahc_first: torch.Tensor + ahc_up: torch.Tensor + mhc_norm: torch.Tensor + mhc_first: torch.Tensor + mhc_up: torch.Tensor + unit_inject: torch.Tensor + q: torch.Tensor # [hidden, 2 * q_width]: per head its q, then its output gate + k: torch.Tensor + v: torch.Tensor + o: torch.Tensor + q_norm: torch.Tensor + k_norm: torch.Tensor + idx: torch.Tensor + idx_qn: torch.Tensor + idx_kn: torch.Tensor + router: torch.Tensor + fc1_w: torch.Tensor + fc1_s: torch.Tensor + fc2_w: torch.Tensor + fc2_s: torch.Tensor + # `[gate | up | gate logit | pad]` (`weights._cat_rows`), as the trunk's, and the down + # projection. + sh_gug: torch.Tensor + sh_dn: torch.Tensor + fhc_norm: torch.Tensor + fhc_first: torch.Tensor + fhc_up: torch.Tensor + # The shared expert's gate|up rows and gate-logit row, row views of `sh_gug`. + sh_gu: torch.Tensor + sh_gate: torch.Tensor + # With `_MTP_MOE_NVFP4_ENV` (one rank), `fc1_w` .. `fc2_s` are the NVFP4 copy in the trunk's + # layout and these its activation global scale and the runner's three scalars. + moe_gscale: torch.Tensor | None = None + out1_scale: torch.Tensor | None = None + out1_gate_scale: torch.Tensor | None = None + out2_scale: torch.Tensor | None = None + + +def _mtp_views(core: Qwen4ExpDefaultModelingV2) -> _MtpWeights: + """Built after `_fuse_hc_first_projections`, whose `"mtp_*"` tensors it reads.""" + w = core.w + p = f"l{core.mtp_layer_idx}_mtp_" + dt = w[f"{p}fc_h"].dtype + return _MtpWeights( + pre_e_norm=w[f"{p}pre_e_norm"], + pre_h_norm=w[f"{p}pre_h_norm"], + fc_e=w[f"{p}fc_e"].t(), + fc_h=w[f"{p}fc_h"].t(), + ahc_norm=w[f"{p}ahc_norm"], + ahc_first=core._hc_first["mtp_ahc"].t(), + ahc_up=w[f"{p}ahc_up"].t(), + mhc_norm=w[f"{p}mhc_norm"], + mhc_first=core._hc_first["mtp_mhc"].t(), + mhc_up=w[f"{p}mhc_up"].t(), + unit_inject=torch.ones( + [core.max_num_tokens, core.hc_count], dtype=dt, device=w[f"{p}fc_h"].device + ), + q=w[f"{p}q_w"].t(), + k=w[f"{p}k_w"].t(), + v=w[f"{p}v_w"].t(), + o=w[f"{p}o_w"].t(), + q_norm=w[f"{p}q_norm"], + k_norm=w[f"{p}k_norm"], + idx=w[f"{p}idx_w"].t(), + idx_qn=w[f"{p}idx_qn"], + idx_kn=w[f"{p}idx_kn"], + router=w[f"{p}router"].t(), + fc1_w=w[f"{p}fc1_w"], + fc1_s=w[f"{p}fc1_s"], + fc2_w=w[f"{p}fc2_w"], + fc2_s=w[f"{p}fc2_s"], + sh_gug=w[f"{p}sh_gu"].t(), + sh_dn=w[f"{p}sh_dn"].t(), + sh_gu=w[f"{p}sh_gu"][: 2 * core.shared_inter].t(), + sh_gate=w[f"{p}sh_gu"][2 * core.shared_inter : 2 * core.shared_inter + 1].t(), + fhc_norm=w[f"{p}fhc_norm"], + fhc_first=core._hc_first["mtp_fhc"].t(), + fhc_up=w[f"{p}fhc_up"].t(), + ) + + +# ------------------------------------------------------------ shared blocks --- + + +def _verify_rows(core, md, spec_metadata, num_seqs: int, num_rows: int) -> tuple[int, bool]: + """`(p, is_verify)` of a step's generation part: `num_rows` rows over `num_seqs` + generation sequences, `p` per sequence -- 1 without speculation, else + `runtime_draft_len + 1`. Host ints of the metadata, constants of a captured step.""" + p = num_rows // num_seqs if num_seqs else 1 + assert p * num_seqs == num_rows, ( + f"{num_rows} generation rows do not divide over {num_seqs} generation sequences; a " + "ragged per-sequence draft length can be expressed neither to the attention call nor " + "to the multi-token recurrence" + ) + assert core.mtp_enabled or p == 1, ( + f"{num_rows} generation rows over {num_seqs} generation sequences with speculation " + "off; the base path implements one query token per generation sequence" + ) + if spec_metadata is not None and num_seqs > 0: + declared = spec_metadata.runtime_draft_len + 1 + assert p == declared, ( + f"the batch carries {p} query rows per generation sequence but the worker's " + f"runtime_draft_len says {declared}; the per-step conv snapshots and the replay " + "history are sized from the second and written from the first" + ) + _pool.check_verify_width(p) + _pool.check_replay_width(core._bound.replay, p) + return p, spec_metadata is not None and num_seqs > 0 and p > 1 + + +def _verify_state(core, mm, num_decodes: int) -> tuple: + """What a verify step's `_gdn_verify` reads besides the layer: the rows + `update_mamba_states` promotes the conv windows from, and the replay's work items. Work + items go to the kernel from `CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE` requests on, and the + kernel leaves the history commit to the manager's all-layer commit exactly when the + manager runs one after the step -- the stock Gated-DeltaNet layer's two conditions.""" + assert core._bound.int_idx is not None, "the cache-manager bind has not run" + int_idx = core._bound.int_idx[:num_decodes] + partitioned = num_decodes >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE + all_layer = partitioned and bool( + getattr(core._bound.cm, "use_gdn_cached_replay_all_layer_commit", False) + ) + work_items = mm.replay_work_items[:num_decodes] if partitioned else None + n_writes = mm.replay_n_writes if partitioned else None + return int_idx, work_items, n_writes, all_layer + + +def _gdn_verify(core, mq, a, b, i, L, si_d, n, p, verify, mq_out=None) -> torch.Tensor: + """The recurrent half of a verify step: `n` generation sequences of `p` rows, written so + that the accepted prefix, unknown until the sampler has run, can be kept. + + The conv writes each row's window to the runtime's per-step buffer + (`update_mamba_states` promotes the accepted one) besides advancing the pool; the cached + replay rebuilds each slot's state from its checkpoint and pending history, runs the `p` + rows and appends them to the history, so rejected rows cost nothing to undo. `mq` is the + `[n * p, conv_dim]` q|k|v column view of the input projection, `a` and `b` its raw gate + columns (column views too): the kernel applies `A_log`, `dt_bias` and the q/k L2 norm + itself. Over a bf16 recurrent pool the replay is the CuTe kernel (`gdn_cached_replay_cute`), + which reads `a` / `b` through their row stride; over an fp32 one the Triton kernel, which is + handed contiguous copies. Both take the same work items and leave the same commits to the + manager. Returns `[1, n * p, num_v_heads, head_v_dim]`. + + `mq_out` is the conv output when the input projection's GEMV ran the conv step as its + epilogue (`fp8_gemv_conv1d_verify`, `nvfp4_split_gemv_conv1d_verify`; `mq` is then None).""" + int_idx, work_items, n_writes, all_layer = verify + pools = core._bound.pools + replay = core._bound.replay + ssm = pools.ssm[i] + cute = ssm.dtype == torch.bfloat16 + pdl = _GDN_REPLAY_PDL and cute + if mq_out is None: + # [n*p, C] -> [n, C, p]: the multi-token conv reads channel-major. + x3 = torch.transpose(torch.reshape(mq, [n, p, core.conv_dim]), 1, 2) + conv_out = causal_conv1d_triton_update( + x=x3, + conv_state=pools.conv[i], + weight=L.conv_w, + bias=None, + activation="silu", + conv_state_indices=si_d, + pad_slot_id=_PAD_SLOT_ID, + # Snapshots, not the accepted-token roll. + num_accepted_tokens=None, + intermediate_conv_window=core._bound.int_conv[i], + intermediate_state_indices=int_idx, + launch_dependent_kernels=pdl, + ) + # Back to token-major: one contiguous copy, which the kernel reads q/k/v from directly. + mq_out = torch.reshape(torch.transpose(conv_out, 1, 2), [n * p, core.conv_dim]) + old_x, old_B, old_dt, old_dA = replay.old[i] + qv, kv, vv = torch.split(mq_out, [core.key_dim, core.key_dim, core.value_dim], 1) + if cute: + replay_op = gdn_cached_replay_cute + else: + replay_op = fused_recurrent_gated_delta_rule_cached_replay_update + a, b = a.contiguous(), b.contiguous() + out = replay_op( + q=torch.reshape(qv, [n, p, core.num_k_heads, core.head_k_dim]), + k=torch.reshape(kv, [n, p, core.num_k_heads, core.head_k_dim]), + v=torch.reshape(vv, [n, p, core.num_v_heads, core.head_v_dim]), + g=torch.reshape(a, [n, p, core.num_v_heads]), + beta=torch.reshape(b, [n, p, core.num_v_heads]), + ssm_states=ssm, + state_indices=si_d, + old_u=old_x, + old_k=old_B, + old_G=old_dt, + old_beta=old_dA, + cache_buf_idx=replay.cache_buf_idx, + prev_num_accepted_tokens=replay.prev_num_accepted, + history_size=replay.history, + scale=None, + use_qk_l2norm_in_kernel=True, + A_log=L.a_log, + dt_bias=L.dt_bias, + launch_with_pdl=pdl, + replay_work_items=work_items, + n_writes=n_writes, + block_v=None, + num_warps=None, + use_l2_state_cache=None, + use_l2_streaming_inputs=None, + use_l2_shared_inputs=None, + main_num_stages=None, + packed_qkv=mq_out, + use_all_layer_commit=all_layer, + output=None, + ) + return torch.reshape(out, [1, n * p, core.num_v_heads, core.head_v_dim]) + + +def _position_key(position_ids: torch.Tensor) -> tuple: + """What identifies a step's positions buffer: its storage, dtype and `[3, N]` layout.""" + view = torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]) + return (view.data_ptr(), view.dtype, tuple(view.shape), view.stride()) + + +def _draft_positions(position_ids: torch.Tensor) -> tuple: + """`(pos3, pos, coords)` of a draft step's mRoPE positions (`[3, 1, N]` from the trunk, + `[3, N]` from the worker's later steps): a fresh contiguous int32 `[3, N]` copy, its + coordinate-0 row (the index rope's) and the `[N, 3]` transpose the position plane + stores.""" + assert position_ids.shape[0] == _QSA_POSITION_AXES, ( + f"expected the mRoPE position layout with {_QSA_POSITION_AXES} coordinate rows, got " + f"{tuple(position_ids.shape)}" + ) + pos3 = _copy_positions(position_ids) + return pos3, pos3[0], torch.transpose(pos3, 0, 1) + + +def _draft_embedding(input_ids: torch.Tensor, embed_tokens: torch.Tensor) -> torch.Tensor: + """The MTP front end's token embedding. A draft loop's first call embeds every slot of + each generation request's accepted tokens, and under rejection sampling + `chain_speculative_sampling` fills the slots past a request's accepted count with -1 + (the in-tree `Embedding` masks such ids for the same reason). Those rows feed nothing the + step keeps: attention is causal, the worker continues from each request's last accepted + row, and the next step's first draft call rewrites their cache positions. So a negative + id reads row 0 instead of past the table; every valid id embeds as before.""" + return nn.functional.embedding(input_ids.clamp_min(0), embed_tokens) + + +def _mtp_front(core, M: _MtpWeights, embed_tokens, input_ids, hidden_states) -> tuple: + """The draft layer's front end and its attention-side hyper-connection. `fc_embedding` + over the next token's normalized embedding is added to every stream of `fc_hidden` + (one `[hidden, hidden]` matrix applied to each normalized stream of the trunk's + residual): the trunk's site with that sum as its injection at unit weight (bitwise the + bf16 add), then the trunk's mixer. Returns `(x, mixed, inject)`: the layer's residual, + the mixed attention input and its injection weights.""" + rows = hidden_states.shape[0] + assert rows <= M.unit_inject.shape[0], ( + f"a draft step of {rows} rows exceeds max_num_tokens {M.unit_inject.shape[0]}" + ) + e = _draft_embedding(input_ids, embed_tokens) + x_e = mtp_cublas_mm(flashinfer_gemma_rmsnorm(e, M.pre_e_norm, core.eps), M.fc_e) + _, h_n = qwen4_exp_hc_inject_norm( + hidden_states, None, None, M.pre_h_norm, core.eps, core.hc_count + ) + x_h = torch.reshape( + mtp_cublas_mm(torch.reshape(h_n, [rows * core.hc_count, core.hidden]), M.fc_h), + [rows, core.stream], + ) + x, normed = qwen4_exp_hc_inject_norm( + x_h, x_e, M.unit_inject[:rows], M.ahc_norm, core.eps, core.hc_count + ) + mixed, inject = _hc_mixer(core, normed, M.ahc_first, M.ahc_up, True) + return x, mixed, inject + + +def _mtp_qkv(core, M: _MtpWeights, mixed, pos3) -> tuple: + """The draft layer's q (with its interleaved output gate), k and v projections at this + rank's heads, packed into one compact `[T, qkv_width]` row and normalized and rotated in + place at the module's rope base. Returns `(qkv, gate, k, v)`: `gate` is the + `[T, heads_q, head_dim]` gate half of the q projection, `k` and v` views of `qkv`.""" + rows = mixed.shape[0] + qg = mtp_cublas_mm(mixed, M.q) + qkv = torch.empty([rows, core.qkv_width], dtype=mixed.dtype, device=mixed.device) + q_dst, k_dst, v_dst = torch.split( + torch.reshape(qkv, [rows, core.heads_q + 2 * core.heads_kv, core.head_dim]), + [core.heads_q, core.heads_kv, core.heads_kv], + 1, + ) + q_src, gate = torch.split( + torch.reshape(qg, [rows, core.heads_q, 2 * core.head_dim]), + [core.head_dim, core.head_dim], + 2, + ) + q_dst.copy_(q_src) + k_dst.copy_(torch.reshape(mtp_cublas_mm(mixed, M.k), [rows, core.heads_kv, core.head_dim])) + v_dst.copy_(torch.reshape(mtp_cublas_mm(mixed, M.v), [rows, core.heads_kv, core.head_dim])) + mtp_fused_qk_norm_rope( + qkv, + num_heads_q=core.heads_q, + num_heads_k=core.heads_kv, + num_heads_v=core.heads_kv, + head_dim=core.head_dim, + rotary_dim=core.rotary_dim, + eps=core.eps, + q_weight=M.q_norm, + k_weight=M.k_norm, + base=core.mtp_theta, + is_neox=True, + position_ids=pos3, + is_qk_norm=True, + use_gemma=True, + use_mrope=True, + mrope_section1=core.mrope_section1, + mrope_section2=core.mrope_section2, + ) + return qkv, gate, k_dst, v_dst + + +def _mtp_index_branch(core, M: _MtpWeights, mixed, draft: _DraftBound, qsa) -> torch.Tensor: + """The draft layer's QSA index branch over its own planes, as a trunk attention layer + runs it: the raw key stored before anything normalizes it, the query's norm and rotation + (coordinate 0 alone, at the module's rope base), then the compression of every group this + step completes. Returns the normalized, rotated index query.""" + req_idx, logical, block_table, tokens_per_block, pos, coords = qsa + rows = mixed.shape[0] + qk_index = mtp_cublas_mm(mixed, M.idx) + token_k = torch.split( + torch.reshape(qk_index, [rows, core.index_heads + core.index_kv_heads, -1]), + [core.index_heads, core.index_kv_heads], + 1, + )[1] + qsa_paged_index_store( + token_k, + coords, + draft.index_cache, + draft.position_cache, + req_idx, + logical, + block_table, + tokens_per_block, + ) + index_scratch = torch.empty( + [rows, core.index_scratch_width], dtype=mixed.dtype, device=mixed.device + ) + index_scratch[:, : core.index_width].copy_(qk_index) + fused_qk_norm_rope( + index_scratch, + num_heads_q=core.index_heads, + num_heads_k=core.index_kv_heads, + num_heads_v=core.index_kv_heads, + head_dim=core.index_head_dim, + rotary_dim=core.rotary_dim, + eps=core.eps, + q_weight=M.idx_qn, + k_weight=M.idx_kn, + base=core.mtp_theta, + is_neox=True, + position_ids=pos, + is_qk_norm=True, + use_gemma=True, + use_mrope=False, + mrope_section1=0, + mrope_section2=0, + ) + qsa_prefill_compress( + logical_positions=logical, + request_indices=req_idx, + block_table=block_table, + index_cache=draft.index_cache, + position_cache=draft.position_cache, + k_norm_weight=M.idx_kn, + cos_sin=core._mtp_cos_sin, + eps=core.eps, + tokens_per_block=tokens_per_block, + compress_ratio=core.compress_ratio, + mrope_section=_QSA_MROPE_SECTION, + ) + return _index_query(core, index_scratch) + + +def _mtp_moe_block(core, M: _MtpWeights, x, pad_rows: int = 0, routed_only=False) -> torch.Tensor: + """The draft layer's MoE: renormalized top-k over the whole routing space, the FP8 + block-scale routed experts over this rank's window (1x128 activation scales, pre-routed), + and the replicated bf16 shared expert in the trunk's form -- one GEMM forms gate|up and + the per-token gate logit -- combined as `routed + sigmoid(gate logit) * shared`. The + activation scales come back transposed `[hidden / 128, tokens]` and go straight to the + runner. + + The window's partial is completed as the trunk's is: under attention TP summed over the + TP group; under attention DP by the round trip over the group's gathered rows + (`core.expert_sum`), `pad_rows` padding this rank's (`_draft_pad_rows`). The routing, + the quantization and the runner are row-wise and run on the rows the window sees, in + chunks of at most `_MTP_MOE_MAX_ROWS`: one chunk at a TP rank's rows, several over the + rows an attention-DP group gathers.""" + + def run(rows: torch.Tensor) -> torch.Tensor: + parts = [] + for lo in range(0, rows.shape[0], _MTP_MOE_MAX_ROWS): + chunk = rows[lo : lo + _MTP_MOE_MAX_ROWS] + logits = cublas_mm(chunk, M.router) + topk_ids, topk_w = renorm_moe_routing_op( + logits, topk=core.topk, output_dtype=chunk.dtype + ) + xq, xsf = fp8_quantize_1x128(chunk, use_ue8m0=_MTP_MOE_USE_UE8M0) + parts.append( + fp8_block_scale_moe_runner( + routing_logits=None, + routing_bias=None, + hidden_states=xq, + hidden_states_scale=xsf, + gemm1_weights=M.fc1_w, + gemm1_weights_scale=M.fc1_s, + gemm2_weights=M.fc2_w, + gemm2_weights_scale=M.fc2_s, + num_experts=core.num_experts, + top_k=core.topk, + num_fused_shared_experts=None, + n_group=None, + topk_group=None, + intermediate_size=core.moe_inter, + local_expert_offset=core.mtp_expert_offset, + local_num_experts=core.mtp_local_experts, + routed_scaling_factor=None, + routing_method_type=_ROUTING_METHOD_RENORM, + topk_weights=topk_w, + topk_ids=topk_ids, + act_type=_ACT_TYPE_SWIGLU, + gemm1_clamp_limit=None, + ) + ) + return parts[0] if len(parts) == 1 else torch.cat(parts, 0) + + if M.moe_gscale is not None: + # At one rank under `_MTP_MOE_NVFP4_ENV` the experts are their NVFP4 copy. + routed = _mtp_routed_nvfp4(core, M, x) + elif core.dp_size == 1: + routed = core.tp_sum(run(x)) + else: + routed = core.expert_sum(run, x, pad_rows) + if routed_only: + return routed + gug = cublas_mm(x, M.sh_gug) + shared = cublas_mm(silu_and_mul_rows(gug, core.shared_inter), M.sh_dn) + gate_logits = gug[:, 2 * core.shared_inter : 2 * core.shared_inter + 1] + return fused_sigmoid_gate_mul_add(routed, gate_logits, shared) + + +def _mtp_routed_nvfp4(core, M: _MtpWeights, x) -> torch.Tensor: + """The draft layer's routed experts on their NVFP4 copy (`_mtp_moe_nvfp4`, one rank): + the activation quantized at the static FC1 input global scale, then the trunk's + logits-routed trtllm-gen runner over every expert, finalized.""" + xq, xsf = fp4_quantize(x.contiguous(), M.moe_gscale, _SF_VEC, False, False) + return fp4_block_scale_moe_renorm_routed( + cublas_mm(x, M.router), + xq, + torch.reshape(xsf, [-1]).view(torch.float8_e4m3fn), + M.fc1_w, + M.fc1_s, + M.fc2_w, + M.fc2_s, + M.out1_scale, + M.out1_gate_scale, + M.out2_scale, + core.num_experts, + core.topk, + core.moe_inter, + ) + + +def _mtp_moe_nvfp4(core) -> dict: + """The draft layer's routed experts requantized to NVFP4 (`_MTP_MOE_NVFP4_ENV`), as the + `_MtpWeights` fields `_mtp_routed_nvfp4` reads; frees the FP8 copies. The runner's scalars + follow `_moe_layer_scalars` with `1 / g_in` and `1 / g_mid` standing for the checkpoint's + `input_scale`s and `1 / g_w` for its `weight_scale_2`: `out1_gate = 1 / (g_in * g_w1[e])`, + `out1 = out1_gate * g_mid`, `out2 = 1 / (g_mid * g_w2[e])`. The activation scales are static + (`_MTP_MOE_NVFP4_FC1_MARGIN` / `_FC2_MARGIN`).""" + w = core.w + p = f"l{core.mtp_layer_idx}_mtp_" + keys = [f"{p}{k}" for k in ("fc1_w", "fc1_s", "fc2_w", "fc2_s")] + fc1_w, fc1_s, fc2_w, fc2_s, g_w1, g_w2 = _weights.mtp_experts_to_nvfp4(*(w[k] for k in keys)) + for k in keys: + w[k].data = w[k].data.new_empty(0) + torch.cuda.empty_cache() + trunk = range(core.num_layers) + to_global = _weights.NVFP4_AMAX_TO_GLOBAL + in_amax = max(float(w[f"l{i}_e_isc1"].max()) for i in trunk) * to_global + mid_amax = max(float(w[f"l{i}_e_isc2"].max()) for i in trunk) * to_global + inv_g_in = _MTP_MOE_NVFP4_FC1_MARGIN * in_amax / to_global + inv_g_mid = _MTP_MOE_NVFP4_FC2_MARGIN * mid_amax / to_global + gate1 = (inv_g_in / g_w1).contiguous() + return dict( + fc1_w=fc1_w, + fc1_s=fc1_s.view(torch.float8_e4m3fn), + fc2_w=fc2_w, + fc2_s=fc2_s.view(torch.float8_e4m3fn), + moe_gscale=torch.full((1,), 1.0 / inv_g_in, dtype=torch.float32, device=fc1_w.device), + out1_scale=(gate1 / inv_g_mid).contiguous(), + out1_gate_scale=gate1, + out2_scale=(inv_g_mid / g_w2).contiguous(), + ) + + +def _mtp_tail(core, M: _MtpWeights, x, attn, inject, pad_rows: int) -> torch.Tensor: + """The draft layer after its attention core: the gated output projection's partial summed + over the TP group and injected, the MoE-side hyper-connection, the MoE and its injection. + `attn` is the `[T, q_width]` gated attention output; `pad_rows` is the MoE round trip's + padding under attention DP (`_draft_pad_rows`). The MoE-side site is the trunk's + (`qwen4_exp_hc_inject_norm`, `_hc_mixer`).""" + x, normed = qwen4_exp_hc_inject_norm( + x, core.tp_sum(cublas_mm(attn, M.o)), inject, M.mhc_norm, core.eps, core.hc_count + ) + mixed, inject = _hc_mixer(core, normed, M.mhc_first, M.mhc_up, True) + return qwen4_exp_post_inject(x, _mtp_moe_block(core, M, mixed, pad_rows), inject) + + +def _decode_projection(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor: + """A generation-only step's dense bf16 projection at every layout but tp1 (the attention + input and output projections and the GDN output projection of `DecodeTarget`, which the + engine captures in CUDA graphs; tp1 takes `_one_rank_projection`). + + Up to `_GEMV_MAX_ROWS` rows of the o / GDN out projection at a tp4/tep4 rank (K=1536) run + `bf16_gemv_pdl`: its weight streams into registers before it waits on the producer. + Everything else runs `low_m_bf16_mm`, whose table picks FlashInfer's low-M kernels at the + rows and shapes it names (a tp4/tep4 rank's, and the whole layer's that a dep4 / dep2 rank + holds) and cuBLAS elsewhere. Both launch with PDL per `_DECODE_PROJ_PDL` and read the + activation only after `griddepcontrol.wait`.""" + if mat_a.shape[0] <= _GEMV_MAX_ROWS and mat_b.shape[0] == _GEMV_TP4_K: + return bf16_gemv_pdl(mat_a, mat_b, _DECODE_PROJ_PDL) + return low_m_bf16_mm(mat_a, mat_b, _DECODE_PROJ_PDL) + + +def _one_rank_projection(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor: + """`_decode_projection` at tp1, where every decode projection's K is one `bf16_gemv_pdl` + tiles: one row on that GEMV, as up to `_GEMV_MAX_ROWS` rows of a projection at most + `_GEMV_NARROW_N` wide (the shared expert's gate logit, which `low_m_bf16_mm` would run on one + CTA); every other call on `low_m_bf16_mm`. Both launch with PDL.""" + rows = mat_a.shape[0] + if mat_b.shape[0] in _GEMV_TILING and ( + rows == 1 or (rows <= _GEMV_MAX_ROWS and mat_b.shape[1] <= _GEMV_NARROW_N) + ): + return bf16_gemv_pdl(mat_a, mat_b, _DECODE_PROJ_PDL) + return low_m_bf16_mm(mat_a, mat_b, _DECODE_PROJ_PDL) + + +class _DraftRowCounts(NamedTuple): + """The group's per-rank row counts of one draft step, in the field `_dp_pad_rows` reads.""" + + all_rank_num_tokens: list + + +def _draft_pad_rows(core, all_rank_num_tokens, rows: int) -> int: + """Rows a draft step pads its MoE input with before the attention-DP round trip: + `_dp_pad_rows` over the counts the spec worker hands this draft step -- every rank's + trunk tokens on the loop's first step, its request count on the later ones (attention-DP + dummy requests included) -- which it passes alongside the metadata rather than writing + into it. Zero without attention DP, and under a decode CUDA graph, which pads every rank + to the captured batch.""" + if core.dp_size == 1: + return 0 + assert all_rank_num_tokens is not None, ( + "a draft step under attention DP came without the group's row counts; the MoE round " + "trip cannot pad this rank's rows to the group's largest count" + ) + return core._dp_pad_rows(_DraftRowCounts(list(all_rank_num_tokens)), rows) + + +def _hc_first(core, normed, first, with_inject): + """A hyper-connection mixer's first half: the first projection over + `[down; block_inject; 0]` (fp32 logits), then the SiLU and injection activations. + Returns `(low, inject)`; `inject` is None without `with_inject`. Up to + `_HC_TC_FIRST_MAX_ROWS` rows (a per-capture constant) it is one tcgen05 launch; larger + steps run cuBLAS's fp32 logits and the activation kernel.""" + if normed.shape[0] <= _HC_TC_FIRST_MAX_ROWS: + return qwen4_exp_hc_decode_tc_first( + normed, first, core.hc_lowrank, core.hc_count, with_inject + ) + return qwen4_exp_hc_lowrank_act( + cublas_mm(normed, first[:, : core.hc_first_cols], None, torch.float32), + core.hc_lowrank, + core.hc_count, + with_inject, + normed.dtype, + ) + + +def _hc_mixer(core, normed, first, up, with_inject): + """One hyper-connection mixer from the site's normed stream: `_hc_first`, then the up + projection with the sigmoid gate and the stream mean in its epilogue. Returns + `(mixed, inject)`; `inject` is None without `with_inject` (the stack-final mixer).""" + low, inject = _hc_first(core, normed, first, with_inject) + return qwen4_exp_hc_up_gate_mean(low, up, normed, core.hc_count), inject + + +def _hc_moe_mixer(core, normed, L: _LayerWeights): + """The MoE-site mixer (`_hc_mixer` with injection). Under attention TP its up + projection's epilogue also writes the routed experts' NVFP4 input -- linear block + scales at the layer's FC1 input global scale -- which the expert window reads as is. + Under attention DP the window runs on the group's gathered bf16 rows and quantizes + them itself (`_routed_experts`). Returns `(mixed, inject, fp4)`, `fp4` the + `(xq, xsf)` pair or None.""" + if core.dp_size > 1: + mixed, inject = _hc_mixer(core, normed, L.mhc_first, L.mhc_up, True) + return mixed, inject, None + low, inject = _hc_first(core, normed, L.mhc_first, True) + mixed, xq, xsf = qwen4_exp_hc_up_gate_mean_fp4( + low, L.mhc_up, normed, core.hc_count, L.moe_gscale + ) + return mixed, inject, (xq, xsf) + + +def _hc_mixer_rstd(core, hc, rstd, norm, norm_w1, first, up, with_inject): + """`_hc_mixer` from the site's residual `hc` and its `[T, C]` rstd + (`qwen4_exp_hc_inject_rstd`) instead of a stored normed stream: the first projection forms + `bf16(hc * rstd * (1 + norm))` on its A tiles, the up projection's epilogue + `hc * rstd * norm_w1` in fp32. Returns `(mixed, inject)`.""" + low, inject = qwen4_exp_hc_first_rescale_tc( + hc, rstd, norm, first, core.hc_count, core.hc_lowrank, with_inject + ) + return qwen4_exp_hc_up_gate_mean_rstd(low, up, hc, rstd, norm_w1, core.hc_count), inject + + +def _hc_moe_mixer_rstd(core, hc, rstd, L: _LayerWeights): + """`_hc_moe_mixer` on the rstd form (see `_hc_mixer_rstd`): under attention TP the up + projection's epilogue also writes the routed experts' NVFP4 input; under attention DP the + window quantizes the gathered rows itself. Returns `(mixed, inject, fp4)`.""" + if core.dp_size > 1: + mixed, inject = _hc_mixer_rstd( + core, hc, rstd, L.mhc_norm, L.mhc_norm_w1, L.mhc_first, L.mhc_up, True + ) + return mixed, inject, None + low, inject = qwen4_exp_hc_first_rescale_tc( + hc, rstd, L.mhc_norm, L.mhc_first, core.hc_count, core.hc_lowrank, True + ) + mixed, xq, xsf = qwen4_exp_hc_up_gate_mean_rstd_fp4( + low, L.mhc_up, hc, rstd, L.mhc_norm_w1, core.hc_count, L.moe_gscale + ) + return mixed, inject, (xq, xsf) + + +def _decode_router_mid(core, L: _LayerWeights): + """A decode step's router projection for layer `L`: `mid_rows_mm` over its view in the + dense FP8 mode's rows (`_MID_MM_ROWS["router"]`), else at tp1 the step's projection entry + (`core.decode_mm`) and at the other layouts `_decode_router`.""" + fallback = core.decode_mm if core.mapping.world_size == 1 else _decode_router + if L.router_m is None: + return fallback + return lambda x, router: _dense_mm(x, router, None, fallback, "router", mid=L.router_m) + + +def _decode_router(x, router): + """A decode step's router projection: up to `_ROUTER_LOW_M_MAX_ROWS` rows FlashInfer's + low-M kernels with PDL (`low_m_bf16_mm`), more rows cuBLAS. Decode only: those + kernels compile on the first call per row count, which the engine's warmup forward at + each captured batch makes before capture; a prefill step's row count is any, so prefill + keeps `cublas_mm`. Both allocate the logits fresh and dense, as the runner's routing + reads them from `data_ptr()`.""" + if x.shape[0] <= _ROUTER_LOW_M_MAX_ROWS: + return low_m_bf16_mm(x, router, _DECODE_PROJ_PDL) + return cublas_mm(x, router) + + +def _routed_experts(core, x, L: _LayerWeights, fp4=None, router_mm=cublas_mm) -> torch.Tensor: + """This rank's partial of the routed experts over every row of `x`: the router logits + over the whole routing space, and the trtllm-gen runner over this rank's expert window + entered with them -- its own Renormalize routing (bitwise `renorm_moe_routing_op` then + the pre-routed runner) -- on the NVFP4 activation with linear block scales at the + layer's FC1 input global scale, finalized. `fp4` is that activation as the MoE-site + mixer wrote it (`_hc_moe_mixer`, the rows of `x`), or None to quantize `x` here; + `router_mm` the step's router projection (`_decode_router` at decode). A slot + routed outside the window adds nothing, so the window outputs sum to the layer's. Under + the intermediate split (tp4) the window is every expert at this rank's + `expert_inter`-wide slice, and the slices' outputs sum to the layer's the same way. + + Every step is row-wise, so `x` runs in chunks of at most the runner's certified rows (a + host decision on `T`): one chunk at a TP rank's rows, several over the rows an + attention-DP group gathers. + + At one rank with the CuteDSL-layout copy (`moe_prefill_copy`) a launch of + `_CUTE_DSL_MOE_MIN_ROWS` to `_CUTE_DSL_MOE_MAX_ROWS` rows runs the CuteDSL MoE over it + instead, in one call. On the CuteDSL backend (one rank) every launch is + `_cute_dsl_routed`, in one call.""" + rows = x.shape[0] + if core.trunk_moe_backend == "CUTEDSL": + xq, xsf = fp4_quantize(x, L.moe_gscale, _SF_VEC, False, False) if fp4 is None else fp4 + routed_out = torch.empty([rows, core.hidden], dtype=x.dtype, device=x.device) + return _cute_dsl_routed(core, L, router_mm(x, L.router), xq, xsf, routed_out) + if L.fc1_w_cd is not None and _CUTE_DSL_MOE_MIN_ROWS <= rows <= _CUTE_DSL_MOE_MAX_ROWS: + xq, xsf = fp4_quantize(x, L.moe_gscale, _SF_VEC, False, False) if fp4 is None else fp4 + return fp4_moe_renorm_cute_dsl( + router_mm(x, L.router), + xq, + xsf, + L.fc1_w_cd, + L.fc1_s_cd, + L.alpha1_cd, + L.fc2_w_cd, + L.fc2_s_cd, + L.alpha2_cd, + L.global_sf_cd, + core.num_experts, + core.topk, + memset_stream=core._moe_memset_stream, + ) + sf_cols = x.shape[1] // _SF_VEC + parts = [] + for lo in range(0, rows, _MOE_MAX_ROWS): + hi = min(rows, lo + _MOE_MAX_ROWS) + chunk = x[lo:hi] + logits = router_mm(chunk, L.router) + if fp4 is None: + xq, xsf = fp4_quantize(chunk, L.moe_gscale, _SF_VEC, False, False) + else: + xq, xsf = fp4[0][lo:hi], fp4[1][lo * sf_cols : hi * sf_cols] + parts.append( + fp4_block_scale_moe_renorm_routed( + logits, + xq, + # The runner reads the chunk's linear block scales flat. + torch.reshape(xsf, [-1]).view(torch.float8_e4m3fn), + L.fc1_w, + L.fc1_s, + L.fc2_w, + L.fc2_s, + L.out1_scale, + L.out1_gate_scale, + L.out2_scale, + core.num_experts, + core.topk, + core.expert_inter, + core.expert_offset, + core.local_experts, + ) + ) + return parts[0] if len(parts) == 1 else torch.cat(parts, 0) + + +def _cute_dsl_routed(core, L, logits, xq, xsf, routed_out) -> torch.Tensor: + """The trunk's routed experts on the CuteDSL pipeline (one rank, every expert): + Renormalize routing, the sort into per-expert row tiles, FC1 gathering its rows from the + unpermuted NVFP4 activation `xq` / `xsf` (linear scales) with SwiGLU and the requantize to + the layer's one FC2 input scale fused, FC2, and the routing-weighted unpermute in a fixed + slot order, which writes every row of `routed_out` `[T, hidden]` and returns it.""" + # fp32 weights: the dtype the unpermute scales by. + topk_ids, topk_w = renorm_moe_routing_op(logits, topk=core.topk, output_dtype=torch.float32) + # The sort pads each expert's (token, expert) pairs to whole tiles and emits the tile + # tables both GEMMs walk. + tile_expert, tile_limit, expanded_to_permuted, permuted_to_expanded, _, live_tiles = moe_sort( + token_selected_experts=topk_ids, + token_final_scales=topk_w, + num_experts=core.num_experts, + top_k=core.topk, + local_expert_offset=0, + local_num_experts=core.num_experts, + tile_tokens_dim=_CUTEDSL_TILE_SIZE, + ) + h, h_sf = cute_dsl_nvfp4_gather_grouped_gemm_act_fusion( + input=xq.view(torch.float4_e2m1fn_x2), + weight=L.fc1_w, + input_scale=xsf.view(torch.uint8).view(xq.shape[0], core.hidden // _SF_VEC), + weight_scale=L.fc1_s, + alpha=L.alpha1, + tile_idx_to_group_idx=tile_expert, + tile_idx_to_mn_limit=tile_limit, + permuted_idx_to_expanded_idx=permuted_to_expanded, + num_non_exiting_tiles=live_tiles, + global_sf=L.global_sf, + num_experts=core.num_experts, + top_k=core.topk, + num_local_experts=core.num_experts, + local_expert_offset=0, + tile_size=_CUTEDSL_TILE_SIZE, + scaling_vector_size=_SF_VEC, + activation_type=int(ActivationType.Swiglu), + ) + # FC2 writes one bf16 row per permuted row. + y = cute_dsl_nvfp4_grouped_gemm( + input=h, + weight=L.fc2_w, + input_scale=h_sf, + weight_scale=L.fc2_s, + alpha=L.alpha2, + tile_idx_to_group_idx=tile_expert, + num_non_exiting_tiles=live_tiles, + num_experts=core.num_experts, + top_k=core.topk, + num_local_experts=core.num_experts, + local_expert_offset=0, + tile_size=_CUTEDSL_TILE_SIZE, + output_dtype=torch.bfloat16, + scaling_vector_size=_SF_VEC, + ) + return moe_unpermute_inplace( + permuted_input=y, + output=routed_out, + expanded_idx_to_permuted_idx=expanded_to_permuted, + topk_scales=topk_w, + ) + + +def _shared_expert(core, x, L: _LayerWeights, decode_gemm=None) -> tuple: + """The shared expert and its per-token gate logit over `x`: `(shared, gate_logits)`. A + decode step under the dense FP8 mode (`decode_gemm` its bf16 projection entry, the copies + present) runs gate|up and down on their FP8 / mid-rows copies and the gate logit on its own + bf16 row; otherwise one GEMM forms gate|up and the gate logit (`sh_gu`'s extra rows).""" + if decode_gemm is not None and L.sh_gu_g is not None: + gu = _dense_mm(x, L.sh_gu, None, decode_gemm, "sh_gu", L.sh_gu_g, mid=L.sh_gu_m) + shared = _dense_mm( + silu_and_mul_rows(gu, core.shared_inter), + L.sh_dn, + None, + decode_gemm, + "sh_dn", + L.sh_dn_g, + mid=L.sh_dn_m, + ) + return shared, decode_gemm(x, L.sh_gate) + gug = cublas_mm(x, L.sh_gug) + shared = cublas_mm(silu_and_mul_rows(gug, core.shared_inter), L.sh_dn) + return shared, gug[:, 2 * core.shared_inter : 2 * core.shared_inter + 1] + + +def _moe_block( + core, x, L: _LayerWeights, pad_rows: int, fp4=None, router_mm=cublas_mm, decode_gemm=None +) -> tuple: + """One MoE layer: NVFP4 routed experts plus the replicated bf16 shared expert, left + uncombined. Returns `(routed, gate_logits, shared)`: the next `qwen4_exp_hc_inject_norm` + forms `routed + sigmoid(gate_logits) * shared` inside its kernel. `routed` is every + window's contribution to this rank's rows (`core.expert_sum`): under attention TP the + window partials' sum over the TP group, the same bits on every rank; under attention DP + the round trip over the group's gathered rows, `pad_rows` padding this rank's. + + The shared expert and its per-token gate logit run on the core's side stream, + concurrently with the routed path, and join it before this returns. `x` is the bf16 + MoE-site mixer output and `fp4` its NVFP4 form (`_hc_moe_mixer`): present only under + attention TP, where the window runs on `x`'s own rows. `router_mm` is the step's + router projection (`_routed_experts`). `decode_gemm` is a decode step's bf16 projection + entry, under which the shared expert reads the dense FP8 mode's copies when they exist.""" + main = torch.cuda.current_stream() + side = core._side_stream + # Fork: the side stream reads `x` only after everything queued on `main` so far. What it + # allocates is consumed on `main` after the join below. + side.wait_stream(main) + with torch.cuda.stream(side): + shared, gate_logits = _shared_expert(core, x, L, decode_gemm) + routed = core.expert_sum( + lambda rows: _routed_experts(core, rows, L, fp4, router_mm), x, pad_rows + ) + main.wait_stream(side) + return routed, gate_logits, shared + + +def _attn_in_heads(core, proj): + """The main q, k and v heads of the fused attention projection's output `proj` + `[T, attn_in_width]`, as `[T, heads, head_dim]` views.""" + q, k, v, _ = torch.split( + torch.reshape(proj, [proj.shape[0], core.attn_in_width // core.head_dim, core.head_dim]), + [core.heads_q, core.heads_kv, core.heads_kv, core.attn_in_tail_heads - core.heads_kv], + 1, + ) + return q, k, v + + +def _packed_q(core, qkv): + """The query heads `[T, heads_q, head_dim]` of the compact `[T, qkv_width]` packed QKV.""" + heads = core.heads_q + 2 * core.heads_kv + return torch.reshape(qkv, [qkv.shape[0], heads, core.head_dim])[:, : core.heads_q] + + +def _attention_rope(core, proj, index_scratch, L: _LayerWeights, pos3, pos) -> None: + """Both planes' Gemma norm + rope. The main q/k heads in place on the fused projection + row (interleaved mRoPE over the three position rows; everything after k is passed as v + heads, which only sets the row stride). The index projection's query and key are copied + into the index scratch and normalized and rotated there (coordinate 0 alone), so the raw + index key stays in `proj` for the index plane's store.""" + fused_qk_norm_rope( + proj, + num_heads_q=core.heads_q, + num_heads_k=core.heads_kv, + num_heads_v=core.attn_in_tail_heads, + head_dim=core.head_dim, + rotary_dim=core.rotary_dim, + eps=core.eps, + q_weight=L.q_norm, + k_weight=L.k_norm, + base=core.theta, + is_neox=True, + position_ids=pos3, + is_qk_norm=True, + use_gemma=True, + use_mrope=True, + mrope_section1=core.mrope_section1, + mrope_section2=core.mrope_section2, + ) + index_scratch[:, : core.index_width].copy_( + proj[:, core.index_col : core.index_col + core.index_width] + ) + fused_qk_norm_rope( + index_scratch, + num_heads_q=core.index_heads, + num_heads_k=core.index_kv_heads, + num_heads_v=core.index_kv_heads, + head_dim=core.index_head_dim, + rotary_dim=core.rotary_dim, + eps=core.eps, + q_weight=L.idx_qn, + k_weight=L.idx_kn, + base=core.theta, + is_neox=True, + position_ids=pos, + is_qk_norm=True, + use_gemma=True, + use_mrope=False, + mrope_section1=0, + mrope_section2=0, + ) + + +def _index_query(core, index_scratch): + """The normalized, rotated index query `[T, index_heads, index_head_dim]`: the q slots + of the index scratch.""" + rows = index_scratch.shape[0] + return torch.split( + torch.reshape( + index_scratch, + [rows, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim], + ), + [core.index_heads, core.index_kv_heads, core.index_kv_heads], + 1, + )[0] + + +def _index_plane(core, proj, i, L: _LayerWeights, qsa) -> None: + """The QSA index branch of attention layer `i`: store the raw index key and the + positions of `proj`'s rows into the paged index plane, then compress every group this + step completes. The raw key is stored before anything normalizes or rotates it, because + the compression applies both itself.""" + req_idx, logical, block_table, tokens_per_block, coords = qsa + rows = proj.shape[0] + token_k = torch.split( + torch.reshape( + proj[:, core.index_col : core.index_col + core.index_scratch_width], + [rows, core.index_scratch_width // core.index_head_dim, core.index_head_dim], + ), + [core.index_heads, core.index_kv_heads, 1], + 1, + )[1] + pools = core._bound.pools + index_cache = pools.index_cache[i] + qsa_paged_index_store( + token_k, + coords, + index_cache, + pools.position_cache, + req_idx, + logical, + block_table, + tokens_per_block, + ) + qsa_prefill_compress( + logical_positions=logical, + request_indices=req_idx, + block_table=block_table, + index_cache=index_cache, + position_cache=pools.position_cache, + k_norm_weight=L.idx_kn, + cos_sin=core._cos_sin, + eps=core.eps, + tokens_per_block=tokens_per_block, + compress_ratio=core.compress_ratio, + mrope_section=_QSA_MROPE_SECTION, + ) + + +def _select(core, scores, row_starts, row_ends, topk_indices, live_cols=None) -> None: + """Each row's canonical top-`block_topk` groups of its window into `topk_indices`: the + groups strictly above the budget's last score, then the highest-index ties, ascending, + `-1` last -- a function of the scores alone, which the expansion and the sparse GQA + reduce in. QSA's row starts are zero, so window-local and table indices agree. + + The exact top-k and the canonical rewrite, two launches, which serve a draft step whose + rows are the rejection kernel's accepted-token runs at every width the engine's graph + buckets use. + + The rewrite reads only the columns of each row's window, and `scores` is as wide as the + block table; a caller that knows a host-side bound on the windows passes it as + `live_cols` so the rewrite reads only those columns and returns the same rows.""" + qsa_topk_select(scores, row_starts, row_ends, topk_indices, index_topk=core.block_topk) + if live_cols is not None: + scores = scores[:, :live_cols] + qsa_canonical_selection(scores, row_starts, row_ends, topk_indices) + + +def _live_score_columns(core, md, batch: int) -> int: + """A host-side bound on every row's selection window in a step with context rows: the + longest sequence's compressed length, with the verify width as slack for the draft rows + that sit past the host's KV lengths.""" + longest = int(md.kv_lens_runtime[:batch].max()) if batch > 0 else 0 + return (longest + _MAX_VERIFY_WIDTH) // core.compress_ratio + 1 + + +def _selection( + core, scores, visible_blocks, logical, seq_lens_q, topk_indices, row_starts, out, live_cols=None +): + """The selected token indices of `scores`' rows into `out` `[rows, token_topk + + compress_ratio - 1]`: the canonical block top-k and its expansion in one launch, or, for + a score table that kernel does not serve, the selection and the expansion apart (the + rewrite bounded by `live_cols`, `_select`).""" + if _topk_expand_supports(scores): + qsa_topk_expand_decode( + scores, + visible_blocks, + logical, + seq_lens_q, + topk_indices, + core.compress_ratio, + core.token_topk, + out=out, + ) + return + _select(core, scores, row_starts, visible_blocks, topk_indices, live_cols) + out.copy_( + qsa_expand_block_indices( + topk_indices, + logical, + seq_lens_q, + compress_ratio=core.compress_ratio, + token_topk=core.token_topk, + ) + ) + + +def _store_kv(core, kv, k, v, request_indices, logical_positions, block_table, tokens_per_block): + """Append this step's K/V rows to the layer's paged pool: a plain move over a bf16 or fp8 + pool (an fp8 pool takes e4m3 unscaled, as the dense op does at the unit KV scales), the + quantizing store over an NVFP4 one, at the dense op's second-level scales.""" + if kv.k_scale is None: + qsa_paged_kv_store( + k, v, kv.k, kv.v, request_indices, logical_positions, block_table, tokens_per_block + ) + return + qsa_paged_kv_store_nvfp4( + k, + v, + kv.k, + kv.v, + kv.k_scale, + kv.v_scale, + request_indices, + logical_positions, + block_table, + tokens_per_block, + core.kv_scale_orig_quant, + ) + + +def _nvfp4_context_groups(pages: list, rows: list, max_pages: int) -> list: + """Consecutive context requests in groups of at most `max_pages` pages, a request with more + pages alone: `(first request, end request, first row, end row)` a group, where request `r` + has `pages[r]` pages and the next `rows[r]` rows.""" + groups = [] + first = first_row = row = listed = 0 + for r, (n_pages, n_rows) in enumerate(zip(pages, rows)): + if r > first and listed + n_pages > max_pages: + groups.append((first, r, first_row, row)) + first, first_row, listed = r, row, 0 + listed += n_pages + row += n_rows + if pages: + groups.append((first, len(pages), first_row, row)) + return groups + + +def _nvfp4_context_plan(core, md, req_idx, block_table, tpb: int) -> tuple: + """The pages a step's context requests read, for the bf16 row-tile pass over an NVFP4 pool: + request `r` lists its first `ceil(kv_len / tpb)` table pages. The requests come in groups + that fit `_NVFP4_CONTEXT_SCRATCH_BYTES` (`_nvfp4_context_groups`); within a group request `r` + takes scratch pages `offsets[r]` on. Returns `(groups, table, k_scratch, v_scratch)`: per + group `(first row, end row, slots, counts, offsets)`, the operands `qsa_nvfp4_pages_to_bf16` + fills per layer and group; `table` (all -1 here) gets the requests' slot rows, the scratch + pools `[pages, Hkv, tpb, D]` bf16 the largest group.""" + num_contexts = md.num_contexts + device = block_table.device + pages = ((md.kv_lens_runtime[:num_contexts] + (tpb - 1)) // tpb).tolist() + rows = md.seq_lens[:num_contexts].tolist() + page_bytes = 2 * core.heads_kv * tpb * core.head_dim * 2 # K and V, bf16 + groups = _nvfp4_context_groups(pages, rows, max(1, _NVFP4_CONTEXT_SCRATCH_BYTES // page_bytes)) + lens = md.seq_lens_cuda[:num_contexts].long() + slots = req_idx[torch.cumsum(lens, 0) - lens] + counts = torch.div( + md.kv_lens_cuda_runtime[:num_contexts].long() + (tpb - 1), tpb, rounding_mode="floor" + ) + planned = [] + for first, end, first_row, end_row in groups: + group_counts = counts[first:end] + planned.append( + ( + first_row, + end_row, + slots[first:end], + group_counts, + torch.cumsum(group_counts, 0) - group_counts, + ) + ) + listed = max(sum(pages[first:end]) for first, end, _, _ in groups) + table = torch.full_like(block_table, -1) + shape = (listed, core.heads_kv, tpb, core.head_dim) + k_scratch = torch.empty(shape, dtype=torch.bfloat16, device=device) + v_scratch = torch.empty(shape, dtype=torch.bfloat16, device=device) + return planned, table, k_scratch, v_scratch + + +def _sparse_gqa_nvfp4_context( + core, kv, plan, q, block_table, selected, req_idx, logical, tpb, num_ctx, out, rows_per_request +) -> torch.Tensor: + """Sparse GQA of a step with context rows over an NVFP4 pool at one rank, into `out` `[rows, + q_width]`. With a plan (`_nvfp4_context_plan`) the first `num_ctx` rows, the context rows, + take the row-tile pass over the plan's scratch, filled here from this layer's pool a group of + requests at a time; the other rows, and every row without a plan, go split-K over the pool + itself. The generation rows after the context rows come `rows_per_request` a request: more + than one (a verify step's) takes the request-grouped entry for them on their own.""" + rows = q.shape[0] + head = 0 + if plan is not None: + groups, table, k_scratch, v_scratch = plan + assert groups[-1][1] == num_ctx, (groups[-1][1], num_ctx) + for first_row, end_row, slots, counts, offsets in groups: + qsa_nvfp4_pages_to_bf16( + kv.k, + kv.v, + kv.k_scale, + kv.v_scale, + core.kv_scale_quant_orig, + block_table, + slots, + counts, + offsets, + table, + k_scratch, + v_scratch, + ) + context = one_rank_sparse_gqa_rowtile( + q=q[first_row:end_row], + k_cache=k_scratch, + v_cache=v_scratch, + block_table=table, + selected_tokens=selected[first_row:end_row], + request_indices=req_idx[first_row:end_row], + tokens_per_block=tpb, + softmax_scale=core.softmax_scale, + query_positions=logical[first_row:end_row], + compress_ratio=core.compress_ratio, + ) + out[first_row:end_row].copy_( + torch.reshape(context, [end_row - first_row, core.q_width]) + ) + head = num_ctx + elif ( + rows_per_request > 1 + and 0 < num_ctx < rows + and _verify_takes_grouped( + q[num_ctx:], kv.k, block_table, selected[num_ctx:], tpb, rows_per_request + ) + ): + context = _sparse_gqa( + core, kv, q[:num_ctx], block_table, selected[:num_ctx], req_idx[:num_ctx], tpb, None + ) + out[:num_ctx].copy_(torch.reshape(context, [num_ctx, core.q_width])) + head = num_ctx + if head < rows: + rest = _sparse_gqa( + core, + kv, + q[head:], + block_table, + selected[head:], + req_idx[head:], + tpb, + None, + rows_per_request=rows_per_request if head == num_ctx else 1, + ) + out[head:rows].copy_(torch.reshape(rest, [rows - head, core.q_width])) + return out + + +def _fp8_context_pages(md, tpb: int) -> int: + """Pool pages a step's context requests hold, `ceil(kv_len / tpb)` each, from the host's + KV lengths: the bf16 scratch `_sparse_gqa_rowtile` widens them into over an fp8 pool.""" + lens = md.kv_lens_runtime[: md.num_contexts] + return int(((lens + (tpb - 1)) // tpb).sum()) + + +def _sparse_gqa_rowtile( + core, kv, q, block_table, selected, request_indices, tpb, query_positions, fp8_pages +): + """The row-tile pass over a context launch, or None where it does not serve it. A bf16 + pool is read in place. Over an fp8 pool, request slot `s`'s first `ceil((p_s + 1) / tpb)` + pages (`p_s` the highest position the launch's rows hold there) are widened, exactly, into + a bf16 scratch pool addressed through a table of `block_table`'s shape; `fp8_pages` is + their host count (`_fp8_context_pages`), read back from the device when not given. A + scratch past `_FP8_ROWTILE_MAX_SCRATCH` bytes declines (None).""" + k_cache, v_cache, table = kv.k, kv.v, block_table + if kv.k.dtype == torch.float8_e4m3fn: + # Each slot's highest position among the launch's rows, -1 for a slot it does not + # hold: a compare-and-max over `[slots, rows]` (a scatter max would serialize every + # row of a request on one address). + slots = torch.arange(block_table.shape[0], dtype=request_indices.dtype, device=q.device) + held = request_indices[None, :] == slots[:, None] + last = torch.where(held, query_positions[None, :].long(), -1).amax(1) + counts = torch.div(last + tpb, tpb, rounding_mode="floor") + if fp8_pages is None: + fp8_pages = int(counts.sum()) + # K and V pages, bf16. + if fp8_pages * kv.k[0].numel() * 4 > _FP8_ROWTILE_MAX_SCRATCH: + return None + k_cache = torch.empty([fp8_pages, *kv.k.shape[1:]], dtype=q.dtype, device=q.device) + v_cache = torch.empty_like(k_cache) + table = torch.full_like(block_table, -1) + qsa_fp8_pages_to_bf16( + kv.k, + kv.v, + block_table, + slots.int(), + counts, + torch.cumsum(counts, 0) - counts, + table, + k_cache, + v_cache, + ) + elif kv.k.dtype != q.dtype: + return None + return qsa_sparse_gqa_rowtile( + q=q, + k_cache=k_cache, + v_cache=v_cache, + block_table=table, + selected_tokens=selected, + request_indices=request_indices, + tokens_per_block=tpb, + softmax_scale=core.softmax_scale, + query_positions=query_positions, + compress_ratio=core.compress_ratio, + ) + + +def _sparse_gqa( + core, + kv, + q, + block_table, + selected, + request_indices, + tpb, + query_positions, + fp8_pages=None, + output_gate=None, + rows_per_request=1, +): + """Sparse GQA over the layer's paged pool. `query_positions` is a context row's bound on + its scan (the prefill pair with `compress_ratio`), or None for decode rows. A context + launch of `_ROWTILE_MIN_ROWS` rows or more over a bf16 or fp8 pool takes the row-tile + pass (`fp8_pages`: see `_sparse_gqa_rowtile`). Over an NVFP4 pool context rows take the + prefill kernel and decode rows the split-K decode kernel.""" + if kv.k_scale is None: + if query_positions is not None and q.shape[0] >= _ROWTILE_MIN_ROWS: + out = _sparse_gqa_rowtile( + core, kv, q, block_table, selected, request_indices, tpb, query_positions, fp8_pages + ) + if out is not None: + return out + return qsa_paged_sparse_gqa( + q=q, + k_cache=kv.k, + v_cache=kv.v, + block_table=block_table, + selected_tokens=selected, + request_indices=request_indices, + tokens_per_block=tpb, + softmax_scale=core.softmax_scale, + query_positions=query_positions, + compress_ratio=None if query_positions is None else core.compress_ratio, + output_gate=None, + ) + if query_positions is None: + # `output_gate` (`[rows, Hq * D]`) returns the gated output; `rows_per_request` > 1 is a + # verify step's rows, consecutive per request, whose K/V the verify kernel loads once. + if rows_per_request > 1: + return qsa_paged_sparse_gqa_nvfp4_verify( + q, + kv.k, + kv.v, + kv.k_scale, + kv.v_scale, + core.kv_scale_quant_orig, + block_table, + selected, + request_indices, + tpb, + core.softmax_scale, + rows_per_request, + output_gate=output_gate, + ) + return qsa_paged_sparse_gqa_nvfp4_decode( + q, + kv.k, + kv.v, + kv.k_scale, + kv.v_scale, + core.kv_scale_quant_orig, + block_table, + selected, + request_indices, + tpb, + core.softmax_scale, + output_gate, + ) + return qsa_paged_sparse_gqa_nvfp4( + q, + kv.k, + kv.v, + kv.k_scale, + kv.v_scale, + core.kv_scale_quant_orig, + block_table, + selected, + request_indices, + tpb, + core.softmax_scale, + query_positions=query_positions, + compress_ratio=core.compress_ratio, + ) + + +def _copy_positions(position_ids: torch.Tensor) -> torch.Tensor: + """A fresh contiguous int32 `[3, N]` copy of the step's mRoPE positions (the engine hands + a slice of a larger buffer, whose row stride the rope rejects).""" + pos3 = torch.empty( + [_QSA_POSITION_AXES, position_ids.numel() // _QSA_POSITION_AXES], + dtype=torch.int32, + device=position_ids.device, + ) + pos3.copy_(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1])) + return pos3 + + +def _step_args(md: TrtllmAttentionMetadata, num_contexts, num_ctx_tokens) -> dict: + """thop_attention's batch state off the prepared metadata. Tensors are engine-owned + persistent buffers refreshed in place; Python ints are per-capture constants.""" + return dict( + sequence_length=md.kv_lens_cuda_runtime, + host_past_key_value_lengths=md.kv_lens_runtime, + host_total_kv_lens=md.host_total_kv_lens, + context_lengths=md.prompt_lens_cuda_runtime, + host_context_lengths=md.prompt_lens_cpu_runtime, + host_request_types=md.host_request_types_runtime, + kv_cache_block_offsets=md.kv_cache_block_offsets, + host_kv_cache_pool_pointers=md.host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping=md.host_kv_cache_pool_mapping, + workspace=md.effective_workspace, + tokens_per_block=md.tokens_per_block, + max_num_requests=md.max_num_requests, + max_context_length=md.max_context_length, + max_seq_len=md.max_seq_len, + # No sliding window: the full admitted length on every attention layer. + attention_window_size=md.max_seq_len, + num_contexts=num_contexts, + num_ctx_tokens=num_ctx_tokens, + trtllm_gen_jit_warmup=md.trtllm_gen_jit_warmup, + use_paged_context_fmha=md.use_paged_context_fmha, + beam_width=md.effective_beam_width, + cache_indirection=md.cache_indirection, + block_ids_per_seq=md.block_ids_per_seq, + max_context_q_len_override=md.max_context_q_len_override, + is_cross=md.is_cross, + is_spec_decoding_enabled=md.is_spec_decoding_enabled, + use_spec_decoding=md.use_spec_decoding, + is_spec_dec_tree=md.is_spec_dec_tree, + spec_decoding_generation_lengths=md.spec_decoding_generation_lengths, + spec_decoding_position_offsets=get_spec_decoding_position_offsets(md), + spec_decoding_packed_mask=md.spec_decoding_packed_mask, + spec_decoding_bl_tree_mask_offset=md.spec_decoding_bl_tree_mask_offset, + spec_decoding_bl_tree_mask=md.spec_decoding_bl_tree_mask, + spec_decoding_target_max_draft_tokens=md.max_total_draft_tokens, + spec_bl_tree_first_sparse_mask_offset_kv=md.spec_bl_tree_first_sparse_mask_offset_kv, + num_sparse_topk=md.num_sparse_topk, + flash_mla_tile_scheduler_metadata=md.flash_mla_tile_scheduler_metadata, + flash_mla_num_splits=md.flash_mla_num_splits, + max_num_sequences=md.max_num_sequences, + force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask, + ) + + +class _KvRangeProbe: + """`TRTLLM_MODELING_V2_VALIDATE` over an fp8 KV pool: reports, once, the largest K/V + magnitude the first eager forward's attention layers store. + + The pool takes K/V at the unit scale, so anything past the e4m3 maximum saturates + silently (a bf16 inf included). One host sync per attention layer of that one forward; + a forward under CUDA graph capture is passed over without a read.""" + + def __init__(self, last_layer: int) -> None: + self._last = last_layer + self._worst = 0.0 + + @classmethod + def for_target(cls, core, enabled: bool) -> "_KvRangeProbe | None": + """A probe over `core`'s attention layers when `enabled` and the pool is fp8.""" + if not enabled or core.quant_mode != QM_FP8_KV_CACHE: + return None + return cls(core.attn_layers[-1]) + + def observe(self, layer: int, k: torch.Tensor, v: torch.Tensor) -> bool: + """Fold one attention layer's K/V in; False once the forward's last layer has + reported.""" + if torch.cuda.is_current_stream_capturing(): + return True + peak = float(torch.maximum(k.abs().amax(), v.abs().amax())) + # A NaN stays: `max` would drop it. + self._worst = math.nan if math.isnan(peak + self._worst) else max(self._worst, peak) + if layer != self._last: + return True + limit = torch.finfo(torch.float8_e4m3fn).max + # Written so that NaN fails too. + if not self._worst <= limit: + logger.warning( + f"FP8 KV cache: the first forward stores a K/V magnitude of {self._worst:g}, " + f"past the e4m3 maximum {limit:g}; the pool is written at the unit scale, " + "so those values saturate" + ) + return False + + +# ------------------------------------------------------------------ prefill --- + + +class _PrefillStep(NamedTuple): + """One step's host-side state, built by `PrefillTarget.prepare` and read by the forward's + breaks. `num_tokens` is the batch's row count `n` and `physical_tokens` the `T` the + forward runs over (`n`, or the prefill CUDA graph bucket it was padded to), and + `dp_pad_rows` the rows the MoE blocks pad them with under attention DP (zero otherwise, + and on every step a prefill CUDA graph runs). A new one replaces it every step, so no + break reads a previous step's.""" + + num_tokens: int + physical_tokens: int + dp_pad_rows: int + input_ids: torch.Tensor + attn: dict + gdn: tuple + ple: tuple + qsa: tuple + vision: tuple | None + # A verify step: the generation rows' side-cache writes are staged for the rollback. + stage: bool = False + # At one rank over an NVFP4 pool, `_nvfp4_context_plan`'s row-tile plan or None. + nvfp4_context: tuple | None = None + # Over an fp8 pool, the context requests' pages the row-tile pass widens + # (`_fp8_context_pages`); None otherwise. + fp8_pages: int | None = None + # `_live_score_columns`: a bound on every row's selection window. + live_cols: int | None = None + + +class PrefillTarget: + """Any batch with context rows (`num_contexts > 0`), possibly with generation rows + beside them.""" + + def __init__(self, core: Qwen4ExpDefaultModelingV2) -> None: + self.core = core + # Off unless TRTLLM_MODELING_V2_VALIDATE asks. + self._contract_pending = step_contract_enabled() + # Dropped once it has reported. + self._kv_probe = _KvRangeProbe.for_target(core, self._contract_pending) + # The rstd sites' `(1 + w)` operands (`_HC_RSTD_MIN_ROWS`), formed once here. + self._layers = [ + L._replace(ahc_norm_w1=_norm_w1(L.ahc_norm), mhc_norm_w1=_norm_w1(L.mhc_norm)) + for L in _layer_views(core) + ] + stack = _stack_views(core) + self._stack = stack._replace(fhc_norm_w1=_norm_w1(stack.fhc_norm)) + # The step `prepare` built, read by the breaks. + self._step: _PrefillStep | None = None + # Whether this rank reported the GDN in-place short-context path yet. + self._gdn_inplace_reported = False + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): calling `step_args` + surfaces an upstream metadata rename here, not mid-forward.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md: TrtllmAttentionMetadata) -> dict: + """Project the prepared metadata onto thop_attention's batch state.""" + return _step_args(md, md.num_contexts, md.num_ctx_tokens) + + # --------------------------------------------------------------- blocks --- + + def _attention_block(self, x, i, L, proj_qkv, index_scratch, attn_out, pos3, pos): + """One gated GQA layer plus its QSA index branch, at this rank's heads. + + The four input projections are one GEMM against the fused operand; the norm + rope + of both planes is row-wise and runs over all `T` rows here; the cache writes, the + selection and the attention core read the step's batch and run in the + `_attention_core` break, which leaves the attention output in `attn_out`. The gate + is read through its row stride, and the output projection's per-rank partial is + summed over the TP group.""" + core = self.core + proj = _dense_mm(x, L.attn_in, L.attn_in8, cublas_mm, "in") + gate = torch.split(proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1)[1] + _attention_rope(core, proj, index_scratch, L, pos3, pos) + eager_on_graph(proj, proj_qkv, index_scratch, attn_out, i, inner=self._attention_core) + fused_sigmoid_mul_inplace(attn_out, gate) + return core.tp_sum(_dense_mm(attn_out, L.o, L.o8, cublas_mm, "out")) + + def _attention_core(self, proj, proj_qkv, index_scratch, attn_out, i) -> None: + """Break: the index plane's store and compression, the dense/sparse decision and the + attention core over the step's `n` rows. Writes rows `[:n]` of `attn_out` + `[T, q_width]` in place and zeroes rows `[n:]`. + + The index plane advances on every step, including below the threshold, so a request + growing past it finds its prefix compressed; above it the main K/V append becomes + this layer's own call. On a verify step the generation rows' index slots (and, on + the first attention layer, their position slots) are staged for `SpeculativeRollback` + before the store and the compression overwrite them.""" + core = self.core + L = self._layers[i] + step = self._step + ( + req_idx, + logical, + seq_lens_q, + visible_blocks, + topk_indices, + topk_row_starts, + block_table, + tokens_per_block, + coords, + sparse, + ) = step.qsa + num_tokens = step.num_tokens + proj = proj[:num_tokens] + _, k_dst, v_dst = _attn_in_heads(core, proj) + if self._kv_probe is not None and not self._kv_probe.observe(i, k_dst, v_dst): + self._kv_probe = None + # The packed `[q | k | v]` as one compact copy: the dense op reads it flat, and the + # sparse GQA reads its query heads from it. + qkv = proj_qkv[:num_tokens] + qkv.copy_(proj[:, : core.qkv_width]) + if step.stage: + core.spec_rollback.stage_index( + core._bound.pools.index_cache[i], core._bound.pools.position_cache + ) + _index_plane(core, proj, i, L, (req_idx, logical, block_table, tokens_per_block, coords)) + + if sparse: + kv = core._bound.pools.kv[i] + # Append the normalized, rotated K/V: nothing else writes the pool on this path, + # and the selection reads the current token from it. + _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) + q_index = _index_query(core, index_scratch[:num_tokens]) + # The step's context rows come first, consecutive rows of a request, and score as + # query tiles; its generation rows, one a request, score row by row. Each part + # selects into its own rows. + num_ctx = step.attn["num_ctx_tokens"] + parts = [(0, num_ctx, True)] + if num_ctx < num_tokens: + parts.append((num_ctx, num_tokens, False)) + selected = torch.empty( + [num_tokens, core.token_topk + core.compress_ratio - 1], + dtype=torch.int32, + device=proj.device, + ) + for lo, hi, context_rows in parts: + scores = qsa_paged_index_scores( + q=q_index[lo:hi], + index_cache=core._bound.pools.index_cache[i], + block_table=block_table, + query_positions=logical[lo:hi], + request_indices=req_idx[lo:hi], + tokens_per_block=tokens_per_block, + compress_ratio=core.compress_ratio, + only_visible_blocks=True, + context_rows=context_rows, + ) + _selection( + core, + scores, + visible_blocks[lo:hi], + logical[lo:hi], + seq_lens_q[lo:hi], + topk_indices[lo:hi], + topk_row_starts[lo:hi], + selected[lo:hi], + step.live_cols, + ) + # Context rows with the prefill pair (a bound on a row's scan), generation rows + # without. + if core.mapping.world_size == 1 and kv.k_scale is not None: + _sparse_gqa_nvfp4_context( + core, + kv, + step.nvfp4_context, + _packed_q(core, qkv), + block_table, + selected, + req_idx, + logical, + tokens_per_block, + num_ctx, + attn_out[:num_tokens], + step.attn["predicted_tokens_per_seq"], + ) + parts = [] + for lo, hi, context_rows in parts: + out = _sparse_gqa( + core, + kv, + _packed_q(core, qkv)[lo:hi], + block_table, + selected[lo:hi], + req_idx[lo:hi], + tokens_per_block, + logical[lo:hi] if context_rows else None, + step.fp8_pages, + ) + attn_out[lo:hi].copy_(torch.reshape(out, [hi - lo, core.q_width])) + else: + thop_attention( + q=qkv, + output=attn_out[:num_tokens], + local_layer_idx=core._bound.pools.pool_idx[i], + num_heads=core.heads_q, + num_kv_heads=core.heads_kv, + head_size=core.head_dim, + quant_mode=core.quant_mode, + kv_scale_orig_quant=core.kv_scale_orig_quant, + kv_scale_quant_orig=core.kv_scale_quant_orig, + **step.attn, + **_CALL_CONSTANTS, + ) + attn_out[num_tokens:].zero_() + + def _gdn_block(self, x, i, L, core_out) -> torch.Tensor: + """One Gated-DeltaNet layer over a batch with a context prefix, at this rank's heads: + the one input projection GEMM (q|k|v, z, a and b are column views of its output) and + the gated-norm epilogue run over all `T` rows here, the recurrence in the `_gdn_core` + break, which leaves its output in `core_out` `[1, T, num_v_heads, head_v_dim]`.""" + core = self.core + proj = _dense_mm(x, L.in_proj, L.in_proj8, cublas_mm, "in") + z = torch.split(proj, core.in_proj_split, 1)[1] + eager_on_graph(proj, core_out, i, inner=self._gdn_core) + return self._gdn_epilogue(core_out, z, L, x.shape[0]) + + def _gdn_core(self, proj, out, i) -> None: + """Break: the recurrence over the step's `n` rows. Writes rows `[:n]` of `out` in place + and zeroes rows `[n:]`. Every branch is on a host int of the step. + + The context rows take one fused conv + post-conv pass (`_gdn_prefill_prep`) and + FlashInfer's chunked recurrence, and a mixed step's one-token generation rows the + decode step's conv and recurrence. A context sequence carrying nothing starts from + zeros without its pool slots being cleared: the conv pass reads no carried window for + it, and the recurrence (`has_initial_state`) gathers zeros for it. + + On a verify step the generation rows are draft chains of unknown accepted length: + they take `_gdn_verify` (per-step conv windows and the cached replay) instead.""" + core = self.core + L = self._layers[i] + step = self._step + ( + num_prefill_tokens, + num_decode_tokens, + si_p, + si_all, + si_d, + hi_p, + hi_all, + qsl_p, + qsl_all, + qsl_long, + qsl_long_p, + num_decodes, + gen_p, + verify, + ) = step.gdn + num_tokens = step.num_tokens + pools = core._bound.pools + ssm_states = pools.ssm[i] + mixed_qkv, _, a, b = torch.split(proj[:num_tokens], core.in_proj_split, 1) + + if verify is not None: + mq_p, mq_d = torch.split(mixed_qkv, [num_prefill_tokens, num_decode_tokens], 0) + a_p, a_d = torch.split(a, [num_prefill_tokens, num_decode_tokens], 0) + b_p, b_d = torch.split(b, [num_prefill_tokens, num_decode_tokens], 0) + # The context rows as on a step without verify rows. + q, k, v, g, beta = self._gdn_prefill_prep(mq_p, a_p, b_p, L, i, qsl_p, si_p, hi_p) + flashinfer_chunk_gated_delta_rule( + q, + k, + v, + g, + beta, + scale=None, + initial_state=ssm_states, + initial_state_indices=si_p, + inplace_indexed_state_update=True, + output_final_state=False, + cu_seqlens=qsl_long_p, + head_first=False, + use_qk_l2norm_in_kernel=False, + output=out[:, :num_prefill_tokens], + state_workspace=None, + has_initial_state=hi_p, + linear_gate=True, + ) + out[:, num_prefill_tokens:num_tokens].copy_( + _gdn_verify( + core, + mq_d, + a_d, + b_d, + i, + L, + si_d, + num_decodes, + gen_p, + verify, + ) + ) + elif num_decode_tokens > 0: + mq_p, mq_d = torch.split(mixed_qkv, [num_prefill_tokens, num_decode_tokens], 0) + a_p, b_p = a[:num_prefill_tokens], b[:num_prefill_tokens] + q, k, v, g, beta = self._gdn_prefill_prep(mq_p, a_p, b_p, L, i, qsl_p, si_p, hi_p) + causal_conv1d_update( + x=mq_d, + conv_state=pools.conv[i], + weight=L.conv_w, + bias=None, + silu_activation=True, + cache_seqlens=None, + conv_state_indices=si_d, + pad_slot_id=_PAD_SLOT_ID, + ) + # The two halves run separate recurrences over disjoint slots and write disjoint + # row ranges of the output. + flashinfer_chunk_gated_delta_rule( + q, + k, + v, + g, + beta, + scale=None, + initial_state=ssm_states, + initial_state_indices=si_p, + inplace_indexed_state_update=True, + output_final_state=False, + cu_seqlens=qsl_long_p, + head_first=False, + use_qk_l2norm_in_kernel=False, + output=out[:, :num_prefill_tokens], + state_workspace=None, + has_initial_state=hi_p, + linear_gate=True, + ) + # The generation rows: the decode step's recurrence (without the gated norm, + # which runs with the context rows' below) over their rows of the projection, + # whose q/k/v columns the conv step rewrote; it computes both gates and the q/k + # norm itself. + gdn_decode_rows( + proj[num_prefill_tokens:num_tokens], + ssm_states, + si_d, + L.a_log, + L.dt_bias, + core.num_k_heads, + core.head_k_dim, + out[0, num_prefill_tokens:num_tokens].view(num_decode_tokens, core.value_dim), + ) + else: + q, k, v, g, beta = self._gdn_prefill_prep( + mixed_qkv, a, b, L, i, qsl_all, si_all, hi_all + ) + # q and k arrive L2-normalized; the recurrence must not repeat it. + flashinfer_chunk_gated_delta_rule( + q, + k, + v, + g, + beta, + scale=None, + initial_state=ssm_states, + initial_state_indices=si_all, + inplace_indexed_state_update=True, + output_final_state=False, + cu_seqlens=qsl_long, + head_first=False, + use_qk_l2norm_in_kernel=False, + output=out[:, :num_tokens], + state_workspace=None, + has_initial_state=hi_all, + linear_gate=True, + ) + out[:, num_tokens:].zero_() + + def _gdn_prefill_prep(self, mq, a, b, L, i, qsl, si, hi): + """The context rows' recurrence inputs in one pass over `mq`, the token-major q|k|v + column view of the projection: the causal conv against each sequence's carried + window (published back to the pool), SiLU, the q|k|v split with q/k L2-normalized, + and the decay/beta gates (the decay as the factor `exp(g)`, `linear_gate=True`).""" + core = self.core + return gdn_prefill_conv_post( + x=mq, + weight=L.conv_w, + conv_states=core._bound.pools.conv[i], + query_start_loc=qsl, + cache_indices=si, + has_initial_state=hi, + a=a, + b=b, + A_log=L.a_log, + dt_bias=L.dt_bias, + num_k_heads=core.num_k_heads, + head_k_dim=core.head_k_dim, + num_v_heads=core.num_v_heads, + head_v_dim=core.head_v_dim, + pad_slot_id=_PAD_SLOT_ID, + l2_norm_eps=_L2_NORM_EPS, + softplus_threshold=_SOFTPLUS_THRESHOLD, + linear_gate=True, + ) + + def _gdn_epilogue(self, out, z, L, num_tokens) -> torch.Tensor: + """The gated RMS norm (sigmoid gate, plain non-gemma weight) and the output + projection, whose per-rank partial is summed over the TP group.""" + core = self.core + o = rms_norm_gated_token_major( + torch.reshape(out, [-1, core.head_v_dim]), + torch.reshape(z, [num_tokens, core.num_v_heads, core.head_v_dim]), + weight=L.gnorm, + eps=core.eps, + fp8_scale=None, + gate_activation="sigmoid", + ) + o = torch.reshape(o, [num_tokens, core.value_dim]) + return core.tp_sum(_dense_mm(o, L.out, L.out8, cublas_mm, "out")) + + def _ple_block(self, hc, hc_out, normed_out) -> None: + """Break: the n-gram side path on its decoder layer, over the step's `n` rows, added + into the stream and normalized for the layer's attention-side site. Writes rows + `[:n]` of `hc_out` (the stream) and `normed_out` (its grouped norm) in place and zeroes + rows `[n:]`. + + Per-request stateful: ids become one `ngram_size`-wide window per token + (predecessors from the paged id history); each window hashes to one row per head; + rows gather fp8 vectors from one table (`ple_ngram_embed`); the key/value projections, + the gate, a dilated depthwise short conv over a paged state, the residual add and the + next site's grouped norm are `ple_gate_conv_inject_norm`. A sequence carrying nothing + is seeded with `eos_token_id`. + + On a verify step the generation rows' pre-step pool rows and ids are staged before + the first launch publishes, and the conv input the fused launch writes out after it, + for `SpeculativeRollback` to rewrite both pools to the accepted prefix.""" + core = self.core + S = self._stack + step = self._step + conv_states, ngram_pool, si_all, carry_all, qsl = step.ple + num_tokens = step.num_tokens + next_norm = self._layers[core.ple_layer].ahc_norm + conv_input = None + if step.stage: + core.spec_rollback.stage_ple_verify(ngram_pool, conv_states, step.input_ids) + conv_input = torch.empty_like(hc[:num_tokens]) + ngram = ple_ngram_embed( + step.input_ids, + ngram_pool, + si_all, + carry_all, + qsl, + core.w["ple_table"], + S.ple_tscale, + S.ple_mult, + S.ple_vocab, + S.ple_offs, + core.eos_token_id, + num_tokens, + ) + ple_gate_conv_inject_norm( + hc[:num_tokens], + cublas_mm(ngram, S.ple_key), + cublas_mm(ngram, S.ple_val), + S.ple_nk_w, + S.ple_nq_w, + S.ple_nc_w, + S.ple_conv_w, + conv_states, + si_all, + qsl, + next_norm, + core.eps, + core.hc_count, + _PLE_CLAMP_MIN, + core.ple_dilation, + _PAD_SLOT_ID, + num_tokens, + carry_all, + conv_input=conv_input, + out=(hc_out[:num_tokens], normed_out[:num_tokens]), + ) + if step.stage: + core.spec_rollback.stage_ple_conv_input(conv_input) + hc_out[num_tokens:].zero_() + normed_out[num_tokens:].zero_() + + def _splice(self, h) -> None: + """Break: write the step's vision rows over their placeholder rows of the token + embedding `h` `[T, hidden]` in place, and zero rows `[n:]`.""" + step = self._step + if step.vision is not None: + rows, mm_token_indices = step.vision + h.index_copy_(0, mm_token_indices, rows) + h[step.num_tokens :].zero_() + + # -------------------------------------------------------------- forward --- + + def prepare( + self, + attn_metadata: AttentionMetadata, + input_ids: torch.IntTensor, + position_ids: torch.IntTensor, + vision: tuple | None, + spec_metadata=None, + ) -> None: + """Everything this step reads from the host, kept for the forward's breaks as one + `_PrefillStep`: the contract check, the attention op's per-step arguments and the + batch's bookkeeping, built once and shared by every layer. + + Called by the core's `begin_step` on every step with context rows, before the + forward. On a step a breakable prefill graph replays, this and the breaks are all of + this target that runs.""" + core = self.core + md = attn_metadata + self._check_step_contract(md) + attn_args = self.step_args(md) + mm = md.mamba_metadata + + # The activation rows are the batch's rows, or a prefill CUDA graph bucket's. + physical_tokens = input_ids.shape[0] + assert physical_tokens in (md.num_tokens, md.padded_num_tokens), ( + f"the engine handed {physical_tokens} activation rows for a batch its metadata " + f"describes as {md.num_tokens} tokens (padded_num_tokens={md.padded_num_tokens})" + ) + num_tokens = md.num_tokens + # Under attention DP the MoE blocks pad the `T` rows to the group's largest count. A + # prefill CUDA graph pads every rank to one bucket, so the body it captures pads + # nothing; that zero is what makes the count a per-capture constant. + dp_pad_rows = core._dp_pad_rows(md, physical_tokens) + assert physical_tokens == num_tokens or dp_pad_rows == 0, ( + f"a prefill CUDA graph bucket of {physical_tokens} rows with {dp_pad_rows} attention-DP " + "padding rows: the graph would freeze the group's row counts" + ) + batch = md.num_seqs + num_prefills = md.num_contexts + num_decodes = batch - num_prefills + num_prefill_tokens = md.num_ctx_tokens + num_decode_tokens = num_tokens - num_prefill_tokens + if not self._gdn_inplace_reported and num_prefill_tokens <= _GDN_INPLACE_MAX_TOKENS: + # The recurrence of a context this short runs in place over the pool rows + # (`flashinfer_chunk_gated_delta_rule`'s short path). + self._gdn_inplace_reported = True + print( + f"modeling_v2 qwen4_exp default target: rank {core.mapping.rank} runs short GDN " + "context chunks in place over the recurrent pool" + ) + # Rows per generation sequence: 1 without speculation, else `runtime_draft_len + 1`. + gen_p, is_verify = _verify_rows(core, md, spec_metadata, num_decodes, num_decode_tokens) + attn_args["predicted_tokens_per_seq"] = gen_p + + si_all = mm.state_indices[:batch] + hi_all = mm.has_initial_states[:batch] + si_p, si_d = torch.split(si_all, [num_prefills, num_decodes], 0) + # The single-token decode kernel takes its slot vector through a 16-byte-aligned + # pointer; the generation half of the int32 slot vector starts `num_prefills` rows + # in, so one copy per step makes it aligned for every layer. + si_d = torch.clone(si_d) + hi_p = hi_all[:num_prefills] + # `query_start_loc` covers the context prefix; `_long` the whole batch, and its prefix + # the context rows' chunked recurrence when generation rows run beside them. + qsl_all = mm.query_start_loc + qsl_p = None if qsl_all is None else qsl_all[: num_prefills + 1] + qsl_long = mm.query_start_loc_long + qsl_long_p = qsl_long[: num_prefills + 1] + gdn_state = ( + num_prefill_tokens, + num_decode_tokens, + si_p, + si_all, + si_d, + hi_p, + hi_all, + qsl_p, + qsl_all, + qsl_long, + qsl_long_p, + num_decodes, + gen_p, + _verify_state(core, mm, num_decodes) if is_verify else None, + ) + + # PLE carry flags must cover the whole batch, but `has_initial_states` is filled for + # the context prefix only; every generation sequence carries, so pad with True. + ple_carry = nn.functional.pad(hi_p, [0, num_decodes], mode="constant", value=1.0) + # Under speculation `query_start_loc_long` has one row per generation sequence, so + # with `gen_p > 1` both PLE stateful calls would be mis-segmented: rebuilt from + # `seq_lens_cuda`, which holds the true counts during the trunk forward. + ple_qsl = ( + qsl_long + if gen_p == 1 + else cu_seqlens_from_seq_lens(md.seq_lens_cuda[:batch], torch.int64) + ) + pools = core._bound.pools + ple_state = (pools.ple_state[0], pools.ple_state[1], si_all, ple_carry, ple_qsl) + + # Dense/sparse decision once per forward, from a host-side metadata read: a + # per-capture constant. + sparse = batch > 0 and int(md.kv_lens_runtime[:batch].max()) > core.dense_threshold + # The `[T, 3]` transpose of the `[3, 1, T]` position buffer, which the QSA position + # plane stores. + coords = torch.transpose(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]), 0, 1) + qsa = ( + md.qsa_req_idx_per_token[:num_tokens], + md.qsa_logical_positions[:num_tokens], + md.qsa_sequence_lengths[:num_tokens], + md.qsa_visible_blocks[:num_tokens], + md.qsa_topk_indices[:num_tokens], + md.qsa_topk_row_starts[:num_tokens], + md.qsa_block_table, + md.tokens_per_block, + coords[:num_tokens], + sparse, + ) + + # A verify step's generation rows write the PLE pools and the QSA side caches for + # every draft row, accepted or not; `SpeculativeRollback` keeps what its commit needs + # to undo the rejected ones. Reopened on every forward, so a step whose commit never + # ran leaves nothing for the next one. Context rows are never staged. + rollback = core.spec_rollback + if rollback is not None: + rollback.discard() + if is_verify: + rollback.begin( + num_contexts=num_prefills, + prefix_tokens=num_prefill_tokens, + num_gens=num_decodes, + rows_per_gen=gen_p, + slots=si_d, + request_indices=md.qsa_req_idx_per_token[num_prefill_tokens:num_tokens], + logical_positions=md.qsa_logical_positions[num_prefill_tokens:num_tokens], + block_table=md.qsa_block_table, + tokens_per_block=md.tokens_per_block, + ) + fp8_pages = None + if ( + sparse + and core.quant_mode == QM_FP8_KV_CACHE + and num_prefill_tokens >= _ROWTILE_MIN_ROWS + ): + fp8_pages = _fp8_context_pages(md, md.tokens_per_block) + nvfp4_context = None + if ( + core.mapping.world_size == 1 + and sparse + and core.quant_mode == QM_NVFP4_KV_CACHE + and num_prefill_tokens > _NVFP4_ROWTILE_MIN_ROWS + ): + nvfp4_context = _nvfp4_context_plan( + core, md, qsa[0], md.qsa_block_table, md.tokens_per_block + ) + self._step = _PrefillStep( + num_tokens=num_tokens, + physical_tokens=physical_tokens, + dp_pad_rows=dp_pad_rows, + input_ids=input_ids[:num_tokens], + attn=attn_args, + gdn=gdn_state, + ple=ple_state, + qsa=qsa, + vision=vision, + stage=is_verify, + fp8_pages=fp8_pages, + nvfp4_context=nvfp4_context, + live_cols=_live_score_columns(core, md, batch), + ) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: torch.IntTensor, + position_ids: torch.IntTensor, + **kwargs, + ): + """The step `prepare` built, over the `T` activation rows the engine handed in: the + batch's `n` rows, or the prefill CUDA graph bucket they were padded to. + + **This is the body a breakable prefill CUDA graph captures**, once per bucket, and + replays for every step it pads to that bucket without running this Python. So nothing + here reads the step: every op is row-wise and sized by `T` alone, and what depends on + the step's batch runs in one of the four breaks (`_splice`, `_ple_block`, + `_attention_core`, `_gdn_core`), each writing rows `[:n]` of a fixed output allocated + here and zeroing rows `[n:]`. The tensor-parallel sums run on all `T` rows, which + every rank holds alike; under attention DP the MoE round trip pads them by the step's + `dp_pad_rows`, zero whenever a graph captures this body.""" + core = self.core + num_tokens = input_ids.shape[0] + assert num_tokens == self._step.physical_tokens, ( + f"the forward got {num_tokens} activation rows and prepare saw " + f"{self._step.physical_tokens}" + ) + pos3 = _copy_positions(position_ids) + pos = pos3[0] + + # The token embedding (replicated) with this step's vision rows spliced in. + h = nn.functional.embedding(input_ids, core.w["embed"]) + eager_on_graph(h, inner=self._splice) + dt = h.dtype + dev = h.device + + # Scratch shared by every layer that uses it; each break rewrites its whole output. + proj_qkv = torch.empty([num_tokens, core.qkv_width], dtype=dt, device=dev) + index_scratch = torch.empty([num_tokens, core.index_scratch_width], dtype=dt, device=dev) + attn_out = torch.empty([num_tokens, core.q_width], dtype=dt, device=dev) + gdn_out = torch.empty( + [1, num_tokens, core.num_v_heads, core.head_v_dim], dtype=dt, device=dev + ) + ple_hc = torch.empty([num_tokens, core.stream], dtype=dt, device=dev) + ple_normed = torch.empty([num_tokens, core.stream], dtype=dt, device=dev) + + # The embedding repeated across the hc_count residual streams. + hc_count = core.hc_count + hc = torch.cat([h] * hc_count, -1) + layers, S = self._layers, self._stack + pad_rows = self._step.dp_pad_rows + # A host decision on `T` alone (each captured bucket takes one branch): at one rank, from + # `_HC_RSTD_MIN_ROWS` rows, every site but the PLE layer's attention-side one is + # `qwen4_exp_hc_inject_rstd` and the rstd mixer. + by_rstd = core.hc_fast_sites and num_tokens >= _HC_RSTD_MIN_ROWS + normed = rstd = None + if by_rstd: + _, rstd = qwen4_exp_hc_inject_rstd(hc, None, None, core.eps, hc_count) + else: + _, normed = qwen4_exp_hc_inject_norm( + hc, None, None, layers[0].ahc_norm, core.eps, hc_count + ) + for i, L in enumerate(layers): + if i == core.ple_layer: + eager_on_graph(hc, ple_hc, ple_normed, inner=self._ple_block) + hc, normed = ple_hc, ple_normed + # The PLE kernel stores this site's normed stream: the stored-normed mixer + # reads it (cheaper than a second pass for its rstd). + mixed, inject = _hc_mixer(core, normed, L.ahc_first, L.ahc_up, True) + elif by_rstd: + mixed, inject = _hc_mixer_rstd( + core, hc, rstd, L.ahc_norm, L.ahc_norm_w1, L.ahc_first, L.ahc_up, True + ) + else: + mixed, inject = _hc_mixer(core, normed, L.ahc_first, L.ahc_up, True) + if core.is_attn[i]: + block_out = self._attention_block( + mixed, i, L, proj_qkv, index_scratch, attn_out, pos3, pos + ) + else: + block_out = self._gdn_block(mixed, i, L, gdn_out) + if by_rstd: + hc, rstd = qwen4_exp_hc_inject_rstd(hc, block_out, inject, core.eps, hc_count) + mixed, inject, fp4 = _hc_moe_mixer_rstd(core, hc, rstd, L) + else: + hc, normed = qwen4_exp_hc_inject_norm( + hc, block_out, inject, L.mhc_norm, core.eps, hc_count + ) + mixed, inject, fp4 = _hc_moe_mixer(core, normed, L) + next_norm = layers[i + 1].ahc_norm if i + 1 < len(layers) else S.fhc_norm + # The MoE block output is formed inside the injection: routed + sigmoid(gate) * + # shared, rounded once. + routed, gate_logits, shared = _moe_block(core, mixed, L, pad_rows, fp4) + if by_rstd: + hc, rstd = qwen4_exp_hc_inject_rstd( + hc, routed, inject, core.eps, hc_count, shared, gate_logits + ) + else: + hc, normed = qwen4_exp_hc_inject_norm( + hc, routed, inject, next_norm, core.eps, hc_count, shared, gate_logits + ) + del routed, gate_logits, shared, fp4 + # The stack-final mixer with no injection weight is the final norm. + if by_rstd: + final, _ = _hc_mixer_rstd( + core, hc, rstd, S.fhc_norm, S.fhc_norm_w1, S.fhc_first, S.fhc_up, False + ) + else: + final, _ = _hc_mixer(core, normed, S.fhc_first, S.fhc_up, False) + if not core.mtp_enabled: + return final + # `lm_head` reads the `[T, hidden]` mix; the MTP module reads the `[T, stream]` stream + # it was collapsed from. They leave packed in one tensor, because a breakable prefill + # graph's body returns one tensor (`split_spec_output`). + return torch.cat([final, hc], -1) + + # ------------------------------------------------------------- the draft --- + + def _capture_selection(self, md, topk_indices, visible_blocks, compact=False) -> None: + """IndexShare: stash, per request, the selection of the row the draft loop continues + from, with that row's complete-group count, in the draft layer's row of the shared + buffers: a context request's last row of this chunk, and a generation request's last + accepted row (`num_accepted - 1` past its first). A `compact` selection holds each + context request's last row, then every generation row (`mtp_forward`'s `keep`).""" + batch = md.num_seqs + assert md.qsa_mtp_num_accepted is not None, ( + "IndexShare needs the worker's accepted counts to pick each request's capture row; " + "the draft loop did not supply them" + ) + seq_lens = md.seq_lens_cuda[:batch].to(torch.int64) + starts = torch.cumsum(seq_lens, 0) - seq_lens + accepted = md.qsa_mtp_num_accepted[:batch].to(torch.int64) + offsets = torch.minimum(torch.clamp(accepted - 1, min=0), seq_lens - 1) + offsets[: md.num_contexts] = seq_lens[: md.num_contexts] - 1 + rows = starts + offsets + if compact: + num_contexts = md.num_contexts + rows[:num_contexts] = torch.arange(num_contexts, device=rows.device) + rows[num_contexts:] -= md.num_ctx_tokens - num_contexts + slot = self.core._draft_bound.share_slot + md.qsa_shared_topk_indices[slot, :batch].copy_(torch.index_select(topk_indices, 0, rows)) + md.qsa_shared_topk_visible_blocks[slot, :batch].copy_( + torch.index_select(visible_blocks, 0, rows) + ) + # Host-side readiness, recorded in program order. + md.qsa_shared_topk_captured_slots.add(slot) + + def mtp_forward( + self, + embed_tokens: torch.Tensor, + all_rank_num_tokens, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + **kwargs, + ) -> torch.Tensor: + """One draft step whose metadata still carries context rows: the draft loop's first + (draft-extend) step of a batch with context rows, eager. Every row writes the layer's + KV and index planes; the worker gathers one row per request from the result. Under + attention TP only `keep` -- each context request's last row, then every generation + row -- goes on past those writes (the query side included on the sparse path; the + dense call's metadata describes every row), and the other returned rows keep the + layer's input stream. Under attention DP every row runs the whole layer. + + The draft layer's attention is the trunk prefill target's: the context rows score + as query tiles and bound their sparse scan by their positions, the generation rows + (`p` a request) score row by row. `all_rank_num_tokens` is the group's row counts of + this draft step, which pad the MoE round trip under attention DP (`_draft_pad_rows`).""" + core = self.core + assert isinstance(attn_metadata, TrtllmAttentionMetadata) + assert core._bound is not None, "the trunk's cache-manager bind has not run" + md = attn_metadata + rows = hidden_states.shape[0] + assert input_ids.shape[0] == rows and hidden_states.shape[1] == core.stream, ( + f"the draft step got {input_ids.shape[0]} token ids and {tuple(hidden_states.shape)} " + f"hidden states; the MTP module reads one {core.stream}-wide stream row per token" + ) + # Projected on every call: the worker rewrites `attn_metadata` in place between + # draft steps. + step = self.step_args(md) + num_ctx = md.num_ctx_tokens + gen_seqs = md.num_seqs - md.num_contexts + p, _ = _verify_rows(core, md, None, gen_seqs, rows - num_ctx) + step["predicted_tokens_per_seq"] = p + draft, layer_idx = core._draft_step_bind(md) + assert rows == md.num_tokens, ( + f"the draft step carries {rows} rows for a batch its metadata describes as " + f"{md.num_tokens} tokens; the QSA row mapping covers the second" + ) + pad_rows = _draft_pad_rows(core, all_rank_num_tokens, rows) + # The trunk's dense/sparse decision from the same host buffer: every draft step of + # one engine forward takes the branch its trunk took. + batch = md.num_seqs + sparse = batch > 0 and int(md.kv_lens_runtime[:batch].max()) > core.dense_threshold + # IndexShare: the only draft step with context rows is the loop's first, which always + # runs the indexer and, with sharing on, captures the selection the later steps reuse. + share = core.mtp_index_share and md.qsa_in_mtp_draft_loop + assert not (share and md.qsa_indexer_skip_topk), ( + "the worker asked a draft step with context rows to reuse a selection" + ) + pos3, pos, coords = _draft_positions(position_ids) + M = core._mtp + req_idx = md.qsa_req_idx_per_token[:rows] + logical = md.qsa_logical_positions[:rows] + block_table = md.qsa_block_table + tokens_per_block = md.tokens_per_block + + x, mixed, inject = _mtp_front(core, M, embed_tokens, input_ids, hidden_states) + qkv, gate, k_dst, v_dst = _mtp_qkv(core, M, mixed, pos3) + q_index = _mtp_index_branch( + core, M, mixed, draft, (req_idx, logical, block_table, tokens_per_block, pos, coords) + ) + keep = None + if core.dp_size == 1 and num_ctx > md.num_contexts: + ctx_lens = md.seq_lens_cuda[: md.num_contexts].to(torch.long) + keep = torch.cat( + [ + torch.cumsum(ctx_lens, 0) - 1, + torch.arange(num_ctx, rows, dtype=torch.long, device=x.device), + ] + ) + if sparse and keep is not None: + kv = draft.kv + _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) + q_rows = keep.shape[0] + q_req, q_logical = req_idx[keep], logical[keep] + visible_blocks = md.qsa_visible_blocks[:rows][keep] + seq_lens_q = md.qsa_sequence_lengths[:rows][keep] + row_starts = md.qsa_topk_row_starts[:rows][keep] + topk_indices = md.qsa_topk_indices[:q_rows] + scores = qsa_paged_index_scores( + q=q_index[keep], + index_cache=draft.index_cache, + block_table=block_table, + query_positions=q_logical, + request_indices=q_req, + tokens_per_block=tokens_per_block, + compress_ratio=core.compress_ratio, + only_visible_blocks=True, + # One row per context request: no run of rows to share a key gather. + context_rows=False, + ) + selected = torch.empty( + [q_rows, core.token_topk + core.compress_ratio - 1], + dtype=torch.int32, + device=qkv.device, + ) + _selection( + core, + scores, + visible_blocks, + q_logical, + seq_lens_q, + topk_indices, + row_starts, + selected, + _live_score_columns(core, md, batch), + ) + if share: + self._capture_selection(md, topk_indices, visible_blocks, compact=True) + q_att = _packed_q(core, qkv)[keep] + attn = torch.empty([q_rows, core.q_width], dtype=qkv.dtype, device=qkv.device) + if kv.k_scale is not None and core.mapping.world_size == 1: + # One row per context request, then `p` per generation request: the split-K + # kernel for the first, the verify kernel for the rest where it serves. + _sparse_gqa_nvfp4_context( + core, + kv, + None, + q_att, + block_table, + selected, + q_req, + q_logical, + tokens_per_block, + md.num_contexts, + attn, + p, + ) + else: + out = _sparse_gqa( + core, kv, q_att, block_table, selected, q_req, tokens_per_block, q_logical + ) + attn.copy_(torch.reshape(out, [q_rows, core.q_width])) + gate = torch.reshape(gate[keep], [q_rows, core.q_width]) + elif sparse: + kv = draft.kv + _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) + topk_indices = md.qsa_topk_indices[:rows] + visible_blocks = md.qsa_visible_blocks[:rows] + seq_lens_q = md.qsa_sequence_lengths[:rows] + row_starts = md.qsa_topk_row_starts[:rows] + parts = [(0, num_ctx, True)] + if num_ctx < rows: + parts.append((num_ctx, rows, False)) + selected = torch.empty( + [rows, core.token_topk + core.compress_ratio - 1], + dtype=torch.int32, + device=qkv.device, + ) + for lo, hi, context_rows in parts: + scores = qsa_paged_index_scores( + q=q_index[lo:hi], + index_cache=draft.index_cache, + block_table=block_table, + query_positions=logical[lo:hi], + request_indices=req_idx[lo:hi], + tokens_per_block=tokens_per_block, + compress_ratio=core.compress_ratio, + only_visible_blocks=True, + context_rows=context_rows, + ) + _selection( + core, + scores, + visible_blocks[lo:hi], + logical[lo:hi], + seq_lens_q[lo:hi], + topk_indices[lo:hi], + row_starts[lo:hi], + selected[lo:hi], + _live_score_columns(core, md, batch), + ) + if share: + self._capture_selection(md, topk_indices, visible_blocks) + attn = torch.empty([rows, core.q_width], dtype=qkv.dtype, device=qkv.device) + for lo, hi, context_rows in parts: + out = _sparse_gqa( + core, + kv, + _packed_q(core, qkv)[lo:hi], + block_table, + selected[lo:hi], + req_idx[lo:hi], + tokens_per_block, + logical[lo:hi] if context_rows else None, + ) + attn[lo:hi].copy_(torch.reshape(out, [hi - lo, core.q_width])) + else: + attn = thop_attention( + q=qkv, + output=torch.empty([rows, core.q_width], dtype=qkv.dtype, device=qkv.device), + local_layer_idx=layer_idx, + num_heads=core.heads_q, + num_kv_heads=core.heads_kv, + head_size=core.head_dim, + quant_mode=core.quant_mode, + kv_scale_orig_quant=core.kv_scale_orig_quant, + kv_scale_quant_orig=core.kv_scale_quant_orig, + **step, + **_CALL_CONSTANTS, + ) + mtp_fused_sigmoid_mul_inplace(attn, gate) + out = x + if keep is not None: + # The sparse path attended the kept rows only; the dense call wrote every row. + if not sparse: + attn = attn[keep] + x, inject = x[keep], inject[keep] + x = _mtp_tail(core, M, x, attn, inject, pad_rows) + return x if keep is None else out.index_copy_(0, keep, x) + + +# ------------------------------------------------------------------- decode --- + + +class DecodeTarget: + """Generation-only batches (no context rows): one token per sequence, every op inside + the captured decode graph.""" + + def __init__(self, core: Qwen4ExpDefaultModelingV2) -> None: + self.core = core + # See PrefillTarget.__init__. + self._contract_pending = step_contract_enabled() + self._kv_probe = _KvRangeProbe.for_target(core, self._contract_pending) + self._layers = _layer_views(core) + self._stack = _stack_views(core) + # The trunk step's positions buffer (`_position_key`) and its `pos3`, which the first + # draft step of the same forward reads instead of a copy of its own. + self._trunk_pos3: tuple | None = None + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast; see PrefillTarget._check_step_contract.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md: TrtllmAttentionMetadata) -> dict: + """See PrefillTarget.step_args. Context counts are stated, not read: routing + guarantees them, and a decode-only CUDA graph would freeze whatever it read.""" + return _step_args(md, num_contexts=0, num_ctx_tokens=0) + + # --------------------------------------------------------------- blocks --- + + def _attention_prologue(self, proj, q_index, index_view, index_cache, L, qsa, side, stage): + """At one rank, without the fused prologue: the norm + rope of both planes on this stream + and the QSA index branch (the verify step's staging, the raw-key store, then compression) + on the side stream, which reads only the projection's index-key columns and the index + pools; the attention core joins it before reading either.""" + core = self.core + pools = core._bound.pools + num_tokens = proj.shape[0] + req_idx, logical, _, _, _, _, block_table, tokens_per_block, pos3, _, coords, _ = qsa + side.wait_stream(torch.cuda.current_stream()) + qwen4_exp_qk_norm_rope( + proj, + q_index, + pos3, + core._cos_sin, + L.q_norm, + L.k_norm, + L.idx_qn, + num_heads_q=core.heads_q, + num_heads_k=core.heads_kv, + head_dim=core.head_dim, + index_heads=core.index_heads, + index_head_dim=core.index_head_dim, + index_q_col=core.qkv_width + core.q_width, + rotary_dim=core.rotary_dim, + eps=core.eps, + mrope_section1=core.mrope_section1, + mrope_section2=core.mrope_section2, + index_mrope_section1=0, + index_mrope_section2=0, + ) + with torch.cuda.stream(side): + token_k = torch.split( + torch.reshape( + index_view, + [ + num_tokens, + core.index_scratch_width // core.index_head_dim, + core.index_head_dim, + ], + ), + [core.index_heads, core.index_kv_heads, 1], + 1, + )[1] + if stage: + core.spec_rollback.stage_index(index_cache, pools.position_cache) + qsa_paged_index_store( + token_k, + coords, + index_cache, + pools.position_cache, + req_idx, + logical, + block_table, + tokens_per_block, + ) + qsa_prefill_compress( + logical_positions=logical, + request_indices=req_idx, + block_table=block_table, + index_cache=index_cache, + position_cache=pools.position_cache, + k_norm_weight=L.idx_kn, + cos_sin=core._cos_sin, + eps=core.eps, + tokens_per_block=tokens_per_block, + compress_ratio=core.compress_ratio, + mrope_section=_QSA_MROPE_SECTION, + ) + + def _attention_one_rank(self, x, i, L, qkv, index_scratch, attn_out, qsa, step, stage): + """`_attention_block` at one rank, where a layer holds every head: the input + projection's output row is `[q | k | v | gate | index q, k | v slot]`. Over a bf16 or + NVFP4 pool above the dense threshold everything up to the index scores is one launch, + `one_rank_attn_decode_prologue` (over an NVFP4 pool it quantizes the K/V append; on a + single-token step of up to `_PROLOGUE_SCORE_ROWS` rows it writes the index scores too); + otherwise `_attention_prologue`. The sparse path reads q/k/v as views of the + projection, selects with `qsa_topk_expand_decode` where it serves, and fuses the + output gate into the sparse GQA; the dense path takes one compact copy of `[q | k | v]`. + A verify step stages the index slots for `SpeculativeRollback` first.""" + core = self.core + pools = core._bound.pools + ( + req_idx, + logical, + seq_lens_q, + visible_blocks, + topk_indices, + topk_row_starts, + block_table, + tokens_per_block, + pos3, + _, + coords, + sparse, + ) = qsa + num_tokens = x.shape[0] + main = torch.cuda.current_stream() + side = core._side_stream + index_cache = pools.index_cache[i] + kv = pools.kv[i] + rows_per_request = step["predicted_tokens_per_seq"] + fused_prologue = sparse and (kv.k_scale is not None or kv.k.dtype == x.dtype) + if fused_prologue and stage: + side.wait_stream(main) + with torch.cuda.stream(side): + core.spec_rollback.stage_index(index_cache, pools.position_cache) + proj = _dense_mm( + x, L.attn_in, L.attn_in8, core.decode_mm, "in", L.attn_in_g, mid=L.attn_in_m + ) + qkv_view, gate, index_view = torch.split( + proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1 + ) + q_dst, k_dst, v_dst = _attn_in_heads(core, proj) + q_index = _index_query(core, index_scratch) + scores = None + if fused_prologue and not stage and num_tokens <= _PROLOGUE_SCORE_ROWS: + scores = torch.empty( + (num_tokens, block_table.shape[1] * tokens_per_block // core.compress_ratio), + dtype=torch.float32, + device=proj.device, + ) + if fused_prologue: + if stage: + main.wait_stream(side) + one_rank_attn_decode_prologue( + proj, + q_index, + pos3, + coords, + core._cos_sin, + L.q_norm, + L.k_norm, + L.idx_qn, + L.idx_kn, + kv.k, + kv.v, + index_cache, + pools.position_cache, + req_idx, + logical, + block_table, + tokens_per_block, + core.compress_ratio, + num_heads_q=core.heads_q, + num_heads_k=core.heads_kv, + head_dim=core.head_dim, + index_heads=core.index_heads, + index_head_dim=core.index_head_dim, + index_q_col=core.qkv_width + core.q_width, + rotary_dim=core.rotary_dim, + eps=core.eps, + mrope_section1=core.mrope_section1, + mrope_section2=core.mrope_section2, + index_mrope_section1=0, + index_mrope_section2=0, + compress_mrope_section=_QSA_MROPE_SECTION, + k_scale_cache=kv.k_scale, + v_scale_cache=kv.v_scale, + kv_scale_orig_quant=None if kv.k_scale is None else core.kv_scale_orig_quant, + rows_per_request=rows_per_request, + index_scores=scores, + ) + else: + self._attention_prologue(proj, q_index, index_view, index_cache, L, qsa, side, stage) + if self._kv_probe is not None and not self._kv_probe.observe(i, k_dst, v_dst): + self._kv_probe = None + gated = False + if sparse: + if not fused_prologue: + main.wait_stream(side) + _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) + if scores is None: + scores = qsa_paged_index_scores( + q=q_index, + index_cache=index_cache, + block_table=block_table, + query_positions=logical, + request_indices=req_idx, + tokens_per_block=tokens_per_block, + compress_ratio=core.compress_ratio, + only_visible_blocks=True, + context_rows=False, + use_pdl=fused_prologue, + prefetch=fused_prologue and not stage and num_tokens <= _SPARSE_DECODE_MAX_ROWS, + rows_per_request=rows_per_request, + ) + if _topk_expand_supports(scores): + selected = qsa_topk_expand_decode( + scores, + visible_blocks, + logical, + seq_lens_q, + topk_indices, + core.compress_ratio, + core.token_topk, + ) + else: + selected = torch.empty( + [num_tokens, core.token_topk + core.compress_ratio - 1], + dtype=torch.int32, + device=proj.device, + ) + _selection( + core, + scores, + visible_blocks, + logical, + seq_lens_q, + topk_indices, + topk_row_starts, + selected, + ) + if kv.k_scale is not None: + attn = _sparse_gqa( + core, + kv, + q_dst, + block_table, + selected, + req_idx, + tokens_per_block, + None, + output_gate=gate, + rows_per_request=rows_per_request, + ) + gated = True + else: + gated = qsa_paged_sparse_gqa.fuses_output_gate(q_dst, kv.k, selected, gate) + if ( + gated + and num_tokens <= _SPARSE_DECODE_MAX_ROWS + and _sparse_decode_supports(q_dst, kv.k, selected) + ): + attn = qsa_paged_sparse_gqa_decode( + q_dst, + kv.k, + kv.v, + block_table, + selected, + req_idx, + tokens_per_block, + core.softmax_scale, + gate, + ) + else: + attn = qsa_paged_sparse_gqa( + q=q_dst, + k_cache=kv.k, + v_cache=kv.v, + block_table=block_table, + selected_tokens=selected, + request_indices=req_idx, + tokens_per_block=tokens_per_block, + softmax_scale=core.softmax_scale, + query_positions=None, + compress_ratio=None, + output_gate=gate if gated else None, + ) + out = torch.reshape(attn, [num_tokens, core.q_width]) + else: + main.wait_stream(side) + qkv.copy_(qkv_view) + out = thop_attention( + q=qkv, + output=attn_out, + local_layer_idx=pools.pool_idx[i], + num_heads=core.heads_q, + num_kv_heads=core.heads_kv, + head_size=core.head_dim, + quant_mode=core.quant_mode, + kv_scale_orig_quant=core.kv_scale_orig_quant, + kv_scale_quant_orig=core.kv_scale_quant_orig, + **step, + **_CALL_CONSTANTS, + ) + if not gated: + fused_sigmoid_mul_inplace(out, gate) + return _dense_mm(out, L.o, L.o8, core.decode_mm, "out", L.o_g, mid=L.o_m) + + def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step, stage): + """One gated GQA layer plus its QSA index branch for a generation-only batch, at this + rank's heads. + + The four input projections are one GEMM against the fused operand; the norm + rope of + both planes runs on it; one compact copy of its leading `[q | k | v]` feeds the dense + op and the sparse GQA's query heads, the sparse K/V append reads k and v as views of + it, and the index branch reads the raw key from it. The index plane is maintained on + every step, at both ends of the threshold; above it the dense core is replaced and the + main K/V append becomes this layer's own call; over a bf16 or e4m3 pool all of this up + to the index scores is one launch (`qwen4_exp_attn_decode_prologue`), which writes the + packed buffer and the index query itself. A verify step's rows (`stage`) come + `predicted_tokens_per_seq` a request, each scored and attended on its own, and their + index slots are staged for `SpeculativeRollback` before the store. At one rank + `_attention_one_rank` runs instead.""" + core = self.core + if core.mapping.world_size == 1: + return self._attention_one_rank(x, i, L, qkv, index_scratch, attn_out, qsa, step, stage) + ( + req_idx, + logical, + seq_lens_q, + visible_blocks, + topk_indices, + topk_row_starts, + block_table, + tokens_per_block, + pos3, + pos, + coords, + sparse, + ) = qsa + num_tokens = x.shape[0] + proj = _dense_mm( + x, L.attn_in, L.attn_in8, core.decode_mm, "in", L.attn_in_g, mid=L.attn_in_m + ) + qkv_view, gate, _ = torch.split( + proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1 + ) + pools = core._bound.pools + # Over a bf16 or e4m3 pool the sparse path's norm + rope, packed copy, K/V append, + # index store and compression are one launch; it writes the packed buffer itself. + prologue = sparse and pools.kv[i].k_scale is None + if stage: + core.spec_rollback.stage_index(pools.index_cache[i], pools.position_cache) + if prologue: + kv = pools.kv[i] + qwen4_exp_attn_decode_prologue( + proj, + qkv, + _index_query(core, index_scratch), + pos3, + coords, + core._cos_sin, + L.q_norm, + L.k_norm, + L.idx_qn, + L.idx_kn, + kv.k, + kv.v, + pools.index_cache[i], + pools.position_cache, + req_idx, + logical, + block_table, + tokens_per_block, + core.compress_ratio, + core.heads_q, + core.heads_kv, + core.head_dim, + core.index_heads, + core.index_head_dim, + core.index_col, + core.rotary_dim, + core.eps, + core.mrope_section1, + core.mrope_section2, + rows_per_request=step["predicted_tokens_per_seq"], + ) + else: + _, k_dst, v_dst = _attn_in_heads(core, proj) + _attention_rope(core, proj, index_scratch, L, pos3, pos) + _index_plane( + core, proj, i, L, (req_idx, logical, block_table, tokens_per_block, coords) + ) + # The packed `[q | k | v]` as one compact copy: the dense op reads it flat, and + # the sparse GQA reads its query heads from it. + qkv.copy_(qkv_view) + if self._kv_probe is not None: + # Both branches leave the normalized, rotated K/V in the packed copy. + k_rows, v_rows = torch.split( + torch.reshape(qkv, [qkv.shape[0], -1, core.head_dim])[:, core.heads_q :], + core.heads_kv, + 1, + ) + if not self._kv_probe.observe(i, k_rows, v_rows): + self._kv_probe = None + + if sparse: + kv = pools.kv[i] + if not prologue: + _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) + scores = qsa_paged_index_scores( + q=_index_query(core, index_scratch), + index_cache=core._bound.pools.index_cache[i], + block_table=block_table, + query_positions=logical, + request_indices=req_idx, + tokens_per_block=tokens_per_block, + compress_ratio=core.compress_ratio, + only_visible_blocks=True, + context_rows=False, + ) + selected = torch.empty( + [num_tokens, core.token_topk + core.compress_ratio - 1], + dtype=torch.int32, + device=proj.device, + ) + _selection( + core, + scores, + visible_blocks, + logical, + seq_lens_q, + topk_indices, + topk_row_starts, + selected, + ) + attn = _sparse_gqa( + core, + kv, + _packed_q(core, qkv), + block_table, + selected, + req_idx, + tokens_per_block, + None, + ) + out = torch.reshape(attn, [num_tokens, core.q_width]) + else: + out = thop_attention( + q=qkv, + output=attn_out, + local_layer_idx=core._bound.pools.pool_idx[i], + num_heads=core.heads_q, + num_kv_heads=core.heads_kv, + head_size=core.head_dim, + quant_mode=core.quant_mode, + kv_scale_orig_quant=core.kv_scale_orig_quant, + kv_scale_quant_orig=core.kv_scale_quant_orig, + **step, + **_CALL_CONSTANTS, + ) + fused_sigmoid_mul_inplace(out, gate) + return core.tp_sum(_dense_mm(out, L.o, L.o8, core.decode_mm, "out", L.o_g, mid=L.o_m)) + + def _hc_site(self, hc, block_out, inject, norm, first, up, with_inject, normed=None, moe=None): + """One hyper-connection site: the previous block's injection (none at the stack's first + site and after the PLE add), this site's grouped norm and its mixer. Returns `(hc, mixed, + inject)`. At one rank (`hc_fast_sites`) the row count, a per-capture constant, picks the + form: one row runs the CuTe GEMV pair, up to `_HC_DECODE_ROWS_MAX` rows the tcgen05 rows + pair, larger batches `qwen4_exp_hc_inject_norm` and `_hc_mixer`, which every batch takes + at the other layouts. `normed`, when given, is this site's grouped + norm of `hc` already computed (the PLE tail's, no injection pending), which only the + inject-norm route reads (the first kernels of the other two compute it themselves). + + After a MoE block `block_out` is its routed output and `moe` the `(gate_logits, shared, + topk_weights, topk_rows)` `_moe` left uncombined: the site's first kernel forms the + block output itself; with the top-k pair `block_out` is the runner's permuted FC2 rows + and that kernel also forms their finalize.""" + core = self.core + gate_logits, shared, topk_w, topk_rows = moe if moe is not None else (None,) * 4 + n_cols = core.hc_lowrank + (core.hc_count if with_inject else 0) + fast = core.hc_fast_sites + if fast and _HC_DECODE_GEMV_MAX_ROWS < hc.shape[0] <= _HC_DECODE_ROWS_MAX: + hc, low, inject, rstd = qwen4_exp_hc_decode_rows_first( + hc, + block_out, + inject, + norm, + first, + core.hc_count, + n_cols, + core.eps, + shared, + gate_logits, + topk_w, + topk_rows, + ) + mixed, _, _ = qwen4_exp_hc_decode_rows_mix(low, rstd, up, hc, norm, core.hc_count) + return hc, mixed, inject + if fast and hc.shape[0] <= _HC_DECODE_GEMV_MAX_ROWS: + hc, part, ss = qwen4_exp_hc_decode_first_cute( + hc, + block_out, + inject, + norm, + first, + core.hc_count, + n_cols, + shared, + gate_logits, + topk_w, + topk_rows, + ) + mixed, inject = qwen4_exp_hc_decode_mix_cute( + part, ss, up, hc, norm, core.eps, core.hc_count, with_inject + ) + return hc, mixed, inject + if normed is None: + hc, normed = qwen4_exp_hc_inject_norm( + hc, + block_out, + inject, + norm, + core.eps, + core.hc_count, + shared, + gate_logits, + topk_w, + topk_rows, + ) + mixed, inject = _hc_mixer(core, normed, first, up, with_inject) + return hc, mixed, inject + + def _mtp_layer( + self, + embed_tokens, + input_ids, + hidden_states, + md, + step, + draft, + layer_idx, + rows, + gen_seqs, + p, + sparse, + share, + reuse, + pos3, + coords, + ) -> torch.Tensor: + """The draft layer at one rank on the decode machinery: the front end's two projections, + both hyper-connection sites as `_hc_site` (the front end's per-stream `x_h + x_e` is the + first site's injection at unit weight, the attention output the second's), the attention + as `_mtp_attention`, the MoE as `_mtp_moe`. Returns the stream with the MoE output + injected.""" + core = self.core + M, G = core._mtp, core._mtp_gemv8 + gemm = core.decode_mm if md.is_cuda_graph else cublas_mm + e = _draft_embedding(input_ids, embed_tokens) + x_e = _draft_mm( + flashinfer_gemma_rmsnorm(e, M.pre_e_norm, core.eps), M.fc_e, G.get("fc_e"), gemm + ) + _, h_n = qwen4_exp_hc_inject_norm( + hidden_states, None, None, M.pre_h_norm, core.eps, core.hc_count + ) + x_h = torch.reshape( + _draft_mm( + torch.reshape(h_n, [rows * core.hc_count, core.hidden]), M.fc_h, G.get("fc_h"), gemm + ), + [rows, core.stream], + ) + x, mixed, inject = self._hc_site( + x_h, x_e, M.unit_inject[:rows], M.ahc_norm, M.ahc_first, M.ahc_up, True + ) + attn = self._mtp_attention( + mixed, + md, + step, + draft, + layer_idx, + rows, + gen_seqs, + p, + sparse, + share, + reuse, + pos3, + coords, + gemm, + ) + x, mixed, inject = self._hc_site(x, attn, inject, M.mhc_norm, M.mhc_first, M.mhc_up, True) + return qwen4_exp_post_inject(x, self._mtp_moe(mixed), inject) + + def _mtp_attention( + self, + x, + md, + step, + draft, + layer_idx, + rows, + gen_seqs, + p, + sparse, + share, + reuse, + pos3, + coords, + gemm, + ) -> torch.Tensor: + """The draft layer's gated GQA and QSA index branch over the draft manager's planes + (`draft`), as `_attention_one_rank` runs a trunk layer: one input projection, the norm + + rope of both planes in one launch at the module's rope base, the index store and + compression on the side stream, the K/V append, the index scores and the one-launch + top-k + expansion, the gated sparse GQA, the output projection. Under IndexShare a draft + decode step (`reuse`) skips the index branch and rebuilds its selection from the + draft-extend step's capture; the draft-extend step (`share`) captures it.""" + core = self.core + M = core._mtp + dt, dev = x.dtype, x.device + proj = _draft_mm(x, core._mtp_attn_in, core._mtp_gemv8.get("attn_in"), gemm) + qkv_view, gate, index_view = torch.split( + proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1 + ) + q_dst, k_dst, v_dst = _attn_in_heads(core, proj) + index_scratch = torch.empty([rows, core.index_scratch_width], dtype=dt, device=dev) + q_index = _index_query(core, index_scratch) + req_idx = md.qsa_req_idx_per_token[:rows] + logical = md.qsa_logical_positions[:rows] + block_table = md.qsa_block_table + tokens_per_block = md.tokens_per_block + main = torch.cuda.current_stream() + side = core._side_stream + side.wait_stream(main) + qwen4_exp_qk_norm_rope( + proj, + q_index, + pos3, + core._mtp_cos_sin, + M.q_norm, + M.k_norm, + M.idx_qn, + num_heads_q=core.heads_q, + num_heads_k=core.heads_kv, + head_dim=core.head_dim, + index_heads=core.index_heads, + index_head_dim=core.index_head_dim, + index_q_col=core.qkv_width + core.q_width, + rotary_dim=core.rotary_dim, + eps=core.eps, + mrope_section1=core.mrope_section1, + mrope_section2=core.mrope_section2, + index_mrope_section1=0, + index_mrope_section2=0, + ) + if not reuse: + with torch.cuda.stream(side): + token_k = torch.split( + torch.reshape( + index_view, + [ + rows, + core.index_scratch_width // core.index_head_dim, + core.index_head_dim, + ], + ), + [core.index_heads, core.index_kv_heads, 1], + 1, + )[1] + qsa_paged_index_store( + token_k, + coords, + draft.index_cache, + draft.position_cache, + req_idx, + logical, + block_table, + tokens_per_block, + ) + qsa_prefill_compress( + logical_positions=logical, + request_indices=req_idx, + block_table=block_table, + index_cache=draft.index_cache, + position_cache=draft.position_cache, + k_norm_weight=M.idx_kn, + cos_sin=core._mtp_cos_sin, + eps=core.eps, + tokens_per_block=tokens_per_block, + compress_ratio=core.compress_ratio, + mrope_section=_QSA_MROPE_SECTION, + ) + main.wait_stream(side) + gated = False + if sparse: + kv = draft.kv + _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) + topk_indices = md.qsa_topk_indices[:rows] + visible_blocks = md.qsa_visible_blocks[:rows] + seq_lens_q = md.qsa_sequence_lengths[:rows] + if reuse: + slot = draft.share_slot + assert slot in md.qsa_shared_topk_captured_slots, ( + "a draft decode step would reuse a selection that no draft-extend step of " + "this loop captured" + ) + qsa_mtp_reuse_topk( + shared_indices=md.qsa_shared_topk_indices[slot], + shared_visible_blocks=md.qsa_shared_topk_visible_blocks[slot], + request_indices=req_idx, + visible_blocks=visible_blocks, + output=topk_indices, + ) + selected = qsa_expand_block_indices( + topk_indices, + logical, + seq_lens_q, + compress_ratio=core.compress_ratio, + token_topk=core.token_topk, + ) + else: + scores = qsa_paged_index_scores( + q=q_index, + index_cache=draft.index_cache, + block_table=block_table, + query_positions=logical, + request_indices=req_idx, + tokens_per_block=tokens_per_block, + compress_ratio=core.compress_ratio, + only_visible_blocks=True, + context_rows=False, + rows_per_request=p, + ) + if _topk_expand_supports(scores): + selected = qsa_topk_expand_decode( + scores, + visible_blocks, + logical, + seq_lens_q, + topk_indices, + core.compress_ratio, + core.token_topk, + ) + else: + selected = torch.empty( + [rows, core.token_topk + core.compress_ratio - 1], + dtype=torch.int32, + device=dev, + ) + _selection( + core, + scores, + visible_blocks, + logical, + seq_lens_q, + topk_indices, + md.qsa_topk_row_starts[:rows], + selected, + ) + if share: + self._capture_selection(md, topk_indices, visible_blocks, gen_seqs, p) + if kv.k_scale is not None: + attn = _sparse_gqa( + core, + kv, + q_dst, + block_table, + selected, + req_idx, + tokens_per_block, + None, + output_gate=gate, + ) + gated = True + else: + gated = qsa_paged_sparse_gqa.fuses_output_gate(q_dst, kv.k, selected, gate) + attn = qsa_paged_sparse_gqa( + q=q_dst, + k_cache=kv.k, + v_cache=kv.v, + block_table=block_table, + selected_tokens=selected, + request_indices=req_idx, + tokens_per_block=tokens_per_block, + softmax_scale=core.softmax_scale, + query_positions=None, + compress_ratio=None, + output_gate=gate if gated else None, + ) + out = torch.reshape(attn, [rows, core.q_width]) + else: + qkv = torch.empty([rows, core.qkv_width], dtype=dt, device=dev) + qkv.copy_(qkv_view) + out = thop_attention( + q=qkv, + output=torch.empty([rows, core.q_width], dtype=dt, device=dev), + local_layer_idx=layer_idx, + num_heads=core.heads_q, + num_kv_heads=core.heads_kv, + head_size=core.head_dim, + quant_mode=core.quant_mode, + kv_scale_orig_quant=core.kv_scale_orig_quant, + kv_scale_quant_orig=core.kv_scale_quant_orig, + **step, + **_CALL_CONSTANTS, + ) + if not gated: + fused_sigmoid_mul_inplace(out, gate) + return _draft_mm(out, M.o, core._mtp_gemv8.get("o"), gemm) + + def _mtp_moe(self, x) -> torch.Tensor: + """The draft layer's MoE at one rank: the FP8 block-scale routed experts (or with + `_MTP_MOE_NVFP4_ENV` the NVFP4 ones on the trunk's logits-routed runner) and the bf16 + shared expert (its dense projections on `_draft_mm`), combined as `routed + + sigmoid(gate logit) * shared`; the NVFP4 path runs the shared expert on the side + stream.""" + core = self.core + M, G = core._mtp, core._mtp_gemv8 + + def shared_expert(): + gu = _draft_mm(x, M.sh_gu, G.get("sh_gu"), cublas_mm) + shared = _draft_mm( + silu_and_mul_rows(gu, core.shared_inter), M.sh_dn, G.get("sh_dn"), cublas_mm + ) + return shared, cublas_mm(x, M.sh_gate) + + if M.moe_gscale is None: + routed = _mtp_moe_block(core, M, x, routed_only=True) + shared, gate_logits = shared_expert() + return fused_sigmoid_gate_mul_add(routed, gate_logits, shared) + main = torch.cuda.current_stream() + side = core._side_stream + side.wait_stream(main) + with torch.cuda.stream(side): + shared, gate_logits = shared_expert() + routed = _mtp_routed_nvfp4(core, M, x) + main.wait_stream(side) + return fused_sigmoid_gate_mul_add(routed, gate_logits, shared) + + def _hc_moe_site(self, hc, block_out, inject, L: _LayerWeights): + """`_hc_site` at an MoE site (always with injection). Under attention TP (or at one + rank) it also returns the routed experts' NVFP4 input (linear scales at the layer's FC1 + input global scale), written by the mixer's own last launch; under attention DP the + window quantizes the group's gathered rows itself and `fp4` is None. Returns `(hc, + mixed, inject, fp4)`.""" + core = self.core + if core.dp_size > 1: + hc, mixed, inject = self._hc_site( + hc, block_out, inject, L.mhc_norm, L.mhc_first, L.mhc_up, True + ) + return hc, mixed, inject, None + n_cols = core.hc_lowrank + core.hc_count + fast = core.hc_fast_sites + if fast and _HC_DECODE_GEMV_MAX_ROWS < hc.shape[0] <= _HC_DECODE_ROWS_MAX: + hc, low, inject, rstd = qwen4_exp_hc_decode_rows_first( + hc, block_out, inject, L.mhc_norm, L.mhc_first, core.hc_count, n_cols, core.eps + ) + mixed, xq, xsf = qwen4_exp_hc_decode_rows_mix( + low, rstd, L.mhc_up, hc, L.mhc_norm, core.hc_count, L.moe_gscale + ) + return hc, mixed, inject, (xq, xsf) + if fast and hc.shape[0] <= _HC_DECODE_GEMV_MAX_ROWS: + hc, part, ss = qwen4_exp_hc_decode_first_cute( + hc, block_out, inject, L.mhc_norm, L.mhc_first, core.hc_count, n_cols + ) + mixed, inject, xq, xsf = qwen4_exp_hc_decode_mix_fp4_cute( + part, ss, L.mhc_up, hc, L.mhc_norm, core.eps, core.hc_count, True, L.moe_gscale + ) + return hc, mixed, inject, (xq, xsf) + hc, normed = qwen4_exp_hc_inject_norm( + hc, block_out, inject, L.mhc_norm, core.eps, core.hc_count + ) + mixed, inject, fp4 = _hc_moe_mixer(core, normed, L) + return hc, mixed, inject, fp4 + + def _moe(self, x, L: _LayerWeights, pad_rows: int, fp4, fold: bool) -> tuple: + """One decode MoE block, left uncombined for the next site: `(routed, (gate_logits, + shared, topk_weights, topk_rows))`. With `fold` (one rank, where the window is the whole + layer and needs no sum) the runner skips its finalize: `routed` is its permuted FC2 rows + and the top-k pair its routing weights and index map, which the next site's first kernel + gathers; otherwise the pair is None and `routed` the finalized rows (`_moe_block`).""" + core = self.core + router_mm = _decode_router_mid(core, L) + if not fold: + routed, gate_logits, shared = _moe_block( + core, x, L, pad_rows, fp4, router_mm, core.decode_mm + ) + return routed, (gate_logits, shared, None, None) + main = torch.cuda.current_stream() + side = core._side_stream + side.wait_stream(main) + # From `_ROUTER_FIRST_MIN_ROWS` rows the router launches before the shared expert (the + # fork above precedes it, so the side stream does not wait for it). + router_first = x.shape[0] >= _ROUTER_FIRST_MIN_ROWS + if router_first: + logits = router_mm(x, L.router) + with torch.cuda.stream(side): + shared, gate_logits = _shared_expert(core, x, L, core.decode_mm) + if not router_first: + logits = router_mm(x, L.router) + xq, xsf = fp4 + rows, topk_w, topk_rows = fp4_block_scale_moe_renorm_permuted( + logits, + xq, + torch.reshape(xsf, [-1]).view(torch.float8_e4m3fn), + L.fc1_w, + L.fc1_s, + L.fc2_w, + L.fc2_s, + L.out1_scale, + L.out1_gate_scale, + L.out2_scale, + core.num_experts, + core.topk, + core.expert_inter, + ) + main.wait_stream(side) + return rows, (gate_logits, shared, topk_w, topk_rows) + + def _gdn_block(self, x, i, L, si_all, verify, gen_p, captured=False) -> torch.Tensor: + """One Gated-DeltaNet layer over a generation-only batch, one token per sequence, at + this rank's heads: one GEMM for all four input projections (q|k|v, z, a and b are + column views of its output), the conv step over the q|k|v columns in place, then the + recurrence and the gated norm as one streaming launch that computes each (row, + v-head) on its own, so over a bf16 pool a row's output and state do not depend on its + batch. The output projection's per-rank partial is summed over the TP group. + + A verify step (`verify`, `gen_p` rows a sequence) takes `_gdn_verify` and the gated + norm epilogue instead. + + With the in_proj's NVFP4 copy (`_GDN_INPROJ_NVFP4_ENV`) every step of up to + `_INPROJ_NVFP4_MAX_ROWS` rows reads it; with its packed FP8 copy (`_DENSE_FP8_ENV`) the + rows `_dense_mm` gives it. In a captured step (`captured`) those GEMVs run the conv step + as their epilogue -- a single-token step's and, up to `_VERIFY_CONV_FUSED_MAX_ROWS` + rows, a verify step's with its window snapshots and the pool advance -- reading the + pool, taps and slots before their PDL wait, bitwise the projection + conv they + replace.""" + core = self.core + pools = core._bound.pools + num_tokens = x.shape[0] + is_verify = verify is not None + nv = ( + L.in_proj_nv + if L.in_proj_nv is not None and num_tokens <= _INPROJ_NVFP4_MAX_ROWS and _gemv_ready(x) + else None + ) + nv_conv = nv is not None and captured and not is_verify + fp8_conv = ( + not nv_conv + and captured + and not is_verify + and L.in_proj_g is not None + and _gemv_ready(x) + and num_tokens <= _FP8_CONV_FUSED_MAX_ROWS + ) + verify_fused = ( + captured + and is_verify + and (L.in_proj_g is not None or nv is not None) + and _gemv_ready(x) + and num_tokens <= min(_FP8_GEMV_ROWS["in"][1], _VERIFY_CONV_FUSED_MAX_ROWS) + ) + if verify_fused: + verify_op = ( + fp8_gemv_conv1d_verify_op if nv is None else nvfp4_split_gemv_conv1d_verify_op + ) + mq_out, zab = verify_op( + x, + L.in_proj_g if nv is None else nv, + pools.conv[i], + L.conv_w, + si_all, + core._bound.int_conv[i], + verify[0], + core.conv_dim, + gen_p, + _PAD_SLOT_ID, + ) + z, a, b = torch.split(zab, core.in_proj_split[1:], 1) + out = _gdn_verify( + core, None, a, b, i, L, si_all, num_tokens // gen_p, gen_p, verify, mq_out + ) + return self._gdn_epilogue(out, z, L, num_tokens) + if nv_conv: + proj = nvfp4_split_gemv_conv1d_op( + x, nv, pools.conv[i], L.conv_w, si_all, core.conv_dim, _PAD_SLOT_ID + ) + elif fp8_conv: + proj = fp8_gemv_conv1d_op( + x, L.in_proj_g, pools.conv[i], L.conv_w, si_all, core.conv_dim, _PAD_SLOT_ID + ) + elif nv is not None: + proj = nvfp4_split_gemv_op(x, nv) + else: + # FlashInfer's low-M kernels at the rows its table names (a verify step's few rows + # a request), cuBLAS past them; the dense FP8 copies where `_dense_mm` takes them. + proj = _dense_mm( + x, + L.in_proj, + L.in_proj8, + core.decode_in_mm, + "in", + L.in_proj_g, + mid=L.in_proj_m, + ) + if is_verify: + mixed_qkv, z, a, b = torch.split(proj, core.in_proj_split, 1) + out = _gdn_verify( + core, + mixed_qkv, + a, + b, + i, + L, + si_all, + num_tokens // gen_p, + gen_p, + verify, + ) + return self._gdn_epilogue(out, z, L, num_tokens) + if not (nv_conv or fp8_conv): + mixed_qkv = torch.split(proj, core.in_proj_split, 1)[0] + causal_conv1d_update( + x=mixed_qkv, + conv_state=pools.conv[i], + weight=L.conv_w, + bias=None, + silu_activation=True, + cache_seqlens=None, + conv_state_indices=si_all, + pad_slot_id=_PAD_SLOT_ID, + ) + if ( + captured + and core.mapping.world_size == 1 + and pools.ssm[i].dtype == torch.float32 + and num_tokens <= _GDN_DECODE_FUSED_MAX_ROWS + ): + # At one rank over an fp32 pool, the recurrence and the gated norm as one launch + # that reads this layer's pool rows before its PDL wait: in a captured + # single-token step nothing ahead of it writes them. + o = gdn_decode_gated_norm_pdl( + proj, + pools.ssm[i], + si_all, + L.a_log, + L.dt_bias, + L.gnorm, + core.eps, + core.num_k_heads, + core.head_k_dim, + ) + return _dense_mm(o, L.out, L.out8, core.decode_mm, "out", L.out_g, mid=L.out_m) + # It computes both gates and the q/k L2 norm itself, at the softplus pair (1, 20) and + # epsilon 1e-6 this target's constants name. + o = gdn_decode_rows_gated_norm( + proj, + pools.ssm[i], + si_all, + L.a_log, + L.dt_bias, + L.gnorm, + core.eps, + core.num_k_heads, + core.head_k_dim, + ) + return core.tp_sum(_dense_mm(o, L.out, L.out8, core.decode_mm, "out", L.out_g, mid=L.out_m)) + + def _gdn_epilogue(self, out, z, L, num_tokens) -> torch.Tensor: + """A verify step's gated RMS norm (sigmoid gate, plain non-gemma weight) and the + output projection, whose per-rank partial is summed over the TP group.""" + core = self.core + o = rms_norm_gated_token_major( + torch.reshape(out, [-1, core.head_v_dim]), + torch.reshape(z, [num_tokens, core.num_v_heads, core.head_v_dim]), + weight=L.gnorm, + eps=core.eps, + fp8_scale=None, + gate_activation="sigmoid", + ) + o = torch.reshape(o, [num_tokens, core.value_dim]) + return core.tp_sum(_dense_mm(o, L.out, L.out8, core.decode_mm, "out", L.out_g, mid=L.out_m)) + + def _ple_block(self, hc, input_ids, ple_state, next_norm, stage) -> tuple: + """The n-gram side path on its decoder layer, added into the stream and normalized + for the layer's attention-side site. Returns `(hc, normed)`. The same steps as the + prefill target's `_ple_block`, over up to `rows` rows a sequence; every generation + sequence carries its history by construction, so no carry flags are passed. A verify + step (`stage`) stages what `SpeculativeRollback` rewrites the two pools from: both + pools' pre-step rows before either launch publishes, and the conv input the second + launch writes out.""" + core = self.core + S = self._stack + conv_states, ngram_pool, si_all, qsl, rows = ple_state + if stage: + core.spec_rollback.stage_ple_verify(ngram_pool, conv_states, input_ids) + ngram = ple_ngram_embed( + input_ids, + ngram_pool, + si_all, + None, + qsl, + core.w["ple_table"], + S.ple_tscale, + S.ple_mult, + S.ple_vocab, + S.ple_offs, + core.eos_token_id, + rows, + ) + conv_input = torch.empty_like(hc) if stage else None + out = ple_gate_conv_inject_norm( + hc, + cublas_mm(ngram, S.ple_key), + cublas_mm(ngram, S.ple_val), + S.ple_nk_w, + S.ple_nq_w, + S.ple_nc_w, + S.ple_conv_w, + conv_states, + si_all, + qsl, + next_norm, + core.eps, + core.hc_count, + _PLE_CLAMP_MIN, + core.ple_dilation, + _PAD_SLOT_ID, + rows, + conv_input=conv_input, + ) + if stage: + core.spec_rollback.stage_ple_conv_input(conv_input) + return out + + # -------------------------------------------------------------- forward --- + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: torch.IntTensor | None = None, + position_ids: torch.IntTensor | None = None, + **kwargs, + ): + core = self.core + md = attn_metadata + self._check_step_contract(md) + step = self.step_args(md) + spec_metadata = kwargs.get("spec_metadata") + mm = md.mamba_metadata + # `coords` is the `[N, 3]` view of the engine's positions the position plane stores; + # `pos3` the contiguous `[3, N]` copy, built below (by `verify_step_setup` on a + # captured verify step). + coords = torch.transpose(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]), 0, 1) + + h = nn.functional.embedding(input_ids, core.w["embed"]) + num_tokens = h.shape[0] + dt = h.dtype + dev = h.device + batch = md.num_seqs + # Decode CUDA graphs pad with real dummy requests, so they are in the metadata. + assert num_tokens == md.num_tokens, ( + f"the engine handed {num_tokens} activation rows for a generation batch its " + f"metadata describes as {md.num_tokens} tokens" + ) + step["predicted_tokens_per_seq"] = 1 + # Zero rows under a decode CUDA graph, which pads every attention-DP rank to the + # captured batch. + pad_rows = core._dp_pad_rows(md, num_tokens) + # Rows per generation sequence: 1 without speculation, else `runtime_draft_len + 1`. + gen_p, is_verify = _verify_rows(core, md, spec_metadata, batch, num_tokens) + step["predicted_tokens_per_seq"] = gen_p + si_all = mm.state_indices[:batch] + verify = _verify_state(core, mm, batch) if is_verify else None + pools = core._bound.pools + # A captured verify step builds `pos3`, the PLE segmentation vector and the rollback's + # staging in one launch (`verify_step_setup`). + rollback = core.spec_rollback + fused_setup = ( + is_verify + and md.is_cuda_graph + and batch < _VERIFY_SETUP_MAX_SEQS + and position_ids.stride(-1) == 1 + and rollback.index_planes is not None + ) + if fused_setup: + pos3, ple_qsl, staged_slots, staged_pages, staged_within, staged_index = ( + verify_step_setup( + position_ids, + md.seq_lens_cuda[:batch], + si_all, + md.qsa_req_idx_per_token[:num_tokens], + md.qsa_logical_positions[:num_tokens], + md.qsa_block_table, + md.tokens_per_block, + *rollback.index_planes, + ) + ) + else: + pos3 = _copy_positions(position_ids) + # Under speculation `query_start_loc_long` has one row per generation sequence: + # with `gen_p > 1` the PLE segmentation is rebuilt from `seq_lens_cuda`. + ple_qsl = ( + mm.query_start_loc_long + if gen_p == 1 + else cu_seqlens_from_seq_lens(md.seq_lens_cuda[:batch], torch.int64) + ) + pos = pos3[0] + self._trunk_pos3 = (_position_key(position_ids), pos3) + ple_state = (pools.ple_state[0], pools.ple_state[1], si_all, ple_qsl, gen_p) + + # Staging for the speculative rollback (see PrefillTarget.prepare); every row here is + # a generation row. A captured step stages and commits the PLE pools in one launch + # each. + if rollback is not None: + rollback.discard() + if fused_setup: + rollback.begin_staged( + num_gens=batch, + rows_per_gen=gen_p, + slots=staged_slots, + pages=staged_pages, + within=staged_within, + index_all=staged_index, + ) + elif is_verify: + rollback.begin( + num_contexts=0, + prefix_tokens=0, + num_gens=batch, + rows_per_gen=gen_p, + slots=si_all, + request_indices=md.qsa_req_idx_per_token[:num_tokens], + logical_positions=md.qsa_logical_positions[:num_tokens], + block_table=md.qsa_block_table, + tokens_per_block=md.tokens_per_block, + fused=md.is_cuda_graph, + ) + + # Dense/sparse decision once per forward, from a host-side metadata read: a + # per-capture constant matching the engine's short/long decode graph families. + sparse = batch > 0 and int(md.kv_lens_runtime[:batch].max()) > core.dense_threshold + qsa = ( + md.qsa_req_idx_per_token[:num_tokens], + md.qsa_logical_positions[:num_tokens], + md.qsa_sequence_lengths[:num_tokens], + md.qsa_visible_blocks[:num_tokens], + md.qsa_topk_indices[:num_tokens], + md.qsa_topk_row_starts[:num_tokens], + md.qsa_block_table, + md.tokens_per_block, + pos3, + pos, + coords, + sparse, + ) + + # Scratch shared by all attention layers: thop_attention overwrites rows + # [:num_tokens] of `output` in place, and the packed buffers are fully rewritten every + # layer (the QKV one on the dense path only). + qkv = torch.empty([num_tokens, core.qkv_width], dtype=dt, device=dev) + index_scratch = torch.empty([num_tokens, core.index_scratch_width], dtype=dt, device=dev) + attn_out = torch.empty([num_tokens, core.q_width], dtype=dt, device=dev) + + hc_count = core.hc_count + hc = torch.cat([h] * hc_count, -1) + layers, S = self._layers, self._stack + # Each block's output and injection go into the next site's first kernel (`_hc_site`), + # after an MoE block with its combine (and at one rank its finalize, `fold`); before + # the PLE layer the pending block is added here, the PLE tail taking the site's norm. + block_out = inject = moe = None + for i, L in enumerate(layers): + normed = None + if i == core.ple_layer: + if block_out is not None: + if moe is not None: + block_out = fused_sigmoid_gate_mul_add(block_out, moe[0], moe[1]) + hc = qwen4_exp_post_inject(hc, block_out, inject) + block_out = inject = moe = None + hc, normed = self._ple_block(hc, input_ids, ple_state, L.ahc_norm, is_verify) + hc, mixed, inject = self._hc_site( + hc, block_out, inject, L.ahc_norm, L.ahc_first, L.ahc_up, True, normed, moe + ) + if core.is_attn[i]: + block_out = self._attention_block( + mixed, i, L, qkv, index_scratch, attn_out, qsa, step, is_verify + ) + else: + block_out = self._gdn_block(mixed, i, L, si_all, verify, gen_p, md.is_cuda_graph) + hc, mixed, inject, fp4 = self._hc_moe_site(hc, block_out, inject, L) + block_out, moe = self._moe( + mixed, L, pad_rows, fp4, core.moe_fold and i + 1 != core.ple_layer + ) + # The stack-final site with no injection weight: its mix *is* the final norm. + hc, final, _ = self._hc_site( + hc, block_out, inject, S.fhc_norm, S.fhc_first, S.fhc_up, False, moe=moe + ) + if spec_metadata is None: + return final + # `lm_head` reads the `[T, hidden]` mix; the MTP module reads the `[T, stream]` stream + # it was collapsed from. + return final, hc + + # ------------------------------------------------------------- the draft --- + + def _capture_selection(self, md, topk_indices, visible_blocks, batch, p) -> None: + """IndexShare: stash, per request, the selection of its last accepted draft-extend + row with that row's complete-group count, in the draft layer's row of the shared + buffers. Every request carries `p` rows here, and the worker continues the draft loop + from row `num_accepted - 1` of each (the golden token counts as accepted).""" + assert md.qsa_mtp_num_accepted is not None, ( + "IndexShare needs the worker's accepted counts to pick each request's capture row; " + "the draft loop did not supply them" + ) + accepted = md.qsa_mtp_num_accepted[:batch].to(torch.int64) + rows = torch.arange(batch, dtype=torch.int64, device=topk_indices.device) * p + rows = rows + torch.clamp(accepted - 1, min=0, max=p - 1) + slot = self.core._draft_bound.share_slot + md.qsa_shared_topk_indices[slot, :batch].copy_(torch.index_select(topk_indices, 0, rows)) + md.qsa_shared_topk_visible_blocks[slot, :batch].copy_( + torch.index_select(visible_blocks, 0, rows) + ) + # Host-side readiness, recorded in program order: under capture this runs while the + # graph is recorded, so a graph holds the reuse only after the capture it reads. + md.qsa_shared_topk_captured_slots.add(slot) + + def mtp_forward( + self, + embed_tokens: torch.Tensor, + all_rank_num_tokens, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + **kwargs, + ) -> torch.Tensor: + """One draft step over generation rows only: every draft step past the first in an + engine forward (one row a request), and the first (`p` rows a request) when the + batch had no context. Inside the captured decode graph when the trunk step is. + + Under IndexShare the loop's draft-extend step (`p` rows) runs the draft indexer and + captures each request's selection; every draft decode step after it skips the index + branch and rebuilds its selection from that capture with `qsa_mtp_reuse_topk`, which + appends the groups completed since (the expansion adds the incomplete tail, so the + tokens drafted since stay visible). What a skipped step would have written to the + index plane sits at or past the newly committed length, and the next engine step's + draft-extend rewrites it before anything reads it. `all_rank_num_tokens` is the + group's row counts of this draft step, which pad the MoE round trip under attention DP + (`_draft_pad_rows`).""" + core = self.core + assert isinstance(attn_metadata, TrtllmAttentionMetadata) + assert core._bound is not None, "the trunk's cache-manager bind has not run" + md = attn_metadata + rows = hidden_states.shape[0] + assert input_ids.shape[0] == rows and hidden_states.shape[1] == core.stream, ( + f"the draft step got {input_ids.shape[0]} token ids and {tuple(hidden_states.shape)} " + f"hidden states; the MTP module reads one {core.stream}-wide stream row per token" + ) + # Projected on every call: the worker rewrites `attn_metadata` in place between + # draft steps. + step = self.step_args(md) + gen_seqs = md.num_seqs + p, _ = _verify_rows(core, md, None, gen_seqs, rows) + step["predicted_tokens_per_seq"] = p + draft, layer_idx = core._draft_step_bind(md) + assert rows == md.num_tokens, ( + f"the draft step carries {rows} rows for a batch its metadata describes as " + f"{md.num_tokens} tokens; the QSA row mapping covers the second" + ) + pad_rows = _draft_pad_rows(core, all_rank_num_tokens, rows) + # The trunk's dense/sparse decision from the same host buffer: the worker advances + # only the device KV lengths through the draft loop, so every draft step of one engine + # forward takes the branch its trunk took, in the graph family the runner chose. + sparse = gen_seqs > 0 and int(md.kv_lens_runtime[:gen_seqs].max()) > core.dense_threshold + share = core.mtp_index_share and md.qsa_in_mtp_draft_loop + reuse = share and md.qsa_indexer_skip_topk + assert not share or reuse == (p == 1), ( + f"the worker's reuse flag ({md.qsa_indexer_skip_topk}) disagrees with the step's " + f"shape ({p} rows per request): only the draft decode steps after the draft-extend " + "step reuse its selection" + ) + # The first draft step reads the trunk's positions buffer, whose int32 copy the trunk + # made (taken once); a contiguous int32 `position_ids` (every later draft step) is read + # in place. The step's ops only read positions. + trunk, self._trunk_pos3 = self._trunk_pos3, None + if trunk is not None and trunk[0] == _position_key(position_ids): + pos3 = trunk[1] + pos, coords = pos3[0], torch.transpose(pos3, 0, 1) + elif position_ids.dtype == torch.int32 and position_ids.is_contiguous(): + pos3 = torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]) + pos, coords = pos3[0], torch.transpose(pos3, 0, 1) + else: + pos3, pos, coords = _draft_positions(position_ids) + M = core._mtp + req_idx = md.qsa_req_idx_per_token[:rows] + logical = md.qsa_logical_positions[:rows] + block_table = md.qsa_block_table + tokens_per_block = md.tokens_per_block + + if core.mapping.world_size == 1: + return self._mtp_layer( + embed_tokens, + input_ids, + hidden_states, + md, + step, + draft, + layer_idx, + rows, + gen_seqs, + p, + sparse, + share, + reuse, + pos3, + coords, + ) + x, mixed, inject = _mtp_front(core, M, embed_tokens, input_ids, hidden_states) + qkv, gate, k_dst, v_dst = _mtp_qkv(core, M, mixed, pos3) + q_index = None + if not reuse: + q_index = _mtp_index_branch( + core, + M, + mixed, + draft, + (req_idx, logical, block_table, tokens_per_block, pos, coords), + ) + if sparse: + kv = draft.kv + _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) + topk_indices = md.qsa_topk_indices[:rows] + visible_blocks = md.qsa_visible_blocks[:rows] + seq_lens_q = md.qsa_sequence_lengths[:rows] + selected = torch.empty( + [rows, core.token_topk + core.compress_ratio - 1], + dtype=torch.int32, + device=qkv.device, + ) + if reuse: + slot = draft.share_slot + assert slot in md.qsa_shared_topk_captured_slots, ( + "a draft decode step would reuse a selection that no draft-extend step of " + "this loop captured" + ) + qsa_mtp_reuse_topk( + shared_indices=md.qsa_shared_topk_indices[slot], + shared_visible_blocks=md.qsa_shared_topk_visible_blocks[slot], + request_indices=req_idx, + visible_blocks=visible_blocks, + output=topk_indices, + ) + selected.copy_( + qsa_expand_block_indices( + topk_indices, + logical, + seq_lens_q, + compress_ratio=core.compress_ratio, + token_topk=core.token_topk, + ) + ) + else: + scores = qsa_paged_index_scores( + q=q_index, + index_cache=draft.index_cache, + block_table=block_table, + query_positions=logical, + request_indices=req_idx, + tokens_per_block=tokens_per_block, + compress_ratio=core.compress_ratio, + only_visible_blocks=True, + context_rows=False, + ) + _selection( + core, + scores, + visible_blocks, + logical, + seq_lens_q, + topk_indices, + md.qsa_topk_row_starts[:rows], + selected, + ) + if share: + self._capture_selection(md, topk_indices, visible_blocks, gen_seqs, p) + attn = torch.reshape( + _sparse_gqa( + core, + kv, + _packed_q(core, qkv), + block_table, + selected, + req_idx, + tokens_per_block, + None, + ), + [rows, core.q_width], + ) + else: + attn = thop_attention( + q=qkv, + output=torch.empty([rows, core.q_width], dtype=qkv.dtype, device=qkv.device), + local_layer_idx=layer_idx, + num_heads=core.heads_q, + num_kv_heads=core.heads_kv, + head_size=core.head_dim, + quant_mode=core.quant_mode, + kv_scale_orig_quant=core.kv_scale_orig_quant, + kv_scale_quant_orig=core.kv_scale_quant_orig, + **step, + **_CALL_CONSTANTS, + ) + mtp_fused_sigmoid_mul_inplace(attn, gate) + return _mtp_tail(core, M, x, attn, inject, pad_rows) + + +# ------------------------------------------------------------ the MTP module --- + + +class MTPLayer: + """The checkpoint's MTP module, replayed once per draft step by the spec worker: one more + decoder layer (two hyper-connections, gated GQA over its own QSA indexer, a 512-expert + FP8 MoE) behind a two-branch front end (`fc_embedding` over the next token's embedding, + `fc_hidden` over the trunk's residual per stream). It returns the residual un-collapsed; + `shared_head` owns the module's mixer. The step runs on the targets (`mtp_forward`), + routed by the phase each call's metadata shows: the worker rewrites `attn_metadata` + between draft steps, so one engine step's draft calls cross from prefill to decode. + + Not an `nn.Module`: the weights live in the core's `ParameterDict`, and a second + registration would break the weight load's both-directions coverage.""" + + def __init__(self, core: Qwen4ExpDefaultModelingV2, logits_processor) -> None: + self.core = core + self.logits_processor = logits_processor + + def __call__(self, *args, **kwargs) -> torch.Tensor: + return self.forward(*args, **kwargs) + + def shared_head(self, hidden_states, lm_head, attn_metadata, return_context_logits): + """The module's own head: its stack-final hyper-connection mixer (a gated mixer with + its own parameters, not a plain RMS norm) over the grouped norm, then the trunk's + `lm_head` through the shell's logits processor (vocab-parallel and gathered under + attention TP). The norm and the mixer are the trunk's stack-final site's.""" + core = self.core + M = core._mtp + _, normed = qwen4_exp_hc_inject_norm( + hidden_states, None, None, M.fhc_norm, core.eps, core.hc_count + ) + final, _ = _hc_mixer(core, normed, M.fhc_first, M.fhc_up, False) + draft_head = getattr(core, "draft_head", None) + if draft_head is None: + return self.logits_processor.forward( + final, lm_head, attn_metadata, return_context_logits + ) + # At tp1 the FP8 draft head, its logits bf16 when only the draft sampler's argmax reads + # them. + return self.logits_processor.forward( + final, + draft_head, + attn_metadata, + return_context_logits, + widen=not self.logits_processor.draft_argmax_only, + ) + + def forward( + self, + embed_tokens: torch.Tensor, + all_rank_num_tokens, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + **kwargs, + ) -> torch.Tensor: + return self.core._select_target(attn_metadata).mtp_forward( + embed_tokens, + all_rank_num_tokens, + input_ids, + position_ids, + hidden_states, + attn_metadata, + **kwargs, + ) + + +class DraftModel: + """The container the spec worker reaches the drafter through: `mtp_layers` (`[0]`, + MTP-Eagle replays one layer), `embed_tokens` and `lm_head`. `embed_tokens` resolves + through the core on every read: this container is built while parameters are meta + tensors, and materialization replaces the tensor objects.""" + + def __init__(self, core: Qwen4ExpDefaultModelingV2, lm_head, logits_processor) -> None: + self.core = core + self.mtp_layers = [MTPLayer(core, logits_processor)] + self.lm_head = lm_head + + @property + def embed_tokens(self) -> torch.Tensor: + return self.core.w["embed"] + + +class SpeculativeRollback: + """The target's auxiliary speculative-state handler: it undoes what a verify step's + rejected draft rows wrote into the side-cache planes the runtime's own state promotion + does not cover: the PLE n-gram id history and dilated-conv state, each attention layer's + QSA index-K plane and the shared QSA position plane. The Gated-DeltaNet state is not + here: the runtime promotes it. + + For a generation request that accepted `a` rows (golden token included, `1 <= a <= p`) + the commit rewrites the n-gram pool row to the id window of row `a - 1` minus its oldest + id, the conv pool row to columns `[a, a + state_len)` of `[pre-step state | conv input of + the request's p rows]`, and every index-K and position slot written by a row of ordinal + `>= a` to its pre-verify contents. + + The targets stage during the forward (`begin`, then `stage_ple_verify` before the fused PLE + launches and `stage_ple_conv_input` after them, and `stage_index` before each attention + layer's store) and the spec worker calls `commit_speculative_states` after its draft loop, + inside the same model forward. Everything staged is allocated by the forward and consumed + by that forward's commit, so a captured graph holds both ends. A step opened with `fused` + (the decode target's captured steps) stages the PLE pools with one `spec_rollback_stage` + launch and commits with one `spec_rollback_commit` launch.""" + + def __init__(self) -> None: + self._pending: dict | None = None + # (base addresses, page strides, plane count) of every attention layer's index-K + # plane, once the cache manager is bound. + self._planes: tuple | None = None + + def set_index_planes(self, planes: list) -> None: + """Bind the attention layers' index-K planes (static for the manager's lifetime).""" + ptrs, strides = _plane_table(planes) + self._planes = (ptrs, strides, len(planes)) + + # ------------------------------------------------------------- staging --- + + def discard(self) -> None: + """Drop whatever an earlier forward staged.""" + self._pending = None + + @property + def index_planes(self) -> tuple | None: + """`(base addresses, page strides, plane count)` of the bound index-K planes.""" + return self._planes + + def begin_staged( + self, + *, + num_gens: int, + rows_per_gen: int, + slots: torch.Tensor, + pages: torch.Tensor, + within: torch.Tensor, + index_all: torch.Tensor, + ) -> None: + """Open a generation-only verify step's staging from what `verify_step_setup` + resolved: the requests' int64 recurrent slots, the rows' int64 QSA pages and offsets, + and every bound index-K plane at them. The PLE staging and the commit are the + one-launch ones (`begin(..., fused=True)`).""" + n, p = num_gens, rows_per_gen + assert n > 0 and p > 1, ( + f"a verify step stages generation requests carrying drafts; got {n} requests of " + f"{p} rows" + ) + assert pages.shape[0] == n * p and index_all.shape[1] == n * p, ( + f"{n} generation requests of {p} rows are {n * p} rows; got {pages.shape[0]} slots" + ) + self._pending = dict( + num_contexts=0, + prefix=0, + n=n, + p=p, + slots=slots, + pages=pages, + within=within, + staged_layers=0, + position=None, + index_all=index_all, + fused=True, + ngram_stream=None, + ) + + def begin( + self, + *, + num_contexts: int, + prefix_tokens: int, + num_gens: int, + rows_per_gen: int, + slots: torch.Tensor, + request_indices: torch.Tensor, + logical_positions: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + fused: bool = False, + ) -> None: + """Open a verify step's staging. The generation rows follow the `prefix_tokens` + context rows as `num_gens` runs of `rows_per_gen`, so row `s * p + t` is ordinal `t` + of generation request `s`. `slots` are those requests' recurrent slots; + `request_indices` and `logical_positions` the generation rows' own. Each row's QSA + slot is resolved here, as `qsa_paged_index_store` resolves it, and every index plane's + slots are gathered before the trunk writes any of them.""" + n, p = num_gens, rows_per_gen + assert n > 0 and p > 1, ( + f"a verify step stages generation requests carrying drafts; got {n} requests of " + f"{p} rows" + ) + assert logical_positions.shape[0] == n * p, ( + f"{n} generation requests of {p} rows are {n * p} rows; got " + f"{logical_positions.shape[0]} positions" + ) + assert self._planes is not None, "the attention layers' index planes are not bound" + positions = logical_positions.to(torch.long) + pages = block_table[request_indices.to(torch.long), positions // tokens_per_block] + pages = pages.to(torch.long) + within = positions % tokens_per_block + ptrs, strides, count = self._planes + self._pending = dict( + num_contexts=num_contexts, + prefix=prefix_tokens, + n=n, + p=p, + slots=slots.to(torch.long), + pages=pages, + within=within, + staged_layers=0, + position=None, + index_all=qsa_rollback_gather(ptrs, strides, pages, within, count), + fused=fused, + ngram_stream=None, + ) + + def stage_ple_verify( + self, ngram_pool: torch.Tensor, conv_pool: torch.Tensor, input_ids: torch.Tensor + ) -> None: + """A verify step's PLE staging, before either pool is written: both pools, the + generation slots' pre-step conv rows, and the n-gram candidates rebuilt from the slots' + pre-step id history and the rows' ids. Row `t`'s history after the step is the last + `context` ids of `[history | ids of rows 0..t]`.""" + pending = self._pending + prefix, n, p = pending["prefix"], pending["n"], pending["p"] + slots = pending["slots"] + pending["ngram_pool"] = ngram_pool + pending["conv_pool"] = conv_pool + if pending["fused"]: + pending["ngram_stream"], pending["conv_prev"] = spec_rollback_stage( + ngram_pool, conv_pool, slots, input_ids[prefix : prefix + n * p], p + ) + return + pending["conv_prev"] = conv_pool[slots] + context = ngram_pool.shape[1] + ids = torch.reshape(input_ids[prefix : prefix + n * p], [n, p]).to(ngram_pool.dtype) + stream = torch.cat([ngram_pool[slots], ids], 1) + # Window k of the unfold is the history after k rows. + pending["ngram_candidates"] = stream.unfold(1, context, 1)[:, 1:] + + def stage_ple_conv_input(self, conv_input: torch.Tensor) -> None: + """The generation rows' conv input (`[T, C]`, token-major), as the fused PLE launch + wrote it into a buffer the forward allocated for this staging alone. A mixed batch's + rows are copied out, so its context prefix is not kept alive through the draft loop.""" + pending = self._pending + prefix, n, p = pending["prefix"], pending["n"], pending["p"] + conv_x = torch.reshape(conv_input[prefix:], [n, p, -1]) + pending["conv_x"] = conv_x.clone() if prefix else conv_x + + def stage_index(self, index_cache: torch.Tensor, position_cache: torch.Tensor) -> None: + """Before an attention layer's index store: on the step's first attention layer (the + first to write the shared plane) the generation rows' position slots. The index-K + slots of every layer were gathered when the step was opened.""" + pending = self._pending + if pending["position"] is None: + pages, within = pending["pages"], pending["within"] + pending["position"] = (position_cache, position_cache[pages, within]) + pending["staged_layers"] += 1 + + # -------------------------------------------------------------- commit --- + + def commit_speculative_states( + self, + num_accepted_tokens: torch.Tensor, + state_indices: torch.Tensor, + num_contexts: int, + ) -> None: + """Rewrite the staged planes to the accepted prefix. Called by the spec worker.""" + pending, self._pending = self._pending, None + if pending is None: + return + n, p = pending["n"], pending["p"] + assert num_contexts == pending["num_contexts"], ( + f"the worker commits with {num_contexts} context requests and the forward staged " + f"{pending['num_contexts']}" + ) + assert state_indices.shape[0] == num_contexts + n, ( + f"the worker hands {state_indices.shape[0]} slots for a batch the forward staged " + f"as {num_contexts} + {n} requests" + ) + assert pending.get("conv_x") is not None and pending["staged_layers"] > 0, ( + "a verify step staged neither its PLE rows nor any attention layer; the forward " + "did not run the staging it opened" + ) + ptrs, strides, _ = self._planes + position, position_old = pending["position"] + pages, within = pending["pages"], pending["within"] + if pending["fused"]: + spec_rollback_commit( + num_accepted_tokens[num_contexts : num_contexts + n], + state_indices[num_contexts:], + pending["ngram_pool"], + pending["ngram_stream"], + pending["conv_pool"], + pending["conv_prev"], + pending["conv_x"], + position, + position_old, + pages, + within, + ptrs, + strides, + pending["index_all"], + ) + return + accepted = num_accepted_tokens[num_contexts : num_contexts + n].to(torch.long) + slots = state_indices[num_contexts:].to(torch.long) + rows = torch.arange(n, device=accepted.device) + ngram_pool, conv_pool = pending["ngram_pool"], pending["conv_pool"] + state_len = conv_pool.shape[-1] + ngram_pool[slots] = pending["ngram_candidates"][rows, accepted - 1] + # [n, C, state_len + p]; window k of the unfold is the state after k tokens. + history = torch.cat([pending["conv_prev"], torch.transpose(pending["conv_x"], 1, 2)], 2) + conv_pool[slots] = history.unfold(2, state_len, 1)[rows, :, accepted] + # [n * p]: True at every ordinal the sampler rejected. + rejected = torch.reshape( + torch.arange(p, device=accepted.device).unsqueeze(0) >= accepted.unsqueeze(1), [-1] + ) + qsa_rollback_restore(ptrs, strides, pages, within, rejected, pending["index_all"]) + position[pages, within] = torch.where( + rejected[:, None], position_old, position[pages, within] + ) + + +# -------------------------------------------------------------------- shell --- + + +class _WidenedHead: + """`lm_head` whose bf16 logits come back as fp32 from `bf16_to_fp32` (one vectorized + kernel; torch's `.float()` is an unvectorized copy at under half the bandwidth), so the + logits processor's own `.float()` returns them as they are. The same values bit for + bit.""" + + def __init__(self, lm_head) -> None: + self.lm_head = lm_head + + def __call__(self, hidden_states: torch.Tensor) -> torch.Tensor: + logits = self.lm_head(hidden_states) + if logits.dtype == torch.bfloat16 and logits.is_contiguous(): + return bf16_to_fp32(logits) + return logits + + +def _fp8_rowwise(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """`weight` [rows, cols] bf16 -> (e4m3 [rows, cols], fp32 [1, rows]): one scale a row, + amax / 448, so each row's largest magnitude maps to e4m3's largest finite value.""" + scale = torch.empty(weight.shape[0], dtype=torch.float32, device=weight.device) + w8 = torch.empty_like(weight, dtype=torch.float8_e4m3fn) + for lo in range(0, weight.shape[0], 16384): + hi = min(lo + 16384, weight.shape[0]) + chunk = weight[lo:hi].float() + scale[lo:hi] = (chunk.abs().amax(dim=1) / 448.0).clamp(min=1e-12) + w8[lo:hi] = (chunk / scale[lo:hi, None]).to(torch.float8_e4m3fn) + return w8, scale.reshape(1, -1).contiguous() + + +def _head_mid(weight_fp8: torch.Tensor, weight_scale: torch.Tensor) -> MidWeight: + """`mid_rows_mm`'s view of a per-row-scaled e4m3 head copy (`_fp8_rowwise`).""" + return MidWeight(weight_fp8, weight_scale.reshape(-1)) + + +class _Fp8Head: + """A vocabulary projection over an FP8 copy of `lm_head`: up to `_FP8_GEMV_HEAD_ROWS` rows on + `fp8_gemv` over the packed block-scaled copy -- or, given the NVFP4 copy `nv`, `nvfp4_gemv` + over it --, up to `_MID_HEAD_ROWS` on `mid_rows_mm` over the per-row-scaled copy, more rows as + per-token e4m3 activations times that copy, scaled in cuBLASLt's epilogue; bf16 out + (`_WidenedHead` widens it as it does the bf16 head's). The draft head always (at tp1), the + verify head under the dense FP8 mode.""" + + def __init__(self, weight_fp8: torch.Tensor, weight_scale: torch.Tensor, packed, nv) -> None: + self.weight_fp8, self.weight_scale, self.packed, self.nv = ( + weight_fp8, + weight_scale, + packed, + nv, + ) + self.mid = _head_mid(weight_fp8, weight_scale) + + def __call__(self, hidden_states: torch.Tensor) -> torch.Tensor: + rows = hidden_states.shape[0] + if rows <= _FP8_GEMV_HEAD_ROWS and _gemv_ready(hidden_states): + if self.nv is not None: + return nvfp4_gemv_op(hidden_states, self.nv) + if self.packed is not None: + return fp8_gemv_op(hidden_states, self.packed) + if rows <= _MID_HEAD_ROWS and _gemv_ready(hidden_states): + return mid_rows_mm_op(hidden_states, self.mid) + x8, xs = torch.ops.tensorrt_llm.quantize_e4m3_activation(hidden_states) + return torch._scaled_mm( + x8, + self.weight_fp8.t(), + scale_a=xs.float().reshape(-1, 1), + scale_b=self.weight_scale, + out_dtype=hidden_states.dtype, + ) + + +class _WideningLogitsProcessor(LogitsProcessor): + """The shell's logits processor (its row gather, its fp32 logits) over `_WidenedHead`, at + tp1; with `verify_head` set (the dense FP8 mode) the target's own head is that `_Fp8Head`. + + `widen=False` (with `return_context_logits`, the speculative paths) returns the head's own + bf16 logits: for a consumer that only takes their row argmax, which is the same token over + bf16 as over their exact fp32 widening (same order, ties, NaN and +-0).""" + + verify_head = None + #: Set by the shell's forward each engine step: the step's draft logits feed only the + #: greedy draft sampler's argmax (`_spec_argmax_only`). + draft_argmax_only = False + + def forward( + self, + hidden_states, + lm_head, + attn_metadata, + return_context_logits: bool = False, + widen: bool = True, + ) -> torch.Tensor: + if self.verify_head is not None and not isinstance(lm_head, _Fp8Head): + lm_head = self.verify_head + if not widen and return_context_logits: + return lm_head(hidden_states) + return super().forward( + hidden_states, _WidenedHead(lm_head), attn_metadata, return_context_logits + ) + + +def _spec_argmax_only(spec_worker, spec_metadata) -> tuple[bool, bool]: + """`(verify, draft)`: whether this step's verify / draft logits are consumed only by a row + argmax, so they can stay bf16. + + Both need no guided decoder (it edits the logits). The draft logits then feed only the draft + sampler, an argmax unless rejection sampling runs on a non-greedy batch. The verify logits + are also the step's `logits` output and the occurrence penalties' and relaxed acceptance's + input; the one-model speculative sampler admits no request asking for logits or log + probabilities, so they stay bf16 only for an all-greedy batch with the penalties off (a + deploy-time flag) and strict acceptance. Every input is a deploy-time constant or + `is_all_greedy_sample`, part of the CUDA graph key.""" + if getattr(spec_worker, "guided_decoder", None) is not None: + return False, False + draft = not spec_metadata.wants_advanced_draft_sampling + verify = ( + spec_metadata.is_all_greedy_sample + and not getattr(spec_metadata, "enable_penalty", False) + and not getattr( + getattr(spec_worker, "spec_config", None), "use_relaxed_acceptance_for_thinking", False + ) + ) + return verify, draft + + +class ModelingV2Qwen38FlashNextNvfp4Sm103Default( + DecoderModelForCausalLM[Qwen4ExpDefaultModelingV2, PretrainedConfig] +): + """Qwen3.8-Flash-Next-NVFP4 at tp1, tp4, tep4, dep4 or dep2 on sm_103.""" + + # At tp1 one-engine speculative decoding takes the engine's fast paths (greedy and + # temperature-only sampling CUDA graph variants, the overlap loop on the execution stream, + # hoisted position ids, the MTP-Eagle draft loop's step-0 epilogue, the recorded steady + # decode-step prepare), validated there; off at every other layout. + spec_decode_fast_paths: bool = False + runtime_hooks: RuntimeHooks | None = None + + @classmethod + def get_preferred_kv_cache_manager_version(cls, pretrained_config=None) -> str: + """Prefer KV cache manager V2: PLE state (`ple_layer_cache`) and QSA's index and + position planes exist only there.""" + return "V2" + + def __init__(self, model_config: ModelConfig): + cfg = model_config.pretrained_config + assert cfg is not None + # Adopt the engine-normalized text tower (`Qwen4ExpTextConfig`) as the model's config: + # everything that sizes the pools reads `model_engine.model.model_config. + # pretrained_config`. The swap is in place because the engine keeps its own reference + # to this `ModelConfig`. + text = getattr(cfg, "text_config", None) + assert text is not None, ( + "this target is the language tower of the composite Qwen4-Exp export; " + "pretrained_config carries no text_config" + ) + assert getattr(text, "architectures", None) == ["Qwen4ExpForCausalLM"], ( + f"the text tower must be the normalized Qwen4-Exp text config; got architectures " + f"{getattr(text, 'architectures', None)!r}" + ) + composite = cfg + model_config.pretrained_config = text + super().__init__( + Qwen4ExpDefaultModelingV2(model_config, composite), + config=model_config, + hidden_size=text.hidden_size, + vocab_size=text.vocab_size, + ) + # `lm_head` is outside the quantization scheme: a dense bf16 `LMHead`. + assert self.lm_head.weight.dtype == torch.bfloat16, ( + f"expected a bf16 lm_head; the shell built {self.lm_head.weight.dtype}" + ) + self.one_rank = model_config.mapping.world_size == 1 + if self.one_rank: + self.logits_processor = _WideningLogitsProcessor() + # No separate draft checkpoint: the MTP module's weights are the target's own. + self.draft_config = None + # The speculative branch, built only when `speculative_config` asked for it. + self.spec_config = getattr(model_config, "spec_config", None) + self.draft_model = None + self.spec_worker = None + if self.spec_config is not None: + self.draft_model = DraftModel(self.model, self.lm_head, self.logits_processor) + # Must match what the resource manager decided: left at `get_spec_worker`'s + # default, the draft step would be handed a manager that maps no draft layer. + self.use_separate_draft_kv_cache = should_use_separate_draft_kv_cache(self.spec_config) + if self.one_rank: + self.spec_decode_fast_paths = True + self.runtime_hooks = RuntimeHooks(loop_on_execution_stream=True) + self.spec_worker = get_spec_worker( + self.spec_config, + model_config, + model_config.mapping, + use_separate_draft_kv_cache=self.use_separate_draft_kv_cache, + spec_decode_fast_paths=self.spec_decode_fast_paths, + ) + assert self.spec_worker is not None, ( + f"the engine resolved no spec worker for spec_dec_mode " + f"{self.spec_config.spec_dec_mode!r}" + ) + self.spec_worker.set_draft_model(self.draft_model) + # The PLE pools and the QSA side caches are written per verify row and nothing the + # runtime owns rolls them back; the worker calls this after its draft loop, inside + # the same forward. + self.spec_worker.register_auxiliary_state_handler(self.model.spec_rollback) + # Media placeholder ids for the executor; built lazily (`__init__` runs under meta + # init). + self._mm_token_ids_value = (self.model.image_token_id, self.model.video_token_id) + self._mm_token_ids_cache: torch.Tensor | None = None + # Replicated on every rank; built once the weights are real (`post_load_weights`). + self.vision_tower: _vision.VisionTower | None = None + + @property + def multimodal_token_ids(self) -> torch.Tensor: + """Placeholder ids whose rows vision embeddings replace. This checkpoint's + placeholders are in vocabulary, so the executor's out-of-vocabulary fallback would + select nothing.""" + if self._mm_token_ids_cache is None: + self._mm_token_ids_cache = torch.tensor(self._mm_token_ids_value, dtype=torch.int32) + return self._mm_token_ids_cache + + @property + def multimodal_data_device_paths(self) -> list[str]: + """Tensors under `multimodal_data` the executor moves to the device. Absent, it moves + everything, putting the `(t, h, w)` grids the tower reads on the host there too.""" + return list(_vision.DEVICE_PATHS) + + def validate_request(self, request) -> None: + """Admission check, run by the executor on every new request: a media request the + vision tower would assert on fails here as its own error (`vision.validate_request`); + a text request returns at once.""" + _vision.validate_request(self.vision_tower, request) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: torch.IntTensor | None = None, + position_ids: torch.IntTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + return_context_logits: bool = False, + spec_metadata=None, + lora_params: dict | None = None, + **kwargs, + ): + """One engine step: the core's `begin_step`, then the base forward with every + argument it was handed (`resource_manager` stays in `**kwargs`). + + With a spec worker the shell owns the logits gather (one-model modes gather the rows + the sampler verifies, `spec_metadata.gather_ids`) and hands the worker the trunk's + ungathered stream rows, which the MTP module reads; `position_ids` goes in the + engine's layout.""" + self.model.begin_step( + attn_metadata, + input_ids, + position_ids, + inputs_embeds=inputs_embeds, + lora_params=lora_params, + spec_metadata=spec_metadata, + vision=_vision.multimodal_embeds(self.vision_tower, kwargs), + ) + if self.spec_worker is None: + assert spec_metadata is None, ( + "spec_metadata arrived without a spec worker; this engine was built without a " + "speculative_config" + ) + return super().forward( + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + return_context_logits, + spec_metadata, + lora_params, + **kwargs, + ) + assert spec_metadata is not None, "the spec worker is built but no spec_metadata came" + final, hc = self.model.split_spec_output( + self.model( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + lora_params=lora_params, + spec_metadata=spec_metadata, + **kwargs, + ) + ) + if isinstance(self.logits_processor, _WideningLogitsProcessor): + verify_argmax_only, self.logits_processor.draft_argmax_only = _spec_argmax_only( + self.spec_worker, spec_metadata + ) + logits = self.logits_processor.forward( + nn.functional.embedding(spec_metadata.gather_ids, final), + self.lm_head, + attn_metadata, + True, + widen=not verify_argmax_only, + ) + else: + logits = self.logits_processor.forward( + nn.functional.embedding(spec_metadata.gather_ids, final), + self.lm_head, + attn_metadata, + True, + ) + # A prefill CUDA graph pads the step's rows; the worker reshapes generation rows per + # request, so it takes the batch's own. + if attn_metadata.padded_num_tokens is not None: + n = attn_metadata.num_tokens + input_ids, hc = input_ids[:n], hc[:n] + position_ids = position_ids[..., :n] + return self.spec_worker( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hc, + logits=logits, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=self.draft_model, + resource_manager=kwargs.get("resource_manager"), + ) + + def load_weights(self, weights, *args, **kwargs): + _weights.MODEL_WEIGHTS.load(self, weights) + + def post_load_weights(self): + super().post_load_weights() + self.model.post_load_weights() + self.vision_tower = _vision.VisionTower(self.model.w, self.model.vision_config) + if self.one_rank: + self._build_heads(self.lm_head.weight.data) + + def _build_heads(self, weight: torch.Tensor) -> None: + """The tp1 heads' copies of `weight`: the draft head's FP8 copies (per row, and packed + for `fp8_gemv` unless its NVFP4 copy replaces it, `_MTP_HEAD_NVFP4_ENV`), and under the + dense FP8 mode the verify head over the same per-row copy with its packed copy or its + NVFP4 one (`_TARGET_HEAD_NVFP4_ENV`); every decode variant compiled now.""" + dense_fp8 = _switch(_DENSE_FP8_ENV) + draft = self.model.mtp_enabled + if not (draft or dense_fp8): + return + rowwise = _fp8_rowwise(weight) + draft_nv = draft and _switch(_MTP_HEAD_NVFP4_ENV) + target_nv = dense_fp8 and _switch(_TARGET_HEAD_NVFP4_ENV) + packed = ( + pack_from_bf16(weight) + if (draft and not draft_nv) or (dense_fp8 and not target_nv) + else None + ) + nv = None + if draft_nv or target_nv: + nv = nvfp4_pack_from_bf16(weight) + if weight.is_cuda: + nvfp4_gemv_precompile(nv) + if weight.is_cuda: + mid_rows_precompile(_head_mid(*rowwise), _FP8_GEMV_HEAD_ROWS + 1, _MID_HEAD_ROWS) + if draft: + self.model.draft_head = _Fp8Head(*rowwise, packed, nv if draft_nv else None) + if dense_fp8: + self.logits_processor.verify_head = _Fp8Head( + *rowwise, packed, nv if target_nv else None + ) + if target_nv: + print( + f"target head NVFP4 ({_TARGET_HEAD_NVFP4_ENV}=1): decode rows <= " + f"{_FP8_GEMV_HEAD_ROWS} on nvfp4_gemv over lm_head {tuple(nv.data.shape)}", + flush=True, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/__init__.py new file mode 100644 index 000000000000..777b6b725c00 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ops used only by this target; the family's shared ops are in `models/qwen4_exp/ops`.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/__init__.py similarity index 74% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/__init__.py index 8fd2a7f0c45b..52a7a9daf028 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/__init__.py @@ -1,3 +1,2 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Qwen3.8-Flash-Next-NVFP4 / sm_103 / tp1.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_fp8_pages_to_bf16__prefill.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_fp8_pages_to_bf16__prefill.py new file mode 100644 index 000000000000..e993fd24eae7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_fp8_pages_to_bf16__prefill.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A batch's pages of an e4m3 paged K/V pool, widened into a bf16 scratch pool, and the block +table that addresses them there. + +Tier 5 (Triton written for this target): the target's context launches over a kv_fp8 pool +read their requests' pages through this scratch with the bf16 row-tile kernel. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Elements a program widens per loop step. +_CHUNK = 4096 + +#: One rank's K/V geometry and the engine's QSA table at max_seq_len 262144 +#: (`ceil(262145 / 32)` columns padded to a multiple of 4): a dep4 rank's two KV heads and +#: 32-row table; `slot_pages` role pages of one lifecycle slot, layer `layer`'s K and V +#: planes among them. +_CHECKPOINT = dict( + kv_heads=2, + head_dim=256, + tokens_per_block=32, + width=8196, + slots=32, + slot_pages=24, + layer=11, + holes=False, + straddle=False, +) + + +def _spec(kv_lens, **overrides) -> dict: + """`kv_lens` is one KV length per listed request, in order.""" + return {**_CHECKPOINT, "kv_lens": tuple(kv_lens), **overrides} + + +@triton.jit +def _fp8_pages_to_bf16_kernel( + k_cache, + v_cache, + block_table, + slots, + counts, + offsets, + table, + k_out, + v_out, + num_pool_pages, + listed, + stride_page, + bt_stride, + table_stride, + PAGE_ELEMS: tl.constexpr, + CHUNK: tl.constexpr, +): + request = tl.program_id(0) + count = tl.load(counts + request) + base = tl.load(offsets + request).to(tl.int64) + slot = tl.load(slots + request).to(tl.int64) + for page in range(tl.program_id(1), count, tl.num_programs(1)): + dst = base + page + if dst < listed: + src = tl.load(block_table + slot * bt_stride + page).to(tl.int64) + resident = (src >= 0) & (src < num_pool_pages) + tl.store(table + slot * table_stride + page, tl.where(resident, dst, -1).to(tl.int32)) + safe = tl.where(resident, src, 0) + for start in tl.static_range(0, PAGE_ELEMS, CHUNK): + elems = start + tl.arange(0, CHUNK) + k = tl.load(k_cache + safe * stride_page + elems, mask=resident, other=0.0) + v = tl.load(v_cache + safe * stride_page + elems, mask=resident, other=0.0) + tl.store(k_out + dst * PAGE_ELEMS + elems, k.to(tl.bfloat16)) + tl.store(v_out + dst * PAGE_ELEMS + elems, v.to(tl.bfloat16)) + + +#: Programs per request; each strides the request's pages. +_PAGE_PROGRAMS = 1024 # capacity: fixed: a launch width; a request may list any page count + + +class QsaFp8PagesToBf16(OpWrapper): + """K and V of listed requests' pages of an e4m3 HND paged pool, as bf16 scratch pools, + and the block table that addresses them there. + + Request `r` lists its first `counts[r]` table pages, `block_table[slots[r], j]` for `j < + counts[r]`, at scratch pages `offsets[r] + j` (those below `listed`, the scratch's page + count):: + + k_out[offsets[r] + j] = bf16(k_cache[block_table[slots[r], j]]) + + and `v_out` likewise. Each e4m3 code widens exactly to bf16 (no scale: the pool holds + K/V at unit scale, as `qsa_paged_sparse_gqa` reads it), so a bf16 kernel over the scratch + reads the values the e4m3 kernels read. `table[slots[r], j]` becomes `offsets[r] + j`, + or -1 where the pool's entry is unmapped or out of the pool (whose scratch page is then + zeros); the call writes no other entry of `table`, which the caller fills with -1. + Returns `(k_out, v_out, table)`. + + Fusion boundary: the listing's widening copy and table; choosing the requests and their + page counts (and `offsets`, their exclusive prefix sum) is the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, one page", spec=_spec((17,))), + Cell(why="dep4 cold context: one 8192-token request", spec=_spec((8192,))), + Cell( + why="dep4 deep chunk: one request at 198160 tokens (6193 pages) over a pool " + "straddling 2^31 elements of the trunk plane", + spec=_spec((198160,), straddle=True), + ), + Cell( + why="three requests, a chunked continuation and two short ones", + spec=_spec((40000 + 3192, 3000, 900)), + ), + Cell(why="holes: unmapped and out-of-pool entries", spec=_spec((4000, 2000), holes=True)), + Cell( + why="a tp4 / tep4 rank's single KV head, [128, 8196] table", + spec=_spec((9000, 2100), kv_heads=1, slots=128), + ), + ) + + note = """ + Metadata consumed: none -- every input is an argument. Call site: the target's + context/mixed forward over a kv_fp8 pool, once per sparse attention layer, for the + step's context requests (`_fp8_context_plan`: slots from the context rows' request + indices, counts `ceil(kv_len / tpb)` from the device KV lengths, `listed` from the host + ones); the row-tile kernel then reads the scratch through `table`. + + One Triton kernel, `(requests, min(width, 1024))` programs, each striding its request's + pages; a page (all KV heads, contiguous in the role plane) is copied in 4096-element + steps, K and V together. Page ids are widened to int64 before the stride product: the + engine's trunk planes reach 6.7e10 elements (dep4 kv_fp8 171362 pages x 24 role pages). + + Preconditions: e4m3 role planes `[pages, Hkv, tpb, D]` with a page's heads, tokens and + dims contiguous (`stride(1) == tpb * D`, `stride(2) == D`, unit last stride; any page + stride), K's and V's at equal strides; `k_out` / `v_out` contiguous bf16 `[listed, Hkv, + tpb, D]`; `table` int32 of `block_table`'s shape. Bitwise deterministic (a copy). + """ + + def __call__( + self, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + slots: torch.Tensor, + counts: torch.Tensor, + offsets: torch.Tensor, + table: torch.Tensor, + k_out: torch.Tensor, + v_out: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + requests, listed = slots.shape[0], k_out.shape[0] + _, heads, tpb, dim = k_cache.shape + if k_cache.dtype != torch.float8_e4m3fn or v_cache.dtype != k_cache.dtype: + raise ValueError("the widening copy reads e4m3 K/V planes") + if k_cache.stride() != v_cache.stride() or k_cache.stride()[1:] != (tpb * dim, dim, 1): + raise ValueError("K/V planes must share strides with contiguous head pages") + if not (k_out.is_contiguous() and v_out.is_contiguous()): + raise ValueError("the bf16 scratch pools must be contiguous") + if k_out.dtype != torch.bfloat16 or k_out.shape[1:] != k_cache.shape[1:]: + raise ValueError("the scratch pools must be bf16 pages of the pool's geometry") + page_elems = heads * tpb * dim + if requests and listed: + _fp8_pages_to_bf16_kernel[(requests, min(block_table.shape[1], _PAGE_PROGRAMS))]( + k_cache, + v_cache, + block_table, + slots, + counts, + offsets, + table, + k_out, + v_out, + k_cache.shape[0], + listed, + k_cache.stride(0), + block_table.stride(0), + table.stride(0), + PAGE_ELEMS=page_elems, + CHUNK=min(_CHUNK, page_elems), + num_warps=8, + ) + return k_out, v_out, table + + def reference( + self, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + slots: torch.Tensor, + counts: torch.Tensor, + offsets: torch.Tensor, + table: torch.Tensor, + k_out: torch.Tensor, + v_out: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + k_ref, v_ref, table_ref = k_out.clone(), v_out.clone(), table.clone() + listed, pool_pages = k_out.shape[0], k_cache.shape[0] + for r in range(slots.shape[0]): + slot, base = int(slots[r]), int(offsets[r]) + count = max(0, min(int(counts[r]), block_table.shape[1], listed - base)) + if count == 0: + continue + src = block_table[slot, :count].long() + resident = (src >= 0) & (src < pool_pages) + dst = torch.arange(base, base + count, device=src.device) + table_ref[slot, :count] = torch.where(resident, dst, -1).to(table.dtype) + safe = torch.where(resident, src, 0) + keep = resident[:, None, None, None] + k_ref[base : base + count] = torch.where(keep, k_cache[safe].float(), 0.0).to( + k_out.dtype + ) + v_ref[base : base + count] = torch.where(keep, v_cache[safe].float(), 0.0).to( + v_out.dtype + ) + return k_ref, v_ref, table_ref + + +qsa_fp8_pages_to_bf16 = QsaFp8PagesToBf16() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_rowtile_attention_cute__prefill.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_rowtile_attention_cute__prefill.py new file mode 100644 index 000000000000..4fa6fef27d45 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_rowtile_attention_cute__prefill.py @@ -0,0 +1,1386 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The row-tile pass of `qsa_sparse_gqa_rowtile` on SM103 tensor cores (CuTe DSL, tcgen05). + +Inputs and output: per group of `rows_per_tile` rows of one request, the plan's packed +slot list (`slots[group]`, `listed` entries), each row's bitmap re-indexed onto it +(`packed_bits`, as int32 halves) and `group_plan[group] = (tiles, listed)`; groups with +`tiles < 0` belong to the lone-row pass. + +A persistent kernel: one CTA per SM, each taking work items -- a (group, KV head) pair, 128 +MMA rows = `rows_per_tile` query rows x `group_size` heads -- from a global ticket counter, +longest-first (the plan's last group first). 9 warps: + +* warp 5 lane 0 draws the tickets and publishes the items through a shared-memory ring that + every role reads; warps 5-8 gather each item's Q rows (into a 60 KB Q buffer) and its list's + K and V rows (64 a tile) with `cp.async` into one ring of `kv_slots` 32 KB SW128 slots, filled + in the MMA warp's order K0, K1, V0, K2, V1, ...; the next item's Q is fetched while the + current item's last tiles load; +* warp 4 issues the tcgen05 MMAs: S[t % 2] = Q K(t)^T (Q from tensor memory), then + O += P(t - 1) V(t - 1) (P from tensor memory, overwriting the S it came from), `t` counting + the CTA's tiles across items; +* warps 0-3, one thread per MMA row, per tile load S with its row max (`tcgen05.ld.red`), + rescale O when that max grew by more than `rescale_threshold` in log2 units (lazily; + otherwise exponentiate against the stale max), write P = 2^(s - max) on the row's own bits + and 0 elsewhere. After an item's last tile they copy the next item's Q from the Q buffer + into tensor memory (so its first QK MMAs overlap this epilogue), then normalize O, stage it + in the Q buffer as four 64-dim SW128 boxes of the item's G heads x R rows and store it with + TMA (thread 0 hands the buffer back to the loaders once the store has read it). + +The Q buffer is shared by Q and the staged O: a softmax thread's Q row and its O row occupy +the same bytes, so the next item's Q is read out before this item's O is staged over it. An +item followed by a lone (or empty) group stores O row by row from registers instead, as the +loaders may then fetch the following item's Q into the buffer at any time. + +Tensor memory: O in columns 0-255, S/P buffers at 256 and 320, Q at 384 (512 columns). +The tile's row max covers every listed entry, an upper bound of the row's own: a valid +softmax offset. P is rounded to bf16 for the PV MMA and summed in fp32 unrounded. +""" + +from __future__ import annotations + +import functools + +import cuda.bindings.driver as cuda_drv +import cutlass +import cutlass.cute as cute +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +import cutlass.pipeline as pipeline +import cutlass.utils.blackwell_helpers as sm100_utils +import torch +from cutlass._mlir.dialects import nvvm as _nvvm +from cutlass.cute.nvgpu import OperandMajorMode, cpasync +from cutlass.cute.runtime import make_ptr +from cutlass.cute.typing import Float32, Int32, Int64 +from cutlass.cutlass_dsl import dsl_user_op + +try: + from cutlass.memory import SmemAllocator +except ImportError: # DSL releases before the memory namespace split + from cutlass.utils import SmemAllocator + +BF16 = cutlass.BFloat16 +_LOG2_E = 1.4426950408889634 +#: The head dimension the tensor-memory plan is laid out for. +HEAD_DIM = 256 +#: MMA rows per CTA and tokens per tile (one 64-bit bitmap word). +BLOCK_M = 128 +BLOCK_N = 64 +#: K/V ring slots (32 KB each); with the 60 KB Q buffer, about 220 KB of shared memory. +KV_SLOTS = 5 +#: Work items in flight in the shared-memory schedule ring. +SCHED_SLOTS = 8 +#: The next item's Q is fetched once this many tiles of the current item remain to load. +Q_LEAD = 3 + + +@dsl_user_op +def tcgen05_fence_before(*, loc=None, ip=None): + """tcgen05.fence::before_thread_sync: order this thread's prior tcgen05 ops before a thread sync.""" + _nvvm.tcgen05_fence(_nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC, loc=loc, ip=ip) + + +@dsl_user_op +def tcgen05_fence_after(*, loc=None, ip=None): + """tcgen05.fence::after_thread_sync: order later tcgen05 ops after a thread sync.""" + _nvvm.tcgen05_fence(_nvvm.Tcgen05FenceKind.AFTER_THREAD_SYNC, loc=loc, ip=ip) + + +class _RowtileAttentionKernel: + def __init__( + self, + group_size: int, + rows_per_tile: int, + head_dim: int = 256, + tokens_per_block: int = 32, + kv_slots: int = KV_SLOTS, + rescale_threshold: float = 8.0, + use_pdl: bool = True, + ): + assert head_dim == 256, ( + "the TMEM plan (O 256 + 2 x S 64 + Q 128 columns) is laid out for D = 256" + ) + assert group_size * rows_per_tile <= 128 + # The O staging's four SW128 boxes sit `64 G R` elements apart: each must start on + # a 1024-byte swizzle atom, so the live MMA rows come in whole 8-row atoms. + assert group_size * rows_per_tile % 8 == 0, (group_size, rows_per_tile) + self.group_size = group_size + self.rows_per_tile = rows_per_tile + self.d = head_dim + self.tpb = tokens_per_block + self.bm = 128 + self.bn = 64 + self.kv_slots = kv_slots + self.sched_slots = SCHED_SLOTS + self.q_lead = Q_LEAD + self.tau = rescale_threshold + self.use_pdl = use_pdl + self.mma_warp = 4 + self.load_warp0 = 5 + self.threads = 32 * 9 + self.tmem_o = 0 + self.tmem_s0 = 256 + self.tmem_q = 384 # Q as the QK MMA's A operand: 256 bf16 a row, two a column + self.tmem_cols = 512 + self.chunk_cols = 32 # O columns a TMEM load/store moves per thread + + @cute.jit + def __call__( + self, + q: cute.Pointer, + k: cute.Pointer, + v: cute.Pointer, + slots: cute.Pointer, + bits: cute.Pointer, + plan: cute.Pointer, + out: cute.Pointer, + work: cute.Pointer, + rows: Int32, + num_groups: Int32, + num_kv_heads: Int32, + num_ctas: Int32, + q_stride_row: Int32, + k_stride_page: Int32, + k_stride_head: Int32, + k_stride_token: Int32, + v_stride_page: Int32, + v_stride_head: Int32, + v_stride_token: Int32, + slots_stride_group: Int32, + bits_stride_row: Int32, + out_stride_row: Int32, + scale_log2: Float32, + stream, + ): + def qk(): + return sm100_utils.make_trivial_tiled_mma( + BF16, + BF16, + OperandMajorMode.K, + OperandMajorMode.K, + Float32, + tcgen05.CtaGroup.ONE, + (self.bm, self.bn), + tcgen05.OperandSource.TMEM, + ) + + def pv(): + return sm100_utils.make_trivial_tiled_mma( + BF16, + BF16, + OperandMajorMode.K, + OperandMajorMode.MN, + Float32, + tcgen05.CtaGroup.ONE, + (self.bm, self.d), + tcgen05.OperandSource.TMEM, + ) + + # Overwrite / accumulate twins of each MMA: the ACCUMULATE field is set once at kernel + # scope, so no loop or branch rebinds an MMA value. + qk_mma, qk_mma_acc, pv_mma, pv_mma_acc = qk(), qk(), pv(), pv() + q_layout = sm100_utils.make_smem_layout_a(qk_mma, (self.bm, self.bn, self.d), BF16, 1) + k_layout = sm100_utils.make_smem_layout_b( + qk_mma, (self.bm, self.bn, self.d), BF16, self.kv_slots + ) + v_layout = sm100_utils.make_smem_layout_b( + pv_mma, (self.bm, self.d, self.bn), BF16, self.kv_slots + ) + p_layout = sm100_utils.make_smem_layout_a(pv_mma, (self.bm, self.d, self.bn), BF16, 1) + + @cute.struct + class SharedStorage: + q_full: cute.struct.MemRange[Int64, 1] + kv_full: cute.struct.MemRange[Int64, self.kv_slots] + kv_empty: cute.struct.MemRange[Int64, self.kv_slots] + s_full: cute.struct.MemRange[Int64, 2] + p_full: cute.struct.MemRange[Int64, 2] + o_full: cute.struct.MemRange[Int64, 1] + o_final: cute.struct.MemRange[Int64, 1] + qs_full: cute.struct.MemRange[Int64, 1] + qs_empty: cute.struct.MemRange[Int64, 1] + sch_full: cute.struct.MemRange[Int64, self.sched_slots] + sch_empty: cute.struct.MemRange[Int64, self.sched_slots] + sch_val: cute.struct.MemRange[Int32, self.sched_slots] + tmem_holding: Int32 + sKV: cute.struct.Align[cute.struct.MemRange[BF16, cute.cosize(k_layout)], 1024] + # The Q buffer / O staging: four 64-dim boxes of the live G x R MMA rows. + sQ: cute.struct.Align[ + cute.struct.MemRange[BF16, self.group_size * self.rows_per_tile * self.d], 1024 + ] + + self.shared_storage = SharedStorage + + # The epilogue's TMA store: out viewed as (D, heads, rows), one box = 64 dims of the + # group's G heads x R rows, staged in the Q buffer as 128-byte SW128 rows (MMA row m = + # local row x G + head), four boxes 64 G R elements apart. + G = self.group_size + R = self.rows_per_tile + o_box = cute.make_composed_layout( + cute.make_swizzle(3, 4, 3), 0, cute.make_layout((64, G, R), stride=(1, 64, 64 * G)) + ) + mO = cute.make_tensor( + out, + cute.make_layout((self.d, num_kv_heads * G, rows), stride=(1, self.d, out_stride_row)), + ) + tma_atom_o, tma_tensor_o = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), mO, o_box, (64, G, R) + ) + + self.kernel( + qk_mma, + qk_mma_acc, + pv_mma, + pv_mma_acc, + q_layout, + k_layout, + v_layout, + p_layout, + tma_atom_o, + tma_tensor_o, + q, + k, + v, + slots, + bits, + plan, + out, + work, + rows, + num_groups, + num_kv_heads, + q_stride_row, + k_stride_page, + k_stride_head, + k_stride_token, + v_stride_page, + v_stride_head, + v_stride_token, + slots_stride_group, + bits_stride_row, + out_stride_row, + scale_log2, + ).launch( + grid=[num_ctas, 1, 1], + block=[self.threads, 1, 1], + stream=stream, + min_blocks_per_mp=1, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + qk_mma: cute.TiledMma, + qk_mma_acc: cute.TiledMma, + pv_mma: cute.TiledMma, + pv_mma_acc: cute.TiledMma, + q_layout: cute.ComposedLayout, + k_layout: cute.ComposedLayout, + v_layout: cute.ComposedLayout, + p_layout: cute.ComposedLayout, + tma_atom_o: cute.CopyAtom, + tma_tensor_o: cute.Tensor, + q: cute.Pointer, + k: cute.Pointer, + v: cute.Pointer, + slots: cute.Pointer, + bits: cute.Pointer, + plan: cute.Pointer, + out: cute.Pointer, + work: cute.Pointer, + rows: Int32, + num_groups: Int32, + num_kv_heads: Int32, + q_stride_row: Int32, + k_stride_page: Int32, + k_stride_head: Int32, + k_stride_token: Int32, + v_stride_page: Int32, + v_stride_head: Int32, + v_stride_token: Int32, + slots_stride_group: Int32, + bits_stride_row: Int32, + out_stride_row: Int32, + scale_log2: Float32, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + qk_mma.set(tcgen05.Field.ACCUMULATE, False) + qk_mma_acc.set(tcgen05.Field.ACCUMULATE, True) + pv_mma.set(tcgen05.Field.ACCUMULATE, False) + pv_mma_acc.set(tcgen05.Field.ACCUMULATE, True) + + smem = SmemAllocator() + storage = smem.allocate(self.shared_storage) + q_full = storage.q_full.data_ptr() + kv_full = storage.kv_full.data_ptr() + kv_empty = storage.kv_empty.data_ptr() + s_full = storage.s_full.data_ptr() + p_full = storage.p_full.data_ptr() + o_full = storage.o_full.data_ptr() + o_final = storage.o_final.data_ptr() + qs_full = storage.qs_full.data_ptr() + qs_empty = storage.qs_empty.data_ptr() + sch_full = storage.sch_full.data_ptr() + sch_empty = storage.sch_empty.data_ptr() + sch_val = storage.sch_val.data_ptr() + if tidx == 0: + cute.arch.mbarrier_init(q_full, 128) + for i in cutlass.range_constexpr(self.kv_slots): + cute.arch.mbarrier_init(kv_full + i, 128) + cute.arch.mbarrier_init(kv_empty + i, 1) + for i in cutlass.range_constexpr(2): + cute.arch.mbarrier_init(s_full + i, 1) + cute.arch.mbarrier_init(p_full + i, 128) + cute.arch.mbarrier_init(o_full, 1) + cute.arch.mbarrier_init(o_final, 1) + cute.arch.mbarrier_init(qs_full, 128) + # One arrival a hand-back: softmax thread 0, after the softmax warps are done with it. + cute.arch.mbarrier_init(qs_empty, 1) + for i in cutlass.range_constexpr(self.sched_slots): + cute.arch.mbarrier_init(sch_full + i, 1) + # One elected lane of each of the 9 warps releases an item. + cute.arch.mbarrier_init(sch_empty + i, 9) + cute.arch.mbarrier_init_fence() + + tmem_users = pipeline.NamedBarrier(barrier_id=1, num_threads=32 * 5) + if warp_idx == self.mma_warp: + cute.arch.alloc_tmem(self.tmem_cols, storage.tmem_holding.ptr) + cute.arch.relinquish_tmem_alloc_permit() + cute.arch.sync_threads() + + # Slices taken before the role branches: the MMA warp's `set(ACCUMULATE)` rebinds the + # tiled MMAs inside its branch, which the other branches must not see. + qk_thr = qk_mma.get_slice(0) + pv_thr = pv_mma.get_slice(0) + sK = storage.sKV.get_tensor(k_layout.outer, swizzle=k_layout.inner) + sV = cute.make_tensor(cute.recast_ptr(sK.iterator, v_layout.inner), v_layout.outer) + sQ = storage.sQ.data_ptr() + sO = storage.sQ.get_tensor( + cute.make_layout( + (64, self.group_size, self.rows_per_tile, self.d // 64), + stride=(1, 64, 64 * self.group_size, 64 * self.group_size * self.rows_per_tile), + ), + swizzle=cute.make_swizzle(3, 4, 3), + ) + + if warp_idx >= self.load_warp0: + self.load_role( + tidx - 32 * self.load_warp0, + sK, + sQ, + q, + k, + v, + slots, + plan, + work, + rows, + num_groups, + num_kv_heads, + q_stride_row, + k_stride_page, + k_stride_head, + k_stride_token, + v_stride_page, + v_stride_head, + v_stride_token, + slots_stride_group, + kv_full, + kv_empty, + qs_full, + qs_empty, + sch_full, + sch_empty, + sch_val, + ) + elif warp_idx == self.mma_warp: + tmem_ptr = cute.arch.retrieve_tmem_ptr(Float32, 16, storage.tmem_holding.ptr) + self.mma_role( + qk_mma, + qk_mma_acc, + pv_mma, + pv_mma_acc, + qk_thr, + pv_thr, + q_layout, + p_layout, + sK, + sV, + tmem_ptr, + plan, + num_kv_heads, + q_full, + kv_full, + kv_empty, + s_full, + p_full, + o_full, + o_final, + sch_full, + sch_empty, + sch_val, + ) + tmem_users.arrive_and_wait() + cute.arch.dealloc_tmem(tmem_ptr, self.tmem_cols) + else: + tmem_ptr = cute.arch.retrieve_tmem_ptr(Float32, 16, storage.tmem_holding.ptr) + self.softmax_role( + qk_thr, + pv_thr, + tmem_ptr, + tidx, + plan, + num_kv_heads, + rows, + bits, + bits_stride_row, + out, + out_stride_row, + scale_log2, + q_full, + s_full, + p_full, + o_full, + o_final, + qs_full, + qs_empty, + sch_full, + sch_empty, + sch_val, + sQ, + sO, + tma_atom_o, + tma_tensor_o, + ) + tmem_users.arrive_and_wait() + + if cutlass.const_expr(self.use_pdl): + # The lone-row pass may still run: finish only after it has, so a later launch sees every row. + cute.arch.griddepcontrol_wait() + + # ------------------------------------------------------------------ schedule + @cute.jit + def read_item(self, n: Int32, sch_full, sch_val): + """Item `n` of this CTA (group * kv_heads + head, or -1 past the last), once published.""" + slot = n % self.sched_slots + cute.arch.mbarrier_wait(sch_full + slot, (n // self.sched_slots) % 2) + return cute.make_tensor(sch_val, cute.make_layout((self.sched_slots,)))[slot] + + @cute.jit + def release_item(self, n: Int32, sch_empty): + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(sch_empty + n % self.sched_slots) + + @cute.jit + def item_info(self, item: Int32, plan: cute.Pointer, num_kv_heads: Int32): + """(group, KV head, tiles, listed) of a published item; tiles = -1 past the last.""" + group = Int32(0) + head = Int32(0) + tiles = Int32(-1) + listed = Int32(0) + if item >= 0: + group = item // num_kv_heads + head = item - group * num_kv_heads + mPlan = cute.make_tensor(plan, cute.make_layout((1 << 30,))) + tiles = mPlan[group * 2] + listed = mPlan[group * 2 + 1] + return group, head, tiles, listed + + # ------------------------------------------------------------------ loaders (warps 5-8) + @cute.jit + def publish(self, n: Int32, ticket: Int32, num_items: Int32, num_groups: Int32, + num_kv_heads: Int32, sch_full, sch_empty, sch_val): # fmt: skip + """Item `n` from a ticket: the plan's groups last to first, both heads of a group in turn.""" + item = Int32(-1) + if ticket < num_items: + back = ticket // num_kv_heads + item = (num_groups - 1 - back) * num_kv_heads + (ticket - back * num_kv_heads) + slot = n % self.sched_slots + cute.arch.mbarrier_wait(sch_empty + slot, ((n // self.sched_slots) % 2) ^ 1) + cute.make_tensor(sch_val, cute.make_layout((self.sched_slots,)))[slot] = item + cute.arch.mbarrier_arrive(sch_full + slot) + + @cute.jit + def load_role( + self, + lt: Int32, + sK, + sQ: cute.Pointer, + q: cute.Pointer, + k: cute.Pointer, + v: cute.Pointer, + slots: cute.Pointer, + plan: cute.Pointer, + work: cute.Pointer, + rows: Int32, + num_groups: Int32, + num_kv_heads: Int32, + q_stride_row: Int32, + k_stride_page: Int32, + k_stride_head: Int32, + k_stride_token: Int32, + v_stride_page: Int32, + v_stride_head: Int32, + v_stride_token: Int32, + slots_stride_group: Int32, + kv_full, + kv_empty, + qs_full, + qs_empty, + sch_full, + sch_empty, + sch_val, + ): + """Per item: its Q rows into the Q buffer (unless fetched during the previous item), then + the ring in the MMA warp's order K0, K1, V0, K2, V1, ..., V(T-1). Thread `lt` gathers + tokens 16 (lt // 32) + 4 t + (lt % 32) // 8 (t < 4), one 16-byte chunk (lt % 8) of each of + their four 128-byte dim blocks; the slots of tile j + 1 are fetched while tile j's copies + issue. Lane 0 of warp 5 publishes item n + 1 when the warps start item n. + """ + D = self.d + atom = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), BF16, num_bits_per_copy=128 + ) + chunk = lt % 8 + # Element-coordinate view of the swizzled ring: (8 elements, (chunk, 64-dim block), token, slot). + sKV_c = cute.make_tensor( + sK.iterator, + cute.make_layout( + (8, (8, D // 64), self.bn, self.kv_slots), + stride=(1, (8, self.bn * 64), 64, self.bn * D), + ), + ) + num_items = num_groups * num_kv_heads + first_token = 16 * (lt // 32) + (lt % 32) // 8 + ticket = Int32(0) + if lt == 0: + ticket = cute.arch.atomic_add(work, Int32(1)) + self.publish( + Int32(0), ticket, num_items, num_groups, num_kv_heads, sch_full, sch_empty, sch_val + ) + ticket = cute.arch.atomic_add(work, Int32(1)) + n = Int32(0) + ring = Int32(0) + ring_phase = Int32(1) + qs_phase = Int32(1) + q_ready = Int32(0) + item = self.read_item(n, sch_full, sch_val) + while item >= 0: + if lt == 0: + self.publish( + n + 1, ticket, num_items, num_groups, num_kv_heads, sch_full, sch_empty, sch_val + ) + ticket = cute.arch.atomic_add(work, Int32(1)) + group, kv_head, tiles, listed = self.item_info(item, plan, num_kv_heads) + nitem = self.read_item(n + 1, sch_full, sch_val) + ngroup, nhead, ntiles, _ = self.item_info(nitem, plan, num_kv_heads) + if tiles > 0: + if q_ready == 0: + cute.arch.mbarrier_wait(qs_empty, qs_phase) + qs_phase = qs_phase ^ 1 + self.issue_q(lt, group, kv_head, rows, q, q_stride_row, sQ) + cute.arch.cp_async_mbarrier_arrive_noinc(qs_full) + q_ready = Int32(0) + q_at = cutlass.max(tiles - self.q_lead, Int32(0)) + mSlots = cute.make_tensor( + slots + Int64(group) * slots_stride_group, cute.make_layout((1 << 30,)) + ) + slot_next = cute.make_rmem_tensor((4,), Int32) + k_cur = cute.make_rmem_tensor((4,), Int64) + v_cur = cute.make_rmem_tensor((4,), Int64) + v_prev = cute.make_rmem_tensor((4,), Int64) + for t in cutlass.range_constexpr(4): + slot_next[t] = mSlots[cutlass.min(first_token + 4 * t, listed - 1)] + for j in cutlass.range(0, tiles, 1, unroll=1): + for t in cutlass.range_constexpr(4): + slot = slot_next[t] + page = slot // self.tpb + tok = slot % self.tpb + k_cur[t] = ( + Int64(page) * k_stride_page + + kv_head * k_stride_head + + tok * k_stride_token + ) + v_cur[t] = ( + Int64(page) * v_stride_page + + kv_head * v_stride_head + + tok * v_stride_token + ) + if j + 1 < tiles: + for t in cutlass.range_constexpr(4): + slot_next[t] = mSlots[ + cutlass.min((j + 1) * self.bn + first_token + 4 * t, listed - 1) + ] + # K(j) + cute.arch.mbarrier_wait(kv_empty + ring, ring_phase) + self.issue_tile(atom, k, k_cur, sKV_c, ring, chunk, first_token) + cute.arch.cp_async_mbarrier_arrive_noinc(kv_full + ring) + ring, ring_phase = self.advance(ring, ring_phase) + if j > 0: + # V(j - 1) + cute.arch.mbarrier_wait(kv_empty + ring, ring_phase) + self.issue_tile(atom, v, v_prev, sKV_c, ring, chunk, first_token) + cute.arch.cp_async_mbarrier_arrive_noinc(kv_full + ring) + ring, ring_phase = self.advance(ring, ring_phase) + for t in cutlass.range_constexpr(4): + v_prev[t] = v_cur[t] + if j == q_at and ntiles > 0: + # The next item's Q, once this item's epilogue no longer needs the buffer. + cute.arch.mbarrier_wait(qs_empty, qs_phase) + qs_phase = qs_phase ^ 1 + self.issue_q(lt, ngroup, nhead, rows, q, q_stride_row, sQ) + cute.arch.cp_async_mbarrier_arrive_noinc(qs_full) + q_ready = Int32(1) + # V(tiles - 1) + cute.arch.mbarrier_wait(kv_empty + ring, ring_phase) + self.issue_tile(atom, v, v_prev, sKV_c, ring, chunk, first_token) + cute.arch.cp_async_mbarrier_arrive_noinc(kv_full + ring) + ring, ring_phase = self.advance(ring, ring_phase) + self.release_item(n, sch_empty) + n = n + 1 + item = nitem + # Tail: every slot released, so no load is in flight when the CTA retires. + for i in cutlass.range_constexpr(self.kv_slots): + cute.arch.mbarrier_wait(kv_empty + ring, ring_phase) + ring, ring_phase = self.advance(ring, ring_phase) + + @cute.jit + def issue_q( + self, + lt: Int32, + group: Int32, + kv_head: Int32, + rows: Int32, + q: cute.Pointer, + q_stride_row: Int32, + sQ: cute.Pointer, + ): + """An item's live Q rows into the Q buffer: MMA row m = local row x G + head at m x D, + its 16-byte chunk c at chunk c ^ (m % 8). Consecutive threads take consecutive chunks of + a row's G heads (contiguous in q). + """ + G = self.group_size + R = self.rows_per_tile + D = self.d + atom = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), BF16, num_bits_per_copy=128 + ) + per_row = G * (D // 8) + total = R * per_row + for it in cutlass.range_constexpr((total + 127) // 128): + x = lt + 128 * it + local_row = x // per_row + rem = x - local_row * per_row + head = rem // (D // 8) + c = rem - head * (D // 8) + row = group * R + local_row + if x < total and row < rows: + m = local_row * G + head + src = cute.make_tensor( + q + + cute.assume( + Int64(row) * q_stride_row + (kv_head * G + head) * D + c * 8, divby=8 + ), + cute.make_layout((8,)), + ) + dst = cute.make_tensor( + sQ + cute.assume(self.box_offset(m, c), divby=8), cute.make_layout((8,)) + ) + cute.copy(atom, src, dst) + + @cute.jit + def box_offset(self, m, c): + """Element offset of row m's 16-byte chunk c (of D / 8) in the Q buffer: the layout of the + epilogue's TMA boxes -- box c // 8 (64 dims of G heads x R rows), its 128-byte row m, + chunk (c % 8) ^ (m % 8) (SW128) -- so a softmax thread's Q row and its O row occupy the + same bytes. + """ + return ( + (c // 8) * (64 * self.group_size * self.rows_per_tile) + + m * 64 + + (((c % 8) ^ (m % 8)) * 8) + ) + + @cute.jit + def issue_tile( + self, + atom, + src_base: cute.Pointer, + bases, + sKV_c, + ring: Int32, + chunk: Int32, + first_token: Int32, + ): + for t in cutlass.range_constexpr(4): + n = first_token + 4 * t + for b in cutlass.range_constexpr(self.d // 64): + offset = cute.assume(bases[t] + b * 64 + chunk * 8, divby=8) + src = cute.make_tensor(src_base + offset, cute.make_layout((8,))) + dst = sKV_c[(None, (chunk, b), n, ring)] + cute.copy(atom, src, dst) + + # ------------------------------------------------------------------ MMA (warp 4) + @cute.jit + def mma_role( + self, + qk_mma: cute.TiledMma, + qk_mma_acc: cute.TiledMma, + pv_mma: cute.TiledMma, + pv_mma_acc: cute.TiledMma, + qk_thr, + pv_thr, + q_layout, + p_layout, + sK, + sV, + tmem_ptr, + plan: cute.Pointer, + num_kv_heads: Int32, + q_full, + kv_full, + kv_empty, + s_full, + p_full, + o_full, + o_final, + sch_full, + sch_empty, + sch_val, + ): + """Per item: QK(0), then QK(j) / PV(j - 1) pairs, then the last PV. S/P buffer b sits at + TMEM columns tmem_s0 + 64 b, b = the CTA's tile count mod 2; QK(t) overwrites S[t % 2] only + after PV(t - 2) was issued (tcgen05 MMAs of one thread execute in order), and P(t - 2) was + its last reader. An item's PV(0) overwrites O: it waits on P(0), which the softmax warps + publish only after reading the previous item's O out. + """ + tQ = cute.make_tensor(tmem_ptr + self.tmem_q, q_layout.outer) + tSrQ = qk_thr.make_fragment_A(tQ)[None, None, None, 0] + tSrQ = cute.make_tensor(cute.recast_ptr(tmem_ptr + self.tmem_q, dtype=BF16), tSrQ.layout) + tSrK = qk_thr.make_fragment_B(sK) + tOrV = pv_thr.make_fragment_B(sV) + tStS = qk_thr.make_fragment_C(qk_thr.partition_shape_C((self.bm, self.bn))) + tOtO_f = pv_thr.make_fragment_C(pv_thr.partition_shape_C((self.bm, self.d))) + tOtO = cute.make_tensor(tmem_ptr + self.tmem_o, tOtO_f.layout) + tP = cute.make_tensor(tmem_ptr + self.tmem_s0, p_layout.outer) + tOrP = pv_thr.make_fragment_A(tP)[None, None, None, 0] + n = Int32(0) + g = Int32(0) + nq = Int32(0) + ring = Int32(0) + ring_phase = Int32(0) + item = self.read_item(n, sch_full, sch_val) + while item >= 0: + group, kv_head, tiles, listed = self.item_info(item, plan, num_kv_heads) + if tiles > 0: + cute.arch.mbarrier_wait(q_full, nq % 2) + tcgen05_fence_after() + # QK(0) + buf0 = g % 2 + cute.arch.mbarrier_wait(kv_full + ring, ring_phase) + self.gemm_qk( + qk_mma, + qk_mma_acc, + tSrQ, + tSrK[None, None, None, ring], + cute.make_tensor(tmem_ptr + (self.tmem_s0 + buf0 * self.bn), tStS.layout), + ) + with cute.arch.elect_one(): + tcgen05.commit(s_full + buf0) + tcgen05.commit(kv_empty + ring) + ring, ring_phase = self.advance(ring, ring_phase) + for j in cutlass.range(1, tiles + 1, 1, unroll=1): + t = g + j + if j < tiles: + # QK(j) into S[t % 2] + buf = t % 2 + cute.arch.mbarrier_wait(kv_full + ring, ring_phase) + self.gemm_qk( + qk_mma, + qk_mma_acc, + tSrQ, + tSrK[None, None, None, ring], + cute.make_tensor( + tmem_ptr + (self.tmem_s0 + buf * self.bn), tStS.layout + ), + ) + with cute.arch.elect_one(): + tcgen05.commit(s_full + buf) + tcgen05.commit(kv_empty + ring) + ring, ring_phase = self.advance(ring, ring_phase) + # PV(j - 1) + prev = t - 1 + pbuf = prev % 2 + cute.arch.mbarrier_wait(kv_full + ring, ring_phase) + cute.arch.mbarrier_wait(p_full + pbuf, (prev // 2) % 2) + tcgen05_fence_after() + tOrP_b = cute.make_tensor( + cute.recast_ptr( + tmem_ptr + (self.tmem_s0 + pbuf * self.bn), dtype=tOrP.element_type + ), + tOrP.layout, + ) + if j == 1: + self.gemm_pv(pv_mma, pv_mma_acc, tOrP_b, tOrV[None, None, None, ring], tOtO) + else: + self.gemm_pv( + pv_mma_acc, pv_mma_acc, tOrP_b, tOrV[None, None, None, ring], tOtO + ) + with cute.arch.elect_one(): + tcgen05.commit(kv_empty + ring) + tcgen05.commit(o_full) + ring, ring_phase = self.advance(ring, ring_phase) + with cute.arch.elect_one(): + tcgen05.commit(o_final) + g = g + tiles + nq = nq + 1 + self.release_item(n, sch_empty) + n = n + 1 + item = self.read_item(n, sch_full, sch_val) + + @cute.jit + def advance(self, ring: Int32, ring_phase: Int32): + ring = ring + 1 + if ring == self.kv_slots: + ring = Int32(0) + ring_phase = ring_phase ^ 1 + return ring, ring_phase + + @cute.jit + def gemm_qk(self, first: cute.TiledMma, rest: cute.TiledMma, tSrQ, tSrK, tStS): + for kb in cutlass.range_constexpr(cute.size(tSrQ, mode=[2])): + if cutlass.const_expr(kb == 0): + cute.gemm(first, tStS, tSrQ[None, None, kb], tSrK[None, None, kb], tStS) + else: + cute.gemm(rest, tStS, tSrQ[None, None, kb], tSrK[None, None, kb], tStS) + + @cute.jit + def gemm_pv(self, first: cute.TiledMma, rest: cute.TiledMma, tOrP, tOrV, tOtO): + for kb in cutlass.range_constexpr(cute.size(tOrP, mode=[2])): + if cutlass.const_expr(kb == 0): + cute.gemm(first, tOtO, tOrP[None, None, kb], tOrV[None, None, kb], tOtO) + else: + cute.gemm(rest, tOtO, tOrP[None, None, kb], tOrV[None, None, kb], tOtO) + + # ------------------------------------------------------------------ softmax + epilogue (warps 0-3) + @cute.jit + def copy_q(self, m: Int32, live, sQ: cute.Pointer, store_p, tPtP0, tPcP, p_cols: int, q_full): + """Row m's Q from the Q buffer into TMEM columns tmem_q.. (the QK MMA's A operand), 64 + elements a store; dead rows store zeros (their S rows are masked out). + """ + D = self.d + for c in cutlass.range_constexpr(D // (2 * p_cols)): + rQ = cute.make_rmem_tensor(tPcP.shape, Float32) + rQ16 = cute.make_tensor( + cute.recast_ptr(rQ.iterator, dtype=BF16), cute.make_layout((2 * p_cols,)) + ) + if live: + for i in cutlass.range_constexpr(2 * p_cols // 8): + cidx = c * (2 * p_cols // 8) + i + off = cute.assume(self.box_offset(m, cidx), divby=8) + cute.autovec_copy( + cute.make_tensor(sQ + off, cute.make_layout((8,))), + cute.make_tensor(rQ16.iterator + i * 8, cute.make_layout((8,))), + ) + else: + rQ.fill(0.0) + cute.copy( + store_p, + rQ, + cute.make_tensor( + tPtP0.iterator + (self.tmem_q - self.tmem_s0 + c * p_cols), tPtP0.layout + ), + ) + cute.arch.fence_view_async_tmem_store() + tcgen05_fence_before() + cute.arch.mbarrier_arrive(q_full) + + @cute.jit + def softmax_role( + self, + qk_thr, + pv_thr, + tmem_ptr, + st: Int32, + plan: cute.Pointer, + num_kv_heads: Int32, + rows: Int32, + bits: cute.Pointer, + bits_stride_row: Int32, + out: cute.Pointer, + out_stride_row: Int32, + scale_log2: Float32, + q_full, + s_full, + p_full, + o_full, + o_final, + qs_full, + qs_empty, + sch_full, + sch_empty, + sch_val, + sQ: cute.Pointer, + sO, + tma_atom_o, + tma_tensor_o, + ): + G = self.group_size + R = self.rows_per_tile + D = self.d + tStS = qk_thr.make_fragment_C(qk_thr.partition_shape_C((self.bm, self.bn))) + tStS0 = cute.make_tensor(tmem_ptr + self.tmem_s0, tStS.layout) + tOtO_f = pv_thr.make_fragment_C(pv_thr.partition_shape_C((self.bm, self.d))) + tOtO = cute.make_tensor(tmem_ptr + self.tmem_o, tOtO_f.layout) + # The softmax warps' own barrier (hand-backs of the Q buffer, the TMA store's staging). + softmax_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=128) + + # S: one 64-column load a tile; P: 32 packed-bf16 columns stored over the S it came from + # (or two 16-column halves). + cS = cute.make_identity_tensor((self.bm, self.bn)) + tScS = qk_thr.partition_C(cS) + # The load also reduces the row max over the tile's 64 columns (tcgen05.ld.red, SM103): an + # upper bound of the row's masked max, which serves as the softmax offset just as well. + load_s = tcgen05.make_tmem_copy( + cute.make_copy_atom( + tcgen05.copy.LdRed32x32bOp( + tcgen05.copy.Repetition(32), redOp=tcgen05.TmemLoadRedOp.MAX + ), + Float32, + ), + tStS0, + ) + thr_load_s = load_s.get_slice(st) + tLtS0 = thr_load_s.partition_S(tStS0) + tLcS = thr_load_s.partition_D(tScS) + p_cols = self.bn // 2 + tStP_layout = cute.composition(tStS0.layout, cute.make_layout((self.bm, p_cols))) + tStP0 = cute.make_tensor(tStS0.iterator, tStP_layout) + tScP = cute.make_tensor( + tScS.iterator, cute.composition(tScS.layout, cute.make_layout((self.bm, p_cols))) + ) + store_p = tcgen05.make_tmem_copy( + cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(p_cols)), Float32), + tStP0, + ) + thr_store_p = store_p.get_slice(st) + tPtP0 = thr_store_p.partition_D(tStP0) + tPcP = thr_store_p.partition_S(tScP) + # O in 32-column chunks (rescale and epilogue). + cO = cute.make_identity_tensor((self.bm, self.d)) + tOcO = pv_thr.partition_C(cO) + tOtO_c = cute.make_tensor( + tOtO.iterator, + cute.composition(tOtO.layout, cute.make_layout((self.bm, self.chunk_cols))), + ) + tOcO_c = cute.make_tensor( + tOcO.iterator, + cute.composition(tOcO.layout, cute.make_layout((self.bm, self.chunk_cols))), + ) + load_o = tcgen05.make_tmem_copy( + cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(self.chunk_cols)), Float32 + ), + tOtO_c, + ) + store_o = tcgen05.make_tmem_copy( + cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(self.chunk_cols)), Float32 + ), + tOtO_c, + ) + thr_load_o = load_o.get_slice(st) + thr_store_o = store_o.get_slice(st) + tLtO = thr_load_o.partition_S(tOtO_c) + tLcO = thr_load_o.partition_D(tOcO_c) + tStO = thr_store_o.partition_D(tOtO_c) + n_chunks = self.d // self.chunk_cols + + # The TMA store's view: per item, box b of (group, KV head). + gO = cute.flat_divide(tma_tensor_o, (64, G, R)) + tOsO, tOgO = cpasync.tma_partition( + tma_atom_o, + 0, + cute.make_layout(1), + cute.group_modes(sO, 0, 3), + cute.group_modes(gO, 0, 3), + ) + + m = st + local_row = m // G + head = m % G + tau = Float32(self.tau) + n = Int32(0) + g = Int32(0) + nq = Int32(0) + qs_phase = Int32(0) + q_ready = Int32(0) + store_pending = Int32(0) + item = self.read_item(n, sch_full, sch_val) + while item >= 0: + group, kv_head, tiles, listed = self.item_info(item, plan, num_kv_heads) + row = group * R + local_row + live = (local_row < R) and (row < rows) + if tiles > 0: + if q_ready == 0: + cute.arch.mbarrier_wait(qs_full, qs_phase) + qs_phase = qs_phase ^ 1 + self.copy_q(m, live, sQ, store_p, tPtP0, tPcP, p_cols, q_full) + # Every row read out: hand the buffer back. + if st < 32: + softmax_bar.arrive_and_wait() + if st == 0: + cute.arch.mbarrier_arrive(qs_empty) + else: + softmax_bar.arrive() + q_ready = Int32(0) + safe_row = group * R + if live: + safe_row = row + mBits = cute.make_tensor( + bits + Int64(safe_row) * bits_stride_row, cute.make_layout((1 << 30,)) + ) + row_max = -Float32.inf + row_sum = Float32(0.0) + lo_next = Int32(0) + hi_next = Int32(0) + if live: + lo_next = mBits[0] + hi_next = mBits[1] + for j in cutlass.range(0, tiles, 1, unroll=1): + t = g + j + lo = lo_next + hi = hi_next + if live and j + 1 < tiles: + lo_next = mBits[2 * j + 2] + hi_next = mBits[2 * j + 3] + buf = t % 2 + cute.arch.mbarrier_wait(s_full + buf, (t // 2) % 2) + tcgen05_fence_after() + rS = cute.make_rmem_tensor(tLcS.shape, Float32) + n_ld = cute.size(rS, mode=[1]) + rMax = cute.make_rmem_tensor(cute.make_layout((1, n_ld)), Float32) + tLtS_b = cute.make_tensor(tLtS0.iterator + buf * self.bn, tLtS0.layout) + for c in cutlass.range_constexpr(n_ld): + cute.copy_atom_call( + load_s, tLtS_b[None, c, 0, 0], (rS[None, c, 0, 0], rMax[None, c]) + ) + cute.arch.fence_view_async_tmem_load() + tile_max = rMax[0] + for c in cutlass.range_constexpr(1, n_ld): + tile_max = cute.arch.fmax(tile_max, rMax[c]) + new_max = cute.arch.fmax(row_max, tile_max) + use_max = row_max + factor = Float32(1.0) + if row_max == -Float32.inf: + use_max = new_max + elif (new_max - row_max) * scale_log2 > tau: + use_max = new_max + factor = cute.math.exp2((row_max - new_max) * scale_log2, fastmath=True) + need = factor != Float32(1.0) + if cute.arch.vote_any_sync(need): + # O holds PV(0..j-1): wait for PV(j-1), then rescale this warp's rows. + cute.arch.mbarrier_wait(o_full, (t - 1) % 2) + tcgen05_fence_after() + for c in cutlass.range_constexpr(n_chunks): + rO = cute.make_rmem_tensor(tLcO.shape, Float32) + cute.copy( + load_o, + cute.make_tensor(tLtO.iterator + c * self.chunk_cols, tLtO.layout), + rO, + ) + for i in cutlass.range_constexpr(cute.size(rO)): + rO[i] = rO[i] * factor + cute.copy( + store_o, + rO, + cute.make_tensor(tStO.iterator + c * self.chunk_cols, tStO.layout), + ) + cute.arch.fence_view_async_tmem_store() + row_max = use_max + neg = -use_max * scale_log2 + rP = cute.make_rmem_tensor(tPcP.shape, Float32) + rP16 = cute.make_tensor(cute.recast_ptr(rP.iterator, dtype=BF16), rS.layout) + # p = 2^(s * scale - max * scale) on the row's own entries, 0 elsewhere; four + # partial sums. + acc0 = (Float32(0.0), Float32(0.0)) + acc1 = (Float32(0.0), Float32(0.0)) + for i in cutlass.range_constexpr(0, self.bn, 4): + word = lo if i < 32 else hi + e = cute.make_rmem_tensor((4,), Float32) + e[0], e[1] = cute.arch.fma_packed_f32x2( + (rS[i], rS[i + 1]), (scale_log2, scale_log2), (neg, neg) + ) + e[2], e[3] = cute.arch.fma_packed_f32x2( + (rS[i + 2], rS[i + 3]), (scale_log2, scale_log2), (neg, neg) + ) + for u in cutlass.range_constexpr(4): + x = cute.math.exp2(e[u], fastmath=True) + rS[i + u] = x if (word & (1 << ((i + u) % 32))) != 0 else Float32(0.0) + acc0 = cute.arch.add_packed_f32x2(acc0, (rS[i], rS[i + 1])) + acc1 = cute.arch.add_packed_f32x2(acc1, (rS[i + 2], rS[i + 3])) + row_sum = row_sum * factor + ((acc0[0] + acc0[1]) + (acc1[0] + acc1[1])) + rP16.store(rS.load().to(BF16)) + cute.copy( + store_p, rP, cute.make_tensor(tPtP0.iterator + buf * self.bn, tPtP0.layout) + ) + cute.arch.fence_view_async_tmem_store() + tcgen05_fence_before() + cute.arch.mbarrier_arrive(p_full + buf) + if j == 0 and store_pending != 0: + # The previous item's TMA store has read the Q buffer: hand it back. + if st == 0: + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.mbarrier_arrive(qs_empty) + store_pending = Int32(0) + + # The last QK of this item has completed (its S arrived): the next item's Q goes + # into tensor memory now, so its first QK MMAs overlap this epilogue. + nitem = self.read_item(n + 1, sch_full, sch_val) + ngroup, nhead, ntiles, _ = self.item_info(nitem, plan, num_kv_heads) + next_q = ntiles > 0 + if next_q: + nrow = ngroup * R + local_row + nlive = (local_row < R) and (nrow < rows) + cute.arch.mbarrier_wait(qs_full, qs_phase) + qs_phase = qs_phase ^ 1 + self.copy_q(m, nlive, sQ, store_p, tPtP0, tPcP, p_cols, q_full) + q_ready = Int32(1) + # The Q buffer is free for staging O unless a later item still has to load its Q + # there (a following lone or empty group): then each thread stores its own row. + staged = next_q or nitem < 0 + + # Epilogue: O / l in bf16. + cute.arch.mbarrier_wait(o_final, nq % 2) + tcgen05_fence_after() + inv = Float32(0.0) + if row_sum > Float32(0.0): + inv = Float32(1.0) / row_sum + if staged: + if nitem < 0: + # No next Q orders this staging after the previous item's TMA store read + # the buffer (thread 0 waited for it at this item's first tile): sync. + softmax_bar.arrive_and_wait() + for c in cutlass.range_constexpr(n_chunks): + rO = cute.make_rmem_tensor(tLcO.shape, Float32) + cute.copy( + load_o, + cute.make_tensor(tLtO.iterator + c * self.chunk_cols, tLtO.layout), + rO, + ) + for i in cutlass.range_constexpr(cute.size(rO)): + rO[i] = rO[i] * inv + rB = cute.make_rmem_tensor((self.chunk_cols,), BF16) + rB.store(rO.load().to(BF16)) + if m < G * R: + for i in cutlass.range_constexpr(self.chunk_cols // 8): + offset = cute.assume( + self.box_offset(m, c * (self.chunk_cols // 8) + i), divby=8 + ) + cute.autovec_copy( + cute.make_tensor(rB.iterator + i * 8, cute.make_layout((8,))), + cute.make_tensor(sQ + offset, cute.make_layout((8,))), + ) + cute.arch.fence_view_async_shared() + if st < 32: + softmax_bar.arrive_and_wait() + if st == 0: + for b in cutlass.range_constexpr(self.d // 64): + cute.copy(tma_atom_o, tOsO[None, b], tOgO[None, b, kv_head, group]) + cute.arch.cp_async_bulk_commit_group() + else: + softmax_bar.arrive() + store_pending = Int32(1) + if nitem < 0: + if st == 0: + cute.arch.cp_async_bulk_wait_group(0, read=True) + store_pending = Int32(0) + else: + for c in cutlass.range_constexpr(n_chunks): + rO = cute.make_rmem_tensor(tLcO.shape, Float32) + cute.copy( + load_o, + cute.make_tensor(tLtO.iterator + c * self.chunk_cols, tLtO.layout), + rO, + ) + for i in cutlass.range_constexpr(cute.size(rO)): + rO[i] = rO[i] * inv + rB = cute.make_rmem_tensor((self.chunk_cols,), BF16) + rB.store(rO.load().to(BF16)) + if live: + for i in cutlass.range_constexpr(self.chunk_cols // 8): + dst_off = cute.assume( + Int64(row) * out_stride_row + + (kv_head * G + head) * D + + c * self.chunk_cols + + i * 8, + divby=8, + ) + cute.autovec_copy( + cute.make_tensor(rB.iterator + i * 8, cute.make_layout((8,))), + cute.make_tensor(out + dst_off, cute.make_layout((8,))), + ) + g = g + tiles + nq = nq + 1 + elif tiles == 0: + # Nothing resident for this group: zeros. + if live: + rZ = cute.make_rmem_tensor((8,), BF16) + rZ.fill(0.0) + for i in cutlass.range_constexpr(D // 8): + dst_off = cute.assume( + Int64(row) * out_stride_row + (kv_head * G + head) * D + i * 8, + divby=8, + ) + cute.autovec_copy( + rZ, cute.make_tensor(out + dst_off, cute.make_layout((8,))) + ) + self.release_item(n, sch_empty) + n = n + 1 + item = self.read_item(n, sch_full, sch_val) + + +@functools.lru_cache(maxsize=None) +def _compiled(group_size: int, rows_per_tile: int, tokens_per_block: int, use_pdl: bool): + kernel = _RowtileAttentionKernel( + group_size, rows_per_tile, HEAD_DIM, tokens_per_block, kv_slots=KV_SLOTS, use_pdl=use_pdl + ) + gmem = cute.AddressSpace.gmem + + def fake(dtype, align): + return make_ptr(dtype, 16 * 1024, gmem, assumed_align=align) + + args = ( + fake(BF16, 16), + fake(BF16, 16), + fake(BF16, 16), + fake(Int32, 4), + fake(Int32, 4), + fake(Int32, 4), + fake(BF16, 16), + fake(Int32, 4), + 1, + 1, + 1, + 1, + 8, + 8, + 8, + 8, + 8, + 8, + 8, + 1, + 1, + 8, + 1.0, + cuda_drv.CUstream(0), + ) + return cute.compile(kernel, *args) + + +@functools.lru_cache(maxsize=None) +def _sm_count(device_index: int) -> int: + return torch.cuda.get_device_properties(device_index).multi_processor_count + + +def num_ctas(num_items: int, device: torch.device) -> int: + """The persistent grid: one CTA per SM, never more than there are items.""" + return max(1, min(num_items, _sm_count(device.index if device.index is not None else 0))) + + +def _aligned(t: torch.Tensor, strides: int) -> bool: + return t.data_ptr() % 16 == 0 and all(t.stride(i) % 8 == 0 for i in range(strides)) + + +def eligible(q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor) -> bool: + """Whether a launch fits this pass: SM103, bf16 q and K/V, D = 256, unit-stride head + dimension and 16-byte-aligned rows. + """ + if q.shape[2] != HEAD_DIM: + return False + if not (q.dtype == k_cache.dtype == v_cache.dtype == torch.bfloat16): + return False + if torch.cuda.get_device_capability(q.device) != (10, 3): + return False + if ( + q.stride(2) != 1 + or q.stride(1) != HEAD_DIM + or k_cache.stride(3) != 1 + or v_cache.stride(3) != 1 + ): + return False + return _aligned(q, 1) and _aligned(k_cache, 3) and _aligned(v_cache, 3) + + +class _Launcher: + """`rowtile_attention_cute[grid](...)`: grid-indexed like the entry's Triton launches; + the grid is `(ctas,)`, the persistent CTAs (`num_ctas`). + """ + + def __getitem__(self, grid): + return functools.partial(self._launch, tuple(grid)) + + @staticmethod + def _launch( + grid, + q, + k_cache, + v_cache, + slots, + packed_bits, + group_plan, + output, + work, + softmax_scale, + rows, + *, + group_size, + rows_per_tile, + tokens_per_block, + use_pdl, + ): + assert len(grid) == 1, grid + assert ( + output.is_contiguous() + and packed_bits.dtype == torch.int64 + and slots.dtype == torch.int32 + and work.dtype == torch.int32 + ) + fn = _compiled(group_size, rows_per_tile, tokens_per_block, bool(use_pdl)) + gmem = cute.AddressSpace.gmem + + def ptr(t, dtype, align): + return make_ptr(dtype, t.data_ptr(), gmem, assumed_align=align) + + fn( + ptr(q, BF16, 16), + ptr(k_cache, BF16, 16), + ptr(v_cache, BF16, 16), + ptr(slots, Int32, 4), + ptr(packed_bits, Int32, 4), + ptr(group_plan, Int32, 4), + ptr(output, BF16, 16), + ptr(work, Int32, 4), + rows, + group_plan.shape[0], + k_cache.shape[1], + grid[0], + q.stride(0), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + slots.stride(0), + 2 * packed_bits.stride(0), + output.stride(0), + softmax_scale * _LOG2_E, + cuda_drv.CUstream(torch.cuda.current_stream(q.device).cuda_stream), + ) + + +rowtile_attention_cute = _Launcher() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_sparse_gqa_rowtile__prefill.py similarity index 86% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_sparse_gqa_rowtile__prefill.py index b91eee94fbeb..2d7cb622703f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qsa_sparse_gqa_rowtile__prefill.py @@ -12,6 +12,10 @@ group, each row masked to its own bits. A group whose rows span requests (generation rows of a mixed batch) walks each row's own selected columns instead, as the per-row kernel does. + +A row group is sized to whole 8-row swizzle atoms (10 rows of 12 heads at a dep4 rank, 20 of 6 +at a tp4 / tep4 rank); scratch is bounded by the group's selections rather than the table's +width. """ from __future__ import annotations @@ -28,14 +32,16 @@ from cutlass.cutlass_dsl import T, dsl_user_op from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa import ( + qsa_paged_sparse_gqa as _per_row_entry, +) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr -from .qsa_paged_sparse_gqa import qsa_paged_sparse_gqa as _per_row_entry -from .qsa_rowtile_attention_cute import BLOCK_M, BLOCK_N -from .qsa_rowtile_attention_cute import eligible as _cute_eligible -from .qsa_rowtile_attention_cute import num_ctas as _rowtile_ctas -from .qsa_rowtile_attention_cute import rowtile_attention_cute as _rowtile_attention_cute +from .qsa_rowtile_attention_cute__prefill import BLOCK_M, BLOCK_N +from .qsa_rowtile_attention_cute__prefill import eligible as _cute_eligible +from .qsa_rowtile_attention_cute__prefill import num_ctas as _rowtile_ctas +from .qsa_rowtile_attention_cute__prefill import rowtile_attention_cute as _rowtile_attention_cute _INTEGER_DTYPES = (torch.int32, torch.int64) # The per-row kernel's scan tile: the prefill bound is rounded up to it, so the @@ -43,12 +49,15 @@ _BOUND_TILE = 64 # Programs of the lone-row pass per KV head; each walks the lone list with this stride. _LONE_PROGRAMS = 256 # capacity: fixed: a launch width; the lone list may be any length -# Bitmap words the plan pass unions and packs per loop step. -_PLAN_WORDS = 16 # capacity: fixed: a loop step; the plan loops over every word +# Bitmap words the plan pass unions and packs per loop step: wider for the 10-row groups of a +# dep4 rank, narrower for the 20-row groups of a tp4 / tep4 rank. +_PLAN_WORDS = {16: 64, 32: 16} # capacity: fixed: a loop step; the plan loops over every word # A one-request group takes the row-tile pass when its list's 64-slot tiles, each costing -# this many quarters of a per-row tile (2.5x), cost no more than its rows' own columns. -_SHARED_TILE_COST = 10 +# this many quarters of a per-row tile (1.5x), cost no more than its rows' own columns. +_SHARED_TILE_COST = 6 _PLAN_WARPS = 4 +#: The O staging's swizzle atom in MMA rows: a group's rows x heads come in whole atoms. +_ROW_ATOM = 8 _LONE_WARPS = 2 _LONE_STAGES = 2 @@ -262,8 +271,8 @@ def _compress_bits(x, m): def _rowtile_plan_kernel( request_indices, row_info, bits, block_table, slots, packed_bits, group_plan, lone_rows, lone_count, rows, - bits_stride_row, block_table_stride_request, block_table_stride_page, slots_stride_group, - NUM_REQUESTS, NUM_CACHE_PAGES, + bits_stride_row, packed_stride_row, block_table_stride_request, block_table_stride_page, + slots_stride_group, NUM_REQUESTS, NUM_CACHE_PAGES, ROWS_PER_TILE: tl.constexpr, ROWS_POW2: tl.constexpr, BLOCK_N: tl.constexpr, SHARED_TILE_COST: tl.constexpr, PAGE_TABLE_WIDTH: tl.constexpr, TOKENS_PER_BLOCK: tl.constexpr, PLAN_WORDS: tl.constexpr, @@ -327,7 +336,7 @@ def _rowtile_plan_kernel( resident = tl.reduce(tl.where(kept, lane_bits[None, :], 0), 1, _or_combine) packed = _compress_bits(row_words, resident[None, :]) shift = (first % 64).to(tl.uint64) - destination = packed_bits + members[:, None] * bits_stride_row + (first // 64)[None, :] + destination = packed_bits + members[:, None] * packed_stride_row + (first // 64)[None, :] writes = live[:, None] & in_range[None, :] & (packed != 0) tl.atomic_or(destination, (packed << shift[None, :]).to(tl.int64, bitcast=True), mask=writes, sem="relaxed") # fmt: skip @@ -527,6 +536,15 @@ def _check_inputs(q, k_cache, v_cache, block_table, selected_tokens, request_ind ) +def rows_per_group(group_size: int) -> int: + """Rows of one request a row-tile item serves: the most whose `rows x group_size` MMA + rows fill whole 8-row swizzle atoms of the 128-row tile (10 at 12 heads, 20 at 6).""" + rows = BLOCK_M // group_size + while rows > 1 and rows * group_size % _ROW_ATOM: + rows -= 1 + return rows + + def qsa_rowtile_sparse_gqa( *, q: torch.Tensor, @@ -555,7 +573,7 @@ def qsa_rowtile_sparse_gqa( bits, row_info = _row_bits( selected_tokens, query_positions, compress_ratio, tokens_per_block, table_width, num_words ) - rows_per_tile = BLOCK_M // group_size + rows_per_tile = rows_per_group(group_size) rows_pow2 = triton.next_power_of_2(rows_per_tile) kv_strides = ( k_cache.stride(0), k_cache.stride(1), k_cache.stride(2), k_cache.stride(3), @@ -571,8 +589,12 @@ def qsa_rowtile_sparse_gqa( ) # fmt: skip out_strides = (output.stride(0), output.stride(1), output.stride(2)) num_groups = triton.cdiv(rows, rows_per_tile) - slots = torch.empty((num_groups, num_words * 64), dtype=torch.int32, device=q.device) - packed_bits = torch.zeros((rows, num_words), dtype=torch.int64, device=q.device) + # A group lists at most its rows' columns: the bound on its slots and re-indexed bits, + # far below the table's tokens at a deep position. + listed_cap = min(num_words * 64, rows_per_tile * selected_tokens.shape[1]) + packed_words = triton.cdiv(listed_cap, 64) + slots = torch.empty((num_groups, packed_words * 64), dtype=torch.int32, device=q.device) + packed_bits = torch.zeros((rows, packed_words), dtype=torch.int64, device=q.device) group_plan = torch.empty((num_groups, 2), dtype=torch.int32, device=q.device) lone_rows = torch.empty((rows,), dtype=torch.int32, device=q.device) # [0]: the lone list's length; [1]: the row-tile pass's work-ticket counter. @@ -581,11 +603,12 @@ def qsa_rowtile_sparse_gqa( _rowtile_plan_kernel[(num_groups,)]( request_indices, row_info, bits, block_table, slots, packed_bits, group_plan, lone_rows, lone_count, rows, - bits.stride(0), block_table.stride(0), block_table.stride(1), slots.stride(0), - block_table.shape[0], k_cache.shape[0], + bits.stride(0), packed_bits.stride(0), block_table.stride(0), block_table.stride(1), + slots.stride(0), block_table.shape[0], k_cache.shape[0], ROWS_PER_TILE=rows_per_tile, ROWS_POW2=rows_pow2, BLOCK_N=BLOCK_N, SHARED_TILE_COST=_SHARED_TILE_COST, PAGE_TABLE_WIDTH=table_width, - TOKENS_PER_BLOCK=tokens_per_block, PLAN_WORDS=_PLAN_WORDS, num_warps=_PLAN_WARPS, + TOKENS_PER_BLOCK=tokens_per_block, PLAN_WORDS=_PLAN_WORDS.get(rows_pow2, 16), + num_warps=_PLAN_WARPS, ) # fmt: skip use_pdl = get_env_enable_pdl() _rowtile_lone_kernel[(min(rows, _LONE_PROGRAMS), num_kv_heads)]( @@ -614,23 +637,30 @@ def qsa_rowtile_sparse_gqa( token_topk=2048, cache_dtype=torch.bfloat16, q_dtype=torch.bfloat16, - # K and V are the two role planes of one [pages, 2, kv, tpb, D] pool. - cache_layout="pool_planes", - # q is the first 24 heads of the packed [T, 24 + 2 + 2, 256] QKV buffer. + # q is the first heads of the packed [T, Hq + 2 Hkv, 256] QKV buffer. q_layout="qkv_slice", request_dtype=torch.int32, - # The engine's table at max_seq_len 10240: 320 pages per slot, -1 past each sequence. - pages=320, + # The engine's table at max_seq_len 262144: `ceil(262145 / 32)` columns padded + # to a multiple of 4, -1 past each sequence. + pages=8196, ) - -def _spec(segments, slots=32, **overrides) -> dict: - """`segments` is `(first position, rows)` per request, in row order.""" - return {**_CHECKPOINT, "segments": tuple(segments), "slots": slots, **overrides} +#: The engine's trunk (`Qwen4ExpHybridCacheManager.get_buffers`): a lifecycle slot holds the 24 +#: role pages of the 12 attention layers, K then V each; the call's K and V are layer 11's two +#: planes, and the requests' pages straddle 2^31 elements of them. +_TRUNK = dict(cache_layout="trunk", slot_pages=24, layer=11, offset_edge=2**31) +#: A tp4 / tep4 rank (attention TP 4): 6 q heads and one of the 2 KV heads, a 128-row table. +_TP4 = dict(q_heads=6, kv_heads=1, slots=128, **_TRUNK) +#: A dep4 rank (attention DP): all 24 / 2 heads, a 32-row table; core's bf16 trunk, or the +#: kv_fp8 variant's pages widened into a contiguous `[pages, 2, 32, 256]` bf16 scratch +#: (`qsa_fp8_pages_to_bf16`) addressed through its own table. +_DEP4 = dict(q_heads=24, kv_heads=2, slots=32, **_TRUNK) +_DEP4_SCRATCH = dict(q_heads=24, kv_heads=2, slots=32, cache_layout="scratch") -def _decodes(count, rows_each=1, first=8200): - return [(first + (37 * g) % 1000, rows_each) for g in range(count)] +def _spec(segments, member, **overrides) -> dict: + """`segments` is `(first position, rows)` per request, in row order.""" + return {**_CHECKPOINT, **member, "segments": tuple(segments), **overrides} class QsaSparseGqaRowtile(OpWrapper): @@ -654,18 +684,19 @@ class QsaSparseGqaRowtile(OpWrapper): fresh, contiguous `[rows, Hq, D]` tensor of `q.dtype`. Four launches. A bitmap pass: one bit per selected logical token per row. - A plan pass per group of 10 consecutive rows of one request: the union of + A plan pass per group of `rows_per_group` consecutive rows of one request + (10 at a dep4 rank's 12 heads a KV head, 20 at a tp4 / tep4 rank's 6): the union of the rows' selected tokens whose page is mapped, listed in logical order as physical cache slots, and each row's bitmap re-indexed onto that list (a per-word parallel bit extract). The group goes to the row-tile pass - when walking its list in 64-slot tiles costs at most 2.5x per-row tiles + when walking its list in 64-slot tiles costs at most 1.5x per-row tiles less than the rows' own columns would; otherwise -- or when its rows span requests -- its rows join the lone list. A lone-row pass (256 programs per KV head striding the list) serves listed rows over their own columns, as the per-row kernel does. The row-tile pass is persistent (one program per SM drawing (group, KV head) items from a ticket counter, the plan's last groups first); per item it loads each listed tile's K/V once for the - 120 (row, head) pairs and masks each row to its own re-indexed bits, and + group's 120 (row, head) pairs and masks each row to its own re-indexed bits, and one item's epilogue overlaps the next item's first MMAs. It is launched programmatically dependent on the lone pass, whose rows are disjoint, so the two overlap. That pass is `qsa_rowtile_attention_cute` (tcgen05 MMAs, @@ -681,16 +712,28 @@ class QsaSparseGqaRowtile(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell( - why="con 1 prefill: 8192 rows of one request from 0, across the budget", - spec=_spec([(0, 8192)]), + why="tp4/tep4 cold prefill: 8192 rows of one request from 0, 6 q / 1 kv heads " + "(20-row groups), trunk pool straddling 2^31, [128, 8196] table", + spec=_spec([(0, 8192)], _TP4), + ), + Cell( + why="tp4/tep4 deep chunk: 8192 rows from 188416 ranking one score vector " + "(a sparse union), 6 q / 1 kv heads, trunk pool straddling 2^31", + spec=_spec([(188416, 8192)], _TP4, shared_scores=True), + ), + Cell( + why="tp4/tep4 2049 rows from 9000: a last group of 9 rows, 6 q / 1 kv heads", + spec=_spec([(9000, 2049)], _TP4), ), Cell( - why="sparse union: 4096 rows from 4096 ranking one score vector", - spec=_spec([(4096, 4096)], shared_scores=True), + why="dep4 kv_fp8 deep chunk: 8160 rows from 188416 over the widened bf16 " + "scratch, 24 q / 2 kv heads (10-row groups), [32, 8196] table", + spec=_spec([(188416, 8160)], _DEP4_SCRATCH, shared_scores=True), ), Cell( - why="con 128 mixed: 3 chunks (16240 rows) + 52 decodes on the lone pass", - spec=_spec([(0, 8192), (2192, 6000), (6144, 2048)] + _decodes(52), slots=128), + why="dep4 core: three context requests (0+3000, 40000+4000, 120000+1160) in one " + "launch, 24 q / 2 kv heads, trunk pool straddling 2^31", + spec=_spec([(0, 3000), (40000, 4000), (120000, 1160)], _DEP4), ), ) @@ -706,8 +749,7 @@ class QsaSparseGqaRowtile(OpWrapper): pass attends to the *set* of tokens a row names; the per-row kernel attends to every column, so a repeated index would weigh twice there and once here. QSA's top-k + expansion never repeats one (distinct - blocks, and a tail outside every complete block): measured zero - repeats on four layers of the target's own 8192-token prefill. + blocks, and a tail outside every complete block). Rows served by the lone-row pass keep per-column semantics. - As for the per-row kernel: each row is prefix-packed with exactly the bound's count of valid entries; entries past the 64-rounded bound @@ -715,29 +757,19 @@ class QsaSparseGqaRowtile(OpWrapper): Work: a row-tile group walks its packed list, so its MMA work grows with the union of its rows' selections rather than with the group's - highest position. On the target's own selections (8192-token prefill, - layers 0/3/6/11) that union is 13-32% shorter than every 64-token tile up - to the highest token, and the walk loads ~5x fewer K/V tiles than the - per-row kernel. Measured on GB300 (eager, whole call) at 8192 context - rows of real selections: 1544-1565 us on layers 0-6 and 1274 us on layer - 11, against 1711-1715 us for a walk of every tile up to the highest token - and 2953 us for `qsa_paged_sparse_gqa`. Uniformly random selections - leave almost no gap in the union, and the call costs what the full walk - does (16292-row mixed step: 3859 vs 3836 us). Below a few hundred rows - it is slower than the per-row kernel (generation-only launches 1.4-2x); - the target's served context launches are thousands of rows. The tensor-core - row-tile pass alone, on six dumped 16384-row context launches of the - target (GB300, layers 0-5 of one step): 581-775 us, against 694-886 us for - the one-program-per-item pass it replaced (the same output, bit for bit). + highest position. Below a few hundred rows it is slower than the per-row + kernel; the target's served context launches are thousands of rows. Compiled once per (group size, rows per tile, tokens per block, PDL): nothing in the row-tile pass's compile key depends on the row count or the plan. - Scratch per call: the bitmap and the re-indexed bitmap, each `rows x - ceil(table_width * tpb / 64)` int64 words (21 MB at 16384 rows and 320 - pages), the slot lists, `ceil(rows / 10) x` that many words' tokens as - int32 (67 MB at 16384 rows), the plan and the lone list. Calls allocate, + Scratch per call: the bitmap, `rows x ceil(table_width * tpb / 64)` int64 + words (267 MB at 8160 rows and an 8196-page table), the re-indexed + bitmap and the slot lists, bounded by a group's selections (`rows x + ceil(rows_per_group * 2051 / 64)` words, 42 MB, and `ceil(rows / + rows_per_group) x` that many words' tokens as int32, 134 MB at 8160 rows), + the plan and the lone list. Calls allocate, so capture in a CUDA graph is not certified; this entry serves the prefill path, which runs eager. Results are deterministic: the bitmap OR, the re-indexed bitmap OR and the lone list order do not change any diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qwen4_exp_attn_decode_prologue__decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qwen4_exp_attn_decode_prologue__decode.py new file mode 100644 index 000000000000..ff83f898bf92 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/qwen4_exp_attn_decode_prologue__decode.py @@ -0,0 +1,719 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A sparse decode attention layer's work between its input projection and its index +scores, one PDL launch, at one rank's heads over a bf16 or e4m3 K/V pool. + +The rotated q and k and the raw v go into the compact packed `[q | k | v]` buffer the sparse +GQA and the dense op read. + +It fuses the composition's seven launches: the main q/k norm + rope in place on the +projection (`fused_qk_norm_rope`), the index-scratch copy and the index plane's norm + rope +(a second `fused_qk_norm_rope`), the index key and coordinate store +(`qsa_paged_index_store`), the group compression (`qsa_prefill_compress`), the packed-QKV +copy and the K/V append (`qsa_paged_kv_store`). Per row: programs of `HPB` main heads +normalize and rotate q and k into the packed buffer, and those holding k heads also store +them into the K pool; one program normalizes and rotates the index query into the index +scratch; one program stores the raw index key and the coordinates and -- at a group's last +token -- pools the group with this token's key taken from the projection row, normalizes, +rotates and stores it at the anchor slot, as the compression kernel does; one program per v +head stores v into the V pool and the packed buffer. + +A verify step (`rows_per_request` > 1, each request's rows consecutive at consecutive +positions) writes several raw keys of one group in the launch: at a group's last token the +earlier members held by the same request's preceding rows are taken from their projection +rows, the others from the index plane. + +With PDL every operand that is static or host-filled -- the page and slot, the rotation +factors at the row's positions, the norm weights, and at a group end the group's earlier +keys from the index plane with their rotation -- is read before `griddepcontrol.wait`; +after it a program reads only projection rows. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + +@triton.jit +def _rope_factors(positions, num_tokens, token, cos_sin, stride_cs, D: tl.constexpr, ROT: tl.constexpr, + S1: tl.constexpr, S2: tl.constexpr): # fmt: skip + """A D-wide head's rotation factors at the row's positions (interleaved mRoPE: pair j + reads coordinate 1 when j % 3 == 1 and j < 3 * S1, coordinate 2 when j % 3 == 2 and + j < 3 * S2, else 0): cos, and sin signed for the NeoX halves; identity past ROT.""" + HALF: tl.constexpr = ROT // 2 + d = tl.arange(0, D) + rot = d < ROT + j = d % HALF + p0 = tl.load(positions + token) + p1 = tl.load(positions + num_tokens + token) + p2 = tl.load(positions + 2 * num_tokens + token) + pos = tl.where((j % 3 == 1) & (j < 3 * S1), p1, tl.where((j % 3 == 2) & (j < 3 * S2), p2, p0)) + row = cos_sin + pos.to(tl.int64) * stride_cs + c = tl.load(row + j, mask=rot, other=1.0) + s = tl.load(row + HALF + j, mask=rot, other=0.0) + s = tl.where(d < HALF, -s, s) + return c, s + + +@triton.jit +def _norm_weights( + heads, weight, weight_k, D: tl.constexpr, ROT: tl.constexpr, WEIGHT_SPLIT: tl.constexpr +): + """`1 + w` for a tile's elements and for their rotation partners; heads at or past + WEIGHT_SPLIT (when > 0) take the second weight (the k heads).""" + HALF: tl.constexpr = ROT // 2 + d = tl.arange(0, D) + rot = d < ROT + partner = tl.where(d < HALF, d + HALF, tl.where(rot, d - HALF, d)) + w = tl.load(weight + d).to(tl.float32) + wp = tl.load(weight + partner).to(tl.float32) + if WEIGHT_SPLIT > 0: + is_k = (heads >= WEIGHT_SPLIT)[:, None] + wk = tl.load(weight_k + d).to(tl.float32) + wkp = tl.load(weight_k + partner).to(tl.float32) + w2 = tl.where(is_k, wk[None, :], w[None, :]) + 1.0 + wp2 = tl.where(is_k, wkp[None, :], wp[None, :]) + 1.0 + else: + w2 = w[None, :] + 1.0 + wp2 = wp[None, :] + 1.0 + return w2, wp2 + + +@triton.jit +def _norm_rope_apply(src, dst, heads, head_mask, w2, wp2, c, s, eps, D: tl.constexpr, ROT: tl.constexpr, + dst_head_stride): # fmt: skip + """rms over each head, `* (1 + w)`, rotate; fp32 throughout, one rounding to `dst`'s + dtype, stored at `dst`. Returns the stored (rounded) tile.""" + HALF: tl.constexpr = ROT // 2 + d = tl.arange(0, D) + rot = d < ROT + partner = tl.where(d < HALF, d + HALF, tl.where(rot, d - HALF, d)) + hm = head_mask[:, None] + x = tl.load(src + heads[:, None] * D + d[None, :], mask=hm, other=0.0).to(tl.float32) + xp = tl.load(src + heads[:, None] * D + partner[None, :], mask=hm & rot[None, :], other=0.0).to( + tl.float32 + ) + rstd = tl.rsqrt(tl.sum(x * x, axis=1) / D + eps) + y = x * rstd[:, None] * w2 + yp = xp * rstd[:, None] * wp2 + out = (y * c[None, :] + yp * s[None, :]).to(dst.dtype.element_ty) + tl.store(dst + heads[:, None] * dst_head_stride + d[None, :], out, mask=hm) + return out + + +@triton.jit +def _compress_norm_rope(x, weight, cosine, sine, eps, HEAD_DIM: tl.constexpr, ROTARY_PAIRS: tl.constexpr): # fmt: skip + """QSA's `_qsa_gemma_norm_rope` for one row with its static factors given (`weight` is + `1 + w`; the fp32 table makes its intermediate rounding to the table's dtype a no-op).""" + x_f32 = x.to(tl.float32) + reciprocal_rms = tl.rsqrt(tl.sum(x_f32 * x_f32, axis=1) / HEAD_DIM + eps) + normalized = x_f32 * reciprocal_rms[:, None] * weight[None, :] + rotated, passthrough = tl.split( + tl.permute(tl.reshape(normalized, (1, 2, HEAD_DIM // 2)), (0, 2, 1)) + ) + first, second = tl.split(tl.permute(tl.reshape(rotated, (1, 2, ROTARY_PAIRS)), (0, 2, 1))) + output_first = first * cosine[None, :] - second * sine[None, :] + output_second = second * cosine[None, :] + first * sine[None, :] + rotated = tl.reshape( + tl.permute(tl.join(output_first, output_second), (0, 2, 1)), (1, HEAD_DIM // 2) + ) + return tl.reshape( + tl.permute(tl.join(rotated, passthrough), (0, 2, 1)), + (1, HEAD_DIM), + ).to(tl.bfloat16) + + +@triton.jit(do_not_specialize=["num_tokens", "num_requests", "table_width", "num_pages"]) +def _attn_decode_prologue_kernel( + proj, + stride_proj, + out, + stride_out, + index_q, + stride_iq_token, + stride_iq_head, + positions, + num_tokens, + cos_sin, + stride_cs, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + eps, + k_cache, + v_cache, + s_kc_page, + s_kc_head, + s_kc_tok, + s_vc_page, + s_vc_head, + s_vc_tok, + request_indices, + logical_positions, + block_table, + s_bt_req, + s_bt_col, + num_requests, + table_width, + num_pages, + index_cache, + s_ic_page, + s_ic_slot, + position_cache, + s_pc_page, + s_pc_slot, + s_pc_axis, + coords, + s_c_row, + s_c_axis, + HQ: tl.constexpr, + HK: tl.constexpr, + D: tl.constexpr, + IH: tl.constexpr, + ID: tl.constexpr, + IQ_COL: tl.constexpr, + ROT: tl.constexpr, + S1: tl.constexpr, + S2: tl.constexpr, + HPB: tl.constexpr, + MAIN_PROGRAMS: tl.constexpr, + TPB: tl.constexpr, + CR: tl.constexpr, + ROTARY_PAIRS: tl.constexpr, + IN_LAUNCH_GROUPS: tl.constexpr, + USE_PDL: tl.constexpr, +): + token = tl.program_id(0).to(tl.int64) + program = tl.program_id(1) + if USE_PDL: + # The consumer (the index scores) reads nothing of this launch before its own wait. + tl.extra.cuda.gdc_launch_dependents() + row = proj + token * stride_proj + out_row = out + token * stride_out + + # The row's page and slot (the K/V pool, the index plane and the coordinates share + # the table); an unmapped row writes no cache. Host-filled: read before the wait. + request = tl.load(request_indices + token).to(tl.int64) + logical = tl.load(logical_positions + token).to(tl.int64) + column = logical // TPB + slot = logical % TPB + live = (request >= 0) & (request < num_requests) & (logical >= 0) & (column < table_width) + page = tl.load(block_table + request * s_bt_req + column * s_bt_col, mask=live, other=-1).to( + tl.int64 + ) + live = live & (page >= 0) & (page < num_pages) + + if program < MAIN_PROGRAMS: + main_heads = program * HPB + tl.arange(0, HPB) + c, s = _rope_factors(positions, num_tokens, token, cos_sin, stride_cs, D, ROT, S1, S2) + w2, wp2 = _norm_weights(main_heads, q_weight, k_weight, D, ROT, HQ) + if USE_PDL: + tl.extra.cuda.gdc_wait() + rotated = _norm_rope_apply( + row, out_row, main_heads, main_heads < HQ + HK, w2, wp2, c, s, eps, D, ROT, D + ) + if program * HPB + HPB > HQ: + # The rounded rows the composition's store read back from the packed buffer. + is_k = (main_heads >= HQ) & (main_heads < HQ + HK) & live + k_heads = (main_heads - HQ).to(tl.int64)[:, None] + d = tl.arange(0, D) + tl.store( + k_cache + page * s_kc_page + k_heads * s_kc_head + slot * s_kc_tok + d[None, :], + rotated.to(k_cache.dtype.element_ty), + mask=is_k[:, None], + ) + elif program == MAIN_PROGRAMS: + index_heads = tl.arange(0, IH) + c, s = _rope_factors(positions, num_tokens, token, cos_sin, stride_cs, ID, ROT, 0, 0) + w2, wp2 = _norm_weights(index_heads, index_q_weight, index_q_weight, ID, ROT, 0) + if USE_PDL: + tl.extra.cuda.gdc_wait() + _norm_rope_apply( + row + IQ_COL, index_q + token * stride_iq_token, index_heads, index_heads < IH, + w2, wp2, c, s, eps, ID, ROT, stride_iq_head, + ) # fmt: skip + elif program >= MAIN_PROGRAMS + 2: + # One v head a program: the V pool, as the K/V append stores it, and the packed buffer. + h = program - MAIN_PROGRAMS - 2 + d = tl.arange(0, D) + if USE_PDL: + tl.extra.cuda.gdc_wait() + v = tl.load(row + (HQ + HK + h) * D + d) + tl.store(out_row + (HQ + HK + h) * D + d, v) + if live: + tl.store( + v_cache + page * s_vc_page + h * s_vc_head + slot * s_vc_tok + d, + v.to(v_cache.dtype.element_ty), + ) + else: + # Before the wait: the coordinates; at a group's last token, the group's earlier + # raw keys (written by earlier steps), the first token's coordinate 0, its rotation + # and the key norm weight. + axes = tl.arange(0, 4) + coordinate = tl.load(coords + token * s_c_row + axes * s_c_axis, mask=axes < 3, other=0) + dims = tl.arange(0, ID) + group_end = live & ((logical + 1) % CR == 0) + offsets = tl.arange(0, CR) + members = logical - (CR - 1) + offsets + member_pages = tl.load( + block_table + request * s_bt_req + (tl.maximum(members, 0) // TPB) * s_bt_col, + mask=group_end, + other=0, + ).to(tl.int64) + earlier = group_end & (offsets < CR - 1) + if IN_LAUNCH_GROUPS: + # A request's rows come consecutive at consecutive positions: the earlier + # member at row `token - (CR - 1) + o`, when that row is this request's at + # that position, is written by this launch -- its raw key is taken from its + # projection row after the wait, not from the index plane, and as the first + # member its coordinates from its row of `coords`. + member_rows = token - (CR - 1) + offsets + maybe = earlier & (member_rows >= 0) + member_requests = tl.load(request_indices + member_rows, mask=maybe, other=-1) + member_logical = tl.load(logical_positions + member_rows, mask=maybe, other=-1) + in_launch = ( + maybe + & (member_requests.to(tl.int64) == request) + & (member_logical.to(tl.int64) == members) + ) + else: + in_launch = offsets < 0 + group = tl.load( + index_cache + + member_pages[:, None] * s_ic_page + + (tl.maximum(members, 0) % TPB)[:, None] * s_ic_slot + + dims[None, :], + mask=(earlier & ~in_launch)[:, None], + other=0.0, + ) + first = tl.maximum(logical - (CR - 1), 0) + first_page = tl.load( + block_table + request * s_bt_req + (first // TPB) * s_bt_col, mask=group_end, other=0 + ).to(tl.int64) + first_base = position_cache + first_page * s_pc_page + (first % TPB) * s_pc_slot + if IN_LAUNCH_GROUPS: + # A first member written by this launch has its coordinates in its row. + first_in_launch = tl.sum(tl.where(offsets == 0, in_launch, False).to(tl.int32)) > 0 + cached = group_end & ~first_in_launch + pos_t = tl.load(coords + (token - (CR - 1)) * s_c_row, mask=first_in_launch, other=0) + pos_t = tl.where(cached, tl.load(first_base, mask=cached, other=0), pos_t) + else: + pos_t = tl.load(first_base, mask=group_end, other=0) + # The compressed key rotates on the first token's coordinate 0 (QSA's plain rope). + pairs = tl.arange(0, ROTARY_PAIRS) + cs_row = cos_sin + pos_t.to(tl.int64) * stride_cs + cosine = tl.load(cs_row + pairs, mask=group_end, other=0.0) + sine = tl.load(cs_row + ROTARY_PAIRS + pairs, mask=group_end, other=0.0) + key_weight = tl.load(index_k_weight + dims).to(tl.float32) + 1.0 + if USE_PDL: + tl.extra.cuda.gdc_wait() + if live: + # The coordinates, as the index store writes them. + tl.store( + position_cache + page * s_pc_page + slot * s_pc_slot + axes * s_pc_axis, + coordinate.to(position_cache.dtype.element_ty), + mask=axes < 3, + ) + # The raw index key -- or, at a group's last token, the group compressed in + # its place (the compression kernel overwrites the same slot). + current_k = tl.load(row + IQ_COL + IH * ID + dims) + stored_k = current_k.to(index_cache.dtype.element_ty) + if (logical + 1) % CR == 0: + if IN_LAUNCH_GROUPS: + launch_k = tl.load( + proj + + member_rows[:, None] * stride_proj + + IQ_COL + + IH * ID + + dims[None, :], + mask=in_launch[:, None], + other=0.0, + ) + group = tl.where( + in_launch[:, None], launch_k.to(index_cache.dtype.element_ty), group + ) + group = tl.where(offsets[:, None] == CR - 1, stored_k[None, :], group) + pooled = (tl.sum(group.to(tl.float32), axis=0) / CR).to(tl.bfloat16) + stored_k = tl.reshape( + _compress_norm_rope( + pooled[None, :], key_weight, cosine, sine, eps, ID, ROTARY_PAIRS + ), + (ID,), + ) + tl.store(index_cache + page * s_ic_page + slot * s_ic_slot + dims, stored_k) + + +def _launch(num_tokens: int) -> Tuple[int, int]: + """Main-plane heads per program and warps per program, by row count: 1/2 at up to 16 + rows, 2/2 to 64, 8/8 past.""" + if num_tokens <= 16: + return 1, 2 + if num_tokens <= 64: + return 2, 2 + return 8, 8 + + +#: One tp4 / tep4 rank (attention TP 4): 6 q heads and one of the 2 KV heads, the fused +#: projection row `[q 6 | k 1 | v 1 | gate 6 | index q 4, k 1 | v slot]` in heads of 256 and +#: 128 (4352 wide, the index query at column 3584), the packed `[q | k | v]` 2048 wide. +_TP4 = dict(geom="tp4") +#: One dep4 rank (attention DP): every head, `[q 24 | k 2 | v 2 | gate 24 | index ...]` +#: (14080 wide, the index query at 13312), the packed buffer 7168 wide. +_DEP4 = dict(geom="dep4") +#: ht_dep4's decode replay: 16 requests at contexts spread 20k..94k. +_HT16 = tuple(int(20000 + r * (94000 - 20000) / 15) for r in range(16)) + + +class Qwen4ExpAttnDecodePrologue(OpWrapper): + """The main and index planes' Gemma norm + rope, the packed-QKV copy, the paged K/V append, + the index key and coordinate store and the group compression of one sparse decode + attention layer, in one launch. + + `proj` `[T, W]` is the layer's fused input projection: q heads `[0, Hq)`, k heads + `[Hq, Hq + Hk)` and v heads `[Hq + Hk, Hq + 2 Hk)` of `head_dim`, the index query heads + at `index_q_col` and the raw index key right after them; it is not written. Writes: + `out` `[T, (Hq + 2 Hk) * head_dim]` (any row stride), q and k normalized (`* (1 + w)`) + and rotated, v copied; `index_q` `[T, index_heads, index_head_dim]` (any token and head + stride), the index query normalized and rotated on coordinate 0; the rotated k and the + raw v into the paged K/V pool (`k_cache` / `v_cache` `[pages, Hk, tpb, head_dim]`, bf16 + or e4m3, the bf16 values converted as the store converts them); the raw index key (or, + at a group's last token, the compressed group) and the coordinates into the index plane + and the position cache, at the row's page and slot. Returns `(out, index_q)`. + + The rotation is NeoX over the first `rotary_dim` dims of a head, the factors read from + `cos_sin` (fp32 `[P, rotary_dim]`, cos block then sin block) at the row's position: + `position_ids` is the contiguous int32 `[3, T]`, pair `j` of the main plane reads + coordinate 1 when `j % 3 == 1 and j < 3 * mrope_section1`, coordinate 2 when + `j % 3 == 2 and j < 3 * mrope_section2`, coordinate 0 otherwise (interleaved mRoPE); the + index plane reads coordinate 0. The compressed key is `qsa_prefill_compress`'s: the + group's four raw keys averaged (fp32, rounded to bf16), normalized with + `index_k_weight` and rotated at its first token's coordinate 0 in the position cache. + + `rows_per_request` > 1 declares a verify step's layout: each request's rows + consecutive, at consecutive positions. + + Fusion boundary: the projection GEMM before it, and the index scores, the selection and + the sparse GQA after it, are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="tp4/tep4 rank, MTP-3 verify: 1 request x 4 rows at 84000..84003 (ll's step; " + "a group end with its three earlier members in the launch), bf16 pool", + spec=dict(_TP4, pos=(84000,), rows_per_request=4), + ), + Cell( + why="tp4/tep4 rank, verify 2 requests x 4 rows: a group end on a page's last slot " + "with one member in the plane, one with all three in the launch; zero, -0.0 and " + "outlier blocks in k and v", + spec=dict(_TP4, pos=(4093, 8696), rows_per_request=4, edge=True), + ), + Cell( + why="tp4/tep4 rank, MTP-7 verify 2 requests x 8 rows, one up to position 262143 " + "(the table's last row)", + spec=dict(_TP4, pos=(1001, 262136), rows_per_request=8), + ), + Cell( + why="tp4/tep4 rank, decode 3 rows: a group end, a mid-group row, an unmapped row", + spec=dict(_TP4, pos=(8699, 1201, 31), unmapped=2), + ), + Cell( + why="tp4/tep4 rank, decode 72 rows (eight-head programs: the one main program " + "masks its eighth head)", + spec=dict(_TP4, pos=tuple(1000 + 113 * r for r in range(72))), + ), + Cell( + why="dep4 rank, decode 16 requests at 20k..94k (ht's step), e4m3 pool", + spec=dict(_DEP4, pos=_HT16, fp8=True), + ), + Cell( + why="dep4 rank, decode 32 requests (the core's bf16 pool, its graph batch)", + spec=dict(_DEP4, pos=tuple(2000 + 151 * r for r in range(32))), + ), + Cell( + why="dep4 rank, e4m3 pool, verify 10 requests x 4 rows (two-head programs)", + spec=dict( + _DEP4, pos=tuple(3000 + 211 * r for r in range(10)), rows_per_request=4, fp8=True + ), # fmt: skip + ), + ) + + note = """ + Metadata consumed: the QSA request indices, logical positions and block table; the + position ids as `[3, T]` (the rope) and as `[T, 3]` (the coordinates stored). + Preconditions: the index plane one KV head of `index_head_dim` (bf16), the position + cache int32 `[pages, tpb, >=3]`, `compress_ratio` dividing `tokens_per_block`; + generation rows: one per request (`rows_per_request` 1: a group's earlier tokens are + already in the index plane), or a verify step's `rows_per_request` consecutive rows a + request at consecutive positions (a group's earlier tokens in the index plane or in the + request's preceding rows). A row whose page lookup fails writes no cache, as the stores + and compression skip it; its packed row and index query are still written. + + Launch: one Triton kernel, grid `(T, cdiv(Hq + Hk, HPB) + 2 + Hk)`, HPB and warps by row + count (`_launch`); PDL when the engine enables it. + + Numerics: fp32 throughout, one rounding at each store. Against the fp32 `reference` + (the same formulas in torch) the rotated q, k, index query and compressed keys differ + only where the reduction order of the rms moves a value across a bf16 rounding + boundary; v, the raw keys and the coordinates are bitwise. Against the composition the + rotation differs: `fused_qk_norm_rope` derives cos/sin with `__sincosf`, this entry reads + the target's fp32 table (the one the compression and the stock implementation use), + which keeps long positions within bf16 rounding of float64. + + Not certified: NVFP4 pools (the forward keeps the composition there); in-launch index + scores; non-zero index-plane mRoPE sections; mRoPE on the compressed key; `T == 0` returns + at once; head geometries other than 6/1 and 24/2 x 256 with a 4 x 128 index plane. + """ + + def __call__( + self, + proj: torch.Tensor, + out: torch.Tensor, + index_q: torch.Tensor, + position_ids: torch.Tensor, + coords: torch.Tensor, + cos_sin: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + index_cache: torch.Tensor, + position_cache: torch.Tensor, + request_indices: torch.Tensor, + logical_positions: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + compress_ratio: int, + num_heads_q: int, + num_heads_k: int, + head_dim: int, + index_heads: int, + index_head_dim: int, + index_q_col: int, + rotary_dim: int, + eps: float, + mrope_section1: int, + mrope_section2: int, + rows_per_request: int = 1, + ) -> Tuple[torch.Tensor, torch.Tensor]: + num_tokens = proj.shape[0] + if num_tokens == 0: + return out, index_q + hpb, warps = _launch(num_tokens) + main_programs = triton.cdiv(num_heads_q + num_heads_k, hpb) + pdl = get_env_enable_pdl() + _attn_decode_prologue_kernel[(num_tokens, main_programs + 2 + num_heads_k)]( + proj, + proj.stride(0), + out, + out.stride(0), + index_q, + index_q.stride(0), + index_q.stride(1), + position_ids, + num_tokens, + cos_sin, + cos_sin.stride(0), + q_weight, + k_weight, + index_q_weight, + index_k_weight, + eps, + k_cache, + v_cache, + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + request_indices, + logical_positions, + block_table, + block_table.stride(0), + block_table.stride(1), + block_table.shape[0], + block_table.shape[1], + k_cache.shape[0], + index_cache, + index_cache.stride(0), + index_cache.stride(1), + position_cache, + position_cache.stride(0), + position_cache.stride(1), + position_cache.stride(2), + coords, + coords.stride(0), + coords.stride(1), + HQ=num_heads_q, + HK=num_heads_k, + D=head_dim, + IH=index_heads, + ID=index_head_dim, + IQ_COL=index_q_col, + ROT=rotary_dim, + S1=mrope_section1, + S2=mrope_section2, + HPB=hpb, + MAIN_PROGRAMS=main_programs, + TPB=tokens_per_block, + CR=compress_ratio, + ROTARY_PAIRS=cos_sin.shape[1] // 2, + IN_LAUNCH_GROUPS=rows_per_request > 1, + USE_PDL=pdl, + num_warps=warps, + launch_pdl=pdl, + ) + return out, index_q + + def reference( + self, + proj: torch.Tensor, + out: torch.Tensor, + index_q: torch.Tensor, + position_ids: torch.Tensor, + coords: torch.Tensor, + cos_sin: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + index_cache: torch.Tensor, + position_cache: torch.Tensor, + request_indices: torch.Tensor, + logical_positions: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + compress_ratio: int, + num_heads_q: int, + num_heads_k: int, + head_dim: int, + index_heads: int, + index_head_dim: int, + index_q_col: int, + rotary_dim: int, + eps: float, + mrope_section1: int, + mrope_section2: int, + rows_per_request: int = 1, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Plain torch, fp32, on the given buffers: the norm + rope of both planes, the + packed buffer, then every row's K/V, raw index key and coordinate stores, then the + compression of every group a row ends (reading the plane those stores left).""" + rows, dev = proj.shape[0], proj.device + half = rotary_dim // 2 + pos = position_ids.long() + table = cos_sin.float() + j = torch.arange(half, device=dev) + + def factors(s1: int, s2: int): + coord = torch.zeros(half, dtype=torch.long, device=dev) + coord[(j % 3 == 1) & (j < 3 * s1)] = 1 + coord[(j % 3 == 2) & (j < 3 * s2)] = 2 + p = pos[coord, :].t() # [T, half] + return table[p, j][:, None, :], table[p, half + j][:, None, :] + + def plane(x: torch.Tensor, w: torch.Tensor, cos, sin) -> torch.Tensor: + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + w) + x1, x2 = x[..., :half], x[..., half:rotary_dim] + return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin, x[..., rotary_dim:]], -1) + + main_heads = num_heads_q + num_heads_k + width = main_heads * head_dim + v_width = num_heads_k * head_dim + weights = torch.cat( + [ + q_weight.float().expand(num_heads_q, head_dim), + k_weight.float().expand(num_heads_k, head_dim), + ] + ) + main = plane( + proj[:, :width].float().reshape(rows, main_heads, head_dim), + weights, + *factors(mrope_section1, mrope_section2), + ).to(out.dtype) + out[:, :width] = main.reshape(rows, width) + v = proj[:, width : width + v_width] + out[:, width : width + v_width] = v + iq = proj[:, index_q_col : index_q_col + index_heads * index_head_dim].float() + iq = plane( + iq.reshape(rows, index_heads, index_head_dim), index_q_weight.float(), *factors(0, 0) + ) + index_q.copy_(iq.to(index_q.dtype)) + key_col = index_q_col + index_heads * index_head_dim + raw = proj[:, key_col : key_col + index_head_dim] + + requests = request_indices.long().tolist() + logical = logical_positions.long().tolist() + num_requests, table_width = block_table.shape + + def page_of(request: int, position: int): + column = position // tokens_per_block + if not (0 <= request < num_requests and position >= 0 and column < table_width): + return None + page = int(block_table[request, column]) + return page if 0 <= page < k_cache.shape[0] else None + + k_rows = main[:, num_heads_q:] + v_rows = v.reshape(rows, num_heads_k, head_dim) + slots = [page_of(r, p) for r, p in zip(requests, logical)] + for t, page in enumerate(slots): + if page is None: + continue + slot = logical[t] % tokens_per_block + k_cache[page, :, slot] = k_rows[t].to(k_cache.dtype) + v_cache[page, :, slot] = v_rows[t].to(v_cache.dtype) + index_cache[page, slot, 0] = raw[t].to(index_cache.dtype) + position_cache[page, slot, :3] = coords[t].to(position_cache.dtype) + + key_weight = 1.0 + index_k_weight.float() + rot = 2 * (cos_sin.shape[1] // 2) + for t, page in enumerate(slots): + if page is None or (logical[t] + 1) % compress_ratio: + continue + members = [logical[t] - compress_ratio + 1 + o for o in range(compress_ratio)] + keys = torch.stack( + [ + index_cache[page_of(requests[t], m), m % tokens_per_block, 0].float() + for m in members + ] + ) + pooled = (keys.sum(0) / compress_ratio).to(torch.bfloat16).float() + first_page = page_of(requests[t], members[0]) + first = int(position_cache[first_page, members[0] % tokens_per_block, 0]) + x = pooled * torch.rsqrt(pooled.pow(2).mean() + eps) * key_weight + cos, sin = table[first, : rot // 2], table[first, rot // 2 : rot] + x1, x2 = x[: rot // 2], x[rot // 2 : rot] + key = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin, x[rot:]]) + index_cache[page, logical[t] % tokens_per_block, 0] = key.to(index_cache.dtype) + return out, index_q + + def is_valid(self, proj, out, index_q, position_ids, coords, cos_sin, *args, **kwargs) -> None: + assert proj.stride(1) == 1 and proj.dtype == torch.bfloat16, "bf16 proj, unit column stride" + assert out.stride(1) == 1 and out.dtype == torch.bfloat16, "bf16 out, unit column stride" + assert index_q.stride(2) == 1, "index_q must have a unit dim stride" + assert position_ids.is_contiguous() and position_ids.shape == (3, proj.shape[0]), ( + "contiguous [3, T] positions (the kernel reads coordinate a of row t at a * T + t)" + ) + assert position_ids.dtype == torch.int32, "int32 positions" + assert coords.shape == (proj.shape[0], 3), "coords [T, 3]" + assert cos_sin.dtype == torch.float32 and cos_sin.stride(1) == 1, ( + "an fp32 cos|sin table with a unit column stride" + ) + + +qwen4_exp_attn_decode_prologue = Qwen4ExpAttnDecodePrologue() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/vision_rope_2d_packed_qk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/vision_rope_2d_packed_qk.py new file mode 100644 index 000000000000..6b8bdc160131 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/attention/vision_rope_2d_packed_qk.py @@ -0,0 +1,255 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025, Tri Dao. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# * Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# +# Adapted from flash-attn 2 (https://github.com/Dao-AILab/flash-attention), +# `flash_attn/ops/triton/rotary.py`: `rotary_kernel` and `apply_rotary`. +"""Vision-tower 2-D rotary, in place on the q|k prefix of the packed qkv projection. + +The rotation half of the vision tower's 2-D axial rotary (`vision_rope_2d__r52` is the +whole op, angle table included). The angle table depends only on the request's +`grid_thw`, so the caller builds it once per tower forward -- spatial-merge `(row, col)` +per token, fp64 angles rounded once to fp32, then `cos` / `sin` -- and every one of the 27 +blocks hands it here with that block's qkv projection output. + +Tier 5: one Triton launch, flash-attn 2's `rotary_kernel` (BSD-3-Clause, header above) +cut down to the path this call takes. flash-attn 2.x is not a pinned dependency +(`requirements.txt` pins only flash-attn 4, and the stock Qwen VL towers guard the 2.x +import), so the kernel is vendored here rather than imported. What changed from the +origin: one sequence (no batch axis, `cu_seqlens` or `seqlen_offsets`), NeoX pairing only +(no `INTERLEAVED` or `CONJUGATE`), int64 token offsets, in place only, a direct +`kernel[grid]` launch instead of `torch.library.wrap_triton`, and 4 heads per program +where the origin takes 2. The arithmetic, `BLOCK_M` 8 and the masked half-width loads are +the origin's. + +Nothing in the tree serves it in one launch: the registered rope ops take 1-D +positions, `fused_dit_qk_norm_rope` always RMS-normalizes q and k, and FlashInfer's +`apply_rope_with_cos_sin_cache_inplace` raises "misaligned address" at head size 72. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: This checkpoint's tower: 1152 / 16 heads = a 72-wide head, spatial_merge_size 2, theta 1e4 +#: (`model.visual.blocks.N.attn.qkv.weight` is [3456, 1152], so a qkv row is 3 * 16 * 72). +NUM_HEADS = 16 +HEAD_DIM = 72 + +#: Heads x tokens per program: the origin's 8 tokens for a rotary width <= 128, and 4 heads +#: (faster than the origin's 2). +_BLOCK_H = 4 +_BLOCK_M = 8 + + +@triton.jit +def _rotary_neox_inplace_kernel( + X, + COS, + SIN, + seqlen, + nheads, + seqlen_ro, + stride_x_seqlen, + stride_x_nheads, + stride_x_headdim, + # ROTARY_DIM is constexpr so the compiler sees the mask is constant every 8 elements + # and issues 128-bit loads (the origin's comment). + ROTARY_DIM: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_M: tl.constexpr, +): + ROTARY_DIM_HALF: tl.constexpr = ROTARY_DIM // 2 + pid_head = tl.program_id(axis=0) + pid_m = tl.program_id(axis=1) + if pid_m * BLOCK_M >= seqlen: + return + + rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + # Token offsets in int64: rows x row stride passes 2^31 at 621379 rows of a 3456-wide row. + rm64 = rm.to(tl.int64) + rk_half = tl.arange(0, BLOCK_K // 2) + + COS = COS + (rm64[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) + SIN = SIN + (rm64[:, None] * ROTARY_DIM_HALF + rk_half[None, :]) + mask_cs = (rm[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) + cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) + sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) + + # The 1st and 2nd halves of each head, rotated as pairs (i, i + ROTARY_DIM_HALF), stored + # back where they were read. + X = X + ( + rh[:, None, None] * stride_x_nheads + + rm64[None, :, None] * stride_x_seqlen + + rk_half[None, None, :] * stride_x_headdim + ) + mask = ( + (rh[:, None, None] < nheads) + & (rm[None, :, None] < seqlen) + & (rk_half[None, None, :] < ROTARY_DIM_HALF) + ) + x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) + x1 = tl.load(X + ROTARY_DIM_HALF * stride_x_headdim, mask=mask, other=0.0).to(tl.float32) + o0 = x0 * cos - x1 * sin + o1 = x0 * sin + x1 * cos + tl.store(X, o0, mask=mask) + tl.store(X + ROTARY_DIM_HALF * stride_x_headdim, o1, mask=mask) + + +class VisionRope2dPackedQk(OpWrapper): + """Rotate every head of `qk` by its token's angles, in place; return `qk`. + + half = head_dim // 2 + qk[t, h, i] <- x[i] cos[t, i] - x[i + half] sin[t, i] + qk[t, h, i + half] <- x[i] sin[t, i] + x[i + half] cos[t, i] i < half + + NeoX half-split pairing `(i, i + half)` over the whole head, in fp32, rounded once on + the store. `qk` is the `[T, 2 * heads, head_dim]` q|k prefix of the tower's packed + `[T, 3 * heads, head_dim]` qkv projection, read and written through its strides; the + v heads that follow in each row are never addressed. `cos` / `sin` are the caller's + per-token fp32 tables `[T, half]`: the height angles, then the width angles. + + Fusion boundary: the rotation. The angle table, the qkv projection, the split into + q / k / v and the attention are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: grid -- `grid_thw` rows (frames, height, width) in patch units, from which + # the test builds the caller's angle table the way the target does; grid_dtype -- + # the integer dtype of that host `grid_thw` (it reaches only the table build); + # dtype -- the qkv projection's dtype. Heads 16 x 72, merge 2, theta 1e4 throughout. + CELLS: tuple[Cell, ...] = ( + Cell( + why="62x90 serving image, int32 host grid: 5580 rows", + spec=dict(grid=((1, 62, 90),), grid_dtype=torch.int32, dtype=torch.bfloat16), + ), + Cell( + why="62x90 with the processor's int64 grid_thw: 5580 rows", + spec=dict(grid=((1, 62, 90),), grid_dtype=torch.int64, dtype=torch.bfloat16), + ), + Cell( + why="video (3,8,8) + image (1,16,24), two grid rows: 576 rows", + spec=dict(grid=((3, 8, 8), (1, 16, 24)), grid_dtype=torch.int64, dtype=torch.bfloat16), + ), + Cell( + why="video (64,16,32): 32768 rows, the tower's per-call ceiling", + spec=dict(grid=((64, 16, 32),), grid_dtype=torch.int64, dtype=torch.bfloat16), + ), + ) + + note = """ + Metadata consumed: none. Call site: the VisionTower's 27 blocks (multimodal unit), + eager only -- the table bakes in the request's grid. Rank-invariant: every member + (tp4, tep4, dep4) runs the tower at these shapes. + + Source: Triton, JIT-compiled at first call for the running device; the cells ran that + path on sm_103 (GB300). Launch: `(cdiv(heads, 4), cdiv(T, 8))` programs, one per 4 + heads x 8 tokens, each loading its 8 tokens' cos / sin once and both half-heads of + its 4 heads with 64-lane masked loads (36 live). Token offsets are int64, so no + address product overflows (the origin's int32 product would pass 2^31 at 621379 rows + of the 3456-wide qkv row); the cos / sin index is `t * half`. + + Preconditions: `qk` 3-D, any strides (the kernel reads all three); `cos` and `sin` + fp32, contiguous, both `[T, head_dim // 2]`. In place: the return value is `qk` + itself, and the caller's qkv buffer holds the rotated q|k on return. + + Numerics: bitwise equal to `reference` (a complex multiply in fp32, rounded once) and + to `vision_rope_2d__r52.reference` (positions by nested loops, fp64 angles, `polar`) + at every cell, q and k alike: relative L2 0, cosine 1. + + Silent, refused by `is_valid` (each demonstrated unguarded in the test): + - a bf16 (or any non-fp32) table: the kernel widens it and rotates by rounded + angles; + - a table with fewer rows than `qk`: rows past it load cos 1 / sin 0, so the tail + tokens come back unrotated; + - a table narrower than half the head: it rotates a prefix as pairs + `(i, i + width)` and leaves the rest of the head untouched (partial rotary); + wider would write across the next head; + - a non-contiguous table (e.g. the first half of a `[T, head_dim]` duplicated + table): it is indexed as `t * half + i` and reads the wrong angles. + + Not certified: float16 / float32 / float8 q|k; head geometries other than 16 x 72; + partial rotary; grids past the cells (sides past 90, calls past 32768 rows); T = 0; + CUDA-graph capture; out-of-place use. + """ + + def __call__(self, qk: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + seqlen, nheads, _ = qk.shape + rotary_dim = 2 * cos.shape[-1] + grid = (triton.cdiv(nheads, _BLOCK_H), triton.cdiv(seqlen, _BLOCK_M)) + _rotary_neox_inplace_kernel[grid]( + qk, + cos, + sin, + seqlen, + nheads, + cos.shape[0], + qk.stride(0), + qk.stride(1), + qk.stride(2), + ROTARY_DIM=rotary_dim, + BLOCK_K=triton.next_power_of_2(rotary_dim), + BLOCK_H=_BLOCK_H, + BLOCK_M=_BLOCK_M, + ) + return qk + + def reference(self, qk: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """A fresh tensor: each head's halves as one complex number, times `cos + i sin`, in + fp32; rounded once to `qk`'s dtype. `qk` is not mutated.""" + half = qk.shape[-1] // 2 + x = qk.float() + rotor = torch.complex(cos.float(), sin.float()).unsqueeze(-2) + rotated = torch.complex(x[..., :half], x[..., half:]) * rotor + return torch.cat((rotated.real, rotated.imag), dim=-1).to(qk.dtype) + + def is_valid(self, qk: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> None: + assert cos.dtype == torch.float32 and sin.dtype == torch.float32, ( + f"cos and sin must be fp32; got {cos.dtype} / {sin.dtype} -- the kernel widens a " + "narrower table and rotates by its rounded angles" + ) + half = qk.shape[-1] // 2 + assert cos.shape == (qk.shape[0], half) and sin.shape == cos.shape, ( + f"cos and sin must be [T, head_dim // 2] = [{qk.shape[0]}, {half}]; got " + f"{tuple(cos.shape)} / {tuple(sin.shape)} -- a short table leaves the tail tokens " + "unrotated, a narrow one rotates only a prefix of each head" + ) + assert cos.is_contiguous() and sin.is_contiguous(), ( + "cos and sin must be contiguous; the kernel indexes them as t * half + i and a " + "strided table reads the wrong angles" + ) + + +vision_rope_2d_packed_qk = VisionRope2dPackedQk() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/gemm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/gemm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc/qwen4_exp_hc_decode_tc_first__decode.py similarity index 89% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc/qwen4_exp_hc_decode_tc_first__decode.py index 834196008834..80e12f692388 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc/qwen4_exp_hc_decode_tc_first__decode.py @@ -1,11 +1,14 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Hyper-connection mixer first half at decode row counts (2..256 rows) as one +"""Hyper-connection mixer first half at decode row counts (1..256 rows) as one tcgen05 split-K GEMM: the first projection `normed @ [down; block_inject]^T` with the SiLU and injection activations in its epilogue, so neither the fp32 logits nor a split-K reduce kernel nor `qwen4_exp_hc_lowrank_act` exists. -Tier 6 (CuTe DSL written for this target). One CTA computes one +The hyper-connection is replicated on every rank, so its shapes are the checkpoint's at +every member (tp4, tep4, dep4); only the rows move. + +Tier 6 (CuTe DSL). One CTA computes one `[128 rows, BN columns]` partial over `K / KS` of the 10240-wide stream (grid z is the 128-row M tile): a TMA warp streams the normed rows (rows past `T` zero-filled by TMA, never read) @@ -547,10 +550,12 @@ def hc_decode_tc_first(normed, first, rank, hc_count, with_inject, bn, ks): #: Tile `(BN, KS)` by row count: `BN` weight rows per CTA (48: 7 tiles cover the #: 324 real rows of 336) and a `KS`-way K split, one cluster per column tile. One #: M tile (<= 128 rows) runs 16-way (7 x 16 = 112 CTAs); two M tiles run 8-way, so -#: the grid (7 x 8 x 2 = 112 CTAs) stays one wave of one CTA an SM. -_TILES = ((128, (48, 16)), (MAX_ROWS, (48, 8))) -#: Every K split `_TILES` uses divides K into 64-wide tiles per CTA when K does by this. -_K_ALIGN = 64 * 16 +#: the grid (7 x 8 x 2 = 112 CTAs) stays one wave of one CTA an SM. Up to 4 rows the +#: split is 10-way (7 x 10 = 70 CTAs). +_TILES = ((4, (48, 10)), (128, (48, 16)), (MAX_ROWS, (48, 8))) +#: Every K split `_TILES` uses (16, 10, 8) divides K into 64-wide tiles per CTA when K +#: does by this. +_K_ALIGN = 64 * 80 def _tile(rows: int) -> tuple[int, int]: @@ -590,68 +595,91 @@ class Qwen4ExpHcDecodeTcFirst(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # spec: rows, inject (True). The checkpoint's geometry throughout: hc_count 4, - # rank 320, hidden 2560, the padded [384, K] weight, bf16. + # rank 320, hidden 2560, the padded [384, K] weight, bf16. HC is replicated, so the + # rows are this rank's tokens at every member. CELLS: tuple[Cell, ...] = ( Cell( - why="checkpoint 4x2560, 324 cols of 384, inject, decode 128 rows (one M tile)", + why="1 row: a dep4 rank's smallest captured decode batch, 1 of a tile's 128 rows", + spec=dict(rows=1), + ), + Cell(why="2 rows: one request's MTP-1 verify step at tep4", spec=dict(rows=2)), + Cell( + why="4 rows: one request's MTP-3 verify step (ll_tp4 / ll_tep4, concurrency 1)", + spec=dict(rows=4), + ), + Cell( + why="16 rows: a dep4 rank's padded decode batch at concurrency 128", + spec=dict(rows=16), + ), + Cell( + why="32 rows: dep4's per-rank graph cap; 4 requests x MTP-7 verify at tep4", + spec=dict(rows=32), + ), + Cell( + why="128 rows: one full M tile (a short prefill step's rows)", spec=dict(rows=128), ), + Cell(why="129 rows: a 1-row second M tile", spec=dict(rows=129)), + Cell(why="256 rows: MAX_ROWS, two full M tiles", spec=dict(rows=256)), Cell( - why="checkpoint 4x2560, 324 cols of 384, inject, 256 rows (MAX_ROWS, two M tiles)", - spec=dict(rows=256), + why="stack-final mixer: 320 cols, no inject, 4 rows (MTP-3 verify)", + spec=dict(rows=4, inject=False), ), Cell( - why="checkpoint 4x2560, 324 cols of 384, inject, 129 rows (a 1-row second M tile)", - spec=dict(rows=129), + why="stack-final mixer: 320 cols, no inject, 16 rows (dep4 decode)", + spec=dict(rows=16, inject=False), ), Cell( why="stack-final mixer: 320 cols, no inject, 200 rows (two M tiles)", spec=dict(rows=200, inject=False), ), Cell( - why="stack-final mixer: 320 cols, no inject, 128 rows", - spec=dict(rows=128, inject=False), + why="MTP draft head mixer: 320 cols, no inject, 1 row (ll draft step)", + spec=dict(rows=1, inject=False), + ), + Cell( + why="MTP draft head mixer: 320 cols, no inject, 32 rows (dep4 rank batch)", + spec=dict(rows=32, inject=False), ), - Cell(why="checkpoint 4x2560, 324 cols of 384, inject, decode 2 rows", spec=dict(rows=2)), ) - note = """ - Tier 6 (CuTe DSL written for this target): tiers 1-5 have no GEMM whose - epilogue applies an activation after a split-K reduction; cuBLAS at - N = 384 runs split-K plus its own reduce kernel and writes fp32 logits - that `qwen4_exp_hc_lowrank_act` reads back (three launches), and the - Triton `tl.dot` K-loop designs measured 36-60 us at 128 rows. - Metadata consumed: none. Call sites: `DecodeTarget`'s trunk - hyper-connection mixers past `_HC_DECODE_GEMV_MAX_ROWS` and up to - `_HC_TC_FIRST_MAX_ROWS` (= MAX_ROWS) rows, 97 per forward, between - `qwen4_exp_hc_inject_norm` (or the PLE tail) and - `qwen4_exp_hc_up_gate_mean(_fp4)`. Preconditions (bf16): `normed` + Tier 6 (CuTe DSL): tiers 1-5 have no GEMM whose epilogue applies an + activation after a split-K reduction; cuBLAS at N = 384 runs split-K plus + its own reduce kernel and writes fp32 logits that `qwen4_exp_hc_lowrank_act` + reads back (three launches). + Metadata consumed: none. Call sites: `_hc_mixer`, every trunk + hyper-connection mixer of a step of at most `_HC_TC_FIRST_MAX_ROWS` + (= MAX_ROWS) rows (every captured decode / MTP verify batch of tp4, + tep4 and dep4, and short prefill steps), 97 per forward, and with MTP + the draft layer's three sites per draft step (`_mtp_front`'s and + `_mtp_tail`'s with the injection, `MTPLayer.shared_head`'s without: + 1 and 4 rows at ll, a rank's batch and its verify rows at dep4), each + between `qwen4_exp_hc_inject_norm` (or the PLE tail) and + `qwen4_exp_hc_up_gate_mean`. Preconditions (bf16): `normed` `[T, K]` contiguous, T <= 256; `first` the `.t()` view of a row-major `[N, K]` weight with N >= R (+ C); 16-byte aligned; K a multiple of - 64 * 16. Only the R (+ C) real weight rows are read (TMA zero-fills + 64 * 80. Only the R (+ C) real weight rows are read (TMA zero-fills the tile past them), never the load-time padding. Launch: grid (KS, cdiv(R + C, BN), cdiv(T, 128)) CTAs of 6 warps, cluster (KS, 1, 1), one CTA per SM (~227 KB shared memory: the K tiles of the slice in flight, as many as fit, plus the padded fp32 partial); - `(BN, KS)` = (48, 16) up to 128 rows (16 x 7 = 112 CTAs), (48, 8) above - (8 x 7 x 2 = 112 CTAs). The partials are reduced through distributed - shared memory between two cluster barriers; no global workspace. + `(BN, KS)` = (48, 10) up to 4 rows (10 x 7 = 70 CTAs), (48, 16) up to 128 + rows (16 x 7 = 112 CTAs), (48, 8) above (8 x 7 x 2 = 112 CTAs). The + partials are reduced through distributed shared memory between two + cluster barriers; no global workspace. Deterministic (fixed K-slice order). Programmatic dependent launch (TRTLLM_ENABLE_PDL=0 turns it off): weight stages before `griddepcontrol.wait`, `normed` after it; dependents released at kernel start. Graph-capturable (current stream, no host sync, no persistent state). Compiles once per (BN, KS, C, R, N, K, inject); `precompile` builds every variant at load. - Measured on GB300 (graph replay, 96 cold weights, us a site, first half - alone): 128 rows 5.5 (48 x 16) against 7.9 for cuBLAS split-K + reduce - + lowrank_act; 256 rows 6.5 (48 x 8) against 9.2. Tile sweep at 129-256 - rows: 48 x 16 (224 CTAs, two waves) 9.6-10.4, 48 x 10 9.5-11.6, 64 x 8 - 6.6-7.4, 112 x 16 7.4-7.9, 32 x 16 16-18; 48 x 8 fastest. - Certified: bf16, 324 columns with the injection at 2, 128, 129 and 256 - rows, 320 columns without it at 128 and 200 rows; graph replay. + Certified: bf16, 324 columns with the injection at 1, 2, 4, 16, 32, + 128, 129 and 256 rows, 320 columns without it at 1, 4, 16, 32 and 200 + rows; + graph replay at 1, 4, 16 and 256 rows. Not certified: hc_count other than 4, an unpadded weight (N = R + C), zero - rows, a single row (below the dispatch range; accepted, no cell drives - them), other row counts without the injection. + rows (no launch; accepted, no cell drives them), other row counts without + the injection. """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc/qwen4_exp_hc_up_gate_mean_fp4__moe.py similarity index 55% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc/qwen4_exp_hc_up_gate_mean_fp4__moe.py index 6d4cfb844b64..5c2968d29a64 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/hc/qwen4_exp_hc_up_gate_mean_fp4__moe.py @@ -4,13 +4,12 @@ (the up projection as a tcgen05 GEMM, sigmoid gate and stream mean in its epilogue) whose epilogue also writes the routed experts' NVFP4 input. -Tier 6 (CuTe DSL written for this target): the `fp4=True` build of -`qwen4_exp_hc_up_gate_mean`'s kernel. The routed experts read the mixer's output -only as NVFP4 (packed e2m1 + one e4m3 scale per 16 elements, linear scale -layout); the bf16 mix is still written for the router and the shared expert. -Each epilogue thread already holds one row's 32 consecutive mix values, i.e. -two whole scale blocks, so the quantization needs no data exchange and removes -the `fp4_quantize` launch and its re-read of the bf16 activation. +It runs the `fp4=True` build of that kernel (`qwen4_exp_hc_up_gate_mean.hc_up_gate_mean_fp4`). +The routed experts read the mixer's output only as NVFP4 (packed e2m1 + one +e4m3 scale per 16 elements, linear scale layout); the bf16 mix is still written +for the router and the shared expert. Each epilogue thread already holds one +row's 32 consecutive mix values, i.e. two whole scale blocks, so the +quantization needs no data exchange. """ from __future__ import annotations @@ -23,8 +22,11 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( fp4_quantize, ) - -from .qwen4_exp_hc_up_gate_mean import _BLOCK_H, hc_up_gate_mean_fp4, qwen4_exp_hc_up_gate_mean +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_up_gate_mean import ( + _BLOCK_H, + hc_up_gate_mean_fp4, + qwen4_exp_hc_up_gate_mean, +) #: NVFP4 block: one e4m3 scale per this many contiguous elements. SF_VEC = 16 @@ -39,51 +41,58 @@ class Qwen4ExpHcUpGateMeanFp4(OpWrapper): `mixed` is bitwise `qwen4_exp_hc_up_gate_mean`'s (same kernel, same epilogue arithmetic); the quantization reads that bf16-rounded value, so - `xq` / `xsf` are bitwise `torch.ops.trtllm.fp4_quantize(mixed, gs, 16, - False, False)`. Returns fresh `(mixed [T, H], xq [T, H/2] uint8, xsf - [T*H/16] uint8)`. `global_scale` is the 1-element fp32 NVFP4 global scale - (`1 / input_scale`). Fusion boundary: the up projection, the sigmoid - gate, the stream mean and the NVFP4 quantization of the mean. + `xq` / `xsf` are `torch.ops.trtllm.fp4_quantize(mixed, gs, 16, False, + False)` outside the reciprocal tie window. Returns fresh `(mixed [T, H], + xq [T, H/2] uint8, xsf [T*H/16] uint8)`. `global_scale` is the 1-element + fp32 NVFP4 global scale (`1 / input_scale`). Fusion boundary: the up + projection, the sigmoid gate, the stream mean and the NVFP4 quantization + of the mean. """ ARCHS = frozenset({Arch.SM_103}) - # spec: rows, hc_count (4), hidden (2560), rank (320); bf16 only. + # spec: rows; hc_count 4, hidden 2560, rank 320 (the checkpoint); bf16 only. Under + # attention TP (tp4, tep4) every rank mixes the full 4 x 2560 residual of every row, + # so only the row count moves per member. CELLS: tuple[Cell, ...] = tuple( Cell(why=f"checkpoint 4x2560, rank 320, {why}", spec=dict(rows=t)) for t, why in ( - (128, "decode 128 rows (max_con)"), - (16384, "prefill chunk 16384 (max_num_tokens)"), + (1, "1 row: one decode token, the smallest captured batch"), + (4, "4 rows: the ll MTP-3 verify step (con1, tp4/tep4)"), + (32, "32 rows: a quarter of one tile"), + (128, "128 rows: one full tile, the largest captured decode batch"), + (1024, "1024 rows: mtp7 verify at tep4, 128 requests x 8 rows"), + (8192, "8192 rows: one prefill chunk at max_num_tokens"), ) - ) + (Cell(why="ragged tail, 16320 + 70 rows (mixed step)", spec=dict(rows=16390)),) + ) + (Cell(why="ragged tail, 8192 + 70 rows (mixed step)", spec=dict(rows=8262)),) note = """ - Tier 6 (CuTe DSL written for this target): `qwen4_exp_hc_up_gate_mean`'s - kernel compiled with its NVFP4 epilogue; tiers 1-5 have no GEMM with this - multi-input epilogue, and the fused up projection itself is tier 6. - Metadata consumed: none. Call sites: the 48 MoE-site mixers (prefill, and - decode past `_HC_DECODE_GEMV_MAX_ROWS` rows); `xq`/`xsf` feed - `fp4_block_scale_moe_renorm_routed` (linear scales), `mixed` the router and - the shared expert. Preconditions: those of `qwen4_exp_hc_up_gate_mean` - (bf16, `R % 64 == 0`, `H % 32 == 0`, `C * 32 <= 256`, 16-byte aligned - operands) and `global_scale` one fp32 element on the device. Compiles once - per `(C, H, R)` on first call; graph-capturable (current stream, no host - sync). + Source: the tier-6 CuTe DSL kernel of `qwen4_exp_hc_up_gate_mean` with its + NVFP4 epilogue; no trtllm op or FlashInfer call fuses this multi-input + epilogue. Metadata consumed: + none. Call sites: the 48 MoE-site mixers under attention TP (tp4, tep4), + at decode, MTP verify and prefill; `xq`/`xsf` feed + `fp4_block_scale_moe_renorm_routed` (linear scales), `mixed` the router + and the shared expert. Under attention DP the routed experts run on the + group's gathered bf16 rows, so this entry is not called there. + Preconditions: those of `qwen4_exp_hc_up_gate_mean` (bf16, `R % 64 == 0`, + `H % 32 == 0`, `C * 32 <= 256`, 16-byte aligned operands) and + `global_scale` one fp32 element on the device. Compiles once per + `(C, H, R)` on first call; graph-capturable (current stream, no host sync). Numerics: the quantization is the trtllm NVFP4 recipe instruction for instruction (`rcp.approx.ftz.f32` reciprocals, `cvt.rn.satfinite` to e4m3 and to e2m1, every multiply `mul.rn.f32`); the GPU test checks it against `fp4_quantize`'s native-torch reference on the returned `mixed` (scales - bitwise, e2m1 codes bitwise outside the reciprocal tie window). Bitwise - `torch.ops.trtllm.fp4_quantize`, ties included, and across tiny and - saturating global scales: measured, not certified. + bitwise, e2m1 codes bitwise outside the reciprocal tie window) and holds + `mixed` bitwise to the non-NVFP4 build. Certified: bf16 at the checkpoint's 4 x 2560, rank 320, a calibrated global - scale: 128 rows, 16384 rows, 16390 rows (a partial last tile); graph replay. + scale: 1, 4, 32, 128, 1024, 8192 and 8262 rows (a partial last tile); graph + replay. - Not certified: zero rows, fewer rows than a tile, hc_count other than 4, - hidden other than 2560, rank other than 320 (the kernel compiles for them; - no cell drives them). + Not certified: zero rows, hc_count other than 4, hidden other than 2560, + rank other than 320 (the kernel compiles for them; no cell drives them). """ def __call__( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/moe/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/moe/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_cached_replay_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/gdn_cached_replay_cute__verify.py similarity index 63% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_cached_replay_cute.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/gdn_cached_replay_cute__verify.py index ff7b5bf2f6a6..9c4ceeef43f3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_cached_replay_cute.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/ssm/gdn_cached_replay_cute__verify.py @@ -2,8 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 """GDN cached-replay verify step on warp tensor cores (CuTe DSL), a drop-in for the Triton `_cached_replay_kernel` (`fused_recurrent_gated_delta_rule_cached_replay_update`) at this -checkpoint's geometry (H=16 q/k heads, HV=48 v heads, K=V=128, -history 16, bf16 checkpoint state, T=1..8 new tokens). +checkpoint's geometry (K=V=128, history 16, bf16 checkpoint state, T=1..8 new tokens), over +H q/k heads and HV = 3 H value heads: a whole layer's 16 / 48, or a tp4/tep4 rank's 4 / 12. + +Tier 5 (CuTe DSL written for this checkpoint); the head counts are compile-time parameters of +`_Replay` (`h`, `hv`). One CTA serves one (request, q/k head, split): the 3 v heads of a q/k head hold 3 x 128 = 384 state rows of 128 (K innermost), cut into 24 m-tiles of 16 rows; a CTA owns `tpc` consecutive @@ -30,7 +33,7 @@ are issued first. Small batches (the launch is latency-bound there) bring every state tile in by bulk copy as soon as the slot is known (`pf=0`: no tile held in registers through the prologue) and take the tiles' state dots before the first barrier, while the history is still in flight -(`early_acc`). `_Replay.kernel_p` is an experimental persistent variant (slower; not used). +(`early_acc`). """ from __future__ import annotations @@ -40,6 +43,7 @@ import cutlass import cutlass.cute as cute import torch +import torch.nn.functional as F from cutlass import const_expr from cutlass._mlir.dialects import llvm from cutlass.cute.runtime import from_dlpack @@ -48,9 +52,13 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -H, HV, K, V, HIST = 16, 48, 128, 128, 16 -GRP = HV // H -QKV_WIDTH = 2 * H * K + HV * V +K, V, HIST = 128, 128, 16 +#: Value heads a q/k head serves. The q/k and value head counts themselves are compile-time +#: parameters (`_Replay(h=..., hv=...)`, from the call's shapes): `GEOMETRIES`. +GRP = 3 +#: `(q/k heads, value heads)` the kernel is compiled for: a whole layer (an attention-DP +#: rank's), and a tp4/tep4 rank's quarter. +GEOMETRIES = ((16, 48), (4, 12)) TILES = GRP * (V // 16) # m-tiles of one (request, q/k head) PU = 272 # shared row pitch (bytes) of 128-element bf16 rows read by ldmatrix (conflict-free) PK = 320 # shared row pitch of the history keys read as permuted 16-byte fragments @@ -790,15 +798,17 @@ def __init__( lo_split=True, state_before_wait=True, smem_tiles=True, - persistent=False, - gmax=0, spec_hist=False, qk_first=True, use_wi=False, evict_first=True, early_acc=False, + h=16, + hv=48, ): assert 1 <= T <= 8 + assert (h, hv) in GEOMETRIES, (h, hv) + self.H, self.HV = h, hv assert tpc in (24, 12, 8, 4, 2, 1) and tpc % warps == 0 self.T = T self.TP = 4 if T <= 4 else 8 @@ -812,40 +822,33 @@ def __init__( # as the slot is known, and no state tile is held in registers through the prologue self.ld, self.pf, self.minb = ld, max(0 if smem_tiles else 1, min(pf, self.tpw)), minb # the tiles' state dots before the first barrier, while the history is in flight - self.use_early = bool(early_acc) and not persistent + self.use_early = bool(early_acc) self.l2norm, self.gating, self.g_f32, self.b_f32 = l2norm, gating, g_f32, b_f32 self.early_trigger, self.lo_split = early_trigger, lo_split self.state_before_wait = state_before_wait # tiles past the first `pf` of a warp: bulk-copied into shared memory at the start - self.persistent, self.gmax = persistent, gmax - if persistent: - assert not (g_f32 or b_f32), "the persistent variant takes bf16 raw g / beta" - self.pf, smem_tiles = 1, True self.nsm = self.tpw - self.pf if smem_tiles else 0 - self.spec_hist = spec_hist and not persistent + self.spec_hist = spec_hist self.qk_first = qk_first # (row, slot, pending count, half) per item from the engine's replay work items - self.use_wi = use_wi and not persistent + self.use_wi = use_wi self.evict_first = evict_first - nset = 2 if (persistent or self.spec_hist) else 1 + nset = 2 if self.spec_hist else 1 self.nset = nset self.RP = (2 * self.TP + 4) * 4 # per-warp scratch row pitch (bytes) o = 0 - tset = 2 if persistent else 1 self.o_st = o - o += tset * warps * self.nsm * 4096 + o += warps * self.nsm * 4096 self.o_mb = o - o += -(-(tset * warps * self.nsm * 8) // 16) * 16 + o += -(-(warps * self.nsm * 8) // 16) * 16 self.o_u = o o += nset * self.nhv * 16 * PU self.o_kh = o o += nset * 16 * PK - self.o_meta = o - o += 256 * 16 if persistent else 0 self.o_kp = o o += 16 * PU if commit else 0 self.o_v = o - o += (2 if persistent else 1) * self.nhv * T * 256 + o += self.nhv * T * 256 self.o_x = o o += warps * 16 * self.RP self.o_kk = o @@ -902,79 +905,41 @@ def __call__( scale: cutlass.Float32, stream, ): - if const_expr(self.persistent): - n_items = n_rows * (H * self.S) - self.kernel_p( - sidx_t.iterator.toint(), - q_ptr, - k_ptr, - v_ptr, - q_sn, - q_st, - k_sn, - k_st, - v_sn, - v_st, - g_ptr, - g_st, - b_ptr, - b_st, - alog_ptr, - dtb_ptr, - o_ptr, - h_ptr, - h_ss, - pnat_ptr, - buf_ptr, - ou_ptr, - ok_ptr, - og_ptr, - scale, - n_items, - ).launch( - grid=[cutlass.min(n_items, self.gmax), 1, 1], - block=[32 * self.W, 1, 1], - smem=self.smem, - stream=stream, - use_pdl=self.pdl, - min_blocks_per_mp=self.minb, - ) - else: - self.kernel( - sidx_t.iterator.toint(), - q_ptr, - k_ptr, - v_ptr, - q_sn, - q_st, - k_sn, - k_st, - v_sn, - v_st, - g_ptr, - g_st, - b_ptr, - b_st, - alog_ptr, - dtb_ptr, - o_ptr, - h_ptr, - h_ss, - pnat_ptr, - buf_ptr, - ou_ptr, - ok_ptr, - og_ptr, - wi_ptr, - scale, - ).launch( - grid=[n_rows * (H * self.S), 1, 1], - block=[32 * self.W, 1, 1], - smem=self.smem, - stream=stream, - use_pdl=self.pdl, - min_blocks_per_mp=self.minb, - ) + self.kernel( + sidx_t.iterator.toint(), + q_ptr, + k_ptr, + v_ptr, + q_sn, + q_st, + k_sn, + k_st, + v_sn, + v_st, + g_ptr, + g_st, + b_ptr, + b_st, + alog_ptr, + dtb_ptr, + o_ptr, + h_ptr, + h_ss, + pnat_ptr, + buf_ptr, + ou_ptr, + ok_ptr, + og_ptr, + wi_ptr, + scale, + ).launch( + grid=[n_rows * (self.H * self.S), 1, 1], + block=[32 * self.W, 1, 1], + smem=self.smem, + stream=stream, + use_pdl=self.pdl, + min_blocks_per_mp=self.minb, + ) # -- straight-line helpers (traced inline) ------------------------------------------------ @@ -1170,6 +1135,8 @@ def kernel( S: cutlass.Constexpr = self.S NTHR: cutlass.Constexpr = 32 * self.W RP: cutlass.Constexpr = self.RP + H: cutlass.Constexpr = self.H + HV: cutlass.Constexpr = self.HV tidx, _, _ = cute.arch.thread_idx() bidx, _, _ = cute.arch.block_idx() @@ -1781,639 +1748,6 @@ def kernel( if const_expr(self.pdl and not self.early_trigger): cute.arch.griddepcontrol_launch_dependents() - # -- the persistent variant: CTAs loop over (request, head, split) items, the next item's - # state tiles and history in flight while the current one computes ------------------------- - - @cute.kernel - def kernel_p( - self, - sidx_ptr: cutlass.Int64, - q_ptr: cutlass.Int64, - k_ptr: cutlass.Int64, - v_ptr: cutlass.Int64, - q_sn: cutlass.Int64, - q_st: cutlass.Int64, - k_sn: cutlass.Int64, - k_st: cutlass.Int64, - v_sn: cutlass.Int64, - v_st: cutlass.Int64, - g_ptr: cutlass.Int64, - g_st: cutlass.Int64, - b_ptr: cutlass.Int64, - b_st: cutlass.Int64, - alog_ptr: cutlass.Int64, - dtb_ptr: cutlass.Int64, - o_ptr: cutlass.Int64, - h_ptr: cutlass.Int64, - h_ss: cutlass.Int64, - pnat_ptr: cutlass.Int64, - buf_ptr: cutlass.Int64, - ou_ptr: cutlass.Int64, - ok_ptr: cutlass.Int64, - og_ptr: cutlass.Int64, - scale: cutlass.Float32, - n_items: cutlass.Int32, - ): - T: cutlass.Constexpr = self.T - TP: cutlass.Constexpr = self.TP - NT: cutlass.Constexpr = self.NT - W: cutlass.Constexpr = self.W - TPW: cutlass.Constexpr = self.tpw - TPC: cutlass.Constexpr = self.tpc - NHV: cutlass.Constexpr = self.nhv - S: cutlass.Constexpr = self.S - NTHR: cutlass.Constexpr = 32 * self.W - RP: cutlass.Constexpr = self.RP - NSM: cutlass.Constexpr = self.nsm - SU: cutlass.Constexpr = NHV * 16 * PU - SK: cutlass.Constexpr = 16 * PK - SG: cutlass.Constexpr = NHV * 64 - ST: cutlass.Constexpr = W * NSM * 4096 - SMB: cutlass.Constexpr = W * NSM * 8 - - tidx, _, _ = cute.arch.thread_idx() - bidx, _, _ = cute.arch.block_idx() - gsz, _, _ = cute.arch.grid_dim() - warp = tidx // 32 - lane = tidx % 32 - g = lane // 4 - c = lane % 4 - zf = cutlass.Float32(0.0) - zi = cutlass.Int32(0) - - smem = cutlass.utils.SmemAllocator() - base_t = smem.allocate_tensor( - cutlass.Int8, cute.make_layout((self.smem - 128,), stride=(1,)), byte_alignment=128 - ) - sbase = cutlass.Int32(base_t.iterator.toint()) - s_u0 = sbase + self.o_u - s_kh0 = sbase + self.o_kh - s_gh0 = sbase + self.o_gh - s_meta = sbase + self.o_meta - s_kp = sbase + self.o_kp - s_v = sbase + self.o_v - s_x = sbase + self.o_x + warp * (16 * RP) - s_kk = sbase + self.o_kk - s_cf = sbase + self.o_cf - s_dec = sbase + self.o_dec - s_g = sbase + self.o_g - s_eg = sbase + self.o_eg - s_beta = sbase + self.o_beta - s_a = sbase + self.o_a - s_cq = sbase + self.o_cq - s_egs = sbase + self.o_egs - s_st0 = sbase + self.o_st + warp * (NSM * 4096) - s_mb0 = sbase + self.o_mb + warp * (NSM * 8) - - nk = (n_items - bidx + gsz - 1) // gsz - # ---- the items' slot / pending count / half, once -------------------------------------- - for m_t in cutlass.range(tidx, nk, NTHR): - m_it = bidx + m_t * gsz - m_n = m_it // (H * S) - m_sl = _ldg_i32(sidx_ptr + cutlass.Int64(m_n) * 4) - m_p = zi - m_b = zi - if m_sl >= 0: - m_p = _ldg_i32(pnat_ptr + cutlass.Int64(m_sl) * 4) - m_b = _ldg_i32(buf_ptr + cutlass.Int64(m_sl) * 4) - _sts_i32(s_meta + m_t * 16, m_sl) - _sts_i32(s_meta + m_t * 16 + 4, m_p) - _sts_i32(s_meta + m_t * 16 + 8, m_b) - if const_expr(NSM > 0): - if lane == 0: - for k_ in cutlass.range_constexpr(2 * NSM): - _mbar_init(s_mb0 + (k_ // NSM) * SMB + (k_ % NSM) * 8, cutlass.Int32(1)) - _fence_mbar_init() - cute.arch.sync_threads() - - cur0 = cute.make_rmem_tensor(cute.make_layout((32,), stride=(1,)), cutlass.Int32) - nxt0 = cute.make_rmem_tensor(cute.make_layout((32,), stride=(1,)), cutlass.Int32) - NQ: cutlass.Constexpr = 16 * NT - GN: cutlass.Constexpr = 2 * T - qkc = cute.make_rmem_tensor(cute.make_layout((NQ,), stride=(1,)), cutlass.Int32) - qkn = cute.make_rmem_tensor(cute.make_layout((NQ,), stride=(1,)), cutlass.Int32) - gtc = cute.make_rmem_tensor(cute.make_layout((GN,), stride=(1,)), cutlass.Int32) - gtn = cute.make_rmem_tensor(cute.make_layout((GN,), stride=(1,)), cutlass.Int32) - gfc = cute.make_rmem_tensor(cute.make_layout((2,), stride=(1,)), cutlass.Float32) - gfn = cute.make_rmem_tensor(cute.make_layout((2,), stride=(1,)), cutlass.Float32) - SVB: cutlass.Constexpr = NHV * T * 256 - jrow = (lane % 8) + 8 * (lane // 16) - v8 = 8 * ((lane // 8) % 2) - r = lane % 16 - hf = lane // 16 - godd = (g % 2) == 1 - gi = tidx - (32 * (W - 1) if W > 1 else 0) - gpred = (gi >= 0) & (gi < NHV) - gic = cutlass.min(cutlass.max(gi, zi), NHV - 1) - - for k in cutlass.range(nk): - # ---- issue: item 0 (first pass only), then item k + 1; everything an item reads - # from global memory is in flight one item ahead ---------------------------------- - for q2 in cutlass.range_constexpr(2): - kk = k + q2 - if const_expr(q2 == 0): - do_it = k == 0 - else: - do_it = kk < nk - if do_it: - st_ = kk % 2 - it_ = bidx + kk * gsz - sp_ = it_ % S - nh_ = it_ // S - h_ = nh_ % H - n64_ = cutlass.Int64(nh_ // H) - sl = _lds_i32(s_meta + kk * 16) - pp = _lds_i32(s_meta + kk * 16 + 4) - bb = _lds_i32(s_meta + kk * 16 + 8) - hvf = (sp_ * TPC) // 8 - hv0_ = h_ * GRP + hvf - if const_expr(q2 == 0): - if const_expr(self.pdl): - cute.arch.griddepcontrol_wait() - if const_expr(self.early_trigger): - cute.arch.griddepcontrol_launch_dependents() - # the item's own inputs first (they are needed first), then its history, then its - # state: a lane's requests come back roughly in order - if sl >= 0: - sl64 = cutlass.Int64(sl) - n_v: cutlass.Constexpr = NHV * T * 16 - for it2 in cutlass.range_constexpr(-(-n_v // NTHR)): - e = tidx + it2 * NTHR - if e < n_v: - hl = e // (T * 16) - t = (e // 16) % T - qq = e % 16 - src = ( - v_ptr - + ( - n64_ * v_sn - + cutlass.Int64(t) * v_st - + cutlass.Int64((hv0_ + hl) * V + qq * 8) - ) - * 2 - ) - _cp16z( - s_v + st_ * SVB + (hl * T + t) * 256 + qq * 16, - src, - cutlass.Int32(16), - ) - for nt in cutlass.range_constexpr(NT): - rq = self._qk_raw( - nt, - q_ptr, - k_ptr, - q_sn, - q_st, - k_sn, - k_st, - n64_, - h_, - g, - c, - scale, - tidx >= 0, - ) - for e in cutlass.range_constexpr(16): - if const_expr(q2 == 0): - qkc[nt * 16 + e] = rq[0][e] - else: - qkn[nt * 16 + e] = rq[0][e] - gr = self._gate_raw( - gpred, hv0_ + gic, n64_, g_ptr, g_st, b_ptr, b_st, alog_ptr, dtb_ptr - ) - for t in cutlass.range_constexpr(T): - if const_expr(q2 == 0): - gtc[t] = gr[0][t] - gtc[T + t] = gr[1][t] - else: - gtn[t] = gr[0][t] - gtn[T + t] = gr[1][t] - if const_expr(q2 == 0): - gfc[0] = gr[2] - gfc[1] = gr[3] - else: - gfn[0] = gr[2] - gfn[1] = gr[3] - rsl = sl64 * 2 + cutlass.Int64(bb) - u_rd_ = ou_ptr + rsl * (HIST * HV * V * 2) - k_rd_ = ok_ptr + rsl * (HIST * H * K * 2) - g_rd_ = og_ptr + rsl * (HV * HIST * 4) - n_u: cutlass.Constexpr = NHV * 256 - for it2 in cutlass.range_constexpr(-(-n_u // NTHR)): - e = tidx + it2 * NTHR - if e < n_u: - hl = e // 256 - j = (e // 16) % 16 - qq = e % 16 - jc = cutlass.min(j, cutlass.max(pp - 1, zi)) - src = u_rd_ + cutlass.Int64( - ((jc * HV + hv0_ + hl) * V + qq * 8) * 2 - ) - _cp16z( - s_u0 + st_ * SU + hl * (16 * PU) + j * PU + qq * 16, - src, - cutlass.Int32(cutlass.select_(j < pp, cutlass.Int32(16), zi)), - ) - for it2 in cutlass.range_constexpr(-(-256 // NTHR)): - e = tidx + it2 * NTHR - if e < 256: - j = e // 16 - qq = e % 16 - jc = cutlass.min(j, cutlass.max(pp - 1, zi)) - src = k_rd_ + cutlass.Int64(((jc * H + h_) * K + qq * 8) * 2) - _cp16z( - s_kh0 + st_ * SK + j * PK + qq * 16, - src, - cutlass.Int32(cutlass.select_(j < pp, cutlass.Int32(16), zi)), - ) - if gpred: - # a v head's G row by the thread that gates that head (read back by it) - for qq in cutlass.range_constexpr(4): - src = g_rd_ + cutlass.Int64(((hv0_ + gi) * HIST + qq * 4) * 4) - _cp16z(s_gh0 + st_ * SG + gi * 64 + qq * 16, src, cutlass.Int32(16)) - vals = self._load_tile( - self._tile_addr(h_ptr, sl64, h_ss, h_, sp_, warp, 0, g, c) - ) - for e in cutlass.range_constexpr(32): - if const_expr(q2 == 0): - cur0[e] = vals[e] - else: - nxt0[e] = vals[e] - if const_expr(NSM > 0): - if lane == 0: - for k_ in cutlass.range_constexpr(NSM): - _bulk_g2s( - s_st0 + st_ * ST + k_ * 4096, - self._tile_addr( - h_ptr, sl64, h_ss, h_, sp_, warp, 1 + k_, 0, 0 - ), - cutlass.Int32(4096), - s_mb0 + st_ * SMB + k_ * 8, - ) - else: - if const_expr(NSM > 0): - if lane == 0: - for k_ in cutlass.range_constexpr(NSM): - _mbar_arrive(s_mb0 + st_ * SMB + k_ * 8) - _cp_commit() # item 0 (empty later) / item k + 1 - - # ---- item k --------------------------------------------------------------------- - st = k % 2 - it = bidx + k * gsz - sp = it % S - nh = it // S - h = nh % H - n = nh // H - n64 = cutlass.Int64(n) - slot = _lds_i32(s_meta + k * 16) - p = _lds_i32(s_meta + k * 16 + 4) - rb = _lds_i32(s_meta + k * 16 + 8) - s_u = s_u0 + st * SU - s_kh = s_kh0 + st * SK - s_gh = s_gh0 + st * SG - s_st = s_st0 + st * ST - s_mb = s_mb0 + st * SMB - s_vk = s_v + st * SVB - ph = (k // 2) % 2 - slot64 = cutlass.Int64(cutlass.max(slot, zi)) - wrap = (p + T) > HIST - wb = cutlass.Int32(cutlass.select_(wrap, 1 - rb, rb)) - woff = cutlass.Int32(cutlass.select_(wrap, zi, p)) - hv_first = (sp * TPC) // 8 - hv0 = h * GRP + hv_first - wslot = slot64 * 2 + cutlass.Int64(wb) - u_wr = ou_ptr + wslot * (HIST * HV * V * 2) - k_wr = ok_ptr + wslot * (HIST * H * K * 2) - g_wr = og_ptr + wslot * (HV * HIST * 4) - - live = slot >= 0 - if live: - bq = [] - for nt in cutlass.range_constexpr(NT): - rw = [qkc[nt * 16 + e] for e in range(16)] - if const_expr(TP == 4): - is_q = g >= 4 - mult = cutlass.Float32(cutlass.select_(is_q, scale, cutlass.Float32(1.0))) - elif const_expr(nt == 1): - is_q = g >= 0 - mult = scale - else: - is_q = g < 0 - mult = cutlass.Float32(1.0) - bq.append(self._qk_frag(rw, is_q, mult)) - - # the new keys into the history (one CTA of the head, warp 0) - if sp == 0: - if warp == 0: - if g < T: - kfr = bq[0] - dst = k_wr + cutlass.Int64((((woff + g) * H + h) * K + c * 8) * 2) - for jj in cutlass.range_constexpr(4): - _stg4( - dst + jj * 64, - kfr[4 * jj], - kfr[4 * jj + 1], - kfr[4 * jj + 2], - kfr[4 * jj + 3], - ) - - # new-token dot products: rows (keys | queries) x key columns - if warp == 0: - d = [zf, zf, zf, zf] - for jj in cutlass.range_constexpr(4): - for s in cutlass.range_constexpr(2): - b0 = bq[0][4 * jj + 2 * s] - b1 = bq[0][4 * jj + 2 * s + 1] - if const_expr(NT == 1): - a = (b0, zi, b1, zi) - else: - a = (b0, bq[1][4 * jj + 2 * s], b1, bq[1][4 * jj + 2 * s + 1]) - d = _mma(a, (b0, b1), d) - _sts2_f32(s_kk + (g * 8 + 2 * c) * 4, d[0], d[1]) - _sts2_f32(s_kk + ((g + 8) * 8 + 2 * c) * 4, d[2], d[3]) - - _cp_wait1() - # gating, cumulative decays, the G history (a thread per v head, off warp 0); the - # thread's own G row copy has landed - if gpred: - hl = gi - hv = hv0 + hl - gs = cutlass.Float32( - cutlass.select_( - p > 0, _lds_f32(s_gh + (hl * 16 + cutlass.max(p - 1, zi)) * 4), zf - ) - ) - graws = [_b16_f32(gtc[t]) for t in range(T)] - braws = [_b16_f32(gtc[T + t]) for t in range(T)] - alv = gfc[0] - dtb = gfc[1] - aexp = _exp(alv) - owns = (hv_first + hl) * 8 >= sp * TPC - gl = zf - for t in cutlass.range_constexpr(T): - graw = graws[t] - braw = braws[t] - if const_expr(self.gating): - x = graw + dtb - spl = cutlass.Float32( - cutlass.select_(x <= 20.0, LN2 * _lg2(1.0 + _ex2(x * LOG2E)), x) - ) - gt = -aexp * spl - beta = _div_approx(cutlass.Float32(1.0), 1.0 + _ex2(-braw * LOG2E)) - else: - gt = graw - beta = braw - gl = gl + gt - gabs = gs + gl - _sts_f32(s_g + (hl * 8 + t) * 4, gabs) - _sts_f32(s_eg + (hl * 8 + t) * 4, _exp(gabs)) - _sts_f32(s_beta + (hl * 8 + t) * 4, beta) - if owns: - _stg_f32( - g_wr + cutlass.Int64((hv * HIST + woff + t) * 4), - cutlass.Float32(cutlass.select_(wrap, gl, gabs)), - ) - _sts_f32(s_egs + hl * 4, _exp(gs)) - cute.arch.sync_threads() - - hk_warp: cutlass.Constexpr = 1 if W > 1 else 0 - if warp == hk_warp: - dh = [[zf, zf, zf, zf] for _ in range(NT)] - for jj in cutlass.range_constexpr(4): - rg = _lds4_i32(s_kh + g * PK + jj * 64 + c * 16) - rh = _lds4_i32(s_kh + (g + 8) * PK + jj * 64 + c * 16) - for s in cutlass.range_constexpr(2): - a = (rg[2 * s], rh[2 * s], rg[2 * s + 1], rh[2 * s + 1]) - for nt in cutlass.range_constexpr(NT): - dh[nt] = _mma( - a, (bq[nt][4 * jj + 2 * s], bq[nt][4 * jj + 2 * s + 1]), dh[nt] - ) - for hl in cutlass.range_constexpr(NHV): - gh_g = _lds_f32(s_gh + (hl * 16 + g) * 4) - gh_h = _lds_f32(s_gh + (hl * 16 + g + 8) * 4) - for nt in cutlass.range_constexpr(NT): - for e in cutlass.range_constexpr(2): - col = nt * 8 + 2 * c + e - if const_expr(TP == 4): - tt = (2 * c + e) % 4 - else: - tt = 2 * c + e - tv = tt < T - gtt = _lds_f32(s_g + (hl * 8 + tt) * 4) - vg = cutlass.Float32( - cutlass.select_((g < p) & tv, dh[nt][e] * _exp(gtt - gh_g), zf) - ) - vh = cutlass.Float32( - cutlass.select_( - ((g + 8) < p) & tv, dh[nt][2 + e] * _exp(gtt - gh_h), zf - ) - ) - _sts_bf16(s_cf + ((hl * 16 + col) * 16 + g) * 2, vg) - _sts_bf16(s_cf + ((hl * 16 + col) * 16 + g + 8) * 2, vh) - - if const_expr(self.commit): - if wrap: - for it2 in cutlass.range_constexpr(-(-256 // NTHR)): - e = tidx + it2 * NTHR - if e < 256: - j = e // 16 - jj = (e // 4) % 4 - aa = e % 4 - w4 = _lds4_i32(s_kh + j * PK + jj * 64 + aa * 16) - for s4 in cutlass.range_constexpr(4): - _sts_i32(s_kp + j * PU + (16 * jj + 4 * s4 + aa) * 4, w4[s4]) - - n_ts: cutlass.Constexpr = NHV * T * T - for it2 in cutlass.range_constexpr(-(-n_ts // NTHR)): - e = tidx + it2 * NTHR - if e < n_ts: - hl = e // (T * T) - t = (e // T) % T - s = e % T - gt = _lds_f32(s_g + (hl * 8 + t) * 4) - gsv = _lds_f32(s_g + (hl * 8 + s) * 4) - bt = _lds_f32(s_beta + (hl * 8 + t) * 4) - dec = _exp(gt - gsv) - kkv = _lds_f32(s_kk + (t * 8 + s) * 4) - qkv = _lds_f32(s_kk + ((TP + t) * 8 + s) * 4) - av = cutlass.Float32(cutlass.select_(s < t, bt * kkv * dec, zf)) - cv = cutlass.Float32(cutlass.select_(s <= t, qkv * dec, zf)) - _sts_f32(s_a + (hl * 64 + t * 8 + s) * 4, av) - _sts_f32(s_cq + (hl * 64 + t * 8 + s) * 4, cv) - if const_expr(self.commit): - if tidx < NHV * 16: - hl = tidx // 16 - j = tidx % 16 - gsv = cutlass.Float32(0.0) - if p > 0: - gsv = _lds_f32(s_gh + (hl * 16 + p - 1) * 4) - ghj = _lds_f32(s_gh + (hl * 16 + j) * 4) - _sts_f32( - s_dec + (hl * 16 + j) * 4, - cutlass.Float32(cutlass.select_(j < p, _exp(gsv - ghj), zf)), - ) - cute.arch.sync_threads() - - for i in cutlass.range_constexpr(TPW): - if const_expr(i == 0): - cur = [cur0[e] for e in range(32)] - else: - kslot = i - 1 - _mbar_wait(s_mb + kslot * 8, ph) - tb = s_st + kslot * 4096 + g * 256 + c * 16 - raw = [] - for ii in cutlass.range_constexpr(4): - jsw = (g % 2) ^ ii - raw.append(_lds4_i32(tb + jsw * 64)) - raw.append(_lds4_i32(tb + 2048 + jsw * 64)) - cur = [] - for jj in cutlass.range_constexpr(4): - for rr in cutlass.range_constexpr(2): - for e in cutlass.range_constexpr(4): - cur.append( - cutlass.Int32( - cutlass.select_( - godd, raw[2 * (jj ^ 1) + rr][e], raw[2 * jj + rr][e] - ) - ) - ) - mt = sp * TPC + warp + W * i - hvr = mt // 8 - hv = h * GRP + hvr - hl = hvr - hv_first - v0 = (mt % 8) * 16 - - acc0 = [[zf, zf, zf, zf] for _ in range(NT)] - acc1 = [[zf, zf, zf, zf] for _ in range(NT)] - for jj in cutlass.range_constexpr(4): - for s in cutlass.range_constexpr(2): - a = ( - cur[8 * jj + 2 * s], - cur[8 * jj + 4 + 2 * s], - cur[8 * jj + 2 * s + 1], - cur[8 * jj + 4 + 2 * s + 1], - ) - for nt in cutlass.range_constexpr(NT): - bfr = (bq[nt][4 * jj + 2 * s], bq[nt][4 * jj + 2 * s + 1]) - if const_expr(s == 0): - acc0[nt] = _mma(a, bfr, acc0[nt]) - else: - acc1[nt] = _mma(a, bfr, acc1[nt]) - acc = [[acc0[nt][e] + acc1[nt][e] for e in range(4)] for nt in range(NT)] - - au = _ldsm4t(s_u + hl * (16 * PU) + jrow * PU + (v0 + v8) * 2) - hacc = [] - for nt in cutlass.range_constexpr(NT): - cfa = s_cf + ((hl * 16 + nt * 8 + g) * 16 + 2 * c) * 2 - bc = (_lds_i32(cfa), _lds_i32(cfa + 16)) - hacc.append(_mma(au, bc, [zf, zf, zf, zf])) - - for nt in cutlass.range_constexpr(NT): - col0 = nt * 8 + 2 * c - if const_expr(TP == 4): - t0 = (2 * c) % 4 - else: - t0 = 2 * c - eg0 = _lds_f32(s_eg + (hl * 8 + t0) * 4) - eg1 = _lds_f32(s_eg + (hl * 8 + t0 + 1) * 4) - _sts2_f32( - s_x + g * RP + col0 * 4, - eg0 * acc[nt][0] + hacc[nt][0], - eg1 * acc[nt][1] + hacc[nt][1], - ) - _sts2_f32( - s_x + (g + 8) * RP + col0 * 4, - eg0 * acc[nt][2] + hacc[nt][2], - eg1 * acc[nt][3] + hacc[nt][3], - ) - cute.arch.sync_warp() - - xs = [] - ys = [] - for q4 in cutlass.range_constexpr(TP // 4): - xs.extend(_lds4_f32(s_x + r * RP + q4 * 16)) - ys.extend(_lds4_f32(s_x + r * RP + (TP + 4 * q4) * 4)) - us = [] - for t in cutlass.range_constexpr(T): - vt = _lds_bf16(s_vk + ((hl * T + t) * 128 + v0 + r) * 2) - bt = _lds_f32(s_beta + (hl * 8 + t) * 4) - u = bt * (vt - xs[t]) - for s in cutlass.range_constexpr(t): - u = u - _lds_f32(s_a + (hl * 64 + t * 8 + s) * 4) * us[s] - us.append(u) - if hf == 0: - for t in cutlass.range_constexpr(T): - o = ys[t] - for s in cutlass.range_constexpr(t + 1): - o = o + _lds_f32(s_cq + (hl * 64 + t * 8 + s) * 4) * us[s] - _stg_bf16(o_ptr + (((n64 * T + t) * HV + hv) * V + v0 + r) * 2, o) - else: - for t in cutlass.range_constexpr(T): - _stg_bf16( - u_wr + cutlass.Int64((((woff + t) * HV + hv) * V + v0 + r) * 2), - us[t], - ) - cute.arch.sync_warp() - - if const_expr(self.commit): - if wrap: - dc = [ - _lds_f32(s_dec + (hl * 16 + 2 * c + jx) * 4) for jx in [0, 1, 8, 9] - ] - egs = _lds_f32(s_egs + hl * 4) - a_hi = [] - a_lo = [] - for q4 in cutlass.range_constexpr(4): - dl = dc[0] if q4 < 2 else dc[2] - dhh = dc[1] if q4 < 2 else dc[3] - x0 = _bf2_lo(au[q4]) * dl - x1 = _bf2_hi(au[q4]) * dhh - hw = _bf2_pack(x0, x1) - a_hi.append(hw) - a_lo.append(_bf2_pack(x0 - _bf2_lo(hw), x1 - _bf2_hi(hw))) - taddr = self._tile_addr(h_ptr, slot64, h_ss, h, sp, warp, i, g, c) - kp_row = ( - s_kp + ((lane % 8) + 8 * ((lane // 8) % 2)) * PU + 16 * (lane // 16) - ) - for jj in cutlass.range_constexpr(4): - bb0 = _ldsm4t(kp_row + (4 * jj) * 16) - bb1 = _ldsm4t(kp_row + (4 * jj + 2) * 16) - ng = [] - nh8 = [] - for s4 in cutlass.range_constexpr(4): - bsrc = bb0 if s4 < 2 else bb1 - bf = (bsrc[2 * (s4 % 2)], bsrc[2 * (s4 % 2) + 1]) - dcm = _mma(a_hi, bf, [zf, zf, zf, zf]) - if const_expr(self.lo_split): - dcm = _mma(a_lo, bf, dcm) - wg = cur[8 * jj + s4] - wh = cur[8 * jj + 4 + s4] - ng.append( - _bf2_pack( - egs * _bf2_lo(wg) + dcm[0], egs * _bf2_hi(wg) + dcm[1] - ) - ) - nh8.append( - _bf2_pack( - egs * _bf2_lo(wh) + dcm[2], egs * _bf2_hi(wh) + dcm[3] - ) - ) - _stg4(taddr + jj * 64, ng[0], ng[1], ng[2], ng[3]) - _stg4(taddr + (2048 + jj * 64), nh8[0], nh8[1], nh8[2], nh8[3]) - - for ce in cutlass.range_constexpr(32): - cur0[ce] = nxt0[ce] - for ce in cutlass.range_constexpr(NQ): - qkc[ce] = qkn[ce] - for ce in cutlass.range_constexpr(GN): - gtc[ce] = gtn[ce] - for ce in cutlass.range_constexpr(2): - gfc[ce] = gfn[ce] - cute.arch.sync_threads() - - if const_expr(self.pdl and not self.early_trigger): - cute.arch.griddepcontrol_launch_dependents() - # --------------------------------------------------------------------------------------- # host side @@ -2422,49 +1756,12 @@ def kernel_p( _COMPILED: dict = {} -def parse_config(text: str) -> Optional[dict]: - """'default', positional 'tpc,warps,pf,minb[,smem_tiles[,persistent[,occ]]]' or 'key=value,...'.""" - if text in ("", "default"): - return None - if "=" in text: - out = {} - for kv in text.split(","): - k_, v_ = kv.split("=") - if v_.lstrip("-").isdigit(): - out[k_] = int(v_) - elif v_ in ("True", "False"): - out[k_] = v_ == "True" - else: - out[k_] = v_ - for b_ in ( - "smem_tiles", - "persistent", - "spec_hist", - "qk_first", - "lo_split", - "early_trigger", - "state_before_wait", - "early_acc", - ): - if b_ in out: - out[b_] = bool(out[b_]) - return out - f = [int(x) for x in text.split(",")] - names = ("tpc", "warps", "pf", "minb", "smem_tiles", "persistent", "occ") - out = dict(zip(names, f)) - for b_ in ("smem_tiles", "persistent"): - if b_ in out: - out[b_] = bool(out[b_]) - return out - - def default_config(n: int, t: int, commit: bool = False) -> dict: - """Launch mapping by request count, measured on GB300 at T = 4 (36 cold layers): small - batches split a head's 24 m-tiles over more, smaller CTAs, every state tile by bulk copy - issued as soon as the slot is known and the state dots taken before the history lands (a - latency-bound launch: 1 request 3.3 -> 2.8 us); from 32 requests (or with the in-kernel - commit from 16) CTAs of 8 m-tiles x 4 warps, capped at 5 (4 with the commit) resident per - SM's register file, keep the most state in flight.""" + """Launch mapping by request count: small batches split a head's 24 m-tiles over more, + smaller CTAs, every state tile by bulk copy issued as soon as the slot is known and the + state dots taken before the history lands (a latency-bound launch); from 32 requests (or + with the in-kernel commit from 16) CTAs of 8 m-tiles x 4 warps, capped at 5 (4 with the + commit) resident per SM's register file, keep the most state in flight.""" if n <= 2: return dict(tpc=2, warps=2, pf=0, minb=4, early_acc=True) if n <= 8: @@ -2536,8 +1833,12 @@ def cached_replay_update( The Triton launch knobs (`block_v`, `num_warps`, `use_l2_*`, `main_num_stages`) are ignored. """ del old_beta - N, T, Hq, Kd = k.shape - assert (Hq, Kd) == (H, K) and v.shape[2:] == (HV, V) and history_size == HIST, "geometry" + N, T, H, Kd = k.shape + HV = v.shape[2] + assert (H, HV) in GEOMETRIES and Kd == K and v.shape[3] == V and history_size == HIST, ( + f"geometry: {H} q/k x {Kd}, {HV} v x {v.shape[3]}, history {history_size}" + ) + qkv_width = 2 * H * K + HV * V assert 1 <= T <= 8 assert ssm_states.dtype == torch.bfloat16, "bf16 checkpoint state only" assert ssm_states.stride()[1:] == (V * K, K, 1) @@ -2567,11 +1868,11 @@ def cached_replay_update( if N == 0: return output if packed_qkv is not None: - assert packed_qkv.shape == (N * T, QKV_WIDTH) and packed_qkv.is_contiguous() + assert packed_qkv.shape == (N * T, qkv_width) and packed_qkv.is_contiguous() base = packed_qkv.data_ptr() es = packed_qkv.element_size() qp, kp, vp = base, base + H * K * es, base + 2 * H * K * es - qs = ks = vs = (T * QKV_WIDTH, QKV_WIDTH) + qs = ks = vs = (T * qkv_width, qkv_width) else: for x, nh, d in ((q, H, K), (k, H, K), (v, HV, V)): assert x.stride(3) == 1 and x.stride(2) == d, "head rows must be dense" @@ -2591,23 +1892,13 @@ def cached_replay_update( # with PDL, issue the state / history (not written by the previous kernel) before the wait; # without it, the step's own small inputs first cfg.setdefault("qk_first", not bool(launch_with_pdl)) - occ = cfg.pop("occ", 1) - use_wi = ( - bool(cfg.pop("use_wi", True)) - and replay_work_items is not None - and not cfg.get("persistent") - ) + use_wi = bool(cfg.pop("use_wi", True)) and replay_work_items is not None cfg["use_wi"] = use_wi - if cfg.get("persistent"): - sms = torch.cuda.get_device_properties(q.device).multi_processor_count - cfg["gmax"] = occ * sms - items = N * H * (TILES // cfg["tpc"]) - assert items <= 256 * min(items, cfg["gmax"]), "more than 256 items per CTA" - else: - cfg.pop("gmax", None) pdl = bool(launch_with_pdl) kw = dict( T=T, + h=H, + hv=HV, commit=commit, pdl=pdl, l2norm=bool(use_qk_l2norm_in_kernel), @@ -2653,65 +1944,399 @@ def cached_replay_update( return output +#: A tp4/tep4 rank's GDN heads: 16 key / 48 value split four ways. +RANK_K_HEADS, RANK_V_HEADS = 4, 12 +#: The stock tep4 engine's recurrent-state slots a rank (bf16 state), and the GDN layers whose +#: states share one slot page in the stock manager: the slot stride is 36 layers' states. +TEP4_SLOTS, PAGE_LAYERS = 1874, 36 + + +def _tep4_cell(why, path, n, t, **spec) -> Cell: + """A tp4/tep4 rank's heads over the stock tep4 engine's bf16 state page.""" + return Cell( + why=f"tep4 {why}", + spec=dict( + path=path, + n=n, + t=t, + state_dtype=torch.bfloat16, + H=RANK_K_HEADS, + HV=RANK_V_HEADS, + history_size=HIST, + slots=TEP4_SLOTS, + page_layers=PAGE_LAYERS, + **spec, + ), + ) + + class GdnCachedReplayCute(OpWrapper): """`cached_replay_update` as a catalog entry: one verify step of `N` sequences x `T` tokens - through one GDN layer under the cached-replay state layout, bf16 checkpoint state only, - with the same contract as `fused_recurrent_gated_delta_rule_cached_replay_update` (see that - entry's docstring: history rebuild, gated delta rule over the T new tokens, history append, - and -- `in_kernel_commit` -- the checkpoint advance of a wrapping row written in place). + through one GDN layer under the cached-replay state layout, bf16 checkpoint state only, with + the contract of `fused_recurrent_gated_delta_rule_cached_replay_update` (history rebuild, + gated delta rule over the T new tokens, history append, and the checkpoint advance of a + wrapping row written in place when no work items are given -- or, `in_kernel_commit=True`, + always). With work items and `use_all_layer_commit` it commits nothing (the manager's + all-layer commit does), exactly as the Triton entry. - With `in_kernel_commit=True` every wrapping row commits here at any `N`: the caller must - not also run the manager's all-layer commit (the checkpoint would advance twice). + The q/k and value head counts are compile-time parameters taken from the call's shapes + (`GEOMETRIES`: a whole layer's 16 / 48 and a tp4/tep4 rank's 4 / 12), and `g` / `beta` may + be strided column views of the input projection (any row stride, unit head stride). Fusion boundary: as the Triton entry's. """ ARCHS = frozenset({Arch.SM_103}) - # spec: the Triton entry's test harness (`_Bench`) spec -- path, n sequences x t tokens, - # consecutive steps with random acceptance, pool slots, layers -- run in the commit mode - # the target uses at that size: in-kernel on the small path, the manager's all-layer commit - # on the partitioned one. + # spec: the Triton entry's test harness (`_Bench`) spec -- path, n sequences x t + # tokens, consecutive steps with random acceptance, start, accept, dup, H / HV, pool slots, + # page_layers (the stock manager's coalesced state page) -- plus `strided_gates` (g / beta + # handed over as column views of a `[N*T, conv + z + 2 HV]` projection, as the target does). CELLS: tuple[Cell, ...] = ( - Cell( - why="MTP3 verify at con1: N=1, T=4, in-kernel commit", - spec=dict(path="small", n=1, t=4, steps=6), + _tep4_cell( + "small n=1 T=4: MTP3 verify at con1, in-kernel commit on the page's last slot", + "small", + 1, + 4, + steps=8, + strided_gates=True, + ), + _tep4_cell( + "small n=4 T=4: the largest captured ll batch", + "small", + 4, + 4, + steps=6, + strided_gates=True, + ), + _tep4_cell( + "small n=4 T=4 dup: rows 2..3 share one CUDA-graph padding slot", + "small", + 4, + 4, + steps=6, + dup=(2, 3), + strided_gates=True, ), - Cell( - why="MTP3 verify at con32: N=32, T=4, work items + the manager's all-layer commit", - spec=dict(path="all_layer", n=32, t=4, steps=4, slots=72), + _tep4_cell("small n=2 T=2: MTP1 verify", "small", 2, 2, steps=8, strided_gates=True), + _tep4_cell( + "small n=1 T=8 full accept: MTP7 appends to the history's end, then wraps", + "small", + 1, + 8, + steps=8, + start="fresh", + accept="full", + strided_gates=True, ), - Cell( - why="MTP3 verify at con256: N=256, T=4, work items + the manager's all-layer commit", - spec=dict(path="all_layer", n=256, t=4, steps=3, slots=300, layers=1), + _tep4_cell( + "all_layer n=16 T=4: the smallest partitioned batch, all-layer commit", + "all_layer", + 16, + 4, + steps=6, + strided_gates=True, ), + _tep4_cell("all_layer n=128 T=4: a large partitioned batch", "all_layer", 128, 4, steps=4), ) note = """ - Metadata consumed: as the Triton entry (slots, per-slot pending count and history half), - or the engine's replay work items (row, slot, pending count, half) over every row in one - load each. Call site: `_replay_verify` over a bf16 recurrent pool, under the Triton - entry's contract: no work items below `CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE` requests - (wrapping rows commit here), work items and the manager's all-layer commit from there on. - `g` / `beta` may be the in_proj's column views (any row stride). Compiles one kernel per - (T, launch mapping, commit, PDL, ...) on first use (5-10 s): warm-up runs before CUDA-graph - capture cover it. - Not supported: fp32 state, T > 8, another geometry. Optimized for T = 4; T = 1..8 are - tested for correctness. Every launch mapping is bitwise every other (the per-row - arithmetic does not depend on it). - - GB300, T = 4, in-kernel commit, 25 % of the rows wrapping, PDL off, a CUDA graph of 36 - layers each streaming its own state, us a layer: 1 request 2.8 (2.76 with the work - items), 4: 4.8 (4.4), 32: 17.0 (16.2), 128: 52.3 (51.9), 256: 100.9 (99.6). + Call site: `_gdn_verify` (`DecodeTarget._gdn_block` / `PrefillTarget._gdn_block` of a verify + step) over the bf16 recurrent pool, at a tp4/tep4 rank's 4 / 12 heads: q/k/v from the conv's + fresh contiguous `[N*T, 2560]` rows (`packed_qkv`), `g` / `beta` the input projection's raw + a / b columns read through its row stride (no copies), `A_log` / `dt_bias` fp32, the + layer's view of the stock manager's bf16 state page (slot stride 36 layers' states), the + history buffers `[slots, 2, 16, 12|4, 128]` / `[slots, 2, 12, 16]`, and work items + + `use_all_layer_commit` exactly when the Triton entry gets them (N >= 16 with the manager's + all-layer commit on). `launch_with_pdl=True`: the Triton conv before it signals its + dependents at its start (`launch_dependent_kernels`), so the kernel streams the checkpoint + and the history -- which nothing in the step writes ahead of it -- while the conv runs, and + reads q/k/v, a and b only after `griddepcontrol.wait`. Compiles one kernel per (geometry, + T, launch mapping, commit, PDL, ...) on first use (5-10 s): the warm-up before CUDA-graph + capture covers it. + + Certified at 4 / 12 heads over a bf16 checkpoint (cells: the stock tep4 engine's 36-layer + state page at 1874 slots, rows on both ends of it, several consecutive verify steps from + warm histories with random acceptance): the in-kernel commit at N = 1, 2, 4 and T = 2, 4, + 8 (T = 8 with full acceptance appending to the history's end, then wrapping), CUDA-graph + padding rows sharing one slot, and the all-layer commit at N = 16 and 128 (T = 4); g / beta + as column views of a `[N*T, 4120]` projection. Against the Triton entry's bands (fractions + of each gate): output elements <= 0.50 of 4 bf16 ulps of their token block, relative RMS + <= 0.71 of 1.5 ulps, bf16 checkpoint <= 0.68 of 2 ulps of its head's absmax, `old_u` <= + 0.50 of 4 ulps, `old_G` <= 0.06 -- the Triton entry's own margins (0.52 / 0.64 / 0.84 / + 0.50 / 0.08). A reference with the key heads reversed is refused by the output gate. + CUDA-graph replay (PDL on and off, N = 1, 4 and 16) is bitwise eager, and the launch + mappings are bitwise each other at N = 1, 2, 4 and T = 2, 4, 8. + + Numerics: the Triton kernel's rounding points (bf16 L2-normalized q / k and history + coefficients into the tensor-core dots, fp32 recurrence, fp32 commit rounded once to the + bf16 pool, bf16 history u / k, fp32 G); the dot accumulation order differs, so outputs + match the Triton kernel to a few bf16 ulps, not bitwise. Every launch mapping is bitwise + every other (the per-row arithmetic does not depend on it). + + Not supported: fp32 state, T > 8, a geometry outside `GEOMETRIES`, K or V other than 128, + history other than 16, pre-gated g / beta (`A_log=None` is compiled but not certified). + The Triton launch knobs (`block_v`, `num_warps`, `use_l2_*`, `main_num_stages`) are ignored. """ - def __call__(self, *args, **kwargs) -> torch.Tensor: - return cached_replay_update(*args, **kwargs) + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + ssm_states: torch.Tensor, + state_indices: torch.Tensor, + old_u: torch.Tensor, + old_k: torch.Tensor, + old_G: torch.Tensor, + old_beta: torch.Tensor, + cache_buf_idx: torch.Tensor, + prev_num_accepted_tokens: torch.Tensor, + history_size: int, + scale: Optional[float] = None, + use_qk_l2norm_in_kernel: bool = False, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + launch_with_pdl: bool = False, + replay_work_items: Optional[torch.Tensor] = None, + n_writes: Optional[torch.Tensor] = None, + block_v: Optional[int] = None, + num_warps: Optional[int] = None, + use_l2_state_cache: Optional[bool] = None, + use_l2_streaming_inputs: Optional[bool] = None, + use_l2_shared_inputs: Optional[bool] = None, + main_num_stages: Optional[int] = None, + packed_qkv: Optional[torch.Tensor] = None, + use_all_layer_commit: bool = False, + output: Optional[torch.Tensor] = None, + in_kernel_commit: Optional[bool] = None, + config: Optional[dict] = None, + ) -> torch.Tensor: + return cached_replay_update( + q=q, + k=k, + v=v, + g=g, + beta=beta, + ssm_states=ssm_states, + state_indices=state_indices, + old_u=old_u, + old_k=old_k, + old_G=old_G, + old_beta=old_beta, + cache_buf_idx=cache_buf_idx, + prev_num_accepted_tokens=prev_num_accepted_tokens, + history_size=history_size, + scale=scale, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + A_log=A_log, + dt_bias=dt_bias, + launch_with_pdl=launch_with_pdl, + replay_work_items=replay_work_items, + n_writes=n_writes, + packed_qkv=packed_qkv, + use_all_layer_commit=use_all_layer_commit, + output=output, + in_kernel_commit=in_kernel_commit, + config=config, + ) + + def reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + ssm_states: torch.Tensor, + state_indices: torch.Tensor, + old_u: torch.Tensor, + old_k: torch.Tensor, + old_G: torch.Tensor, + old_beta: torch.Tensor, + cache_buf_idx: torch.Tensor, + prev_num_accepted_tokens: torch.Tensor, + history_size: int, + scale: Optional[float] = None, + use_qk_l2norm_in_kernel: bool = False, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + launch_with_pdl: bool = False, + replay_work_items: Optional[torch.Tensor] = None, + n_writes: Optional[torch.Tensor] = None, + block_v: Optional[int] = None, + num_warps: Optional[int] = None, + use_l2_state_cache: Optional[bool] = None, + use_l2_streaming_inputs: Optional[bool] = None, + use_l2_shared_inputs: Optional[bool] = None, + main_num_stages: Optional[int] = None, + packed_qkv: Optional[torch.Tensor] = None, + use_all_layer_commit: bool = False, + output: Optional[torch.Tensor] = None, + in_kernel_commit: Optional[bool] = None, + config: Optional[dict] = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """`(out, ssm_states, old_u, old_k, old_G)` as the call leaves them; + mutates nothing. The state is rebuilt from the checkpoint and history + in fp32 and the new tokens are stepped one at a time in fp32 -- no + matmul, so no TF32 -- and every store is rounded once to its buffer's + dtype. The history is read at whatever dtype and depth its buffers + have, so the same function driven on fp32 history deeper than any run + ever fills is the exact token-by-token recurrence. + + Rows with a negative slot keep `output`'s prior content, or are NaN + (unspecified) when no `output` is given. `old_beta`, `launch_with_pdl`, + `n_writes` and the launch-tuning knobs are ignored; `replay_work_items` + only decides whether a wrap commits the checkpoint here, unless `in_kernel_commit` + says (True: every wrap commits here; False: none does); `config` is ignored.""" + N, T, H, K = k.shape + HV, V = v.shape[2], v.shape[3] + rep = HV // H + dev = v.device + scale = K**-0.5 if scale is None else scale + if packed_qkv is not None: + rows = packed_qkv.view(N, T, -1) + q = rows[..., : H * K].view(N, T, H, K) + k = rows[..., H * K : 2 * H * K].view(N, T, H, K) + v = rows[..., 2 * H * K :].view(N, T, HV, V) + qf, kf, vf = q.float(), k.float(), v.float() + if use_qk_l2norm_in_kernel: + qf = qf / torch.sqrt(qf.square().sum(-1, keepdim=True) + 1e-6) + kf = kf / torch.sqrt(kf.square().sum(-1, keepdim=True) + 1e-6) + qf = (qf * scale).repeat_interleave(rep, dim=2) # [N, T, HV, K] + kv = kf.repeat_interleave(rep, dim=2) # [N, T, HV, K] + if A_log is not None: + gf = -torch.exp(A_log.float()) * F.softplus(g.float() + dt_bias.float()) + bf = torch.sigmoid(beta.float()) + else: + gf, bf = g.float(), beta.float() - def reference(self, *args, **kwargs): - raise NotImplementedError( - "the multi-step cell harness drives the Triton entry's reference (see the test)" + out_dtype = q.dtype if output is None else output.dtype + out = ( + torch.full((N, T, HV, V), float("nan"), device=dev) + if output is None + else output.float().clone() + ) + pool, u_buf, k_buf, G_buf = (x.clone() for x in (ssm_states, old_u, old_k, old_G)) + live = (state_indices[:N] >= 0).nonzero().flatten() + if live.numel() == 0: + return out.to(out_dtype), pool, u_buf, k_buf, G_buf + slot = state_indices[:N][live].long() + p = prev_num_accepted_tokens[slot].long() + b = cache_buf_idx[slot].long() + M, depth = live.numel(), old_u.shape[2] + + # 1. the pending state: checkpoint plus the p history entries of half b + pending = torch.arange(depth, device=dev)[None, :] < p[:, None] # [M, depth] + Gh = torch.where(pending[:, None, :], old_G[slot, b].float(), 0.0) # [M, HV, depth] + Uh = torch.where(pending[:, :, None, None], old_u[slot, b].float(), 0.0) + Kh = torch.where(pending[:, :, None, None], old_k[slot, b].float(), 0.0) + Kh = Kh.repeat_interleave(rep, dim=2) # [M, depth, HV, K] + last = (p - 1).clamp(min=0)[:, None, None].expand(M, HV, 1) + g_start = torch.where(p[:, None] > 0, Gh.gather(2, last).squeeze(2), 0.0) # [M, HV] + weight = torch.where(pending[:, None, :], torch.exp(g_start[:, :, None] - Gh), 0.0) + h = torch.exp(g_start)[:, :, None, None] * ssm_states[slot].float() # [M, HV, V, K] + for j in range(depth): + h = h + weight[:, :, j, None, None] * Uh[:, j, :, :, None] * Kh[:, j, :, None, :] + checkpoint = h + + # 2. the new tokens, in order + gl, bl = gf[live], bf[live] # [M, T, HV] + ql, kl, vl = qf[live], kv[live], vf[live] + G_local = torch.zeros(M, HV, device=dev) + u_new, G_new = [], [] + for t in range(T): + h = h * torch.exp(gl[:, t])[:, :, None, None] + u = bl[:, t, :, None] * (vl[:, t] - (h * kl[:, t, :, None, :]).sum(-1)) # [M, HV, V] + h = h + u[:, :, :, None] * kl[:, t, :, None, :] + out[live, t] = (h * ql[:, t, :, None, :]).sum(-1) + G_local = G_local + gl[:, t] + u_new.append(u) + G_new.append(G_local.clone()) + + # 3. the history append, and the checkpoint commit on a wrap + wrap = p + T > history_size + w_buf = torch.where(wrap, 1 - b, b) + w_off = torch.where(wrap, 0, p) + for t in range(T): + pos = w_off + t + u_buf[slot, w_buf, pos] = u_new[t].to(u_buf.dtype) + k_buf[slot, w_buf, pos] = kf[live, t].to(k_buf.dtype) + G_store = torch.where(wrap[:, None], G_new[t], g_start + G_new[t]) + G_buf[slot, w_buf, :, pos] = G_store.to(G_buf.dtype) + commit = replay_work_items is None if in_kernel_commit is None else bool(in_kernel_commit) + if commit and bool(wrap.any()): + pool[slot[wrap]] = checkpoint[wrap].to(pool.dtype) + return out.to(out_dtype), pool, u_buf, k_buf, G_buf + + def is_valid( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + ssm_states: torch.Tensor, + state_indices: torch.Tensor, + old_u: torch.Tensor, + old_k: torch.Tensor, + old_G: torch.Tensor, + old_beta: torch.Tensor, + cache_buf_idx: torch.Tensor, + prev_num_accepted_tokens: torch.Tensor, + history_size: int, + scale: Optional[float] = None, + use_qk_l2norm_in_kernel: bool = False, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + launch_with_pdl: bool = False, + replay_work_items: Optional[torch.Tensor] = None, + n_writes: Optional[torch.Tensor] = None, + block_v: Optional[int] = None, + num_warps: Optional[int] = None, + use_l2_state_cache: Optional[bool] = None, + use_l2_streaming_inputs: Optional[bool] = None, + use_l2_shared_inputs: Optional[bool] = None, + main_num_stages: Optional[int] = None, + packed_qkv: Optional[torch.Tensor] = None, + use_all_layer_commit: bool = False, + output: Optional[torch.Tensor] = None, + in_kernel_commit: Optional[bool] = None, + config: Optional[dict] = None, + ) -> None: + N, _, H, K = k.shape + HV, V = v.shape[2], v.shape[3] + assert state_indices.dim() == 1 and state_indices.shape[0] >= N, ( + f"state_indices has shape {tuple(state_indices.shape)} for {N} sequences -- the " + "kernel reads one slot per sequence with no bound, so a shorter tensor hands the " + "trailing rows whatever follows it in memory, and they write that slot's history" + ) + pool_slots, slots = ssm_states.shape[0], old_u.shape[0] + assert slots >= pool_slots, ( + f"the history buffers have {slots} slots for a {pool_slots}-slot pool -- a slot " + "past the history's end is read and written out of bounds" ) + for name, tensor in ( + ("prev_num_accepted_tokens", prev_num_accepted_tokens), + ("cache_buf_idx", cache_buf_idx), + ): + assert tensor.dim() == 1 and tensor.is_contiguous() and tensor.shape[0] >= pool_slots, ( + f"{name} has shape {tuple(tensor.shape)} and stride {tensor.stride()}; the kernel " + f"reads one element per slot of the {pool_slots}-slot pool, at the slot's index" + ) + expected = { + "old_u": (old_u, (slots, 2, history_size, HV, V)), + "old_k": (old_k, (slots, 2, history_size, H, K)), + "old_G": (old_G, (slots, 2, HV, history_size)), + } + for name, (tensor, shape) in expected.items(): + assert tuple(tensor.shape) == shape, ( + f"{name} has shape {tuple(tensor.shape)}; history_size={history_size} and this " + f"geometry need {shape} -- the kernel addresses the history by history_size and " + "the head geometry, never by the buffer's own shape" + ) gdn_cached_replay_cute = GdnCachedReplayCute() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/vision/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/vision/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/vision/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/vision/triton_pos_embed_interpolate.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/vision/triton_pos_embed_interpolate.py new file mode 100644 index 000000000000..5b395305068b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ops/vision/triton_pos_embed_interpolate.py @@ -0,0 +1,210 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Vision position-embedding bilinear interpolation onto one image's patch +grid, emitted in spatial-merge order, by the stock Qwen3-VL tower's fused +Triton kernel. + +`__call__` is one call to `_triton_pos_embed_interpolate` in +`tensorrt_llm/_torch/models/modeling_qwen3vl.py`, which allocates the output +and makes one launch of the Triton JIT kernel `_bilinear_pos_embed_kernel` +beside it: the four corner gathers, the bilinear weights and the merge-order +row permutation in one pass. It is the kernel the stock `modeling_qwen4_exp` +vision tower runs (`Qwen3VisionModel.fast_pos_embed_interpolate`), one call per +image. Concatenating several images' rows is the caller's. +""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.models.modeling_qwen3vl import _triton_pos_embed_interpolate + +#: This checkpoint: num_position_embeddings 2304 (a 48 x 48 grid), hidden 1152, +#: spatial_merge_size 2. +NUM_GRID = 48 +HIDDEN = 1152 +MERGE = 2 + + +def _ckpt(t: int, h: int, w: int) -> dict: + """The checkpoint's bf16 table at one `(t, h, w)` grid.""" + return dict( + t=t, + h=h, + w=w, + num_grid_per_side=NUM_GRID, + m_size=MERGE, + dtype=torch.bfloat16, + hidden=HIDDEN, + ) + + +class TritonPosEmbedInterpolate(OpWrapper): + """Resample a square row-major `[N*N, hidden]` table onto one `(t, h, w)` + patch grid; return a fresh `[t*h*w, hidden]` tensor in `dtype`. + + Corner-aligned bilinear (`align_corners=True`): row `i` of `h` samples the + table at `i * (N - 1) / (h - 1)` (0 when `h == 1`), `lo = floor`, + `hi = min(lo + 1, N - 1)`; `h` indexes table rows (stride `N`), `w` + columns. Rows come out in spatial-merge order -- whole `m_size x m_size` + blocks row-major, row-major inside a block -- which is what the patch + merger's 4-to-1 reshape consumes. `t` repeats the frame's rows. + + Arithmetic is the HF reference's (`fast_pos_embed_interpolate`): the four + corner weights are computed in fp32, rounded to `dtype`, and the + multiply-adds run in `dtype`. Not an fp32 accumulation; see `note`. + + Fusion boundary: one image's position embeddings. Adding them to the patch + embeddings, the 2-D rope, and joining images are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: t, h, w, num_grid_per_side, m_size, dtype, hidden. The tower is + # replicated on every rank, so these are every member's shapes. + CELLS: tuple[Cell, ...] = ( + Cell(why="one 4x8 image: 32 rows, both sides below the table's 48", spec=_ckpt(1, 4, 8)), + Cell(why="the 62x90 serving grid: 5580 rows, one image per call", spec=_ckpt(1, 62, 90)), + Cell(why="a 3-frame 6x10 video: t repeats one frame's 60 rows", spec=_ckpt(3, 6, 10)), + Cell( + why="128x256: 32768 rows, the tower's per-call cap and the envelope's largest", + spec=_ckpt(1, 128, 256), + ), + ) + + note = """ + Source: a Triton JIT kernel, compiled at first call for the device it runs + on. `H` and `W` are int32, the scales fp32, and `NUM_GRID`, `M_SIZE`, + `HIDDEN_DIM` (1152) and `BLOCK_D` (2048) are constexprs, so another table + geometry compiles another kernel. Tier 2: a callable already in the tree (a + private function of the stock Qwen3-VL modeling file), not a registered + `torch.ops.trtllm` op. Importing this entry imports `modeling_qwen3vl`. + + Accuracy, against the exact (fp64) corner-aligned bilinear value, as + row-scale element error / relative RMS in bf16 ulp (`assert_within_ulp`): + at most 3.522 / 0.752 over 50 seeds of the test's smooth table (a bicubic + upsampled normal field, neighbour cosine 0.911 against the checkpoint + table's 0.907), and at most 3.349 / 0.745 on the checkpoint's own + `model.visual.pos_embed.weight`. The test gates 4 / 1. The closest wrong + variant (align_corners=False) lands at 41.3 / 15.2 or more; raster row + order and swapped table axes above 400 / 300. + + Not certified on tables whose neighbouring rows are uncorrelated: a normal + table with independent rows reaches 4.182 / 0.879 over 20 seeds. + Averaging uncorrelated corners shrinks the output's row maximum below the + corners' magnitude, and the bf16 rounding error scales with the corners. A + learned 2-D position table is not such a table. + + Silent, refused by `is_valid` (each measured on GB300): + - a strided table: the kernel indexes `embed_weight` as dense rows of + `hidden`. A column slice of a wider buffer returned, differing from the + same values made contiguous by up to 5.05. + - `num_grid_per_side ** 2 != embed_weight.shape[0]`: the grid side is an + argument, not read off the table. 47 on the 2304-row table returned a + correctly shaped result 564 / 340 ulp off. + - an empty side (`t`, `h` or `w` of 0): returns `[0, hidden]`, no error. + - `h` or `w` not divisible by `m_size`: the wrapper's own check is a bare + `assert`. Under `python -O` a 6x9 grid returned a finite `[54, 1152]` + tensor, its rows read from positions outside the grid and the table. + + Hard limits: the output offset is `program_id * hidden` in int32, so + `t*h*w*hidden` must stay below 2^31 (1,864,135 rows at hidden 1152; the + tower's per-call cap is 32768). Positions are `row * scale` in fp32, + exact up to 2^24 rows per side. + + Not certified: fp16 or fp32 tables, a `dtype` other than the table's, + `m_size` other than 2, a hidden size other than 1152, a CPU table, and + CUDA-graph capture (the call site, the vision tower, is eager only; the + grid is a host argument, frozen at capture). + """ + + def __call__( + self, + embed_weight: torch.Tensor, + t: int, + h: int, + w: int, + num_grid_per_side: int, + m_size: int, + dtype: torch.dtype, + ) -> torch.Tensor: + return _triton_pos_embed_interpolate( + embed_weight, t, h, w, num_grid_per_side, m_size, dtype + ) + + def reference( + self, + embed_weight: torch.Tensor, + t: int, + h: int, + w: int, + num_grid_per_side: int, + m_size: int, + dtype: torch.dtype, + ) -> torch.Tensor: + """The explicit four-corner weighted sum in fp64, returned as fp32. + + Sample positions are the exact `i * (N - 1) / (size - 1)` in fp64. + `dtype` is not applied: the result is the exact value rounded once to + fp32, so a test's band in `dtype` ulp measures the op's error alone + rather than adding the reference's own rounding to `dtype`. + """ + n = num_grid_per_side + hidden = embed_weight.shape[1] + device = embed_weight.device + grid = embed_weight.double().reshape(n, n, hidden) + + def positions(size: int) -> torch.Tensor: + index = torch.arange(size, dtype=torch.float64, device=device) + return index * ((n - 1) / (size - 1)) if size > 1 else torch.zeros_like(index) + + h_pos, w_pos = positions(h), positions(w) + h_lo, w_lo = h_pos.floor().long(), w_pos.floor().long() + h_hi = (h_lo + 1).clamp(max=n - 1) + w_hi = (w_lo + 1).clamp(max=n - 1) + d_h = (h_pos - h_lo)[:, None, None] + d_w = (w_pos - w_lo)[None, :, None] + raster = ( + grid[h_lo][:, w_lo] * ((1 - d_h) * (1 - d_w)) + + grid[h_lo][:, w_hi] * ((1 - d_h) * d_w) + + grid[h_hi][:, w_lo] * (d_h * (1 - d_w)) + + grid[h_hi][:, w_hi] * (d_h * d_w) + ) + rows = ( + raster.reshape(h // m_size, m_size, w // m_size, m_size, hidden) + .permute(0, 2, 1, 3, 4) + .reshape(h * w, hidden) + ) + return (rows.repeat(t, 1) if t > 1 else rows).float() + + def is_valid( + self, + embed_weight: torch.Tensor, + t: int, + h: int, + w: int, + num_grid_per_side: int, + m_size: int, + dtype: torch.dtype, + ) -> None: + assert embed_weight.is_contiguous(), ( + "embed_weight must be contiguous; the kernel indexes it as dense rows " + "and reads a strided table wrongly without raising" + ) + assert num_grid_per_side * num_grid_per_side == embed_weight.shape[0], ( + f"num_grid_per_side {num_grid_per_side} does not describe a " + f"{embed_weight.shape[0]}-row table; the kernel resamples the grid it is " + f"told, in or out of the table" + ) + assert t > 0 and h > 0 and w > 0, ( + f"empty grid side (t, h, w) = ({t}, {h}, {w}); the kernel returns no rows " + f"rather than raising" + ) + assert h % m_size == 0 and w % m_size == 0, ( + f"grid sides h={h}, w={w} are not whole m_size={m_size} merge blocks; " + f"the kernel's bare assert is all that stops an out-of-grid read" + ) + + +triton_pos_embed_interpolate = TritonPosEmbedInterpolate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/pool.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/pool.py new file mode 100644 index 000000000000..4d88a74eaeb0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/pool.py @@ -0,0 +1,488 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The KV cache manager's pools as one rank of this target sees them, and their contract. + +Every rank builds the stock ``Qwen4ExpHybridCacheManager`` (``modules/qwen4_exp/cache_manager.py``, +a ``MambaHybridCacheManagerV2``): QSA is this architecture's default sparse attention, so +``get_kv_cache_manager_cls`` resolves the manager through ``get_sparse_attn_kv_cache_manager``, +which requires KV cache manager V2. The manager sizes each pool from the checkpoint's text config +and the rank's ``Mapping``; `pool_geometry` derives the same per-rank shapes the same way: + +=================== ============================================ ============================= +pool shape after the slot dim split over ranks by +=================== ============================================ ============================= +paged K/V (HND) ``[2, H_kv, tokens_per_block, kv_width]`` attention TP (a) +NVFP4 block scales ``[2, H_kv, tokens_per_block, head_dim/16]`` attention TP (a) +QSA index plane ``[tokens_per_block, 1, index_head_dim]`` replicated (b) +QSA position plane ``[tokens_per_block, 3]`` int32 replicated (b) +GDN conv state ``[2 * H_k * D_k + H_v * D_v, kernel - 1]`` recurrent-state TP (c) +GDN recurrent state ``[H_v, D_v, D_k]`` recurrent-state TP (c) +PLE conv state ``[hc_count * hidden, (ple_kernel-1) * n]`` replicated (d) +PLE n-gram history ``[n - 1]`` int64 replicated (d) +=================== ============================================ ============================= + +``H_kv`` is ``num_key_value_heads``; ``H_k``/``D_k`` and ``H_v``/``D_v`` are the GDN key and value +heads (``linear_num_key_heads``, ``linear_key_head_dim``, ...); ``kernel`` is +``linear_conv_kernel_dim``, ``ple_kernel`` ``ple_conv_kernel_size`` and ``n`` ``ngram_size``. Head +counts are the rank's share. + +(a) ``KVCacheManagerV2.__init__`` splits ``num_key_value_heads`` over ``tp_size``, or over 1 under + attention DP, rounding up: with more TP ranks than KV heads each head is repeated, and rank + ``r`` holds global KV head ``r * H / tp`` (the order ``duplicate_kv_weight`` repeats them in). +(b) ``Qwen4ExpHybridCacheManager._extra_buffers_per_layer`` sizes them from ``QSASparseParams`` + (``index_kv_heads``, ``index_head_dim``), which no mapping divides: the indexer is replicated. + Under an FP8 or NVFP4 pool both are packed into one ``QSA_SIDE_PLANES`` buffer; the views + keep these shapes either way. +(c) ``MambaHybridCacheManagerV2.__init__`` divides the GDN value heads, key heads (``n_groups``) + and ``conv_dim`` by ``mamba_effective_tp_size``: ``tp_size``, or 1 under attention DP. The + conv channels of a rank are ``[q | k | v]`` of its own heads (``conv_state_layout`` "q_k_v", + the manager's ``conv_section_dims``), and rank ``r`` holds key heads ``[r * H_k / tp, ...)`` + and value heads ``[r * H_v / tp, ...)``: the contiguous head ranges a disaggregated transfer + between layouts moves (``get_disagg_role_mapper_kinds``: conv ``SECTIONED``, recurrent state + by head). +(d) ``extract_qwen4_exp_ple_cache_params``: the PLE layer reads the whole hyper-connection + stream on every rank. + +At tp4 and tep4 a rank therefore holds one of the two KV heads, four key and twelve value GDN +heads; at dep4 it holds all of them, and the recurrent pools carry one more slot (the +attention-DP dummy request's). With MTP, attention TP puts the draft layer in a separate draft +KV cache (``should_use_separate_draft_kv_cache``); under attention DP the draft layer is a layer +of the trunk's manager. Its pools have the trunk's per-rank shapes in either. + +The ``check_*`` functions read only ``.shape``, ``.dtype`` and ``.stride()``, so they run on meta +tensors as well as on a manager's views. Each one is a wrong pool its consuming kernel would +accept and mis-address. +""" + +from __future__ import annotations + +from typing import NamedTuple + +import torch + +from tensorrt_llm._torch.pyexecutor.config_utils import mamba_effective_tp_size +from tensorrt_llm.quantization.mode import QuantAlgo + +#: Coordinates per token in the QSA position plane (the mRoPE layout's three rows). +QSA_POSITION_AXES = 3 + +#: Values per e4m3 block scale of an NVFP4 K/V pool. +NVFP4_BLOCK = 16 + +#: The recurrent-state pool dtypes the target serves (`kv_cache_config.mamba_ssm_cache_dtype`): the +#: checkpoint's fp32, or bf16. +SSM_STATE_DTYPES = (torch.float32, torch.bfloat16) + +#: The GDN cached replay's history length the replay entry is certified at. +GDN_REPLAY_HISTORY = 16 + +#: The most query rows a verify step carries per generation request, `max_draft_len + 1`: the +#: replay kernel asserts `T <= 8` and `thop_attention`'s `predicted_tokens_per_seq` is certified +#: up to it. +MAX_VERIFY_WIDTH = 8 + + +class PoolGeometry(NamedTuple): + """The pool shapes one rank binds. Head counts and channel widths are this rank's share.""" + + tokens_per_block: int + heads_kv: int # KV heads in this rank's K/V pages + kv_head_start: int # the global KV head the rank's first page head holds + head_dim: int + index_kv_heads: int + index_head_dim: int + conv_dim: int # this rank's GDN conv channels, [q | k | v] of its own heads + conv_kernel: int + num_k_heads: int # this rank's GDN key heads + k_head_start: int # the global key head the rank's first one is + num_v_heads: int # this rank's GDN value heads + v_head_start: int # the global value head the rank's first one is + head_v_dim: int + head_k_dim: int + stream: int # the hyper-connection stream width, the PLE conv channels + ple_state_len: int + ngram_context_len: int + dtype: torch.dtype # the activation dtype: GDN conv, QSA index and PLE conv pools + kv_dtype: torch.dtype # the paged K/V pool's dtype: bf16, e4m3 (FP8) or uint8 (NVFP4) + kv_width: int # its innermost extent: head_dim, or head_dim // 2 for NVFP4's e2m1 pairs + + @property + def conv_sections(self) -> tuple[int, int, int]: + """The q, k and v channel widths of this rank's conv state, in that order.""" + key = self.num_k_heads * self.head_k_dim + return key, key, self.num_v_heads * self.head_v_dim + + @property + def kv_block_scales(self) -> bool: + """Whether the K/V pool is NVFP4, whose e4m3 block scales sit in a second pool.""" + return self.kv_dtype == torch.uint8 + + +def attention_tp_size(mapping) -> int: + """The ranks the KV heads are split over (`KVCacheManagerV2.__init__`).""" + return 1 if mapping.enable_attention_dp else mapping.tp_size + + +def _share(total: int, ranks: int, rank: int) -> tuple[int, int]: + """(heads on `rank`, its first global head) when `total` heads are split over `ranks`: an even + split, or each head repeated on `ranks // total` neighbouring ranks when there are more + ranks.""" + return -(-total // ranks), rank * total // ranks + + +#: The paged K/V pool's storage per `kv_cache_quant_algo`: (dtype or None for the activation +#: dtype, values per stored element). +_KV_STORAGE = { + None: (None, 1), + QuantAlgo.FP8: (torch.float8_e4m3fn, 1), + QuantAlgo.NVFP4: (torch.uint8, 2), +} + + +def pool_geometry( + cfg, mapping, *, tokens_per_block: int, kv_cache_quant_algo, dtype: torch.dtype +) -> PoolGeometry: + """The pools rank `mapping.rank` binds. + + Args: + cfg: the engine-normalized text config (`Qwen4ExpTextConfig`) the manager is sized from. + mapping: this rank's `Mapping`. + tokens_per_block: the manager's page size (`kv_cache_config.tokens_per_block`). + kv_cache_quant_algo: the KV-cache algorithm the engine resolved: None (the checkpoint's + bf16), `QuantAlgo.FP8` or `QuantAlgo.NVFP4`. + dtype: the activation dtype. + """ + if kv_cache_quant_algo not in _KV_STORAGE: + raise ValueError( + f"kv_cache_quant_algo {kv_cache_quant_algo!r}; the K/V pool stores the checkpoint's " + "bf16, FP8 or NVFP4" + ) + attn_tp = attention_tp_size(mapping) + heads_kv, kv_head_start = _share( + cfg.num_key_value_heads, attn_tp, mapping.tp_rank if attn_tp > 1 else 0 + ) + state_tp = mamba_effective_tp_size(mapping) + state_rank = mapping.tp_rank if state_tp > 1 else 0 + num_k_heads, k_head_start = _share(cfg.linear_num_key_heads, state_tp, state_rank) + num_v_heads, v_head_start = _share(cfg.linear_num_value_heads, state_tp, state_rank) + kv_dtype, per_element = _KV_STORAGE[kv_cache_quant_algo] + return PoolGeometry( + tokens_per_block=tokens_per_block, + heads_kv=heads_kv, + kv_head_start=kv_head_start, + head_dim=cfg.head_dim, + index_kv_heads=cfg.indexer_kv_heads, + index_head_dim=cfg.indexer_head_dim, + conv_dim=2 * num_k_heads * cfg.linear_key_head_dim + + num_v_heads * cfg.linear_value_head_dim, + conv_kernel=cfg.linear_conv_kernel_dim, + num_k_heads=num_k_heads, + k_head_start=k_head_start, + num_v_heads=num_v_heads, + v_head_start=v_head_start, + head_v_dim=cfg.linear_value_head_dim, + head_k_dim=cfg.linear_key_head_dim, + stream=cfg.hc_count * cfg.hidden_size, + ple_state_len=(cfg.ple_conv_kernel_size - 1) * cfg.ngram_size, + ngram_context_len=cfg.ngram_size - 1, + dtype=dtype, + kv_dtype=kv_dtype or dtype, + kv_width=cfg.head_dim // per_element, + ) + + +# ------------------------------------------------------------------ contract --- + + +def require(cm, names, what: str) -> None: + """`cm` exposes every attribute in `names`, which this target reads to `what`.""" + for name in names: + assert hasattr(cm, name), ( + f"the KV cache manager {type(cm).__name__} does not expose {name}; this target {what}" + ) + + +def _check_pool(t, what, lead, tail, dtype=None) -> None: + """`t` has `lead` leading dims of any size, then exactly `tail`, at `dtype`.""" + assert t is not None, f"{what}: the runtime did not allocate it" + expected = ["*"] * lead + [str(n) for n in tail] + assert t.dim() == lead + len(tail) and tuple(t.shape[lead:]) == tuple(tail), ( + f"{what} is {tuple(t.shape)}, expected [{', '.join(expected)}]" + ) + assert dtype is None or t.dtype == dtype, f"{what} is {t.dtype}, expected {dtype}" + + +def check_position_plane(g: PoolGeometry, position) -> None: + shape = (g.tokens_per_block, QSA_POSITION_AXES) + _check_pool(position, "QSA position plane", 1, shape, torch.int32) + + +def check_attention_pools(g: PoolGeometry, layer: int, index_cache, pool) -> None: + """QSA index plane and paged K/V for one attention layer.""" + index_shape = (g.tokens_per_block, g.index_kv_heads, g.index_head_dim) + _check_pool(index_cache, f"layer {layer}: QSA index plane", 1, index_shape, g.dtype) + # The sparse path needs adjacent HND K/V role planes: [pages, 2, H, block, D]. + kv_shape = (g.heads_kv, g.tokens_per_block, g.kv_width) + _check_pool(pool, f"layer {layer}: paged K/V", 2, kv_shape, g.kv_dtype) + assert pool.shape[1] == 2, f"layer {layer}: paged K/V has {pool.shape[1]} role planes, not 2" + + +def check_kv_scale_planes(g: PoolGeometry, layer: int, pool, scales) -> None: + """An NVFP4 layer's e4m3 block scales: one uint8 per `NVFP4_BLOCK` values, laid out like the + data pages (`[slots, 2, H, block, D / 16]`), and holding as many pages per slot as the data + pool does. + + The attention op addresses a layer's data pages and its scale pages with one block-offset + table, so a scale pool whose pages-per-slot differs from the data pool's reads and writes other + slots' scales without any error. Both counts are read off the views' slot strides.""" + scale_shape = (g.heads_kv, g.tokens_per_block, g.head_dim // NVFP4_BLOCK) + _check_pool(scales, f"layer {layer}: NVFP4 block scales", 2, scale_shape, torch.uint8) + assert scales.shape[1] == 2 and scales.shape[0] == pool.shape[0], ( + f"layer {layer}: the block scales are {tuple(scales.shape)} against paged K/V " + f"{tuple(pool.shape)}; they must hold the same slots, K then V" + ) + pages = [] + for what, t in (("paged K/V", pool), ("block scales", scales)): + page = t[0, 0].numel() + assert t.stride(0) % page == 0, ( + f"layer {layer}: the {what} slot stride {t.stride(0)} is not a whole number of " + f"{page}-element pages" + ) + pages.append(t.stride(0) // page) + assert pages[0] == pages[1], ( + f"layer {layer}: the paged K/V pool holds {pages[0]} pages per slot and the block-scale " + f"pool {pages[1]}; the attention op would address the scales of other slots" + ) + + +def check_recurrent_pools(g: PoolGeometry, conv, ssm) -> None: + _check_pool(conv, "conv state pool", 1, (g.conv_dim, g.conv_kernel - 1), g.dtype) + # Key dim innermost. + _check_pool(ssm, "recurrent state pool", 1, (g.num_v_heads, g.head_v_dim, g.head_k_dim)) + assert ssm.dtype in SSM_STATE_DTYPES, ( + f"the recurrent state pool is {ssm.dtype}; this target serves " + f"{' or '.join(str(d) for d in SSM_STATE_DTYPES)} (kv_cache_config.mamba_ssm_cache_dtype)" + ) + + +def check_conv_layout(g: PoolGeometry, section_dims) -> None: + """The manager splits a rank's conv channels as `[q | k | v]` of the rank's own heads, the + order this target's input projection writes them in. Another order (`x_b_c`) or another split + has the same total width, so the conv kernel would read every section from the wrong place.""" + assert tuple(section_dims) == g.conv_sections, ( + f"the manager lays a rank's conv state out as sections {tuple(section_dims)}; this target " + f"writes [q | k | v] = {g.conv_sections} (conv_state_layout 'q_k_v')" + ) + + +def check_ple_pools(g: PoolGeometry, ple_conv, ple_ngram, conv) -> None: + _check_pool(ple_conv, "PLE conv state", 1, (g.stream, g.ple_state_len), g.dtype) + # The id-history entry writes int64 into it; an int32 pool raises. + _check_pool(ple_ngram, "PLE n-gram context pool", 1, (g.ngram_context_len,), torch.int64) + # The forward addresses all three with one scheduler slot vector. + assert ple_conv.shape[0] == conv.shape[0] == ple_ngram.shape[0], ( + "the PLE and Gated-DeltaNet pools must share one slot space" + ) + + +def check_step_buffers(g: PoolGeometry, int_conv, conv, int_idx) -> None: + """MTP verify buffers both recurrent modes use: the per-step conv windows + `[slots, steps, ...]` at the pool's dtype (the promotion copies between them without + converting) and the rows `update_mamba_states` reads its sources from.""" + assert int_conv is not None, ( + "the runtime allocated no per-step conv buffer; a partially accepted draft " + "cannot be rolled back" + ) + _check_pool(int_conv, "per-step conv buffer", 2, (g.conv_dim, g.conv_kernel - 1), conv.dtype) + assert int_idx is not None and int_idx.dtype == torch.int32, ( + "the cache manager exposes no int32 intermediate_state_indices" + ) + + +def check_replay_state(history, width, prev_accepted, cache_buf_idx, ssm) -> None: + """Cached replay, manager-wide: the geometry the replay entry is certified at, and the + per-slot bookkeeping, which must cover every pool slot.""" + assert ssm.dtype in SSM_STATE_DTYPES, ( + f"the recurrent pool is {ssm.dtype}; the replay entry is certified at an fp32 or a " + "bf16 checkpoint" + ) + assert history == GDN_REPLAY_HISTORY, ( + f"replay_history_size {history}; the replay entry is certified at {GDN_REPLAY_HISTORY}" + ) + assert 2 <= width <= min(MAX_VERIFY_WIDTH, history), ( + f"replay_step_width {width}; the replay entry is certified at 2..{MAX_VERIFY_WIDTH} " + f"verify rows and never above the history ({history})" + ) + for name, t in (("prev_num_accepted_tokens", prev_accepted), ("cache_buf_idx", cache_buf_idx)): + assert t.dim() == 1 and t.shape[0] >= ssm.shape[0] and t.dtype == torch.int32, ( + f"{name} is {tuple(t.shape)} {t.dtype}; the kernel reads one int32 per pool slot " + f"({ssm.shape[0]})" + ) + + +def check_replay_history( + g: PoolGeometry, layer: int, history, pool_slots, old_x, old_B, old_dt +) -> None: + """One GDN layer's replay history: `[slots, 2 halves, ...]` over this rank's heads, addressed + by `history` and the head geometry, never by the buffer's own shape.""" + for name, t in (("old_x", old_x), ("old_B", old_B), ("old_dt", old_dt)): + assert t is not None, f"layer {layer}: the manager's layer cache carries no {name}" + assert old_x.shape[0] >= pool_slots, ( + f"layer {layer}: the replay history has {old_x.shape[0]} slots for a {pool_slots}-slot pool" + ) + bf16, fp32 = torch.bfloat16, torch.float32 + x_tail, b_tail = (history, g.num_v_heads, g.head_v_dim), (history, g.num_k_heads, g.head_k_dim) + _check_pool(old_x, f"layer {layer}: old_x", 2, x_tail, bf16) + _check_pool(old_B, f"layer {layer}: old_B", 2, b_tail, bf16) + _check_pool(old_dt, f"layer {layer}: old_dt", 2, (g.num_v_heads, history), fp32) + assert old_x.shape[1] == old_B.shape[1] == old_dt.shape[1] == 2, ( + f"layer {layer}: the replay history is not double-buffered" + ) + + +def check_verify_width(rows: int) -> None: + """The rows per generation request one step carries: `max_draft_len + 1` (`runtime_draft_len + + 1` at run time), a `thop_attention` `predicted_tokens_per_seq` the entry is certified at.""" + assert 1 <= rows <= MAX_VERIFY_WIDTH, ( + f"predicted_tokens_per_seq {rows} is outside the 1..{MAX_VERIFY_WIDTH} range " + "attention/thop_attention is certified at" + ) + + +def check_replay_width(replay, rows: int) -> None: + """Under the cached replay a single-row generation step cannot take the ordinary + recurrence: the pool row is a checkpoint whose pending history only the replay kernel + applies. The manager's step width (`replay.width`) is the one row count per sequence a step + may carry.""" + assert replay is None or rows == replay.width, ( + f"the batch carries {rows} query rows per generation sequence and the GDN cached " + f"replay runs at step width {replay.width}; a runtime_draft_len of {rows - 1} (a " + "draft-length schedule, or skipped drafting) cannot run under it. Export " + "TRTLLM_USE_GDN_REPLAY=0 for variable draft lengths" + ) + + +def check_shared_selection(shared, slots, layer: int, block_topk: int) -> None: + """IndexShare: the metadata's shared selection buffers have a row for `layer` and are + `block_topk` wide.""" + assert shared is not None, "IndexShare is on but the QSA metadata allocated no shared buffers" + assert layer in slots and slots[layer] < shared.shape[0], ( + f"the KV cache manager gives layer {layer} no row of the {shared.shape[0]}-row shared " + f"selection buffers (qsa_shared_topk_slots={slots})" + ) + assert shared.shape[2] == block_topk, ( + f"the shared selection rows are {shared.shape[2]} wide and this target selects " + f"{block_topk} blocks per row" + ) + + +# ------------------------------------------------------------- the views --- + + +class KvPlanes(NamedTuple): + """One layer's paged K/V views, and under an NVFP4 pool its e4m3 block-scale views: + `[slots, H_kv, tokens_per_block, kv_width]` and `[slots, H_kv, tokens_per_block, D / 16]`.""" + + k: torch.Tensor + v: torch.Tensor + k_scale: torch.Tensor | None = None + v_scale: torch.Tensor | None = None + + +class RankPools(NamedTuple): + """The trunk's pool views on one rank, from one manager. A view is a `TensorWrapper` over a + raw device address that keeps no memory alive, so it is valid only for that manager.""" + + pool_idx: dict # attention layer -> pool offset (what thop_attention addresses) + index_cache: dict # attention layer -> QSA index plane + position_cache: torch.Tensor # the shared QSA position plane + kv: dict # attention layer -> `KvPlanes` + conv: dict # GDN layer -> conv state pool + ssm: dict # GDN layer -> recurrent state pool + ple_state: tuple # (PLE conv state, PLE n-gram id history) + + +def kv_planes(cm, g: PoolGeometry, layer: int, index_cache) -> KvPlanes: + """One layer's paged K/V views from `cm`, checked against the contract: the data planes, and + under NVFP4 the block-scale planes that go with them.""" + pool = cm.get_buffers(layer, kv_layout="HND") + check_attention_pools(g, layer, index_cache, pool) + if not g.kv_block_scales: + return KvPlanes(pool[:, 0], pool[:, 1]) + require(cm, ("get_block_scale_buffers",), "reads an NVFP4 pool's e4m3 block-scale planes") + scales = cm.get_block_scale_buffers(layer, kv_layout="HND") + check_kv_scale_planes(g, layer, pool, scales) + return KvPlanes(pool[:, 0], pool[:, 1], scales[:, 0], scales[:, 1]) + + +def take_pools(cm, g: PoolGeometry, attn_layers, gdn_layers, ple_layer: int) -> RankPools: + """Every trunk pool view of `cm` this rank binds, each checked against `g`. + + `attn_layers` and `gdn_layers` are the full-attention and Gated-DeltaNet decoder layers, and + `ple_layer` the zero-based PLE layer. Reads no attention metadata: what depends on the step + (pool ids, windows, the QSA selection width) is the binder's to check.""" + require( + cm, + ("get_index_k_buffer", "get_qsa_position_buffer", "get_buffers", "layer_offsets"), + "reads the QSA index plane, the shared position plane and the paged K/V", + ) + # `ple_layer_cache` exists only on the V2 hybrid manager; under V1 the PLE state silently + # does not exist. + require( + cm, + ("get_conv_states", "get_ssm_states", "conv_section_dims", "ple_layer_cache"), + "reads the recurrent, conv and PLE pools", + ) + # thop_attention addresses the pool by the layer's pool offset, not its ordinal among + # attention layers. + pool_idx = {i: cm.layer_offsets[i] for i in attn_layers} + position_cache = cm.get_qsa_position_buffer() + check_position_plane(g, position_cache) + index_cache, kv = {}, {} + for i in attn_layers: + index_cache[i] = cm.get_index_k_buffer(i) + kv[i] = kv_planes(cm, g, i, index_cache[i]) + + check_conv_layout(g, cm.conv_section_dims) + conv = {i: cm.get_conv_states(i) for i in gdn_layers} + ssm = {i: cm.get_ssm_states(i) for i in gdn_layers} + for i in gdn_layers: + check_recurrent_pools(g, conv[i], ssm[i]) + + ple_cache = cm.ple_layer_cache(ple_layer) + assert ple_cache is not None, ( + f"no PLE lifecycle state for decoder layer {ple_layer}: the runtime allocated neither " + "the dilated conv state nor the n-gram id history" + ) + ple_conv, ple_ngram = ple_cache + check_ple_pools(g, ple_conv, ple_ngram, conv[gdn_layers[0]]) + return RankPools(pool_idx, index_cache, position_cache, kv, conv, ssm, (ple_conv, ple_ngram)) + + +__all__ = [ + "GDN_REPLAY_HISTORY", + "MAX_VERIFY_WIDTH", + "NVFP4_BLOCK", + "QSA_POSITION_AXES", + "SSM_STATE_DTYPES", + "KvPlanes", + "PoolGeometry", + "RankPools", + "attention_tp_size", + "check_attention_pools", + "check_conv_layout", + "check_kv_scale_planes", + "check_ple_pools", + "check_position_plane", + "check_recurrent_pools", + "check_replay_history", + "check_replay_state", + "check_replay_width", + "check_shared_selection", + "check_step_buffers", + "check_verify_width", + "kv_planes", + "pool_geometry", + "require", + "take_pools", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision.py new file mode 100644 index 000000000000..cc2e590770f1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision.py @@ -0,0 +1,694 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The vision tower and the multimodal splice's host bookkeeping. + +The checkpoint's 27-block ViT (no KV cache, deepstack or quantization) and its 2x2 patch +merger turn a request's images and videos into one `[rows, hidden]` embedding row per +placeholder token, which the prefill target's `_splice` break writes over the token +embedding. The tower is replicated: every rank holds all of its weights and runs it on the +media of the requests it serves -- under attention TP every rank the whole batch's, and gets +the same rows back, since every op is a deterministic function of its operands; under +attention DP a rank its own requests'. + +The shell calls `multimodal_embeds` on every step, before the core's `begin_step`, and the +executor calls `validate_request` on every new request. Neither runs inside a captured +graph: the tower is eager and reads the request's grid on the host. +""" + +from __future__ import annotations + +import math +from typing import NamedTuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import ( + thop_attention, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.gelu_erf import ( + gelu_erf, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.gelu_tanh import ( + gelu_tanh, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.norm.layer_norm import ( + layer_norm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.vision.patch_embed_conv3d import ( + patch_embed_conv3d, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.attention.vision_rope_2d_packed_qk import ( # noqa: E501 + vision_rope_2d_packed_qk, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.vision.triton_pos_embed_interpolate import ( # noqa: E501 + triton_pos_embed_interpolate, +) +from tensorrt_llm.inputs.multimodal import get_multimodal_item_metadata + +#: Tower rope base (`vision_config` has no rope field); unrelated to the text 1e7. +VISION_ROPE_THETA = 10000.0 + +#: LayerNorm epsilon of `norm1`, `norm2`, `merger.norm`: a literal in the reference. +VISION_NORM_EPS = 1e-6 + +#: The most patch rows one tower call carries, the ceiling the vision entries' cells reach: +#: 8192 placeholder tokens x `spatial_merge_size`**2 patch rows each. `VisionTower` splits a +#: request's frames into calls under it, so the bound a request meets is per frame: one +#: frame's `h * w` patches. +VISION_CALL_MAX_PATCHES = 32768 + +#: `mask_type` 0 = padding mask: bidirectional within a batch segment. Causal (1) would make +#: a patch depend on raster order. +_VISION_MASK_PADDING = 0 + +# Per-call constants of the tower's attention: the encoder-only mode of `thop_attention`. +# Omitting `kv_cache_block_offsets`, `host_kv_cache_pool_pointers` and +# `host_kv_cache_pool_mapping` selects that mode. `update_kv_cache=False` is rejected ("KV +# cache update cannot be disabled now"), so it stays True. +_VISION_CALL_CONSTANTS = dict( + k=None, + v=None, + output_sf=None, + kv_cache_block_offsets=None, + host_kv_cache_pool_pointers=None, + host_kv_cache_pool_mapping=None, + tokens_per_block=None, + use_paged_context_fmha=False, + is_fused_qkv=True, + update_kv_cache=True, + predicted_tokens_per_seq=1, + # No pool, so no layer to address. + local_layer_idx=0, + max_context_q_len_override=None, + cache_indirection=None, + kv_scale_orig_quant=None, + kv_scale_quant_orig=None, + out_scale=None, + quant_mode=0, + beam_width=1, + mask_type=_VISION_MASK_PADDING, + q_scaling=1.0, + # Rotation is applied before the op (`vision_rope_2d_packed_qk`); rope_* inert. + position_embedding_type=0, + rotary_inv_freq=None, + rotary_cos_sin=None, + rope_dim=0, + rope_base=10000.0, + rope_scale_type=0, + rope_scale=1.0, + rope_short_m_scale=1.0, + rope_long_m_scale=1.0, + rope_max_positions=1024, + rope_original_max_positions=1024, + attention_input_type=0, + is_mla_enable=False, + chunked_prefill_buffer_batch_size=1, + attention_sinks=None, + block_ids_per_seq=None, + latent_cache=None, + q_pe=None, + q_lora_rank=None, + kv_lora_rank=None, + qk_nope_head_dim=None, + qk_rope_head_dim=None, + v_head_dim=None, + rope_append=None, + attention_chunk_size=None, + softmax_stats_tensor=None, + is_spec_decoding_enabled=False, + use_spec_decoding=False, + is_spec_dec_tree=False, + spec_decoding_generation_lengths=None, + spec_decoding_position_offsets=None, + spec_decoding_packed_mask=None, + spec_decoding_bl_tree_mask_offset=None, + spec_decoding_bl_tree_mask=None, + spec_bl_tree_first_sparse_mask_offset_kv=None, + sparse_kv_indices=None, + sparse_kv_offsets=None, + sparse_attn_indices=None, + sparse_attn_offsets=None, + sparse_attn_indices_block_size=0, + mrope_rotary_cos_sin=None, + mrope_position_deltas=None, + helix_position_offsets=None, + helix_is_inactive_rank=None, + is_cross=False, + cross_kv=None, + relative_attention_bias=None, + relative_attention_max_distance=0, + kv_norm_weight=None, + kv_norm_eps=1e-6, + skip_correction_threshold=0.0, +) + +#: Per modality, the keys under `multimodal_data[modality]`: the flat patch rows and the +#: host-side `(t, h, w)` grid. +MEDIA_KEYS = { + "image": ("pixel_values", "image_grid_thw"), + "video": ("pixel_values_videos", "video_grid_thw"), +} + +#: Tensors under `multimodal_data` the executor moves to the device. Everything else stays +#: on the host: the grids are read there, and moved they would cost a synchronization each. +DEVICE_PATHS = ( + "image.pixel_values", + "video.pixel_values_videos", + "mrope_config.mrope_position_ids", + "mrope_config.mrope_position_deltas", +) + +#: The position-ids axes of the mRoPE layout a media request carries. +_POSITION_AXES = 3 + + +# --------------------------------------------------------------- weight views --- + + +class _VisionBlockWeights(NamedTuple): + """One vision block's weights (`.t()` GEMM views). LayerNorm has a bias, unlike every + norm in the language tower.""" + + n1_w: torch.Tensor + n1_b: torch.Tensor + qkv: torch.Tensor + qkv_b: torch.Tensor + proj: torch.Tensor + proj_b: torch.Tensor + n2_w: torch.Tensor + n2_b: torch.Tensor + fc1: torch.Tensor + fc1_b: torch.Tensor + fc2: torch.Tensor + fc2_b: torch.Tensor + + +class _VisionWeights(NamedTuple): + """The tower's weights: patch embedding, position table, blocks and patch merger.""" + + patch_w: torch.Tensor + patch_b: torch.Tensor + pos: torch.Tensor + blocks: list + m_n_w: torch.Tensor + m_n_b: torch.Tensor + m_fc1: torch.Tensor + m_fc1_b: torch.Tensor + m_fc2: torch.Tensor + m_fc2_b: torch.Tensor + + +def _vision_views(w: dict, depth: int) -> _VisionWeights: + """The weight table's `vis_*` entries (`weights.py`) as the views the ops read.""" + return _VisionWeights( + # Not transposed: the entry takes the checkpoint's `[out, in, kt, kh, kw]`. + patch_w=w["vis_patch_w"], + patch_b=w["vis_patch_b"], + pos=w["vis_pos"], + blocks=[ + _VisionBlockWeights( + n1_w=w[f"l{b}_vis_n1_w"], + n1_b=w[f"l{b}_vis_n1_b"], + # Row layout (q heads, k, v) is already thop_attention's packed form. + qkv=w[f"l{b}_vis_qkv_w"].t(), + qkv_b=w[f"l{b}_vis_qkv_b"], + proj=w[f"l{b}_vis_proj_w"].t(), + proj_b=w[f"l{b}_vis_proj_b"], + n2_w=w[f"l{b}_vis_n2_w"], + n2_b=w[f"l{b}_vis_n2_b"], + fc1=w[f"l{b}_vis_fc1_w"].t(), + fc1_b=w[f"l{b}_vis_fc1_b"], + fc2=w[f"l{b}_vis_fc2_w"].t(), + fc2_b=w[f"l{b}_vis_fc2_b"], + ) + for b in range(depth) + ], + m_n_w=w["vis_m_n_w"], + m_n_b=w["vis_m_n_b"], + m_fc1=w["vis_m_fc1_w"].t(), + m_fc1_b=w["vis_m_fc1_b"], + m_fc2=w["vis_m_fc2_w"].t(), + m_fc2_b=w["vis_m_fc2_b"], + ) + + +# ------------------------------------------------------------------- the tower --- + + +def _frame_calls( + grid: list[tuple[int, int, int]], max_rows: int +) -> list[list[tuple[int, int, int]]]: + """Cut `grid` (one `(frames, height, width)` row per image or video) into tower calls of + whole frames, in order, each at most `max_rows` patch rows. A call is its grid rows: the + frames of one item that fall in it are one `(frames, height, width)` row.""" + calls: list[list[tuple[int, int, int]]] = [[]] + room = max_rows + for frames, height, width in grid: + rows = height * width + assert rows <= max_rows, ( + f"one frame is {height}x{width} = {rows} patch rows; the vision ops are " + f"certified up to {max_rows} rows in one call" + ) + while frames: + take = min(frames, room // rows) + if take == 0: + calls.append([]) + room = max_rows + continue + calls[-1].append((take, height, width)) + room -= take * rows + frames -= take + return calls + + +def _merge_order_positions(grid: list[tuple[int, int, int]], merge: int) -> torch.Tensor: + """Each patch row's `(row, col)` in its frame, int64 `[rows, 2]` on the host, in + spatial-merge order: whole `merge x merge` blocks row-major, row-major inside a block, + the order the patch merger's 4-to-1 reshape consumes. A frame group repeats its frame.""" + intra = torch.arange(merge) + parts = [] + for frames, height, width in grid: + blocks_h, blocks_w = height // merge, width // merge + shape = (blocks_h, blocks_w, merge, merge) + rows = torch.arange(blocks_h)[:, None, None, None] * merge + intra[None, None, :, None] + cols = torch.arange(blocks_w)[None, :, None, None] * merge + intra[None, None, None, :] + coords = torch.stack((rows.expand(shape).reshape(-1), cols.expand(shape).reshape(-1)), 1) + parts.append(coords.repeat(frames, 1) if frames > 1 else coords) + return parts[0] if len(parts) == 1 else torch.cat(parts, 0) + + +def _rope_tables( + grid: list[tuple[int, int, int]], merge: int, theta: float, head_dim: int, device +) -> tuple[torch.Tensor, torch.Tensor]: + """The 2-D rotary's per-row fp32 `cos` / `sin`, `[rows, head_dim // 2]` each and + contiguous: the height angles, then the width angles, as `vision_rope_2d_packed_qk` + takes them. Built once per tower call and shared by every block. + + Each axis owns half of the `head_dim // 2` rotation pairs. The angles are formed in + fp64 and rounded once to fp32: the largest is `(side - 1)` radians, and fp32 throughout + rounds twice there.""" + rope_dim = head_dim // 2 + inv_freq = 1.0 / (theta ** (torch.arange(0, rope_dim, 2, dtype=torch.float64) / rope_dim)) + pos = _merge_order_positions(grid, merge).to(torch.float64) + angles = (pos[:, :, None] * inv_freq).flatten(1).to(device=device, dtype=torch.float32) + return angles.cos().contiguous(), angles.sin().contiguous() + + +class VisionTower: + """One request's images or videos through the 27-block ViT and its patch merger. + + Not a target: runs only for a context request carrying media, from the shell, before + the step. Attention is bidirectional per frame (a video's frame group is a frame), + expressed to `thop_attention` as one context "sequence" per frame; running the rows as + one sequence is far outside tolerance. The position-embedding interpolation and the 2-D + rope both emit in spatial-merge order (what the merger's 4-to-1 reshape consumes); a + raster-order emitter gives a correctly shaped, silently wrong tower. + + The frames run in calls of at most `VISION_CALL_MAX_PATCHES` rows, cut only between + frames: attention is per frame, the position embedding and rope are per grid row and + repeat per frame, the norms and MLPs are row-wise, and a frame's `h * w` is a multiple + of `merge**2`, so a cut keeps every merger block whole. One frame past the bound cannot + be served; `validate_request` refuses it at admission.""" + + def __init__(self, w: dict, vision_config) -> None: + """`w` is the core's weight table after load; `vision_config` the composite + config's vision half, which `routing.py` matched against this checkpoint's.""" + vcfg = vision_config + self.depth = vcfg.depth + self.hidden = vcfg.hidden_size + self.heads = vcfg.num_heads + self.head_dim = self.hidden // self.heads + self.merge = vcfg.spatial_merge_size + self.patch = vcfg.patch_size + # One flat pixel row per patch, `(in_channels, kt, kh, kw)` channel-slowest. + self.pixel_width = vcfg.in_channels * vcfg.temporal_patch_size * self.patch * self.patch + # `spatial_merge_size ** 2` patches become one token at the merger. + self.merge_dim = self.hidden * self.merge * self.merge + # The position table is a square grid of this side. + self.grid_side = math.isqrt(vcfg.num_position_embeddings) + self.dtype = w["embed"].dtype + self._w = _vision_views(w, self.depth) + # Persistent scratch: `thop_attention` grows `workspace` in place; never the + # engine's own workspace. + self._ws: torch.Tensor | None = None + + def __call__(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: + """`pixel_values` is `[sum(t*h*w), in_channels*kt*kh*kw]`, one flat patch per row, + channel-slowest. `grid_thw` is one `(frames, height, width)` row per image or video + and stays on the host (no synchronization). Returns + `[sum(t*h*w) / merge**2, out_hidden_size]`, one row per placeholder, in grid row + order.""" + assert pixel_values.dtype == self.dtype, ( + f"the tower runs at the activation dtype {self.dtype}; the runtime handed " + f"pixel_values as {pixel_values.dtype}. The input processor's `dtype` property is " + "what decides this, and no cast happens here" + ) + assert pixel_values.dim() == 2 and pixel_values.shape[1] == self.pixel_width, ( + f"pixel_values must be [patches, {self.pixel_width}] flat patch rows; got " + f"{tuple(pixel_values.shape)}" + ) + grid: list[tuple[int, int, int]] = [] + for row in grid_thw.tolist(): + assert len(row) == 3, f"grid_thw rows are (t, h, w); got {row}" + frames, height, width = (int(v) for v in row) + assert frames >= 1 and height >= 1 and width >= 1, row + assert height % self.merge == 0 and width % self.merge == 0, ( + f"grid sides must divide by spatial_merge_size {self.merge}; got " + f"height={height}, width={width}" + ) + grid.append((frames, height, width)) + assert grid, "grid_thw describes no image or video" + described = sum(t * h * w for t, h, w in grid) + assert described == pixel_values.shape[0], ( + f"grid_thw describes {described} patch rows and pixel_values carries " + f"{pixel_values.shape[0]}" + ) + + outputs = [] + start = 0 + for call in _frame_calls(grid, VISION_CALL_MAX_PATCHES): + rows = sum(t * h * w for t, h, w in call) + outputs.append(self._run(pixel_values[start : start + rows], call)) + start += rows + return outputs[0] if len(outputs) == 1 else torch.cat(outputs, 0) + + def _run(self, pixel_values: torch.Tensor, grid: list[tuple[int, int, int]]) -> torch.Tensor: + """One tower call over whole frames: `grid` describes exactly the rows of + `pixel_values`, at most `VISION_CALL_MAX_PATCHES` of them.""" + dt, dev = pixel_values.dtype, pixel_values.device + assert pixel_values.is_cuda, "pixel_values must reach the tower on the device" + + # ---- the frame segmentation, host-side ------------------------------ + # One attention segment per frame: a video's frames do not attend to each other. + seg_lens: list[int] = [] + for frames, height, width in grid: + seg_lens.extend([height * width] * frames) + num_patches = pixel_values.shape[0] + num_segments = len(seg_lens) + max_seg = max(seg_lens) + + # Batch-state tensors for the encoder-only call; they describe the frames. + host_lens = torch.tensor(seg_lens, dtype=torch.int32) + host_types = torch.zeros([num_segments], dtype=torch.int32) + host_total = torch.tensor([num_patches, 0], dtype=torch.int32) + dev_lens = host_lens.to(dev) + if self._ws is None: + self._ws = torch.empty([0], dtype=torch.int8, device=dev) + # The rotary's angles, once for all blocks. + cos, sin = _rope_tables(grid, self.merge, VISION_ROPE_THETA, self.head_dim, dev) + + # ---- patch embedding + absolute position ---------------------------- + V = self._w + pos = [ + triton_pos_embed_interpolate(V.pos, t, h, w, self.grid_side, self.merge, dt) + for t, h, w in grid + ] + h = torch.add( + patch_embed_conv3d(pixel_values, V.patch_w, V.patch_b), + pos[0] if len(pos) == 1 else torch.cat(pos, 0), + ) + attn_out = torch.empty([num_patches, self.hidden], dtype=dt, device=dev) + + for B in V.blocks: + qkv = cublas_mm(layer_norm(h, B.n1_w, B.n1_b, VISION_NORM_EPS), B.qkv, B.qkv_b) + # The q|k prefix of each packed row, rotated in place; v is never addressed. + qkv_heads = torch.reshape(qkv, [num_patches, 3 * self.heads, self.head_dim]) + vision_rope_2d_packed_qk(qkv_heads[:, : 2 * self.heads], cos, sin) + attn_out = thop_attention( + q=qkv, + output=attn_out, + workspace=self._ws, + sequence_length=dev_lens, + host_past_key_value_lengths=host_lens, + host_total_kv_lens=host_total, + context_lengths=dev_lens, + host_context_lengths=host_lens, + host_request_types=host_types, + max_num_requests=num_segments, + max_context_length=max_seg, + max_seq_len=max_seg, + attention_window_size=max_seg, + num_contexts=num_segments, + num_ctx_tokens=num_patches, + num_heads=self.heads, + num_kv_heads=self.heads, + head_size=self.head_dim, + **_VISION_CALL_CONSTANTS, + ) + h = torch.add(h, cublas_mm(attn_out, B.proj, B.proj_b)) + mlp = cublas_mm(layer_norm(h, B.n2_w, B.n2_b, VISION_NORM_EPS), B.fc1, B.fc1_b) + # Non-gated MLP with the tanh GELU (not the merger's erf). + h = torch.add(h, cublas_mm(gelu_tanh(mlp), B.fc2, B.fc2_b)) + + # ---- the patch merger ----------------------------------------------- + # The norm runs over `hidden` before the 4-to-1 shuffle; the reshape is legal only + # because both order-bearing steps emitted in spatial-merge order. + merged = torch.reshape( + layer_norm(h, V.m_n_w, V.m_n_b, VISION_NORM_EPS), [-1, self.merge_dim] + ) + # The erf GELU (`nn.GELU()`); within one bf16 ulp of tanh, hence a separate entry. + return cublas_mm(gelu_erf(cublas_mm(merged, V.m_fc1, V.m_fc1_b)), V.m_fc2, V.m_fc2_b) + + +# ------------------------------------------------------------ the shell's side --- + + +def multimodal_embeds(tower: VisionTower | None, kwargs: dict): + """Run the vision tower for this batch; returns `(rows, mm_token_indices)` (the tower's + `[R, hidden]` rows in packed-batch order and the `[R]` placeholder rows they replace, + which the prefill target splices into its embedding), or None when the step carries no + media. + + Called by the shell, because the inherited `DecoderModelForCausalLM.forward` drops + `multimodal_params` from `**kwargs` on the no-speculation path. Host-side bookkeeping: + `multimodal_params` is a subset of the context requests, so never zip it positionally; + `multimodal_runtime` (chunked prefill / KV reuse only) gives the cached and in-chunk row + counts over the request's images and videos together, in prompt order; `mm_token_indices` + names every placeholder row. + + The tower runs once per modality a request carries, and `in_prompt_order` merges the + two outputs. Positions are not derived here: every coordinate of a vision span, a + video's temporal one included, is in the payload's `mrope_config.mrope_position_ids`, + which the engine slices into `position_ids`.""" + mm_params = kwargs.get("multimodal_params") or [] + if not mm_params: + return None + rows = [] + for params in mm_params: + data = getattr(params, "multimodal_data", None) or {} + image, video = data.get("image"), data.get("video") + if image is None and video is None: + continue + mrope_ids = (data.get("mrope_config") or {}).get("mrope_position_ids") + assert ( + mrope_ids is not None and mrope_ids.dim() == 3 and mrope_ids.shape[0] == _POSITION_AXES + ), ( + f"a media request needs mrope_config.mrope_position_ids of shape " + f"[{_POSITION_AXES}, 1, L]; got {None if mrope_ids is None else tuple(mrope_ids.shape)}" + ) + runtime = getattr(params, "multimodal_runtime", None) + if runtime is not None and int(runtime.num_mm_tokens_in_chunk) == 0: + # Every placeholder row of this request lies in another chunk. + continue + assert tower is not None, "post_load_weights builds the vision tower before any step" + towers = {} + if image is not None: + towers["image"] = run_tower(tower, image, *MEDIA_KEYS["image"]) + if video is not None: + towers["video"] = run_tower(tower, video, *MEDIA_KEYS["video"]) + embeds = in_prompt_order(tower.merge, params, towers) + if runtime is not None: + # Chunked prefill or reused prefix: splice only this chunk's slice. + start = int(runtime.num_cached_mm_tokens) + count = int(runtime.num_mm_tokens_in_chunk) + total = embeds.shape[0] + assert int(runtime.total_embeds_in_request) == total, ( + f"the executor counted {runtime.total_embeds_in_request} placeholder rows for " + f"this request and the tower produced {total} over its {sorted(towers)}; the " + "two are derived from the same grids and a disagreement means the merger's " + "row count is wrong" + ) + assert 0 <= start <= start + count <= total, (start, count, total) + embeds = embeds[start : start + count] + rows.append(embeds) + if not rows: + return None + vision = rows[0] if len(rows) == 1 else torch.cat(rows, 0) + mm_token_indices = kwargs.get("mm_token_indices") + assert mm_token_indices is not None, ( + "the executor shipped multimodal params but no mm_token_indices; that pairing is what " + "the shell's `multimodal_token_ids` buys, and without it there is no sync-free way to " + "find the placeholder rows" + ) + assert mm_token_indices.shape[0] == vision.shape[0], ( + f"the batch has {mm_token_indices.shape[0]} placeholder rows and the tower produced " + f"{vision.shape[0]} embedding rows" + ) + return vision, mm_token_indices + + +def run_tower(tower, media: dict, pixel_key: str, grid_key: str): + """The tower over one modality's payload: `(rows, grid_thw)`, the rows one per + placeholder, item by item in the order the processor emitted them.""" + pixel_values, grid_thw = media.get(pixel_key), media.get(grid_key) + assert pixel_values is not None and grid_thw is not None, ( + f"the multimodal payload carries {sorted(media)}; this target reads {pixel_key!r} and " + f"{grid_key!r} from it" + ) + assert not grid_thw.is_cuda, f"{grid_key} must stay on the host; see DEVICE_PATHS" + return tower(pixel_values, grid_thw), grid_thw + + +def in_prompt_order(merge: int, params, towers: dict) -> torch.Tensor: + """One request's tower rows in the order its placeholders appear. + + `towers` maps each modality the request carries to `(rows, grid_thw)`; each grid row is + one item of `t*h*w / merge**2` rows. One modality is already in prompt order. Two are + reassembled item by item from `item_metadata.item_refs` (`(modality, index)` per item in + prompt order, `index` counting that modality's items); without it there is no host-side + way to tell which placeholder run is an image and which a video.""" + if len(towers) == 1: + return next(iter(towers.values()))[0] + metadata = params.item_metadata + order = metadata.item_refs if metadata is not None else None + assert order, ( + f"a request carrying {sorted(towers)} needs item metadata to splice them in prompt " + "order; the payload holds each modality's rows apart" + ) + merge_unit = merge * merge + items = {} + for modality, (embeds, grid_thw) in towers.items(): + counts = [t * h * w // merge_unit for t, h, w in grid_thw.tolist()] + assert sum(counts) == embeds.shape[0], ( + f"the {modality} grid describes {sum(counts)} rows and the tower produced " + f"{embeds.shape[0]}" + ) + items[modality] = torch.split(embeds, counts, 0) + listed = {modality: [] for modality in items} + for modality, index in order: + assert modality in items, ( + f"item_refs lists a {modality!r} item; this request carries {sorted(items)}" + ) + listed[modality].append(int(index)) + for modality, parts in items.items(): + assert sorted(listed[modality]) == list(range(len(parts))), ( + f"item_refs lists {modality} items {listed[modality]}; the payload carries {len(parts)}" + ) + return torch.cat([items[modality][int(index)] for modality, index in order], 0) + + +# ---------------------------------------------------------------- admission --- + + +def validate_request(tower, request) -> None: + """Admission check, run by the executor on every new request. A media request the + forward would assert on fails here as that request's own error; an assert inside the + step ends the executor under the overlap scheduler. Reads only the request's host-side + payload, and a request without media returns at once. `tower` supplies the geometry + (`merge`, `patch`, `pixel_width`). + + Refused: a frame past `VISION_CALL_MAX_PATCHES` patches (the one bound the tower keeps), + a malformed grid or pixel payload, media whose rows disagree with the prompt's + placeholder count, missing or mis-sized `mrope_position_ids`, and a request with both + images and videos whose item metadata is missing or does not list its items.""" + data = request.py_multimodal_data + if not data: + return + grids = {} + for modality, (pixel_key, grid_key) in MEDIA_KEYS.items(): + media = data.get(modality) + if media is not None: + grids[modality] = _media_grid(tower, modality, media, pixel_key, grid_key) + if not grids: + return + + mrope_ids = (data.get("mrope_config") or {}).get("mrope_position_ids") + if ( + mrope_ids is None + or mrope_ids.dim() != 3 + or mrope_ids.shape[0] != _POSITION_AXES + or mrope_ids.shape[2] != request.prompt_len + ): + raise ValueError( + f"a media request needs the [{_POSITION_AXES}, 1, {request.prompt_len}] " + "mrope_config.mrope_position_ids from the input processor, one position per prompt " + f"token; got {None if mrope_ids is None else tuple(mrope_ids.shape)}" + ) + + if len(grids) > 1: + metadata = get_multimodal_item_metadata(data) + _check_item_order(grids, metadata.item_refs if metadata is not None else None) + + cumsum = data.get("multimodal_embed_mask_cumsum") + if cumsum is not None: + merge_unit = tower.merge * tower.merge + rows = sum(t * h * w for grid in grids.values() for t, h, w in grid) // merge_unit + placeholders = int(cumsum[-1]) + if placeholders != rows: + raise ValueError( + f"the prompt holds {placeholders} image/video placeholder tokens and its " + f"{' and '.join(sorted(grids))} payload produces {rows} embedding rows" + ) + + +def _media_grid(tower, modality, media, pixel_key, grid_key): + """One modality's `(frames, height, width)` rows, checked on the host.""" + pixel_values, grid_thw = media.get(pixel_key), media.get(grid_key) + if pixel_values is None or grid_thw is None: + raise ValueError( + f"the {modality} payload carries {sorted(media)}; this target reads {pixel_key!r} " + f"and {grid_key!r} from it" + ) + if grid_thw.dim() != 2 or grid_thw.shape[0] == 0 or grid_thw.shape[1] != 3: + raise ValueError( + f"{grid_key} must hold one (t, h, w) row per {modality}; got shape " + f"{tuple(grid_thw.shape)}" + ) + merge = tower.merge + grid = [tuple(int(v) for v in row) for row in grid_thw.tolist()] + for item, (frames, height, width) in enumerate(grid): + if min(frames, height, width) < 1 or height % merge or width % merge: + raise ValueError( + f"{modality} {item} has grid {(frames, height, width)}; every entry must be " + f"positive and the sides divisible by spatial_merge_size {merge}" + ) + if height * width > VISION_CALL_MAX_PATCHES: + raise ValueError( + f"{modality} {item} has frames of {height}x{width} = {height * width} patches " + f"({height * tower.patch}x{width * tower.patch} pixels); the vision tower " + f"serves frames of at most {VISION_CALL_MAX_PATCHES} patches, so send it at a " + "lower resolution" + ) + patches = sum(t * h * w for t, h, w in grid) + if tuple(pixel_values.shape) != (patches, tower.pixel_width): + raise ValueError( + f"{pixel_key} is {tuple(pixel_values.shape)}; its grids describe {patches} patch " + f"rows of {tower.pixel_width}" + ) + return grid + + +def _check_item_order(grids, order) -> None: + """A request with several modalities lists its items in prompt order.""" + if not order: + raise ValueError( + f"a request carrying {sorted(grids)} needs item metadata to splice them in prompt " + "order; the payload holds each modality's rows apart" + ) + listed = {modality: [] for modality in grids} + for entry in order: + try: + modality, index = entry + listed[modality].append(int(index)) + except (KeyError, TypeError, ValueError) as error: + raise ValueError( + f"item_refs entry {entry!r} is not an item of this request, which carries " + f"{sorted(grids)}; entries are (modality, index) pairs" + ) from error + for modality, grid in grids.items(): + if sorted(listed[modality]) != list(range(len(grid))): + raise ValueError( + f"item_refs lists {modality} items {listed[modality]}; the payload carries " + f"{len(grid)}" + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/weights.py new file mode 100644 index 000000000000..ffa1185ca5e0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/weights.py @@ -0,0 +1,1861 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weight table and loader: Qwen3.8-Flash-Next-NVFP4 (hf-fc694b5_orig) / sm_103 / default. + +One `W` per weight role (shape, checkpoint source, the rank's cut, dtype, load-time +transform); `Qwen4ExpDefaultWeights` unrolls the table into declaration, manifest and load. +Storage is HF `[out, in]`; GEMM views, grouped-norm slices, fp32 gating params and NVFP4 +scalars are derived in `modeling.py` after load. + +**Every weight is declared at this rank's share.** The shares are computed once, in +`rank_dims`, from the engine's `Mapping`: + +* **Attention heads** are split over the attention TP group: the TP group, or one rank + under attention DP. Rank `r` of four holds q heads `[6r, 6r + 6)`: their rows of + `q_proj` (q and the output gate interleave per head, so a head's `2 * head_dim` rows + are contiguous) and their input columns of `o_proj`. The two KV heads are repeated on + neighbouring ranks: rank `r` holds head `r * 2 // 4`, the head `duplicate_kv_weight` + gives it and the K/V pool it binds holds (`pool.pool_geometry`). The QSA indexer and + the q/k norms are replicated. The MTP layer's attention is cut the same way. +* **Gated-DeltaNet heads** are split over `mamba_effective_tp_size` (the TP group, or + one rank under attention DP), as the recurrent pools are: rank `r` holds key heads + `[4r, 4r + 4)` and value heads `[12r, 12r + 12)`. Of `in_proj_qkv` and the depthwise + `conv1d` it reads three row blocks, the q and k rows of its key heads and the v rows + of its value heads, into one `[q | k | v]` operand: the channel order of the conv pool + (`pool.check_conv_layout`). `in_proj_z`, `in_proj_a`, `in_proj_b`, `A_log` and + `dt_bias` are cut by value head, `out_proj` by its value-head input columns. +* **Routed experts** are split over the MoE EP group (tep4, dep4): rank `r` holds the + contiguous window `[128r, 128r + 128)`, NVFP4 codes and block scales. tp4 splits each + expert's intermediate instead: rank `r` holds rows `[160r, 160r + 160)` of every expert's + `up_proj` and `gate_proj` and the same columns of its `down_proj`, all 512 experts. 160 + is short of the layout the runner's operands need (the FC1 row shuffle and scale swizzle + work in 128-row blocks), so each half is zero-extended to `expert_inter` = 256 rows (FC2 + to 256 columns) with zero codes and zero block scales before the relayout: SwiGLU of a + zero row pair is zero and a zero FC2 column adds nothing, so the padding is exact. +* **The MTP layer's FP8 block-scale experts** are split by expert over the whole TP + group at every member, `[128r, 128r + 128)`: the FP8 runner takes `H` and `I` in + multiples of 128, which 640 / 4 is not. + +Everything else is replicated: hyper-connections, the PLE block and its n-gram table, +the router, the shared expert, the embedding, the vision tower, the MTP layer's front +end, router, shared expert and mixer, and every routed expert's per-tensor scalars over +the whole routing space (one activation quantization feeds all 512 experts, so the core +checks their FC1 input scale across all of them, not across its window). `lm_head` is +the shell's: `load` reads the rows the shell built, vocab-parallel under attention TP, +whole under attention DP. + +**A rank reads only its share.** Each split weight carries a cut, the rows or columns of +the checkpoint tensor this rank holds, applied to the checkpoint entry before it is +materialized: on a lazy safetensors slice that reads those bytes and nothing else, on an +mmapped tensor it is a view, and either way only the share is copied to the device. An +off-window expert is never read. The kernel relayouts then run on the share, never +before the cut: the 32-row block shuffle and the 128x4 scale swizzle permute within +blocks of rows, so cutting after them would cut through blocks. + +Transforms: + +* trunk attention input projection: `q_proj`, `k_proj`, `v_proj` and the indexer's + `index_qk_proj` read the same activation, so they are one + `[q | k | v | gate | index | pad, hidden]` operand and one GEMM (`_attn_in_w`), at this + rank's heads: 4352 rows at attention TP, 14080 under attention DP. The q halves go + first so the GEMM output's leading columns are the packed `[q | k | v]` the fused + norm+rope rewrites in place, then the gate halves, the replicated index projection and + `index_head_dim` zero rows (a row is a whole number of `head_dim` heads). The draft + (MTP) layer keeps its separate projections; +* fused shared-expert gate|up: `[gate ; up ; gate logit ; pad]` (`_cat_rows`); +* NVFP4 routed-expert stacks, laid out for the trunk's routed backend + (`trunk_moe_backend`): `[up ; gate]`, then 128x4 scale swizzle. trtllm-gen also + interleaves FC1 row by row and 32-row block shuffles weight and scale bytes (FC2: + shuffle only); CuteDSL interleaves FC1 in 64-row up/gate groups and leaves FC2 in + checkpoint row order. Under trtllm-gen a second copy of the window in the CuteDSL + layout (`fc*_cd`) is declared when the core keeps it (`moe_prefill_copy`); +* n-gram table: 128 row chunks concatenated into one `[rows, 160]` buffer (the hash + emits global row indices). It is replicated, and host-resident by default + (`_host_resident_empty`): one copy per rank. + +The n-gram table is not in the table: its height is a checkpoint fact, so `load` +registers it with `torch.empty` (`_declare_ngram_table`) and adds its rows to the +manifest (`ngram_table_sources`). + +Not loaded: `mtp.*` when there is no `speculative_config`. With one, the MTP module is +loaded at this rank's share, including the draft layer's own QSA indexer. MTP experts +are plain row-major FP8 128x128: `[up ; gate]` concat, bf16 `weight_scale_inv` cast to +fp32; `mtp_experts_to_nvfp4` requantizes a rank's window to NVFP4 after load. + +`load(model, weights)` fills every declared parameter once and asserts coverage both +ways: every parameter written, and every checkpoint key consumed (a key read through a +cut is consumed) or predicted by `expected_unconsumed` -- the experts outside this +rank's windows, and the MTP module when there is no speculation. +""" + +import os +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any, Callable + +import torch +from cuda.bindings import runtime as cudart +from torch import nn + +from tensorrt_llm._torch.pyexecutor.config_utils import mamba_effective_tp_size +from tensorrt_llm._utils import CUASSERT + +from .pool import attention_tp_size + +#: `layers=ALL_LAYERS`, the default: the weight is on every layer. +ALL_LAYERS = object() + +#: An index into a checkpoint tensor: a tuple of slices, one per leading axis. +Cut = tuple[slice, ...] + + +@dataclass(frozen=True) +class W: + """One weight role: what it is, where its bytes come from, what happens on the way. + + `shape` takes the dimension bundle and returns the declared shape: this rank's + share, every width derived from the config and the mapping. + + `src` says where the bytes come from, in one of two forms: + + * a checkpoint key template, with `{p}` standing for the layer prefix (a key with no + `{p}` is itself); + * a callable `(d, i) -> [(key, index), ...]` or `[(key, index, cut), ...]`, taking + the dimension bundle and the layer index. `index` addresses the destination + parameter (None for all of it). A `key` may itself be a tuple, meaning the + transform is handed those tensors together. A source that names its own `cut` + reads one block of a tensor several of whose blocks land in the parameter. + + `cut` takes the dimension bundle and returns this rank's rows or columns of the + checkpoint tensor, read from the entry before it is materialized; for a tuple key + it returns one cut (or None) per key. `None` reads the tensor whole. + + `layers` says which layers carry this weight: `ALL_LAYERS` for every one, + `None` for a weight that is not per-layer at all, or a callable + `d -> iterable[int]` when only some layers have it. + + `transform` runs on the cut tensor(s) before the copy, on the destination's + device, with the dimension bundle: the kernels' relayouts, with no slicing of + their own. `None` means the cut is already the declared shape. + """ + + name: str + shape: Callable[[Any], tuple[int, ...]] + src: str | Callable + dtype: torch.dtype | None = None + cut: Callable[[Any], Any] | None = None + transform: Callable | None = None + layers: object | None = ALL_LAYERS + + def layer_indices(self, d: SimpleNamespace) -> list[int] | None: + """The layers this weight exists on, or None if it is not per-layer.""" + if self.layers is None: + return None + if self.layers is ALL_LAYERS: + return list(range(d.num_layers)) + return list(self.layers(d)) + + def sources(self, d: SimpleNamespace, i: int | None, prefix: str) -> list[tuple]: + """`[(checkpoint key or key tuple, index into the param | None, cut | None), ...]`.""" + rows = self.src(d, i) if callable(self.src) else [(self.src.format(p=prefix), None)] + cut = self.cut(d) if self.cut is not None else None + out = [] + for row in rows: + if len(row) == 3: + assert self.cut is None, f"{self.name}: a source with its own cut takes no W.cut" + out.append(tuple(row)) + else: + out.append((*row, cut)) + return out + + +# ----------------------------------------------------------- host residency -- +# +# Allocates the one PLE n-gram table. The table stays a `cuda` tensor of the same +# shape/dtype. Registered, device-mapped pages (not a host gather) because the gather's +# row indices are produced on device inside the captured decode graph. + +#: `cudaMemoryType` values; the core asserts on them after load. +CUDA_MEMORY_TYPE_HOST = 1 +CUDA_MEMORY_TYPE_DEVICE = 2 + + +def memory_kind(ptr: int) -> int: + """The driver's own answer to "what kind of memory is this pointer": a `cudaMemoryType`.""" + (attributes,) = CUASSERT(cudart.cudaPointerGetAttributes(ptr)) + return int(attributes.type) + + +#: `cudaHostRegister` flags (`driver_types.h`): mapped, so a kernel may dereference the +#: pages, and portable, for every context in the process. +_CUDA_HOST_REGISTER_PORTABLE = 0x01 +_CUDA_HOST_REGISTER_MAPPED = 0x02 + + +def _local_cpus(device) -> set | None: + """The CPUs sysfs lists as local to `device`'s PCI function and this process may run on, + or None when that is unknown.""" + try: + props = torch.cuda.get_device_properties(device) + path = ( + f"/sys/bus/pci/devices/{props.pci_domain_id:04x}:{props.pci_bus_id:02x}:" + f"{props.pci_device_id:02x}.0/local_cpulist" + ) + with open(path) as fh: + text = fh.read().strip() + except (OSError, AttributeError, RuntimeError): + return None + cpus = set() + for part in filter(None, text.split(",")): + lo, _, hi = part.partition("-") + cpus.update(range(int(lo), int(hi or lo) + 1)) + return (cpus & os.sched_getaffinity(0)) or None + + +class _RegisteredHostBuffer: + """A pageable CPU allocation registered (pinned and mapped) for the device, and + unregistered before its storage is released: at this size glibc serves the + allocation with `mmap` and returns it with `munmap`, so freeing registered pages + would leave the driver mapping memory the OS may hand out again. + + Not torch's pinned allocator: it rounds the 47.7 GiB table up to 64 GiB, and the + device's random row gathers from its pages are slower than from registered pageable + memory.""" + + def __init__(self, nbytes: int, device=None) -> None: + # The registration faults the pages in on this thread, so under the default local + # allocation policy they land on the NUMA node of the CPU it runs on. Pinned to the + # CPUs local to `device` meanwhile, the table sits next to the GPU that gathers from + # it: with a process per GPU, half the ranks would otherwise gather from the other + # socket's memory. + local = _local_cpus(device) if device is not None else None + previous = os.sched_getaffinity(0) if local else None + if local: + os.sched_setaffinity(0, local) + try: + self._storage = torch.empty(nbytes, dtype=torch.uint8) + self.pointer = self._storage.data_ptr() + CUASSERT( + cudart.cudaHostRegister( + self.pointer, nbytes, _CUDA_HOST_REGISTER_PORTABLE | _CUDA_HOST_REGISTER_MAPPED + ) + ) + finally: + if previous: + os.sched_setaffinity(0, previous) + # Bound now: at interpreter shutdown the module may already be torn down. + self._unregister = cudart.cudaHostUnregister + + def __del__(self) -> None: + pointer, self.pointer = self.pointer, 0 + if pointer: + self._unregister(pointer) + + +def _host_resident_empty( + rows: int, cols: int, dtype, device +) -> tuple[torch.Tensor, _RegisteredHostBuffer]: + """`[rows, cols]` of `dtype` on `device`, whose bytes live in host memory. + + Returns the tensor and the registered CPU allocation that owns the pages. The + caller **must** keep the second alive for as long as the first: the tensor + is a view through a device address and owns nothing. + """ + itemsize = torch.empty(0, dtype=dtype).element_size() + assert itemsize == 1, ( + f"the host-resident allocation reinterprets a `|u1` buffer, so its " + f"element must be one byte; {dtype} is {itemsize}" + ) + assert device.type == "cuda", f"host offload needs a cuda device to map into; got {device}" + + host = _RegisteredHostBuffer(rows * cols, device) + (device_pointer,) = CUASSERT(cudart.cudaHostGetDevicePointer(host.pointer, 0)) + + # No public torch API builds a tensor over a pointer the caching allocator + # did not hand out; `__cuda_array_interface__` is the numba/cupy path. + # `|u1` because the protocol has no float8 typestr; reinterpreted below. + interface = SimpleNamespace( + __cuda_array_interface__={ + "shape": (rows, cols), + "typestr": "|u1", + "data": (int(device_pointer), False), + "version": 3, + "strides": None, + } + ) + table = torch.as_tensor(interface, device=device).view(dtype) + # `torch.as_tensor` copies when it decides the buffer is not already on the + # requested device; a copy here would be a silent 48 GiB device allocation. + assert table.data_ptr() == int(device_pointer), ( + "torch.as_tensor copied the host-mapped buffer instead of viewing it, " + "so the table would be device-resident after all" + ) + assert table.shape == (rows, cols) and table.dtype == dtype, (table.shape, table.dtype) + return table, host + + +# ------------------------------------------------------------------ the table -- + +_LM = "model.language_model" +_VISUAL = "model.visual" +_MTP = "mtp" +_MTP_LAYER = "mtp.layers.0" + +# NVFP4 block size: one e4m3 scale per 16 contiguous elements along K. +SF_VEC = 16 + +# The FP8 block-scale family's block edge, on both axes of the MTP experts' +# weight scales and along the activation row. +FP8_BLOCK = 128 + +#: The PLE n-gram table's storage: the checkpoint's FP8 per-tensor quantization, kept as +#: stored (the gather dequantizes the rows it reads). +PLE_TABLE_DTYPE = torch.float8_e4m3fn + +#: How the trunk's routed experts are split (`rank_dims`'s `expert_split`) that this table +#: declares and `load` fills: by expert (`ep`, a window of whole experts) or by each +#: expert's intermediate (`tp`, every expert at a zero-padded slice). +LOADED_EXPERT_SPLITS = ("ep", "tp") + +#: The multiple a rank's stored routed-expert intermediate is padded up to. It is the +#: engine's own storage rule for an NVFP4 trtllm-gen shard +#: (`NVFP4TRTLLMGenFusedMoEMethod.resolve_alignments`: a padded intermediate that is not a +#: multiple of 128 drops to a 128 weight alignment), which turns tp4's 160 into 256, the +#: width `fp4_block_scale_moe_runner` certifies the padded 512-expert stack at. 640 +#: (tep4, dep4) is already a multiple. +EXPERT_INTER_ALIGN = 128 + +# Parameter-key suffixes after which to release transform scratch. The expert +# relayouts allocate per-expert scratch once per local expert and the table +# fill stages one shard at a time; without this, peak load memory is the whole +# model rather than one parameter. +_RELEASE_AFTER = ( + "_fc2_s", + "_fc1_s", + "_fc2_s_cd", + "_fc1_s_cd", + "ple_table", + "mtp_fc1_w", + "mtp_fc2_w", +) + + +def _read(entry, cut: Cut | None) -> torch.Tensor: + """Realize one checkpoint entry, reading only `cut` of it. + + `entry[cut]` on a lazy safetensors slice reads those rows or columns and nothing + else; on a plain tensor it is a view. Uncut, `[:]` reads it whole -- but a 0-dim + tensor (every NVFP4 per-tensor scalar on this checkpoint) rejects that index, and no + 0-dim entry has a cut.""" + if cut is None: + return entry if getattr(entry, "ndim", 1) == 0 else entry[:] + return entry[cut[0]] if len(cut) == 1 else entry[cut] + + +# --- cuts: this rank's block of rows or columns of a checkpoint tensor --- + + +def _span(start: int, width: int) -> slice: + return slice(start, start + width) + + +def _q_rows(d) -> Cut: + """`q_proj` rows of this rank's q heads: each head's q rows, then its gate rows.""" + return (_span(d.q_head_start * 2 * d.head_dim, 2 * d.q_width),) + + +def _kv_rows(d) -> Cut: + """`k_proj` / `v_proj` rows of this rank's KV heads.""" + return (_span(d.kv_head_start * d.head_dim, d.kv_width),) + + +def _o_cols(d) -> Cut: + """`o_proj` input columns of this rank's q heads.""" + return (slice(None), _span(d.q_head_start * d.head_dim, d.q_width)) + + +def _key_rows(d, section: int) -> Cut: + """Rows of this rank's key heads in section `section` (0 q, 1 k) of a `[q | k | v]` + Gated-DeltaNet tensor.""" + return (_span(section * d.key_dim_total + d.k_head_start * d.head_k_dim, d.key_dim),) + + +def _value_rows(d, base: int = 0) -> Cut: + """Rows of this rank's value heads, `base` rows into the tensor.""" + return (_span(base + d.v_head_start * d.head_v_dim, d.value_dim),) + + +def _value_head_rows(d) -> Cut: + """One row per value head (`in_proj_a`, `in_proj_b`, `A_log`, `dt_bias`).""" + return (_span(d.v_head_start, d.num_v_heads),) + + +def _value_cols(d) -> Cut: + """`out_proj` input columns of this rank's value heads.""" + return (slice(None), _span(d.v_head_start * d.head_v_dim, d.value_dim)) + + +def _qkv_blocks(d) -> tuple[tuple[slice, Cut], ...]: + """(destination rows, checkpoint cut) of the three blocks of a `[q | k | v]` + Gated-DeltaNet tensor this rank holds, in the `[q | k | v]` order of its conv pool.""" + return ( + (slice(0, d.key_dim), _key_rows(d, 0)), + (slice(d.key_dim, 2 * d.key_dim), _key_rows(d, 1)), + (slice(2 * d.key_dim, d.conv_dim), _value_rows(d, 2 * d.key_dim_total)), + ) + + +# --- relayouts, on the rank's cut --- + + +def _block32_perm(rows: int, device) -> torch.Tensor: + """Gather index of the 32-row block shuffle: inside each aligned block of + 32 rows, source row `4u + v` (0<=u<8, 0<=v<4) lands at destination row + `8v + u`.""" + assert rows % 32 == 0, rows + src = torch.arange(32) + dst = (src % 4) * 8 + src // 4 + idx = torch.empty(32, dtype=torch.long) + idx[dst] = src + blocks = rows // 32 + return (idx.repeat(blocks) + torch.arange(blocks).repeat_interleave(32) * 32).to(device) + + +def _interleave_perm(rows: int, device) -> torch.Tensor: + """Gather index that interleaves the `[up | gate]` halves of a `2*I` row + stack into up0, gate0, up1, gate1, ...""" + p = torch.empty(rows, dtype=torch.long) + p[0::2] = torch.arange(0, rows // 2) + p[1::2] = torch.arange(rows // 2, rows) + return p.to(device) + + +#: The group the CuteDSL FC1 epilogue pairs up and gate rows in: rows `[128k, 128k + 64)` +#: of the fused operand are up rows and the next 64 their gate rows. +_CUTEDSL_GATE_UP_GROUP = 64 + + +def _group_interleave_perm(rows: int, group: int, device) -> torch.Tensor: + """Gather index that interleaves the `[up | gate]` halves of a `2*I` row stack in + `group`-row blocks: up[0:g], gate[0:g], up[g:2g], gate[g:2g], ...""" + assert rows % (2 * group) == 0, (rows, group) + idx = torch.arange(rows, device=device).view(2, rows // (2 * group), group) + return idx.transpose(0, 1).reshape(-1) + + +def _fc1_perm(rows: int, device, d) -> torch.Tensor: + """The FC1 row gather of the trunk's routed backend. trtllm-gen: interleave then + 32-row block shuffle, composed into one gather. CuteDSL: the 64-row group + interleave alone.""" + if d.trunk_moe_backend == "CUTEDSL": + return _group_interleave_perm(rows, _CUTEDSL_GATE_UP_GROUP, device) + return _interleave_perm(rows, device)[_block32_perm(rows, device)] + + +def _swizzle(x: torch.Tensor) -> torch.Tensor: + """128x4 block-scale swizzle of an `[M, C]` uint8 scale matrix, returned + flat (`M % 128 == 0` and `C % 4 == 0` hold at this geometry, so the result + is exactly `M * C` bytes).""" + return torch.ops.trtllm.block_scale_interleave(x.contiguous()) + + +#: Rows after the shared expert's gate|up rows in its fused operand: the per-token gate's +#: one row, then zeros, so the GEMM's output row stays a whole number of 16-byte vectors. +SH_GATE_ROWS = 8 + + +def _cat_rows(t: tuple[torch.Tensor, ...], d) -> torch.Tensor: + """Fused shared-expert weight bytes: **gate** rows first, the half order + `flashinfer_silu_and_mul` reads, then up, then the per-token gate's row and + `SH_GATE_ROWS - 1` zero rows -- one GEMM forms gate|up and the gate logit.""" + gate, up, sh_gate = t + pad = sh_gate.new_zeros(SH_GATE_ROWS - 1, sh_gate.shape[-1]) + return torch.cat((gate, up, sh_gate, pad), dim=0).contiguous() + + +def _attn_in_w(t: tuple[torch.Tensor, ...], d) -> torch.Tensor: + """(q_proj, k_proj, v_proj, index_qk_proj), each at this rank's cut -> the fused + `[q | k | v | gate | index | pad]` input projection of one trunk attention layer (see + the module docstring). + + `q_proj`'s rows are `[heads, 2, head_dim]`: head `h`'s q rows, then its gate rows. + """ + q, k, v, idx = t + heads, head_dim, hidden = d.heads_q, d.head_dim, q.shape[-1] + qg = q.reshape(heads, 2, head_dim, hidden) + pad = q.new_zeros(d.index_head_dim, hidden) + rows = (qg[:, 0].reshape(heads * head_dim, hidden), k, v) + rows += (qg[:, 1].reshape(heads * head_dim, hidden), idx, pad) + return torch.cat(rows, dim=0).contiguous() + + +def _expert_fc1_w(t: tuple[torch.Tensor, ...], d) -> torch.Tensor: + """(up, gate) packed nibbles -> one expert's FC1 operand.""" + up, gate = t + w = torch.cat([up, gate], dim=0) + return torch.index_select(w, 0, _fc1_perm(w.shape[0], w.device, d)).contiguous() + + +def _expert_fc1_s(t: tuple[torch.Tensor, ...], d) -> torch.Tensor: + """(up, gate) e4m3 scales -> one expert's FC1 scale operand: the weight's own row + permutation, then the swizzle.""" + up, gate = t + s = torch.cat([up.view(torch.uint8), gate.view(torch.uint8)], dim=0) + rows, cols = s.shape + s = torch.index_select(s, 0, _fc1_perm(rows, s.device, d)) + return _swizzle(s).reshape(rows, cols) + + +def _expert_fc2_w(t: torch.Tensor, d) -> torch.Tensor: + """`down_proj` packed nibbles -> one expert's FC2 operand (no interleave): the + 32-row block shuffle on trtllm-gen, checkpoint row order on CuteDSL.""" + if d.trunk_moe_backend == "CUTEDSL": + return t.contiguous() + return torch.index_select(t, 0, _block32_perm(t.shape[0], t.device)).contiguous() + + +def _expert_fc2_s(t: torch.Tensor, d) -> torch.Tensor: + """`down_proj` e4m3 scales -> one expert's FC2 scale operand: the weight's own row + order, then the swizzle.""" + s = t.view(torch.uint8) + rows, cols = s.shape + if d.trunk_moe_backend != "CUTEDSL": + s = torch.index_select(s, 0, _block32_perm(rows, s.device)) + return _swizzle(s).reshape(rows, cols) + + +def _zero_extend(x: torch.Tensor, dim: int, size: int) -> torch.Tensor: + """`x` extended along `dim` to `size` with zero bytes: a zero packed byte is two +0.0 + e2m1 codes, a zero scale byte is an e4m3 +0.0 block scale.""" + if x.shape[dim] == size: + return x + assert x.shape[dim] < size, (tuple(x.shape), dim, size) + raw = x.view(torch.uint8) + pad = list(raw.shape) + pad[dim] = size - raw.shape[dim] + return torch.cat((raw, raw.new_zeros(pad)), dim).view(x.dtype) + + +def _zero_padded(transform: Callable | None, dim: int, unit: int) -> Callable: + """`transform` run on its source(s) zero-extended along `dim` to the rank's stored + routed-expert intermediate, `d.expert_inter // unit` there: unit 1 for the rows of an + `up_proj` / `gate_proj` and of their scales, 2 for `down_proj`'s packed code columns, + `SF_VEC` for its scale columns. Each half of FC1 is extended before the `[up ; gate]` + concatenation, so the pad rows sit at the end of each half, where the runner's + `intermediate_size` puts them. A no-op where the rank holds whole experts.""" + + def padded(t, d): + size = d.expert_inter // unit + if isinstance(t, tuple): + t = tuple(_zero_extend(x, dim, size) for x in t) + else: + t = _zero_extend(t, dim, size) + return t if transform is None else transform(t, d) + + return padded + + +def _expert_inter_rows(d) -> Cut | None: + """This rank's rows of one expert's `up_proj` or `gate_proj` (or their block scales): + its slice of the intermediate, or None (whole) when the experts are split by expert.""" + if d.expert_inter_share == d.moe_inter: + return None + return (_span(d.expert_inter_start, d.expert_inter_share),) + + +def _fc1_cut(d): + """The `(up, gate)` pair's cuts: the same intermediate rows of both halves.""" + rows = _expert_inter_rows(d) + return None if rows is None else (rows, rows) + + +def _fc2_cut(unit: int) -> Callable: + """`down_proj`'s input columns of this rank's intermediate slice, `unit` intermediate + elements to a stored column (2 for the packed codes, `SF_VEC` for the block scales).""" + + def cut(d): + if _expert_inter_rows(d) is None: + return None + return ( + slice(None), + _span(d.expert_inter_start // unit, d.expert_inter_share // unit), + ) + + return cut + + +#: The layout argument of the second, CuteDSL-layout expert copy (`fc*_cd`). +_CUTEDSL_LAYOUT = SimpleNamespace(trunk_moe_backend="CUTEDSL") + + +def _expert_fc1_w_cd(t: tuple[torch.Tensor, ...], d) -> torch.Tensor: + """(up, gate) packed nibbles -> one expert's FC1 operand in the CuteDSL layout.""" + return _expert_fc1_w(t, _CUTEDSL_LAYOUT) + + +def _expert_fc1_s_cd(t: tuple[torch.Tensor, ...], d) -> torch.Tensor: + """(up, gate) e4m3 scales -> one expert's FC1 scale operand in the CuteDSL layout.""" + return _expert_fc1_s(t, _CUTEDSL_LAYOUT) + + +def _expert_fc2_s_cd(t: torch.Tensor, d) -> torch.Tensor: + """`down_proj` e4m3 scales -> one expert's FC2 scale operand in the CuteDSL layout.""" + return _expert_fc2_s(t, _CUTEDSL_LAYOUT) + + +def _mtp_fc1_w(t: tuple[torch.Tensor, ...], d) -> torch.Tensor: + """(up, gate) e4m3 codes -> one draft expert's FC1 operand, `[up ; gate]` on + the row axis and nothing else. The FP8 block-scale runner reads the + checkpoint's plain row-major layout.""" + up, gate = t + return torch.cat([up, gate], dim=0).contiguous() + + +def _mtp_fc1_s(t: tuple[torch.Tensor, ...], d) -> torch.Tensor: + """(up, gate) 128x128 block scales -> one draft expert's FC1 scale operand: + the same `[up ; gate]` concatenation, cast bf16 -> fp32.""" + up, gate = t + return torch.cat([up, gate], dim=0).to(torch.float32).contiguous() + + +def _mtp_fc2_s(t: torch.Tensor, d) -> torch.Tensor: + """`down_proj.weight_scale_inv` -> the FC2 scale operand: a bf16 -> fp32 + cast, no relayout.""" + return t.to(torch.float32).contiguous() + + +#: The layout argument of the draft experts' NVFP4 copy: the draft MoE runs it on trtllm-gen +#: whatever the trunk's routed backend is. +_TRTLLM_GEN_LAYOUT = SimpleNamespace(trunk_moe_backend="TRTLLM") + +#: `448 * 6`: an NVFP4 global scale is this over the tensor's amax (modelopt's recipe). +NVFP4_AMAX_TO_GLOBAL = 448.0 * 6.0 + + +def _fp8_block_dequant(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """`[E, N, K]` e4m3 codes times their `[E, N/128, K/128]` fp32 block scales, in fp32.""" + s = scales.repeat_interleave(FP8_BLOCK, 1).repeat_interleave(FP8_BLOCK, 2) + return codes.float() * s + + +def mtp_experts_to_nvfp4(fc1_w, fc1_s, fc2_w, fc2_s, chunk: int = 32) -> tuple: + """Requantize the draft layer's FP8 128x128 block-scaled experts (this rank's window) + to NVFP4 in the trunk's trtllm-gen layout (`_expert_fc1_w/_s`, `_expert_fc2_w/_s`). + + Per expert and projection (FC1's up|gate share one, as the fused FC1's alpha + requires): dequantize with the block scales, take the global scale + `448 * 6 / amax`, then `fp4_quantize` (16-element blocks, e4m3 block scales, + linear order), which then gets the trunk's row permutation and 128x4 swizzle. + Returns `(fc1_w [E, 2I, H/2], fc1_s [E, 2I, H/16], fc2_w [E, H, I/2], + fc2_s [E, H, I/16])` uint8 and the fp32 `[E]` weight global scales `(g1, g2)`; + a weight's real value is its NVFP4 reconstruction over its global scale. + `fc1_w` / `fc1_s` are the `[up ; gate]` stacks `_mtp_fc1_w` / `_mtp_fc1_s` build.""" + experts, two_i, hidden = fc1_w.shape + inter = two_i // 2 + dev = fc1_w.device + u8 = torch.uint8 + out = ( + torch.empty(experts, two_i, hidden // 2, dtype=u8, device=dev), + torch.empty(experts, two_i, hidden // SF_VEC, dtype=u8, device=dev), + torch.empty(experts, hidden, inter // 2, dtype=u8, device=dev), + torch.empty(experts, hidden, inter // SF_VEC, dtype=u8, device=dev), + ) + g1 = torch.empty(experts, dtype=torch.float32, device=dev) + g2 = torch.empty(experts, dtype=torch.float32, device=dev) + layout = _TRTLLM_GEN_LAYOUT + for e0 in range(0, experts, chunk): + e1 = min(e0 + chunk, experts) + w1 = _fp8_block_dequant(fc1_w[e0:e1], fc1_s[e0:e1]) + w2 = _fp8_block_dequant(fc2_w[e0:e1], fc2_s[e0:e1]) + g1[e0:e1] = NVFP4_AMAX_TO_GLOBAL / w1.abs().amax(dim=(1, 2)).clamp_min(1e-30) + g2[e0:e1] = NVFP4_AMAX_TO_GLOBAL / w2.abs().amax(dim=(1, 2)).clamp_min(1e-30) + w1, w2 = w1.to(torch.bfloat16), w2.to(torch.bfloat16) + for j, e in enumerate(range(e0, e1)): + q, s = torch.ops.trtllm.fp4_quantize(w1[j], g1[e : e + 1], SF_VEC, False, False) + s = s.view(two_i, hidden // SF_VEC) + out[0][e] = _expert_fc1_w((q[:inter], q[inter:]), layout) + out[1][e] = _expert_fc1_s((s[:inter], s[inter:]), layout) + q, s = torch.ops.trtllm.fp4_quantize(w2[j], g2[e : e + 1], SF_VEC, False, False) + out[2][e] = _expert_fc2_w(q, layout) + out[3][e] = _expert_fc2_s(s.view(hidden, inter // SF_VEC), layout) + del w1, w2 + return (*out, g1, g2) + + +# --- which layers carry a weight --- + + +def _attn_layers(d): + """The full-attention layers (`i % 4 == 3` on this checkpoint).""" + return [i for i, a in enumerate(d.is_attn) if a] + + +def _gdn_layers(d): + """The Gated-DeltaNet (linear-attention) layers.""" + return [i for i, a in enumerate(d.is_attn) if not a] + + +def _trunk_expert_layers(d): + """Every layer when the routed experts are split in a way this table declares + (`LOADED_EXPERT_SPLITS`); none otherwise.""" + return range(d.num_layers) if d.expert_split in LOADED_EXPERT_SPLITS else () + + +def _cute_dsl_copy_layers(d): + """The trunk expert layers when the core keeps the second, CuteDSL-layout copy that + prefill-sized launches read under trtllm-gen (`moe_prefill_copy`); none otherwise.""" + return _trunk_expert_layers(d) if d.moe_prefill_copy else () + + +def _vision_blocks(d): + """The vision tower's blocks. Keyed `l{b}_vis_*`: block `b`, not decoder layer `b`.""" + return range(d.v_depth) + + +def _mtp_only(d): + """The MTP module, at its own layer index `num_hidden_layers`, only when it is live.""" + return [d.num_layers] if d.mtp_enabled else [] + + +# --- where the bytes come from --- + + +def _ple(suffix): + """A PLE weight. The PLE layer's index is config, so the key is computed.""" + return lambda d, i: [(f"{_LM}.layers.{d.ple_layer}.ple.{suffix}", None)] + + +def _ngram_prefix(ple_layer: int) -> str: + """The quantized n-gram table submodule: the 128 shards and the one + per-tensor scale that dequantizes them both hang off this, one level below + `ple_embedding` itself.""" + return f"{_LM}.layers.{ple_layer}.ple.ple_embedding.ngram_embedding" + + +def _visual(suffix): + """A vision-tower weight that is not per-block.""" + return lambda d, i: [(f"{_VISUAL}.{suffix}", None)] + + +def _vision_block(suffix): + """One vision block's weight.""" + return lambda d, i: [(f"{_VISUAL}.blocks.{i}.{suffix}", None)] + + +def _trunk_window(d) -> tuple[int, int]: + """(first expert, count) of this rank's routed-expert window.""" + return d.expert_offset, d.local_experts + + +def _mtp_window(d) -> tuple[int, int]: + """(first expert, count) of this rank's window of the draft layer's experts.""" + return d.mtp_expert_offset, d.mtp_local_experts + + +def _routing_space(d) -> tuple[int, int]: + """Every expert: what the replicated per-tensor scalars cover.""" + return 0, d.num_experts + + +def _experts(prefix, *suffixes, window): + """One source per expert of `window(d)`, written at its local index; a pair of + suffixes hands the transform both tensors together. The checkpoint key depends on + the rank, which is why the source is computed rather than templated.""" + + def build(d, i): + start, count = window(d) + base = prefix(i) + out = [] + for local in range(count): + keys = tuple(f"{base}.experts.{start + local}.{s}" for s in suffixes) + out.append((keys if len(keys) > 1 else keys[0], (local,))) + return out + + return build + + +def _trunk_mlp(i): + return f"{_LM}.layers.{i}.mlp" + + +def _mtp_mlp(i): + return f"{_MTP_LAYER}.mlp" + + +def _in_proj(d, i): + """The Gated-DeltaNet input projections at this rank's heads, in the row order of + the concatenated `in_proj_w` (and so the column order of its GEMM output): + `[q | k | v (conv_dim) | z (value_dim) | a | b]`, `d.in_proj_split`.""" + p = f"{_LM}.layers.{i}.linear_attn" + out = [(f"{p}.in_proj_qkv.weight", (rows,), cut) for rows, cut in _qkv_blocks(d)] + row = d.conv_dim + out.append((f"{p}.in_proj_z.weight", (_span(row, d.value_dim),), _value_rows(d))) + row += d.value_dim + for name in ("in_proj_a", "in_proj_b"): + out.append((f"{p}.{name}.weight", (_span(row, d.num_v_heads),), _value_head_rows(d))) + row += d.num_v_heads + return out + + +def _conv(d, i): + """The depthwise conv's `[q | k | v]` channels of this rank's heads.""" + key = f"{_LM}.layers.{i}.linear_attn.conv1d.weight" + return [(key, (rows,), cut) for rows, cut in _qkv_blocks(d)] + + +#: The six per-tensor NVFP4 scalars every routed expert carries, as +#: (parameter suffix, checkpoint suffix). The checkpoint stores modelopt +#: reciprocals; the core folds them into the runner's scalars after load. +_EXPERT_SCALARS = ( + ("e_isc1", "gate_proj.input_scale"), + ("e_isc1_up", "up_proj.input_scale"), + ("e_ws2_1", "gate_proj.weight_scale_2"), + ("e_ws2_1_up", "up_proj.weight_scale_2"), + ("e_isc2", "down_proj.input_scale"), + ("e_ws2_2", "down_proj.weight_scale_2"), +) + +#: The checkpoint tensors of one routed expert that are not per-tensor scalars: what a +#: rank outside its window leaves unread. The draft layer's FP8 experts carry +#: `weight_scale_inv` in place of the NVFP4 `weight_scale`. +_PROJECTIONS = ("gate_proj", "up_proj", "down_proj") +_TRUNK_EXPERT_TENSORS = tuple(f"{p}.{s}" for p in _PROJECTIONS for s in ("weight", "weight_scale")) +_MTP_EXPERT_TENSORS = tuple( + f"{p}.{s}" for p in _PROJECTIONS for s in ("weight", "weight_scale_inv") +) + + +#: (parameter tag, checkpoint module) of the two per-layer hyper-connection gates. +_HC_MODULES = (("ahc", "attn_hyper_connection"), ("mhc", "mlp_hyper_connection")) + + +def _hyper_connection(tag, src, layers): + """The four weights of one gated-residual hyper-connection; `src(name)` names the source.""" + return ( + W(f"{tag}_norm", shape=lambda d: (d.stream,), src=src("hc_norm"), layers=layers), + W( + f"{tag}_down", + shape=lambda d: (d.hc_lowrank, d.stream), + src=src("input_mix_weight_down"), + layers=layers, + ), + W( + f"{tag}_up", + shape=lambda d: (d.stream, d.hc_lowrank), + src=src("input_mix_weight_up"), + layers=layers, + ), + W( + f"{tag}_inject", + shape=lambda d: (d.hc_count, d.stream), + src=src("block_inject_weight"), + layers=layers, + ), + ) + + +def _mixer(tag, prefix, layers): + """A stack-final hyper-connection mixer: a layer's gate minus the injection weight.""" + return ( + W( + f"{tag}_norm", + shape=lambda d: (d.stream,), + src=f"{prefix}.hc_norm.weight", + layers=layers, + ), + W( + f"{tag}_down", + shape=lambda d: (d.hc_lowrank, d.stream), + src=f"{prefix}.input_mix_weight_down.weight", + layers=layers, + ), + W( + f"{tag}_up", + shape=lambda d: (d.stream, d.hc_lowrank), + src=f"{prefix}.input_mix_weight_up.weight", + layers=layers, + ), + ) + + +def _attn_in_fused(d, i): + """The trunk attention layers' one input projection: q_proj, k_proj and v_proj at this + rank's heads and the QSA indexer's replicated index_qk_proj, re-laid by `_attn_in_w`.""" + p = f"{_LM}.layers.{i}.self_attn" + keys = ( + f"{p}.q_proj.weight", + f"{p}.k_proj.weight", + f"{p}.v_proj.weight", + f"{p}.indexer.index_qk_proj.weight", + ) + return [(keys, None, (_q_rows(d), _kv_rows(d), _kv_rows(d), None))] + + +def _attn_out_and_norms(tag, prefix, layers): + """o_proj at this rank's q heads, and the replicated per-head q/k norms.""" + return ( + W( + f"{tag}o_w", + shape=lambda d: (d.hidden, d.q_width), + src=f"{prefix}.o_proj.weight", + cut=_o_cols, + layers=layers, + ), + W( + f"{tag}q_norm", + shape=lambda d: (d.head_dim,), + src=f"{prefix}.q_norm.weight", + layers=layers, + ), + W( + f"{tag}k_norm", + shape=lambda d: (d.head_dim,), + src=f"{prefix}.k_norm.weight", + layers=layers, + ), + ) + + +def _attn_in_separate(tag, prefix, layers): + """The draft layer's separate q (with its interleaved gate), k and v projections, at + this rank's heads.""" + return ( + W( + f"{tag}q_w", + shape=lambda d: (2 * d.q_width, d.hidden), + src=f"{prefix}.q_proj.weight", + cut=_q_rows, + layers=layers, + ), + W( + f"{tag}k_w", + shape=lambda d: (d.kv_width, d.hidden), + src=f"{prefix}.k_proj.weight", + cut=_kv_rows, + layers=layers, + ), + W( + f"{tag}v_w", + shape=lambda d: (d.kv_width, d.hidden), + src=f"{prefix}.v_proj.weight", + cut=_kv_rows, + layers=layers, + ), + ) + + +def _index_branch(tag, prefix, layers, fused_in=False): + """The QSA index branch's three weights: one replicated, unquantized projection + producing `[index_heads q | 1 k]`, and two per-head gemma norms. With `fused_in`, the + projection is part of `_attn_in_fused`'s one operand and only the norms are declared.""" + projection = W( + f"{tag}idx_w", + shape=lambda d: (d.index_width, d.hidden), + src=f"{prefix}.index_qk_proj.weight", + layers=layers, + ) + return ( + *(() if fused_in else (projection,)), + W( + f"{tag}idx_qn", + shape=lambda d: (d.index_head_dim,), + src=f"{prefix}.q_layernorm.weight", + layers=layers, + ), + W( + f"{tag}idx_kn", + shape=lambda d: (d.index_head_dim,), + src=f"{prefix}.k_layernorm.weight", + layers=layers, + ), + ) + + +def _shared_expert(tag, prefix, layers): + """The replicated bf16 shared expert beside the routed ones, its per-token gate in the + same operand as its gate|up (`_cat_rows`).""" + return ( + W( + f"{tag}sh_gu", + shape=lambda d: (2 * d.shared_inter + SH_GATE_ROWS, d.hidden), + src=lambda d, i: [ + ( + ( + f"{prefix(i)}.shared_expert.gate_proj.weight", + f"{prefix(i)}.shared_expert.up_proj.weight", + f"{prefix(i)}.shared_expert_gate.weight", + ), + None, + ) + ], + transform=_cat_rows, + layers=layers, + ), + W( + f"{tag}sh_dn", + shape=lambda d: (d.hidden, d.shared_inter), + src=lambda d, i: [(f"{prefix(i)}.shared_expert.down_proj.weight", None)], + layers=layers, + ), + ) + + +u8, f8, f32, i64 = torch.uint8, torch.float8_e4m3fn, torch.float32, torch.int64 + + +class Qwen4ExpDefaultWeights: + """Qwen3.8-Flash-Next-NVFP4 / sm_103 / default, at one rank's share. + + The table, and the passes over it: `declare` allocates every weight it names, + `manifest` maps checkpoint keys and cuts onto them, `load` fills them. `declare` and + `manifest` share one "which parameters exist" decision (`_targets`), so they + cannot disagree about a weight. + """ + + WEIGHTS: tuple[W, ...] = ( + # --- the two per-layer hyper-connection gates: every trunk layer, replicated --- + *( + w + for tag, module in _HC_MODULES + for w in _hyper_connection( + tag, lambda name, m=module: f"{{p}}.{m}.{name}.weight", ALL_LAYERS + ) + ), + # --- the gated GQA layers and their QSA index branch --- + # q/k/v at this rank's heads and the indexer's replicated [index_heads q | 1 k] + # projection as one operand, then o_proj at this rank's heads, the q/k norms and + # the indexer's two per-head gemma norms. + W( + "attn_in_w", + shape=lambda d: (d.attn_in_rows, d.hidden), + src=_attn_in_fused, + transform=_attn_in_w, + layers=_attn_layers, + ), + *_attn_out_and_norms("", "{p}.self_attn", _attn_layers), + *_index_branch("", "{p}.self_attn.indexer", _attn_layers, fused_in=True), + # --- the Gated-DeltaNet layers, at this rank's key and value heads --- + # The four input projections read the same activation, so they are one + # weight and one GEMM: rows `[q|k|v (conv_dim) | z (value_dim) | a | b]` + # (`d.in_proj_split`), each checkpoint block copied into its own row range. + W( + "in_proj_w", + shape=lambda d: (sum(d.in_proj_split), d.hidden), + src=_in_proj, + layers=_gdn_layers, + ), + W( + "A_log", + shape=lambda d: (d.num_v_heads,), + src="{p}.linear_attn.A_log", + cut=_value_head_rows, + layers=_gdn_layers, + ), + W( + "dt_bias", + shape=lambda d: (d.num_v_heads,), + src="{p}.linear_attn.dt_bias", + cut=_value_head_rows, + layers=_gdn_layers, + ), + # Stored [conv_dim, 1, width]; the conv entries want [conv_dim, width], which is + # a zero-copy view the targets take once, after load. + W( + "conv_w", + shape=lambda d: (d.conv_dim, 1, d.conv_kernel), + src=_conv, + layers=_gdn_layers, + ), + W( + "gnorm", + shape=lambda d: (d.head_v_dim,), + src="{p}.linear_attn.norm.weight", + layers=_gdn_layers, + ), + W( + "out_w", + shape=lambda d: (d.hidden, d.value_dim), + src="{p}.linear_attn.out_proj.weight", + cut=_value_cols, + layers=_gdn_layers, + ), + # --- MoE on every layer: this rank's window of NVFP4 routed experts, a replicated + # bf16 router and shared expert --- + W("router", shape=lambda d: (d.num_experts, d.hidden), src="{p}.mlp.gate.weight"), + # Each expert at this rank's intermediate share (`_fc1_cut`, `_fc2_cut`), stored + # zero-padded to `expert_inter` (`_zero_padded`); whole experts under an EP window. + W( + "fc1_w", + shape=lambda d: (d.local_experts, 2 * d.expert_inter, d.hidden // 2), + dtype=u8, + src=_experts(_trunk_mlp, "up_proj.weight", "gate_proj.weight", window=_trunk_window), + cut=_fc1_cut, + transform=_zero_padded(_expert_fc1_w, 0, 1), + layers=_trunk_expert_layers, + ), + W( + "fc1_s", + shape=lambda d: (d.local_experts, 2 * d.expert_inter, d.hidden // SF_VEC), + dtype=u8, + src=_experts( + _trunk_mlp, "up_proj.weight_scale", "gate_proj.weight_scale", window=_trunk_window + ), + cut=_fc1_cut, + transform=_zero_padded(_expert_fc1_s, 0, 1), + layers=_trunk_expert_layers, + ), + W( + "fc2_w", + shape=lambda d: (d.local_experts, d.hidden, d.expert_inter // 2), + dtype=u8, + src=_experts(_trunk_mlp, "down_proj.weight", window=_trunk_window), + cut=_fc2_cut(2), + transform=_zero_padded(_expert_fc2_w, 1, 2), + layers=_trunk_expert_layers, + ), + W( + "fc2_s", + shape=lambda d: (d.local_experts, d.hidden, d.expert_inter // SF_VEC), + dtype=u8, + src=_experts(_trunk_mlp, "down_proj.weight_scale", window=_trunk_window), + cut=_fc2_cut(SF_VEC), + transform=_zero_padded(_expert_fc2_s, 1, SF_VEC), + layers=_trunk_expert_layers, + ), + # Under trtllm-gen, the same window again in the CuteDSL layout, when the core keeps + # it for prefill-sized launches. + W( + "fc1_w_cd", + shape=lambda d: (d.local_experts, 2 * d.expert_inter, d.hidden // 2), + dtype=u8, + src=_experts(_trunk_mlp, "up_proj.weight", "gate_proj.weight", window=_trunk_window), + cut=_fc1_cut, + transform=_zero_padded(_expert_fc1_w_cd, 0, 1), + layers=_cute_dsl_copy_layers, + ), + W( + "fc1_s_cd", + shape=lambda d: (d.local_experts, 2 * d.expert_inter, d.hidden // SF_VEC), + dtype=u8, + src=_experts( + _trunk_mlp, "up_proj.weight_scale", "gate_proj.weight_scale", window=_trunk_window + ), + cut=_fc1_cut, + transform=_zero_padded(_expert_fc1_s_cd, 0, 1), + layers=_cute_dsl_copy_layers, + ), + W( + "fc2_w_cd", + shape=lambda d: (d.local_experts, d.hidden, d.expert_inter // 2), + dtype=u8, + src=_experts(_trunk_mlp, "down_proj.weight", window=_trunk_window), + cut=_fc2_cut(2), + transform=_zero_padded(None, 1, 2), + layers=_cute_dsl_copy_layers, + ), + W( + "fc2_s_cd", + shape=lambda d: (d.local_experts, d.hidden, d.expert_inter // SF_VEC), + dtype=u8, + src=_experts(_trunk_mlp, "down_proj.weight_scale", window=_trunk_window), + cut=_fc2_cut(SF_VEC), + transform=_zero_padded(_expert_fc2_s_cd, 1, SF_VEC), + layers=_cute_dsl_copy_layers, + ), + # Replicated over the whole routing space, at every split. + *( + W( + name, + shape=lambda d: (d.num_experts,), + dtype=f32, + src=_experts(_trunk_mlp, suffix, window=_routing_space), + ) + for name, suffix in _EXPERT_SCALARS + ), + *_shared_expert("", _trunk_mlp, ALL_LAYERS), + # --- not per-layer: the stack-final mixer, which is also the final norm --- + *_mixer("fhc", f"{_LM}.hyper_connection_mixer", None), + W( + "embed", + shape=lambda d: (d.vocab, d.hidden), + src=f"{_LM}.embed_tokens.weight", + layers=None, + ), + # --- PLE, on its one decoder layer (the n-gram table itself is declared at load) --- + W( + "ple_key_w", + shape=lambda d: (d.stream, d.hidden), + src=_ple("key_proj.weight"), + layers=None, + ), + W( + "ple_val_w", + shape=lambda d: (d.ple_embed_dim, d.hidden), + src=_ple("value_proj.weight"), + layers=None, + ), + W("ple_nk", shape=lambda d: (d.stream,), src=_ple("norm_key.weight"), layers=None), + W("ple_nq", shape=lambda d: (d.stream,), src=_ple("norm_query.weight"), layers=None), + W("ple_nc", shape=lambda d: (d.stream,), src=_ple("norm_conv.weight"), layers=None), + W( + "ple_conv_w", + shape=lambda d: (d.stream, 1, d.ple_conv_kernel), + src=_ple("conv1d.weight"), + layers=None, + ), + W( + "ple_mult", + shape=lambda d: (d.ngram_size,), + dtype=i64, + src=_ple("ple_embedding.layer_multipliers"), + layers=None, + ), + W( + "ple_vocab", + shape=lambda d: (d.ngram_heads,), + dtype=i64, + src=_ple("ple_embedding.ngram_heads_vocab_sizes"), + layers=None, + ), + W( + "ple_offs", + shape=lambda d: (d.ngram_heads,), + dtype=i64, + src=_ple("ple_embedding.ngram_heads_offsets"), + layers=None, + ), + # The table's per-tensor scale sits with the **table**, one level below the three + # metadata tensors above: the shards and this scale are the quantized + # `ngram_embedding` submodule, while the multipliers, offsets and vocabulary sizes + # describe the id space and stay on `ple_embedding` itself. + W( + "ple_tscale", + shape=lambda d: (1,), + src=lambda d, i: [(f"{_ngram_prefix(d.ple_layer)}.weight_scale", None)], + layers=None, + ), + # --- the vision tower: all bf16, replicated, one-to-one, no transform --- + W( + "vis_patch_w", + shape=lambda d: (d.v_hidden, d.v_channels, d.v_tpatch, d.v_patch, d.v_patch), + src=_visual("patch_embed.proj.weight"), + layers=None, + ), + W( + "vis_patch_b", + shape=lambda d: (d.v_hidden,), + src=_visual("patch_embed.proj.bias"), + layers=None, + ), + W( + "vis_pos", + shape=lambda d: (d.v_pos_rows, d.v_hidden), + src=_visual("pos_embed.weight"), + layers=None, + ), + # LayerNorm **with a bias**, unlike every norm in the language tower. + W( + "vis_n1_w", + shape=lambda d: (d.v_hidden,), + src=_vision_block("norm1.weight"), + layers=_vision_blocks, + ), + W( + "vis_n1_b", + shape=lambda d: (d.v_hidden,), + src=_vision_block("norm1.bias"), + layers=_vision_blocks, + ), + # One fused `[q | k | v]` projection, MHA, with a bias. + W( + "vis_qkv_w", + shape=lambda d: (3 * d.v_hidden, d.v_hidden), + src=_vision_block("attn.qkv.weight"), + layers=_vision_blocks, + ), + W( + "vis_qkv_b", + shape=lambda d: (3 * d.v_hidden,), + src=_vision_block("attn.qkv.bias"), + layers=_vision_blocks, + ), + W( + "vis_proj_w", + shape=lambda d: (d.v_hidden, d.v_hidden), + src=_vision_block("attn.proj.weight"), + layers=_vision_blocks, + ), + W( + "vis_proj_b", + shape=lambda d: (d.v_hidden,), + src=_vision_block("attn.proj.bias"), + layers=_vision_blocks, + ), + W( + "vis_n2_w", + shape=lambda d: (d.v_hidden,), + src=_vision_block("norm2.weight"), + layers=_vision_blocks, + ), + W( + "vis_n2_b", + shape=lambda d: (d.v_hidden,), + src=_vision_block("norm2.bias"), + layers=_vision_blocks, + ), + # A **non-gated** MLP: nothing here is split in half, unlike every SwiGLU in the + # language tower. + W( + "vis_fc1_w", + shape=lambda d: (d.v_inter, d.v_hidden), + src=_vision_block("mlp.linear_fc1.weight"), + layers=_vision_blocks, + ), + W( + "vis_fc1_b", + shape=lambda d: (d.v_inter,), + src=_vision_block("mlp.linear_fc1.bias"), + layers=_vision_blocks, + ), + W( + "vis_fc2_w", + shape=lambda d: (d.v_hidden, d.v_inter), + src=_vision_block("mlp.linear_fc2.weight"), + layers=_vision_blocks, + ), + W( + "vis_fc2_b", + shape=lambda d: (d.v_hidden,), + src=_vision_block("mlp.linear_fc2.bias"), + layers=_vision_blocks, + ), + # The patch merger. Its norm is applied **before** the 4-to-1 shuffle, so it is + # `hidden` wide and not `merge_dim` wide. + W( + "vis_m_n_w", + shape=lambda d: (d.v_hidden,), + src=_visual("merger.norm.weight"), + layers=None, + ), + W("vis_m_n_b", shape=lambda d: (d.v_hidden,), src=_visual("merger.norm.bias"), layers=None), + W( + "vis_m_fc1_w", + shape=lambda d: (d.v_merge_dim, d.v_merge_dim), + src=_visual("merger.linear_fc1.weight"), + layers=None, + ), + W( + "vis_m_fc1_b", + shape=lambda d: (d.v_merge_dim,), + src=_visual("merger.linear_fc1.bias"), + layers=None, + ), + W( + "vis_m_fc2_w", + shape=lambda d: (d.hidden, d.v_merge_dim), + src=_visual("merger.linear_fc2.weight"), + layers=None, + ), + W( + "vis_m_fc2_b", + shape=lambda d: (d.hidden,), + src=_visual("merger.linear_fc2.bias"), + layers=None, + ), + # --- the MTP draft module, at layer index num_hidden_layers, only when live --- + # + # Structurally one more decoder layer of this architecture with a two-branch + # front end; every shape equals a trunk full-attention layer's (cut the same way) + # except the front end and the expert quantization (FP8 128x128 block scales, not + # NVFP4). Its keys carry no `{p}`: the module's prefix is not the trunk's. + W( + "mtp_pre_e_norm", + shape=lambda d: (d.hidden,), + src=f"{_MTP}.pre_fc_norm_embedding.weight", + layers=_mtp_only, + ), + # `[hc_count * hidden]`: the module's input is the trunk's four-stream residual. + W( + "mtp_pre_h_norm", + shape=lambda d: (d.stream,), + src=f"{_MTP}.pre_fc_norm_hidden.weight", + layers=_mtp_only, + ), + W( + "mtp_fc_e", + shape=lambda d: (d.hidden, d.hidden), + src=f"{_MTP}.fc_embedding.weight", + layers=_mtp_only, + ), + W( + "mtp_fc_h", + shape=lambda d: (d.hidden, d.hidden), + src=f"{_MTP}.fc_hidden.weight", + layers=_mtp_only, + ), + *( + w + for tag, module in _HC_MODULES + for w in _hyper_connection( + f"mtp_{tag}", + lambda name, m=module: f"{_MTP_LAYER}.{m}.{name}.weight", + _mtp_only, + ) + ), + *_attn_in_separate("mtp_", f"{_MTP_LAYER}.self_attn", _mtp_only), + *_attn_out_and_norms("mtp_", f"{_MTP_LAYER}.self_attn", _mtp_only), + *_index_branch("mtp_", f"{_MTP_LAYER}.self_attn.indexer", _mtp_only), + W( + "mtp_router", + shape=lambda d: (d.num_experts, d.hidden), + src=f"{_MTP_LAYER}.mlp.gate.weight", + layers=_mtp_only, + ), + # e4m3 codes stored row-major with one fp32 scale per 128x128 weight block, the + # FC1 halves concatenated `[up ; gate]` on the row axis; this rank's window. + W( + "mtp_fc1_w", + shape=lambda d: (d.mtp_local_experts, 2 * d.moe_inter, d.hidden), + dtype=f8, + src=_experts(_mtp_mlp, "up_proj.weight", "gate_proj.weight", window=_mtp_window), + transform=_mtp_fc1_w, + layers=_mtp_only, + ), + W( + "mtp_fc1_s", + shape=lambda d: ( + d.mtp_local_experts, + (2 * d.moe_inter) // FP8_BLOCK, + d.hidden // FP8_BLOCK, + ), + dtype=f32, + src=_experts( + _mtp_mlp, + "up_proj.weight_scale_inv", + "gate_proj.weight_scale_inv", + window=_mtp_window, + ), + transform=_mtp_fc1_s, + layers=_mtp_only, + ), + # FC2 needs no relayout at all: the checkpoint's `[hidden, I]` e4m3 block is + # exactly the operand. + W( + "mtp_fc2_w", + shape=lambda d: (d.mtp_local_experts, d.hidden, d.moe_inter), + dtype=f8, + src=_experts(_mtp_mlp, "down_proj.weight", window=_mtp_window), + layers=_mtp_only, + ), + W( + "mtp_fc2_s", + shape=lambda d: ( + d.mtp_local_experts, + d.hidden // FP8_BLOCK, + d.moe_inter // FP8_BLOCK, + ), + dtype=f32, + src=_experts(_mtp_mlp, "down_proj.weight_scale_inv", window=_mtp_window), + transform=_mtp_fc2_s, + layers=_mtp_only, + ), + *_shared_expert("mtp_", _mtp_mlp, _mtp_only), + # The module's own stack-final mixer, which is its `shared_head`. + *_mixer("mtp_fhc", f"{_MTP}.hyper_connection_mixer", _mtp_only), + ) + + def dims(self, core) -> SimpleNamespace: + """This rank's dimension bundle, from the core's normalized text config, the + composite's vision config (the core holds it as `vision_config` before it + declares anything), the engine's `Mapping` and the core's engine choices.""" + return rank_dims( + core.model_config.pretrained_config, + core.vision_config, + core.model_config.mapping, + mtp_enabled=getattr(core.model_config, "spec_config", None) is not None, + trunk_moe_backend=core.trunk_moe_backend, + moe_prefill_copy=core.moe_prefill_copy, + dtype=core.model_config.torch_dtype, + ) + + def expected_unconsumed(self, d: SimpleNamespace, keys) -> set: + """The checkpoint keys this rank predicts it leaves unread: the tensors of every + routed expert outside its window (its per-tensor scalars are read on every rank; + under the intermediate split the window is every expert, each read at this rank's + slice), and, with MTP live, the draft experts outside its draft window; with MTP off, the + whole `mtp.` subtree. The windows are enumerated from `d`, not read off the + checkpoint, so a key that belongs to a real weight can never land here.""" + out = set() + if d.expert_split in LOADED_EXPERT_SPLITS: + start, count = _trunk_window(d) + for i in range(d.num_layers): + for e in range(d.num_experts): + if not start <= e < start + count: + out.update( + f"{_trunk_mlp(i)}.experts.{e}.{t}" for t in _TRUNK_EXPERT_TENSORS + ) + if not d.mtp_enabled: + return out | {k for k in keys if k.startswith(f"{_MTP}.")} + start, count = _mtp_window(d) + for e in range(d.num_experts): + if not start <= e < start + count: + out.update(f"{_mtp_mlp(None)}.experts.{e}.{t}" for t in _MTP_EXPERT_TENSORS) + return out + + def declare(self, d: SimpleNamespace) -> nn.ParameterDict: + """Allocate every weight the table declares, at this rank's share. + + Meta-init intercepts `torch.empty` here -- real CUDA storage arrives + when the engine materializes the registry. + """ + w = nn.ParameterDict() + for entry in self.WEIGHTS: + dtype = entry.dtype or d.dtype + for key, _, _ in self._targets(entry, d): + w[key] = nn.Parameter( + torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False + ) + return w + + def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None, str]]: + """`[(param key, layer index | None, checkpoint layer prefix), ...]`. + + One place decides which parameters an entry declares, so `declare` and + `manifest` cannot disagree about whether a weight exists. + """ + layers = entry.layer_indices(d) + if layers is None: + return [(entry.name, None, "")] + return [(f"l{i}_{entry.name}", i, f"{_LM}.layers.{i}") for i in layers] + + def manifest(self, d: SimpleNamespace) -> dict[str, list[tuple]]: + """param key -> [(checkpoint key, index into the param | None, cut | None, + transform | None)].""" + rows: dict[str, list[tuple]] = {} + for entry in self.WEIGHTS: + for key, i, prefix in self._targets(entry, d): + rows[key] = [ + (ckpt, index, cut, entry.transform) + for ckpt, index, cut in entry.sources(d, i, prefix) + ] + return rows + + def read(self, weights, ckpt_key, cut, transform, d, device=None) -> torch.Tensor: + """One manifest source -> the tensor its destination receives. + + Each key of the source is read through its cut -- only this rank's rows or + columns leave the checkpoint -- moved to `device`, and handed to `transform`. + The one path `load` takes, so the shard test drives it directly. + """ + multi = isinstance(ckpt_key, tuple) + keys = ckpt_key if multi else (ckpt_key,) + cuts = (cut if cut is not None else (None,) * len(keys)) if multi else (cut,) + assert len(cuts) == len(keys), (ckpt_key, cut) + for key in keys: + assert key in weights, f"checkpoint key missing: {key}" + src_all = tuple(_read(weights[key], c) for key, c in zip(keys, cuts)) + if device is not None: + src_all = tuple(t.to(device, non_blocking=True) for t in src_all) + if transform is not None: + return transform(src_all if multi else src_all[0], d) + assert len(src_all) == 1, "a multi-key row needs a transform" + return src_all[0] + + def load(self, model, weights) -> None: + """Fill every declared parameter. The n-gram table is declared here first, each + source is read through its cut, multi-key rows feed their transform several + tensors at once, and coverage is asserted in both directions.""" + core = model.model + d = self.dims(core) + assert d.expert_split in LOADED_EXPERT_SPLITS, ( + f"routed experts split by {d.expert_split!r}: this table declares the splits " + f"{LOADED_EXPERT_SPLITS}" + ) + # The one load-time declaration, before the manifest is built: the table's + # height comes from the shards, and the manifest addresses it by row range. + rows_per_shard = _declare_ngram_table(core, d, weights) + manifest = self.manifest(d) + manifest["ple_table"] = ngram_table_sources(d, rows_per_shard) + consumed: set = set() + + def fill(param: nn.Parameter, ckpt_key, index, cut, transform) -> None: + dst = param.data if index is None else param.data[index] + # Materialize where the destination lives: the expert relayouts are row + # gathers and the scale swizzle is a kernel, both on the device. + src = self.read(weights, ckpt_key, cut, transform, d, dst.device) + assert dst.shape == src.shape, (ckpt_key, tuple(dst.shape), tuple(src.shape)) + assert src.dtype == dst.dtype, (ckpt_key, src.dtype, dst.dtype) + dst.copy_(src, non_blocking=True) + consumed.update(ckpt_key if isinstance(ckpt_key, tuple) else (ckpt_key,)) + + assert set(manifest.keys()) == set(core.w.keys()), ( + "manifest/parameter drift", + set(manifest.keys()) ^ set(core.w.keys()), + ) + for param_key, sources in manifest.items(): + for ckpt_key, index, cut, transform in sources: + fill(core.w[param_key], ckpt_key, index, cut, transform) + if param_key.endswith(_RELEASE_AFTER): + torch.cuda.empty_cache() + + # The shell owns lm_head (untied, bf16 on this checkpoint); read the rows it built + # (`lm_head_share`) and nothing else. + cut, pad = lm_head_share(model.lm_head.weight.shape[0], d.vocab, d.tp_rank) + assert tuple(model.lm_head.weight.shape[1:]) == (d.hidden,), ( + f"the shell built lm_head {tuple(model.lm_head.weight.shape)}; this checkpoint's " + f"is a dense [{d.vocab}, {d.hidden}] weight" + ) + fill(model.lm_head.weight, "lm_head.weight", None, cut, pad) + # Parameter-side coverage, the other direction: name every registered parameter + # outside the target's own ParameterDict instead of assuming lm_head is the only + # one. + shell = {n for n, _ in model.named_parameters() if not n.startswith("model.w.")} + assert shell == {"lm_head.weight"}, f"unfed shell parameters: {sorted(shell)}" + + torch.cuda.synchronize() + leftover = set(weights.keys()) - consumed + expected = self.expected_unconsumed(d, weights.keys()) + assert leftover == expected, ( + "checkpoint coverage: leftover keys are not exactly this rank's off-window " + f"experts{'' if d.mtp_enabled else ' and the MTP module'}; unexpected " + f"{sorted(leftover - expected)[:8]}, missed {sorted(expected - leftover)[:8]}" + ) + + +MODEL_WEIGHTS = Qwen4ExpDefaultWeights() + + +def _heads(total: int, ranks: int, rank: int) -> tuple[int, int]: + """(heads on `rank`, its first global head) when `total` heads are split over `ranks`: + an even split, or each head repeated on `ranks // total` neighbouring ranks when there + are more ranks than heads -- the rule the K/V and recurrent pools follow + (`pool.pool_geometry`).""" + assert total % ranks == 0 or ranks % total == 0, (total, ranks) + return -(-total // ranks), rank * total // ranks + + +def rank_dims( + cfg, + vcfg, + mapping, + *, + mtp_enabled: bool, + trunk_moe_backend: str, + moe_prefill_copy: bool, + dtype: torch.dtype, +) -> SimpleNamespace: + """Every width the table declares at rank `mapping.rank`, and the shares behind them. + + Args: + cfg: the engine-normalized text config (`Qwen4ExpTextConfig`). + vcfg: the composite's vision config. + mapping: this rank's `Mapping`. + mtp_enabled: whether the engine has a `speculative_config` (the MTP module is live). + trunk_moe_backend: the routed backend the trunk runs, which picks the expert layout. + moe_prefill_copy: whether the core keeps the CuteDSL-layout copy of its window. + dtype: the activation dtype, which every unquantized weight is declared in. + """ + hidden = cfg.hidden_size + head_dim = cfg.head_dim + index_heads, index_kv_heads = cfg.indexer_n_heads, cfg.indexer_kv_heads + + attn_tp = attention_tp_size(mapping) + attn_rank = mapping.tp_rank if attn_tp > 1 else 0 + heads_q, q_head_start = _heads(cfg.num_attention_heads, attn_tp, attn_rank) + heads_kv, kv_head_start = _heads(cfg.num_key_value_heads, attn_tp, attn_rank) + + state_tp = mamba_effective_tp_size(mapping) + state_rank = mapping.tp_rank if state_tp > 1 else 0 + num_k_heads, k_head_start = _heads(cfg.linear_num_key_heads, state_tp, state_rank) + num_v_heads, v_head_start = _heads(cfg.linear_num_value_heads, state_tp, state_rank) + key_dim = num_k_heads * cfg.linear_key_head_dim + value_dim = num_v_heads * cfg.linear_value_head_dim + + num_experts = cfg.num_experts + # The routed experts' MoE group is the TP group: tp4 splits it by intermediate, + # tep4 and dep4 by expert. + assert mapping.moe_tp_size * mapping.moe_ep_size == mapping.tp_size, ( + mapping.moe_tp_size, + mapping.moe_ep_size, + mapping.tp_size, + ) + assert num_experts % mapping.moe_ep_size == 0 and num_experts % mapping.tp_size == 0, ( + num_experts, + mapping.moe_ep_size, + mapping.tp_size, + ) + local_experts = num_experts // mapping.moe_ep_size + mtp_local_experts = num_experts // mapping.tp_size + # Each expert's intermediate over the MoE TP ranks (tp4; one rank, the whole of it, + # otherwise). A share holds whole group-16 scale vectors, so `down_proj`'s scale + # columns cut on a vector boundary; it is stored zero-padded to EXPERT_INTER_ALIGN. + moe_inter = cfg.moe_intermediate_size + assert moe_inter % mapping.moe_tp_size == 0, (moe_inter, mapping.moe_tp_size) + expert_inter_share = moe_inter // mapping.moe_tp_size + assert expert_inter_share % SF_VEC == 0, (expert_inter_share, SF_VEC) + expert_inter = -(-expert_inter_share // EXPERT_INTER_ALIGN) * EXPERT_INTER_ALIGN + + ngram_heads = cfg.heads_per_ngram * (cfg.ngram_size - 1) + v_hidden = vcfg.hidden_size + return SimpleNamespace( + num_layers=cfg.num_hidden_layers, + is_attn=[t == "full_attention" for t in cfg.layer_types], + hidden=hidden, + hc_count=cfg.hc_count, + hc_lowrank=cfg.hc_lowrank, + stream=cfg.hc_count * hidden, + tp_rank=mapping.tp_rank, + # --- attention, at this rank's heads --- + head_dim=head_dim, + heads_q=heads_q, + q_head_start=q_head_start, + heads_kv=heads_kv, + kv_head_start=kv_head_start, + q_width=heads_q * head_dim, + kv_width=heads_kv * head_dim, + index_head_dim=cfg.indexer_head_dim, + index_width=(index_heads + index_kv_heads) * cfg.indexer_head_dim, + # [q | k | v | gate | index q, k | pad of one index head] + attn_in_rows=2 * heads_q * head_dim + + 2 * heads_kv * head_dim + + (index_heads + index_kv_heads + 1) * cfg.indexer_head_dim, + # --- Gated-DeltaNet, at this rank's key and value heads --- + num_k_heads=num_k_heads, + k_head_start=k_head_start, + num_v_heads=num_v_heads, + v_head_start=v_head_start, + head_k_dim=cfg.linear_key_head_dim, + head_v_dim=cfg.linear_value_head_dim, + key_dim=key_dim, + key_dim_total=cfg.linear_num_key_heads * cfg.linear_key_head_dim, + value_dim=value_dim, + conv_dim=2 * key_dim + value_dim, + in_proj_split=(2 * key_dim + value_dim, value_dim, num_v_heads, num_v_heads), + conv_kernel=cfg.linear_conv_kernel_dim, + # --- MoE --- + num_experts=num_experts, + expert_split="ep" if mapping.moe_tp_size == 1 else "tp", + local_experts=local_experts, + expert_offset=local_experts * mapping.moe_ep_rank, + # This rank's intermediate rows `[expert_inter_start, + expert_inter_share)` of + # every expert it holds, stored `expert_inter` wide (the runner's + # `intermediate_size`). + expert_inter_share=expert_inter_share, + expert_inter_start=expert_inter_share * mapping.moe_tp_rank, + expert_inter=expert_inter, + mtp_local_experts=mtp_local_experts, + mtp_expert_offset=mtp_local_experts * mapping.tp_rank, + moe_inter=moe_inter, + shared_inter=cfg.shared_expert_intermediate_size, + vocab=cfg.vocab_size, + # --- PLE --- + ple_layer=list(cfg.ple_layer_ids)[0] - 1, # one-based in the config + ple_embed_dim=cfg.ple_embed_dim, + ple_conv_kernel=cfg.ple_conv_kernel_size, + ngram_size=cfg.ngram_size, + ngram_heads=ngram_heads, + ple_head_dim=cfg.ple_embed_dim // ngram_heads, + ple_shards=cfg.split_ngram_parts, + ple_table_dtype=PLE_TABLE_DTYPE, + # --- the vision tower --- + v_depth=vcfg.depth, + v_hidden=v_hidden, + v_inter=vcfg.intermediate_size, + v_channels=vcfg.in_channels, + v_tpatch=vcfg.temporal_patch_size, + v_patch=vcfg.patch_size, + v_pos_rows=vcfg.num_position_embeddings, + v_merge_dim=v_hidden * vcfg.spatial_merge_size * vcfg.spatial_merge_size, + # --- engine choices --- + mtp_enabled=mtp_enabled, + trunk_moe_backend=trunk_moe_backend, + moe_prefill_copy=moe_prefill_copy, + dtype=dtype, + ) + + +def lm_head_share(rows: int, vocab: int, tp_rank: int) -> tuple[Cut, Callable]: + """The cut and the pad that fill an lm_head the shell built with `rows` rows. + + Whole under attention DP (`rows == vocab`); a vocab-parallel shard under attention TP. + A vocab that does not split evenly leaves the last shard short, and its missing rows + are zero, as the built-in LMHead pads them -- so the cut stops at `vocab` and the pad + fills the rest. + """ + start = 0 if rows == vocab else tp_rank * rows + stop = min(start + rows, vocab) + assert start < stop, (rows, vocab, tp_rank) + + def pad(t: torch.Tensor, d) -> torch.Tensor: + if t.shape[0] == rows: + return t.contiguous() + out = torch.zeros(rows, t.shape[1], dtype=t.dtype, device=t.device) + out[: t.shape[0]] = t + return out + + return (slice(start, stop),), pad + + +def ngram_shard_keys(d: SimpleNamespace) -> list[str]: + """The n-gram table's shard keys, in concatenation order.""" + prefix = _ngram_prefix(d.ple_layer) + return [f"{prefix}.shard_{s}.weight" for s in range(d.ple_shards)] + + +def ngram_table_sources(d: SimpleNamespace, rows_per_shard: int) -> list[tuple]: + """The `ple_table` manifest rows: shard `s` fills rows `[s * rows_per_shard, ...)`, + read whole on every rank (the table is replicated).""" + return [ + (key, (_span(s * rows_per_shard, rows_per_shard),), None, None) + for s, key in enumerate(ngram_shard_keys(d)) + ] + + +def _declare_ngram_table(core, d: SimpleNamespace, weights) -> int: + """Register the n-gram table with the height its shards concatenate to, and return + the rows each shard holds. + + The shards are **contiguous row chunks**, not an interleaving: global row `r` lives in + shard `r // rows_per_shard` at local row `r % rows_per_shard`. That is what makes the + flat buffer indexable by the hash's global rows. + + **This is the whole of PLE host offload.** `core.ple_host_offload` picks between an + ordinary device allocation and host memory the device addresses; shape, dtype, + `ParameterDict` key, the shard-by-shard fill and the forward's gather are the same either + way, so an offloaded table is loaded and covered in both directions like any other. The + core asserts where it landed against the driver. + """ + keys = ngram_shard_keys(d) + for key in keys: + assert key in weights, f"checkpoint key missing: {key}" + shapes = {tuple(weights[key].shape) for key in keys} + assert len(shapes) == 1, ( + f"the n-gram table's {len(keys)} shards do not share one shape: " + f"{sorted(shapes)}; they are concatenated as equal contiguous row " + f"chunks and an uneven one would shift every row above it" + ) + rows_per_shard, head_dim = shapes.pop() + assert head_dim == d.ple_head_dim, ( + f"the n-gram table is {head_dim} wide but {d.ngram_heads} heads must " + f"tile ple_embed_dim {d.ple_embed_dim}, i.e. {d.ple_head_dim}" + ) + dtype = weights[keys[0]].dtype + assert dtype == d.ple_table_dtype, ( + f"the n-gram table is stored {dtype}; this target declares " + f"{d.ple_table_dtype} and does not convert it" + ) + device = core.w["embed"].device + rows = d.ple_shards * rows_per_shard + if core.ple_host_offload: + table, core._ple_table_host = _host_resident_empty(rows, head_dim, dtype, device) + else: + table = torch.empty(rows, head_dim, dtype=dtype, device=device) + core._ple_table_host = None + core.w["ple_table"] = nn.Parameter(table, requires_grad=False) + # `nn.Parameter` (`_make_subclass`) shares storage rather than copying; host offload + # depends on that, so it is asserted. + assert core.w["ple_table"].data_ptr() == table.data_ptr(), ( + "wrapping the n-gram table in nn.Parameter moved its storage" + ) + core.ple_rows_per_shard = rows_per_shard + return rows_per_shard diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py deleted file mode 100644 index 6f48df505098..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py +++ /dev/null @@ -1,8119 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""ModelingV2 target: Qwen3.8-Flash-Next-NVFP4 / sm_103 / tp1 — self-contained modeling code. - -The core holds the static structure (weights, derived operands, per-manager pool -views); `PrefillTarget` (any batch with context rows) and `DecodeTarget` -(generation-only) run a step over it, passing every argument to every op at the -call site. - -The checkpoint is the composite Qwen4-Exp VLM export. The MTP module is built only -when the engine has a ``speculative_config`` (else ``mtp.*`` stays in -``weights.py``'s non-load set). The vision tower is always built; a text-only -request never reaches it. Every speculation-dependent change is gated on a host -int that is 1 when speculation is off. The PLE pools and the QSA side caches are -written for every verify row, so `SpeculativeRollback` (registered with the spec -worker) rewrites them to the accepted prefix after each verify step. - -Differences from other targets: - -* **Hyper-connections.** The inter-block residual is ``4 * 2560 = 10240`` wide. - There is **no** ``model.language_model.norm.weight``: the stack-final - ``hyper_connection_mixer`` is the final norm and ``lm_head`` reads its output. - In the trunk each site is ``ops/hc/qwen4_exp_hc_inject_norm`` (the previous block's - accumulate fused with this site's grouped norm; after a MoE block it also forms that - block's output from the routed and shared outputs and the gate logit) and then the - mixer: ``catalog/gemm/cublas_mm`` over the concatenated ``[down; block_inject]`` - weight, ``ops/hc/qwen4_exp_hc_lowrank_act`` and ``ops/hc/qwen4_exp_hc_up_gate_mean`` - (the up projection as a tcgen05 GEMM whose epilogue applies the sigmoid gate and the - stream mean). ``DecodeTarget`` mixers run the first projection as one tcgen05 launch - (``ops/hc/qwen4_exp_hc_decode_tc_first``, up to ``_HC_TC_FIRST_MAX_ROWS`` rows) or, at - ``_HC_DECODE_GEMV_MAX_ROWS`` rows, the whole site as two GEMV launches - (``ops/hc/qwen4_exp_hc_decode_first_cute``, ``ops/hc/qwen4_exp_hc_decode_mix_fp4_cute``); - at the 48 MoE sites the mixer's last launch also writes the routed experts' NVFP4 - input. The MTP layer's sites run the same fused site; its front end's grouped norm - (``_grouped_norm``) and the accumulates before the PLE block and after the MTP layer's - MoE block (``qwen4_exp_post_inject``) run alone. -* **QSA sparse attention.** The index plane is maintained on every step so a - request can cross ``dense_seq_len_threshold`` later without rebuilding its - prefix; sparse selection replaces the dense core only above it. The MTP draft - layer runs the same branch over its own indexer and its own manager's planes. -* **PLE.** Decoder layer 1 only (``ple_layer_ids: [2]`` is one-based): an n-gram - side path over a 320,001,536-row fp8 table (47.68 GiB of a 123.5 GiB checkpoint). -* **NVFP4 MoE on every layer** (routed experts only); everything else is bf16. The - trunk's routed experts run on trtllm-gen's block-scale runner (``TRTLLM``, and what - ``AUTO`` selects) or on the CuteDSL sort / grouped-GEMM pipeline (``CUTEDSL``), - chosen by ``moe_config.backend``; any other request is refused. The MTP draft - layer's FP8 block-scale experts run on trtllm-gen under either. Under ``TRTLLM`` the - prefill-sized launches (``_CUTE_DSL_MOE_MIN_ROWS``) run the CuteDSL pipeline over a - second copy of the experts in its layout, about 1.3 GiB a layer, kept only when a step - can reach that size and ``TRTLLM_MODELING_V2_MOE_PREFILL_COPY`` does not turn it off - (``_moe_prefill_copy``). A graph-captured decode step routes its block - projections through ``ops/gemm/low_m_bf16_mm`` (FlashInfer low-M kernels at - M <= 32 on the shapes it lists, cuBLAS otherwise); see ``DecodeTarget.forward``. - -Integration facts this file depends on: - -1. The engine rebuilds the composite config and normalizes its text tower - (``config_utils.load_pretrained_config``). The shell adopts that text config as - ``model_config.pretrained_config`` before building the core, because the - recurrent-state and PLE pool geometry are read from it. -2. ``position_ids`` arrives in the mRoPE layout (``[3, 1, T]``) even for text-only - batches, because the normalized text config carries ``rope_scaling['type'] == - 'mrope'``. The main and MTP rope read all three rows; the QSA index plane takes - coordinate 0 only (see ``_QSA_MROPE_SECTION``). -3. The engine owns the multimodal input processor (resolved off the checkpoint's - ``architectures[0]``; modeling_v2 has no hook), so it is the built-in - ``Qwen3VLInputProcessorBase``. This target asserts every ``multimodal_data`` key - where it is read. - -**KV storage.** The paged K/V pool holds the checkpoint's bf16, fp8-e4m3 when the engine -resolves ``kv_cache_quant_algo: FP8`` (``kv_cache_config.dtype: fp8``), or NVFP4 when it resolves -``NVFP4`` (``kv_cache_config.dtype: nvfp4``): a lossy substitution, since the checkpoint's KV is -bf16 and the fp8 and NVFP4 attention kernels run q, K, V and the softmax probabilities through -e4m3. An NVFP4 pool stores e2m1 pairs (uint8, half a head wide) and keeps one e4m3 scale per 16 -values in a second pool; the second-level scales are 1.0. The NVFP4 attention kernels exist for -the paged-context path only, so an NVFP4 engine needs chunked prefill, block reuse or -speculative decoding. The QSA index and position planes keep their own types under all three. - -**Recurrent-state storage.** The Gated-DeltaNet recurrent pool holds the checkpoint's fp32 -(``mamba_ssm_dtype``) or bf16 when ``kv_cache_config.mamba_ssm_cache_dtype`` asks for it, which -is also what ``auto`` resolves to for this checkpoint (`_SSM_STATE_DTYPES`): a lossy -substitution, since the state is rounded to bf16 once per decode token, once per prefill call and -at every cached-replay checkpoint commit. The conv, PLE and replay-history buffers keep their own -types under both. Over a bf16 pool a decode step takes one recurrence kernel at every batch size, -so a row's tokens do not depend on its batch; see ``DecodeTarget._gdn_block``. - -**Breakable prefill CUDA graph.** With ``prefill_cuda_graph_backend: breakable`` the engine -captures the core's ``forward`` (the body, not the shell) for a step with context rows at each of -``prefill_capture_num_tokens``, pads any such step up to the nearest bucket and replays it without -calling the body. The shell therefore calls the core's ``begin_step`` before the body on every -step: the asserts, the cache-manager bind, the target choice and the prefill target's host -bookkeeping (``PrefillTarget.prepare``, one immutable ``_PrefillStep``) all live there. The prefill -``forward`` is row-wise over the ``T`` rows it is handed, with four breaks -(``catalog/graph/eager_on_graph``) that run eagerly between the captured segments and read the -step: the vision splice, the PLE block, and each layer's attention core and Gated-DeltaNet -recurrence. Each writes rows ``[:n]`` of a fixed output and zeroes rows ``[n:]``. Under MTP the -prefill body returns one packed ``[T, hidden + stream]`` tensor that ``split_spec_output`` splits, -and the shell hands the spec worker the batch's own ``n`` rows. With the backend off every step is -eager and unpadded and the breaks are plain calls. - -**Host offload.** By default the PLE table lives in host memory the device -addresses (`weights._host_resident_empty`), and the gather reads it unchanged; -``TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD=0`` restores the resident declaration. Output -is identical, so only the device-memory delta sees it. The gather stays -device-driven because row indices are produced on device inside a captured decode -forward, which cannot express a host readback. -""" - -import math -import os -import re -from typing import NamedTuple - -import torch -from torch import nn -from transformers import PretrainedConfig - -from tensorrt_llm._torch._experimental.modeling_v2._router_index import step_contract_enabled -from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.flashinfer_silu_and_mul import ( # noqa: E501 - flashinfer_silu_and_mul, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( - fused_qk_norm_rope, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import ( - QM_FP8_KV_CACHE, - QM_NVFP4_KV_CACHE, - thop_attention, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm -from tensorrt_llm._torch._experimental.modeling_v2.catalog.graph.eager_on_graph import ( - eager_on_graph, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( - fp4_quantize, -) -from tensorrt_llm._torch.attention.backends.fmha.phased import get_spec_decoding_position_offsets -from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata -from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata -from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM -from tensorrt_llm._torch.modules.fla.cached_replay import CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE -from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor -from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks -from tensorrt_llm._torch.speculative import get_spec_worker, should_use_separate_draft_kv_cache -from tensorrt_llm._torch.utils import ActivationType -from tensorrt_llm.inputs.multimodal import get_multimodal_item_metadata -from tensorrt_llm.logger import logger -from tensorrt_llm.quantization.mode import QuantAlgo - -from . import weights as _weights -from .ops.activation.fused_sigmoid_mul_inplace import fused_sigmoid_mul_inplace -from .ops.activation.gelu_erf import gelu_erf -from .ops.activation.gelu_tanh import gelu_tanh -from .ops.activation.silu_and_mul_rows import silu_and_mul_rows -from .ops.attention.qsa_canonical_selection import qsa_canonical_selection -from .ops.attention.qsa_expand_block_indices import qsa_expand_block_indices -from .ops.attention.qsa_mtp_reuse_topk import qsa_mtp_reuse_topk -from .ops.attention.qsa_nvfp4_pages_to_bf16 import qsa_nvfp4_pages_to_bf16 -from .ops.attention.qsa_paged_index_scores import qsa_paged_index_scores -from .ops.attention.qsa_paged_index_store import qsa_paged_index_store -from .ops.attention.qsa_paged_kv_store import qsa_paged_kv_store -from .ops.attention.qsa_paged_kv_store_nvfp4 import qsa_paged_kv_store_nvfp4 -from .ops.attention.qsa_paged_sparse_gqa import qsa_paged_sparse_gqa -from .ops.attention.qsa_paged_sparse_gqa_decode import MAX_ROWS as _SPARSE_DECODE_MAX_ROWS -from .ops.attention.qsa_paged_sparse_gqa_decode import qsa_paged_sparse_gqa_decode -from .ops.attention.qsa_paged_sparse_gqa_decode import supports as _sparse_decode_supports -from .ops.attention.qsa_paged_sparse_gqa_nvfp4 import qsa_paged_sparse_gqa_nvfp4 -from .ops.attention.qsa_paged_sparse_gqa_nvfp4_decode import qsa_paged_sparse_gqa_nvfp4_decode -from .ops.attention.qsa_paged_sparse_gqa_nvfp4_verify import qsa_paged_sparse_gqa_nvfp4_verify -from .ops.attention.qsa_paged_sparse_gqa_nvfp4_verify import takes_grouped as _verify_takes_grouped -from .ops.attention.qsa_prefill_compress import qsa_prefill_compress -from .ops.attention.qsa_rollback_planes import plane_table as _plane_table -from .ops.attention.qsa_rollback_planes import qsa_rollback_gather, qsa_rollback_restore -from .ops.attention.qsa_sparse_gqa_rowtile import qsa_sparse_gqa_rowtile -from .ops.attention.qsa_topk_expand_decode import qsa_topk_expand_decode -from .ops.attention.qsa_topk_expand_decode import supports as _topk_expand_supports -from .ops.attention.qsa_topk_select import qsa_topk_select -from .ops.attention.qwen4_exp_attn_decode_prologue import SCORE_ROWS as _PROLOGUE_SCORE_ROWS -from .ops.attention.qwen4_exp_attn_decode_prologue import qwen4_exp_attn_decode_prologue -from .ops.attention.qwen4_exp_qk_norm_rope import qwen4_exp_qk_norm_rope -from .ops.attention.spec_rollback_commit import spec_rollback_commit, spec_rollback_stage -from .ops.attention.verify_step_setup import _BLOCK_S as _VERIFY_SETUP_MAX_SEQS -from .ops.attention.verify_step_setup import verify_step_setup -from .ops.attention.vision_rope_2d import vision_rope_2d -from .ops.gemm.bf16_gemv_pdl import MAX_ROWS as _GEMV_MAX_ROWS -from .ops.gemm.bf16_gemv_pdl import TILING as _GEMV_TILED_K -from .ops.gemm.bf16_gemv_pdl import bf16_gemv_pdl -from .ops.gemm.fp8_gemv import ( - fp8_gemv_conv1d_op, - fp8_gemv_conv1d_verify_op, - fp8_gemv_op, - pack_from_bf16, -) -from .ops.gemm.fp8_per_tensor_mm import fp8_per_tensor_mm, fp8_per_tensor_weight -from .ops.gemm.low_m_bf16_mm import low_m_bf16_mm -from .ops.gemm.mid_rows_mm import MidWeight, mid_rows_mm_op -from .ops.gemm.mid_rows_mm import precompile as mid_rows_precompile -from .ops.gemm.nvfp4_gemv import ( - nvfp4_gemv_op, - nvfp4_split_gemv_conv1d_op, - nvfp4_split_gemv_conv1d_verify_op, - nvfp4_split_gemv_op, - precompile_split, -) -from .ops.gemm.nvfp4_gemv import pack_from_bf16 as nvfp4_pack_from_bf16 -from .ops.gemm.nvfp4_gemv import pack_split_from_bf16 as nvfp4_pack_split_from_bf16 -from .ops.gemm.nvfp4_gemv import precompile as nvfp4_gemv_precompile -from .ops.hc.qwen4_exp_hc_decode_first_cute import qwen4_exp_hc_decode_first_cute -from .ops.hc.qwen4_exp_hc_decode_mix_fp4_cute import ( - qwen4_exp_hc_decode_mix_cute, - qwen4_exp_hc_decode_mix_fp4_cute, -) -from .ops.hc.qwen4_exp_hc_decode_rows_tc import MAX_ROWS as _HC_DECODE_ROWS_MAX -from .ops.hc.qwen4_exp_hc_decode_rows_tc import ( - qwen4_exp_hc_decode_rows_first, - qwen4_exp_hc_decode_rows_mix, -) -from .ops.hc.qwen4_exp_hc_decode_tc_first import precompile as hc_tc_first_precompile -from .ops.hc.qwen4_exp_hc_decode_tc_first import qwen4_exp_hc_decode_tc_first -from .ops.hc.qwen4_exp_hc_first_rescale_tc import qwen4_exp_hc_first_rescale_tc -from .ops.hc.qwen4_exp_hc_inject_norm import qwen4_exp_hc_inject_norm -from .ops.hc.qwen4_exp_hc_inject_rstd import qwen4_exp_hc_inject_rstd -from .ops.hc.qwen4_exp_hc_lowrank_act import qwen4_exp_hc_lowrank_act -from .ops.hc.qwen4_exp_hc_up_gate_mean import qwen4_exp_hc_up_gate_mean -from .ops.hc.qwen4_exp_hc_up_gate_mean_fp4 import qwen4_exp_hc_up_gate_mean_fp4 -from .ops.hc.qwen4_exp_hc_up_gate_mean_rstd import ( - qwen4_exp_hc_up_gate_mean_rstd, - qwen4_exp_hc_up_gate_mean_rstd_fp4, -) -from .ops.hc.qwen4_exp_post_inject import qwen4_exp_post_inject -from .ops.moe.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import ( - cute_dsl_nvfp4_gather_grouped_gemm_act_fusion, -) -from .ops.moe.cute_dsl_nvfp4_grouped_gemm import cute_dsl_nvfp4_grouped_gemm -from .ops.moe.fp4_block_scale_moe_renorm_permuted import fp4_block_scale_moe_renorm_permuted -from .ops.moe.fp4_block_scale_moe_renorm_routed import fp4_block_scale_moe_renorm_routed -from .ops.moe.fp4_moe_renorm_cute_dsl import fp4_moe_renorm_cute_dsl -from .ops.moe.fp8_block_scale_moe_runner import fp8_block_scale_moe_runner -from .ops.moe.fused_sigmoid_gate_mul_add import fused_sigmoid_gate_mul_add -from .ops.moe.moe_sort import moe_sort -from .ops.moe.moe_unpermute_inplace import moe_unpermute_inplace -from .ops.moe.renorm_moe_routing_op import renorm_moe_routing_op -from .ops.norm.flashinfer_gemma_rmsnorm import flashinfer_gemma_rmsnorm -from .ops.norm.layer_norm import layer_norm -from .ops.norm.rms_norm_gated_token_major import rms_norm_gated_token_major -from .ops.ple.ple_gate_conv_inject_norm import ple_gate_conv_inject_norm -from .ops.ple.ple_ngram_embed import ple_ngram_embed -from .ops.quantization.bf16_to_fp32 import bf16_to_fp32 -from .ops.quantization.fp8_quantize_1x128 import fp8_quantize_1x128 -from .ops.ssm.causal_conv1d_triton_update import causal_conv1d_triton_update -from .ops.ssm.causal_conv1d_update import causal_conv1d_update -from .ops.ssm.cu_seqlens_from_seq_lens import cu_seqlens_from_seq_lens -from .ops.ssm.flashinfer_chunk_gated_delta_rule import flashinfer_chunk_gated_delta_rule -from .ops.ssm.fused_gdn_gating import fused_gdn_gating -from .ops.ssm.fused_recurrent_gated_delta_rule_cached_replay_update import ( - fused_recurrent_gated_delta_rule_cached_replay_update, -) -from .ops.ssm.fused_recurrent_gated_delta_rule_update import fused_recurrent_gated_delta_rule_update -from .ops.ssm.gdn_cached_replay_cute import gdn_cached_replay_cute -from .ops.ssm.gdn_decode_gated_norm_pdl import MAX_ROWS as _GDN_DECODE_FUSED_MAX_ROWS -from .ops.ssm.gdn_decode_gated_norm_pdl import gdn_decode_gated_norm_pdl -from .ops.ssm.gdn_decode_rows_gated_norm import gdn_decode_rows, gdn_decode_rows_gated_norm -from .ops.ssm.gdn_in_proj_conv1d_pdl import MAX_ROWS as _GDN_CONV_FUSED_MAX_ROWS -from .ops.ssm.gdn_in_proj_conv1d_pdl import gdn_in_proj_conv1d_pdl -from .ops.ssm.gdn_prefill_conv_post import gdn_prefill_conv_post -from .ops.ssm.pack_gdn_decode_qkv import pack_gdn_decode_qkv -from .ops.vision.patch_embed_conv3d import patch_embed_conv3d -from .ops.vision.pos_embed_interpolate import pos_embed_interpolate - -# The GPU architecture this target is certified for; a direct instantiation could -# bypass routing. -_SM = (10, 3) - -# QSA backend metadata buffers rebuilt by prepare(). Their absence means the engine -# did not select the QSA metadata and the index plane would never be written. -_QSA_FIELDS = ( - "qsa_block_table", - "qsa_req_idx_per_token", - "qsa_logical_positions", - "qsa_sequence_lengths", - "qsa_visible_blocks", - "qsa_topk_indices", - "qsa_topk_row_starts", -) - -# Fields read off `attn_metadata.mamba_metadata`. Absence is silent: a smoke run -# passes with no recurrent state at all. -_MAMBA_FIELDS = ( - "state_indices", - "has_initial_states", - "query_start_loc", - "query_start_loc_long", -) - - -# A context launch at or above this many rows takes `qsa_sparse_gqa_rowtile`; -# measured on GB300 it loses at 543 rows and wins from 4096 rows up. -_ROWTILE_MIN_ROWS = 2048 # capacity: fixed: a dispatch threshold, not a buffer size -# Over an NVFP4 pool, a step with more context rows than this reads its context requests' -# pages dequantized into a bf16 scratch through `qsa_sparse_gqa_rowtile`. GB300, a layer: -# 16384 context rows 1.36 ms there against 77-948 ms for the NVFP4 per-row kernel. -_NVFP4_ROWTILE_MIN_ROWS = 256 # capacity: fixed: a dispatch threshold, not a buffer size -# The bf16 K/V scratch those context rows read holds every page of each context request's KV -# history, not just the step's rows. Consecutive context requests are dequantized and attended -# in groups whose pages fit this many bytes; a request with more pages forms a group of its -# own, so the scratch peaks at the larger of this and the longest request's history. -_NVFP4_CONTEXT_SCRATCH_BYTES = 128 << 20 # capacity: fixed: a scratch memory budget - -#: Context-step MoE launches of at least this many rows run on the CuTe DSL MoE -#: (`fp4_moe_renorm_cute_dsl`, from the second expert copy `weights.py` keeps in -#: its layout); smaller ones, every decode one among them, keep the trtllm-gen -#: runner. From GB300 micro-benches, us a layer, CuTe DSL vs trtllm-gen: 2048 -#: rows 275 vs 273, 8192 rows 451 vs 555, 16384 rows 656 vs 882. -_CUTE_DSL_MOE_MIN_ROWS = 4096 # capacity: fixed: a dispatch threshold, not a buffer size - -#: Widest decode projection that `bf16_gemv_pdl` takes at 2..`_GEMV_MAX_ROWS` rows: the shared -#: expert's one-column gate logit, on one CTA under `low_m_bf16_mm` (~3 us a layer at 4 rows). -_GEMV_NARROW_N = 8 # capacity: fixed: a dispatch threshold, not a buffer size - - -def _decode_projection(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor: - """A graph-captured decode step's bf16 projection, always a PDL launch. - - One row runs `bf16_gemv_pdl` on every K it tiles (each decode projection of this - checkpoint): its weight streams into registers before it waits on the producer, and it - triggers its own dependents. So do up to `_GEMV_MAX_ROWS` rows of a projection at most - `_GEMV_NARROW_N` wide (the shared expert's gate logit), which `low_m_bf16_mm` would run - on one CTA. Other calls run `low_m_bf16_mm` with its FlashInfer branches' PDL on (its - cuBLAS branch, above M=32 and off its table, ignores the flag). Every kernel either - path launches reads the activation only after `griddepcontrol.wait`. - """ - rows = mat_a.shape[0] - if mat_b.shape[0] in _GEMV_TILED_K and ( - rows == 1 or (rows <= _GEMV_MAX_ROWS and mat_b.shape[1] <= _GEMV_NARROW_N) - ): - return bf16_gemv_pdl(mat_a, mat_b, True) - return low_m_bf16_mm(mat_a, mat_b, True) - - -# Per-call constants of the dense attention call shape: packed-QKV causal GQA, RoPE -# applied outside (fused_qk_norm_rope), bf16, fp8-e4m3 or NVFP4 pool. The sparse / skip-softmax -# group stays None/0: QSA calls selection and sparse attention as separate kernels and -# `num_sparse_topk` is 0. `quant_mode`, the two KV scales (all three follow the pool's -# storage) and `predicted_tokens_per_seq` are passed beside these. -_CALL_CONSTANTS = dict( - output_sf=None, - k=None, - v=None, - is_fused_qkv=True, - update_kv_cache=True, - attention_input_type=0, - is_mla_enable=False, - mask_type=1, - q_scaling=1.0, - attention_sinks=None, - # In-kernel RoPE disabled; rope_* values are inert placeholders. - position_embedding_type=0, - rotary_inv_freq=None, - rotary_cos_sin=None, - rope_dim=0, - rope_base=10000.0, - rope_scale_type=0, - rope_scale=1.0, - rope_short_m_scale=1.0, - rope_long_m_scale=1.0, - rope_max_positions=1024, - rope_original_max_positions=1024, - out_scale=None, - latent_cache=None, - q_pe=None, - q_lora_rank=None, - kv_lora_rank=None, - qk_nope_head_dim=None, - qk_rope_head_dim=None, - v_head_dim=None, - rope_append=None, - chunked_prefill_buffer_batch_size=1, - attention_chunk_size=None, - softmax_stats_tensor=None, - sparse_kv_indices=None, - sparse_kv_offsets=None, - sparse_attn_indices=None, - sparse_attn_offsets=None, - sparse_attn_indices_block_size=0, - mrope_rotary_cos_sin=None, - mrope_position_deltas=None, - helix_position_offsets=None, - helix_is_inactive_rank=None, - cross_kv=None, - relative_attention_bias=None, - relative_attention_max_distance=0, - # MLA-only surface, unused here. - kv_norm_weight=None, - kv_norm_eps=1e-6, - skip_correction_threshold=0.0, -) - -# ------------------------------------------------------------ vision ------ - -#: Tower rope base (`vision_config` has no rope field); unrelated to the text 1e7. -_VISION_ROPE_THETA = 10000.0 - -#: LayerNorm epsilon of `norm1`, `norm2`, `merger.norm`: a literal in the reference. -_VISION_NORM_EPS = 1e-6 - -#: `mask_type` 0 = padding mask: bidirectional within a batch segment. Causal (1) -#: would make a patch depend on raster order. -_VISION_MASK_PADDING = 0 - -# Per-call constants of the tower's attention: the encoder-only mode of -# `attention/thop_attention`. Omitting `kv_cache_block_offsets`, -# `host_kv_cache_pool_pointers` and `host_kv_cache_pool_mapping` selects that mode. -# `update_kv_cache=False` is rejected ("KV cache update cannot be disabled now"), -# so it stays True. -_VISION_CALL_CONSTANTS = dict( - k=None, - v=None, - output_sf=None, - kv_cache_block_offsets=None, - host_kv_cache_pool_pointers=None, - host_kv_cache_pool_mapping=None, - tokens_per_block=None, - use_paged_context_fmha=False, - is_fused_qkv=True, - update_kv_cache=True, - predicted_tokens_per_seq=1, - # No pool, so no layer to address. - local_layer_idx=0, - max_context_q_len_override=None, - cache_indirection=None, - kv_scale_orig_quant=None, - kv_scale_quant_orig=None, - out_scale=None, - quant_mode=0, - beam_width=1, - mask_type=_VISION_MASK_PADDING, - q_scaling=1.0, - # Rotation is applied outside (`ops/attention/vision_rope_2d`); rope_* inert. - position_embedding_type=0, - rotary_inv_freq=None, - rotary_cos_sin=None, - rope_dim=0, - rope_base=10000.0, - rope_scale_type=0, - rope_scale=1.0, - rope_short_m_scale=1.0, - rope_long_m_scale=1.0, - rope_max_positions=1024, - rope_original_max_positions=1024, - attention_input_type=0, - is_mla_enable=False, - chunked_prefill_buffer_batch_size=1, - attention_sinks=None, - block_ids_per_seq=None, - latent_cache=None, - q_pe=None, - q_lora_rank=None, - kv_lora_rank=None, - qk_nope_head_dim=None, - qk_rope_head_dim=None, - v_head_dim=None, - rope_append=None, - attention_chunk_size=None, - softmax_stats_tensor=None, - is_spec_decoding_enabled=False, - use_spec_decoding=False, - is_spec_dec_tree=False, - spec_decoding_generation_lengths=None, - spec_decoding_position_offsets=None, - spec_decoding_packed_mask=None, - spec_decoding_bl_tree_mask_offset=None, - spec_decoding_bl_tree_mask=None, - spec_bl_tree_first_sparse_mask_offset_kv=None, - sparse_kv_indices=None, - sparse_kv_offsets=None, - sparse_attn_indices=None, - sparse_attn_offsets=None, - sparse_attn_indices_block_size=0, - mrope_rotary_cos_sin=None, - mrope_position_deltas=None, - helix_position_offsets=None, - helix_is_inactive_rank=None, - is_cross=False, - cross_kv=None, - relative_attention_bias=None, - relative_attention_max_distance=0, - kv_norm_weight=None, - kv_norm_eps=1e-6, - skip_correction_threshold=0.0, -) - -#: Head geometries the encoder-only mode is certified at: MHA, `head_size` 72 or -#: 128, context rows only. `head_size` 72 is context-phase only ("unsupported -#: head_size 72" in generation). -_VISION_CERTIFIED_HEADS = {(16, 16, 72), (16, 16, 128)} - -#: Head widths `ops/attention/vision_rope_2d` is certified at. -_VISION_ROPE_HEAD_DIMS = (40, 64, 72, 128) - -#: The most patch rows one tower call carries, the ceiling the vision ops' cells reach: -#: 8192 placeholder tokens x `spatial_merge_size`**2 patch rows each. `VisionTower` splits -#: a request's frames into calls under it, so the bound a request meets is per frame: one -#: frame's `h * w` patches. -_VISION_CALL_MAX_PATCHES = 32768 - -# NVFP4 block size: one e4m3 scale per 16 contiguous elements along K. -_SF_VEC = _weights.SF_VEC -# Routing happens outside (renorm_moe_routing_op), so the runner's routing method is -# inert, but it must be 1 (Renormalize): the tier table covers (512 experts, 16 -# top-k) only for it, and a tier miss is a logged error with no exception and no -# kernel launch. -_ROUTING_METHOD_RENORM = 1 -_ACT_TYPE_SWIGLU = 0 - -# The trunk's routed-expert backend by the requested `moe_config.backend` -# (`ModelConfig.moe_backend_requested`; `moe_backend` itself is already resolved, -# AUTO -> CUTLASS). Any other request is refused at construction. The MTP draft -# layer's FP8 block-scale experts run on trtllm-gen whatever this says. -_TRUNK_MOE_BACKENDS = {"AUTO": "TRTLLM", "TRTLLM": "TRTLLM", "CUTEDSL": "CUTEDSL"} -# The CuteDSL grouped GEMMs' MMA tile M, which is also the tile `moe_sort` pads each -# expert's rows to: the GEMMs walk the sort's tile tables. sm_103 admits 128 and 256. -_CUTEDSL_TILE_SIZE = 128 -# The longest MTP draft the CuteDSL pipeline is certified at: its cells capture up to 128 rows -# (32 requests of 4 verify rows). -_CUTEDSL_MAX_DRAFT_LEN = 3 - -# Skip sentinel the conv ops and paged short conv compare `cache_indices` against -# (the runtime's mamba `PAD_SLOT_ID`). Never a substitution: the scheduler's slot -# vector gives every row a real slot, and `ops/ple/ngram_token_history` reads its -# carry via `index_select`, so a -1 there is a device-side assert that poisons the -# process. -_PAD_SLOT_ID = -1 - -# Gated-DeltaNet gating form: the pair the decode kernels compute with and -# `gdn_prefill_conv_post`'s default, so prefill and decode compute the same `g`. -_SOFTPLUS_BETA = 1.0 -_SOFTPLUS_THRESHOLD = 20.0 - -# Gated-DeltaNet Q/K L2-norm epsilon: a literal, not `rms_norm_eps`. -_L2_NORM_EPS = 1e-6 - -# PLE gate magnitude floor before its signed square root (a model constant). -_PLE_CLAMP_MIN = 1e-6 - -# --------------------------------------------------------------- host offload -# -# Where the PLE n-gram table lives (module docstring, *Host offload*). No engine -# arg owns this, hence an env var, read once at construction (the table is -# allocated once, during `load_weights`). Default on: 47.68 GiB otherwise comes -# out of the KV pool. -_PLE_HOST_OFFLOAD_ENV = "TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD" - -#: Spellings of "off"; anything else, including empty, leaves offload on. -_PLE_HOST_OFFLOAD_OFF = frozenset({"0", "false", "no", "off"}) - - -def _ple_host_offload_enabled() -> bool: - """Whether the n-gram table is allocated in host memory.""" - return os.environ.get(_PLE_HOST_OFFLOAD_ENV, "1").strip().lower() not in _PLE_HOST_OFFLOAD_OFF - - -# Whether trtllm-gen keeps the second, CuteDSL-layout copy of the routed experts for the -# launches of at least `_CUTE_DSL_MOE_MIN_ROWS` rows (module docstring). No engine arg owns -# it, hence an env var read at construction; "off" (the spellings above) drops the copy and -# every launch runs on trtllm-gen. -_MOE_PREFILL_COPY_ENV = "TRTLLM_MODELING_V2_MOE_PREFILL_COPY" - - -def _moe_prefill_copy(trunk_moe_backend: str, max_num_tokens: int) -> bool: - """Whether the trunk keeps the second expert copy: only under trtllm-gen, only when a - step's rows (at most `max_num_tokens`) can reach the CuteDSL dispatch threshold, and not - when the env var turns it off.""" - return ( - trunk_moe_backend == "TRTLLM" - and max_num_tokens >= _CUTE_DSL_MOE_MIN_ROWS - and os.environ.get(_MOE_PREFILL_COPY_ENV, "1").strip().lower() not in _PLE_HOST_OFFLOAD_OFF - ) - - -# QSA's fixed constants (backend `constants.py`): one index KV head, three position -# axes, index head exactly twice the rotary width (128 == 2 * 64). -_QSA_INDEX_KV_HEADS = 1 -_QSA_POSITION_AXES = 3 -_QSA_INDEX_HEAD_TO_ROTARY_RATIO = 2 - -# The QSA index plane takes coordinate 0 alone, a deliberate departure from the -# checkpoint. `qsa_prefill_compress` (key side, `mrope_section`) and -# `fused_qk_norm_rope` on `index_scratch` (query side, `use_mrope`) must agree: -# both read the three coordinates or neither, else scores mix two position -# encodings. Both-on is certified (`mrope_section` (11, 11, 10)); None is a -# modeling choice. No cost on text; on an image or video batch the index plane sees -# a frame at one position, changing block selection only above -# `dense_seq_len_threshold`. -# To flip, set the config's `(s0, s1, s2)` and `_QSA_INDEX_USE_MROPE = True`. -_QSA_MROPE_SECTION = None -_QSA_INDEX_USE_MROPE = False - -# ---------------------------------------------------------------- MTP ------ - -#: The checkpoint's MTP module count (`mtp_num_hidden_layers`), replayed -#: `max_draft_len` times; NOT the draft length. -_MTP_NUM_LAYERS = 1 - -#: Set to 1 to give the FP8 draft head an NVFP4 copy of its lm_head (e2m1 values, one e4m3 -#: scale per 16 hidden elements of a vocabulary row, an fp32 global scale; built at load, 357 -#: MB in place of the 636 MB packed FP8 copy unless the dense FP8 mode's verify head needs that) -#: and run its decode rows -- up to `_FP8_GEMV_HEAD_ROWS` -- on `nvfp4_gemv` (GB300 micro, -#: 248320 x 2560: 54-57 us at 1-8 rows and 66-67 at 9-16, against fp8_gemv's 87-90 and 96-98); -#: more rows keep the per-row FP8 copy through cuBLASLt. The verify head is untouched: outputs -#: are unchanged and the acceptance length is the cost to watch. 0 (default): `fp8_gemv`. -_MTP_HEAD_NVFP4_ENV = "TRTLLM_MODELING_V2_MTP_HEAD_NVFP4" - - -def _mtp_head_nvfp4_enabled() -> bool: - mode = os.environ.get(_MTP_HEAD_NVFP4_ENV, "0") - assert mode in ("0", "1"), f"{_MTP_HEAD_NVFP4_ENV}={mode!r}: expected 0 or 1" - return mode == "1" - - -#: Set to 1 to give each GDN layer an NVFP4 copy of its input projection (e2m1 values, one e4m3 -#: scale per 16 hidden elements of a row, an fp32 global scale; `nvfp4_gemv.pack_split`, 23.7 MB -#: a layer, built at load from the bf16 weight) and run its decode rows -- up to -#: `_INPROJ_NVFP4_MAX_ROWS` -- on `fp8_gemv`'s kernel over it, the conv step as the GEMV's -#: epilogue in a captured step (single-token and verify). Prefill, mixed steps and larger -#: batches keep their entries. Lossy (gsm8k-gated); 0 (default): off. -_GDN_INPROJ_NVFP4_ENV = "TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4" - -#: Rows up to which a decode step's GDN in_proj reads the NVFP4 copy (`_GDN_INPROJ_NVFP4_ENV`). -#: GB300 micro (us, conv included, behind a PDL predecessor): 6.4 / 6.1 at 1 / 4 rows against -#: fp8_gemv_conv1d's 7.9 / 7.7; alone 7.1 against 7.9 at 8 rows, 9.1 against 8.6 at 16. -_INPROJ_NVFP4_MAX_ROWS = 8 # capacity: fixed: a dispatch threshold, at most fp8_gemv's MAX_ROWS - - -def _gdn_inproj_nvfp4_enabled() -> bool: - mode = os.environ.get(_GDN_INPROJ_NVFP4_ENV, "0") - assert mode in ("0", "1"), f"{_GDN_INPROJ_NVFP4_ENV}={mode!r}: expected 0 or 1" - return mode == "1" - - -#: Set to 1 to run the GDN input / output and attention input / output projections in FP8 -#: (`fp8_per_tensor_mm`: e4m3 weights with a per-tensor scale, the activation quantized per -#: call with a dynamic per-tensor scale) wherever a call reaches `_DENSE_FP8_MIN_ROWS` rows. A -#: lossy mode beside the checkpoint's bf16 projections (gsm8k-gated); off by default. -_DENSE_FP8_ENV = "TRTLLM_MODELING_V2_DENSE_FP8" - -#: Rows from which a projection takes the FP8 entry, per projection kind: the input -#: projections from 2 rows (one row stays on the bf16 GEMV / fused launches, whose PDL -#: overlap the two quantization launches would break), the K=6144 output projections only -#: where the GEMM saving outgrows the activation quantization's two launches (GB300 -#: micro-bench: cuBLASLt bf16 8-10 us vs FP8 6-7 us + ~4 us below ~512 rows). -_DENSE_FP8_MIN_ROWS = {"in": 2, "out": 512} # capacity: fixed: dispatch thresholds - - -def _dense_fp8_enabled() -> bool: - mode = os.environ.get(_DENSE_FP8_ENV, "0") - assert mode in ("0", "1"), f"{_DENSE_FP8_ENV}={mode!r}: expected 0 or 1" - return mode == "1" - - -#: Row ranges `(lo, hi)` in which a projection with a packed FP8 decode copy runs `fp8_gemv` -#: (FP8 weights, bf16 activations, one PDL launch), per projection kind, every kind from one row: -#: the input projections are 1.6-1.8x the bf16 GEMV through 16 rows; at one row the output and -#: shared-expert projections on fp8_gemv serve STP concurrency 1 1.5% faster than on -#: `bf16_gemv_pdl` (two same-node pairs), from two rows the bf16 entry is the slower low-M GEMM -#: (micro: out/o 5.8 vs 4.3 us, gate|up 4.4 vs 2.1, down 3.2 vs 1.5 at 4 rows). The upper ends are each -#: kind's crossover (GB300 micro, us, fp8_gemv against the next entry): in/attention-in 8.7 vs -#: ~9 `mid_rows_mm` at 16 rows, 10.9 vs 9.1 at 17; out/o 6.4 vs 8.4 bf16 at 16 rows, 9.3 vs 8.8 -#: at 20; gate|up 2.6 vs 4.0 at 16, 4.6 vs 4.1 at 20; down 1.7 vs 2.9 at 8, 2.3 vs 2.0 at 12. -_FP8_GEMV_ROWS = {"in": (1, 16), "out": (1, 16), "sh_gu": (1, 16), "sh_dn": (1, 8)} # dispatch - -#: Rows up to which a verify step's conv runs in the in_proj GEMV's epilogue -#: (`fp8_gemv_conv1d_verify`, certified at 2..16 rows); more rows take the Triton conv. -_VERIFY_CONV_FUSED_MAX_ROWS = 16 # capacity: fixed: the entry's certified rows - -#: Rows up to which a captured single-token step runs the conv step as `fp8_gemv_conv1d`'s -#: epilogue over the packed FP8 in_proj copy (dense FP8 mode). -_FP8_CONV_FUSED_MAX_ROWS = 8 # capacity: fixed: a dispatch threshold - - -#: Rows up to which the vocabulary projection (the draft head; under the dense FP8 mode the -#: verify head too) runs `fp8_gemv` on the packed copy; more rows take cuBLASLt over the -#: per-row-scaled copy. GB300, 248320 x 2560: fp8_gemv 86-98 us at 1-16 rows against 107-134, -#: then 142 at 20 rows and 172 at 32 against a flat ~108. -_FP8_GEMV_HEAD_ROWS = 16 # capacity: fixed: a dispatch threshold, at most fp8_gemv's MAX_ROWS - -#: Rows up to which the FP8 vocabulary projection's rows past `_FP8_GEMV_HEAD_ROWS` run -#: `mid_rows_mm` over the per-row-scaled copy (activation quantized per row); more rows keep -#: cuBLASLt. GB300, 248320 x 2560, us at 17 / 32 / 64 / 128 rows: 99.7 / 101.5 / 106.0 / 125.8 -#: against 110.0 / 111.9 / 107.7 / 137.5; at 256 rows 158 against 150. -_MID_HEAD_ROWS = 128 # capacity: fixed: a dispatch threshold, at most mid_rows_mm's MAX_ROWS - - -#: Row ranges `(lo, hi)` in which a decode projection runs `mid_rows_mm` (tcgen05, the weight -#: streamed before the PDL wait; at most its MAX_ROWS): the dense FP8 projections over their -#: per-tensor e4m3 codes with the activation quantized per row (`fp8_row_quant`, one launch), -#: the shared expert and the router over their bf16 weights. Each range is where it beat the -#: kind's previous entry in the GB300 micro-bench (CUDA-graph chain over HBM-cold weights, -#: quantization launch included; us at 17 / 32 / 64 / 128 / 256 rows): GDN in_proj 9.1 / 9.3 -#: / 9.9 / 11.4 / 13.8 against 10.9 (fp8_gemv) / 12.4 / 12.8 / 13.7 / 15.4 (per-tensor FP8), -#: attention in-proj 8.9 / 9.1 / 9.2 / 10.3 / 13.4 against 9.7 / 12.0 / 12.4 / 13.2 / 14.0, -#: out/o 5.6 / 6.1 / 7.0 / 8.7 against 6.7 / 6.7 / 9.3 / 9.5 bf16 (10.1 vs 12.5 at 256), the -#: router 2.3 / 2.5 / 2.9 against 2.4 / 2.6 / 4.8 (equal at 128), gate|up 2.9 / 3.3 / 3.7 -#: against 4.4 / 4.5 / 4.5 (behind from 128); the shared expert's down projection never won by -#: more than noise. An empty range `(0, -1)` keeps the kind off it. -_MID_MM_ROWS = { - "in": (17, 256), - "out": (17, 128), - "sh_gu": (17, 64), - "sh_dn": (0, -1), - "router": (17, 64), -} # capacity: fixed: dispatch thresholds - - -#: Decode rows from which `DecodeTarget._moe_block` launches the router before the shared -#: expert's side-stream work: the router heads the routed chain (router, routing, FC1, FC2), -#: and launched after the side stream's gate|up GEMM its CTAs waited for SMs behind that -#: GEMM's (STP con256 nsys: a router call 9.2 us launched second, 5.8 launched first). GB300 -#: micro-bench (`DecodeTarget`'s MoE block, 24 layers of cold weights in one graph), us a -#: block, shared expert first vs router first: 2 rows 21.7 vs 22.4, 4 rows 30.8 vs 30.9, 8 -#: rows 46.5 vs 45.7, 16 rows 76.0 vs 74.6, 32 rows 112.9 vs 112.1, 128 rows 206.7 vs 204.4, -#: 256 rows 275.6 vs 271.7, 1024 rows 366.6 vs 364.7. -_ROUTER_FIRST_MIN_ROWS = 8 # capacity: fixed: a dispatch threshold, not a buffer size - - -def _mid_rows(kind: str, rows: int) -> bool: - lo, hi = _MID_MM_ROWS[kind] - return lo <= rows <= hi - - -def _gemv_ready(x: torch.Tensor) -> bool: - """`fp8_gemv`'s activation layout: unit column stride, rows 16-byte aligned.""" - return x.stride(1) == 1 and x.stride(0) % 8 == 0 and x.data_ptr() % 16 == 0 - - -def _dense_mm( - x: torch.Tensor, - weight_t: torch.Tensor, - fp8, - gemm, - kind: str, - packed=None, - mid=None, -) -> torch.Tensor: - """A dense projection under the dense FP8 mode: `fp8_gemv` over the packed decode copy - (`packed`) within `_FP8_GEMV_ROWS[kind]` rows, `mid_rows_mm` over `mid` (a `MidWeight`) - within `_MID_MM_ROWS[kind]` rows, the per-tensor FP8 GEMM (`fp8` = `(codes_t, scale)`) from - `_DENSE_FP8_MIN_ROWS[kind]` rows, else the step's bf16 entry `gemm` -- which is every call - with the mode off (all copies None).""" - rows = x.shape[0] - lo, hi = _FP8_GEMV_ROWS.get(kind, (1, 0)) - if packed is not None and lo <= rows <= hi and _gemv_ready(x): - return fp8_gemv_op(x, packed) - if mid is not None and _mid_rows(kind, rows) and (not mid.fp8 or _gemv_ready(x)): - return mid_rows_mm_op(x, mid) - if fp8 is not None and rows >= _DENSE_FP8_MIN_ROWS[kind]: - return fp8_per_tensor_mm(x, *fp8) - return gemm(x, weight_t) - - -#: `_mid` weight-name patterns -> the `_MID_MM_ROWS` kind their projections dispatch as. -_MID_KINDS = ( - (r"l\d+_(attn_in_w|in_proj_w)", "in"), - (r"l\d+_(o_w|out_w)", "out"), - (r"l\d+_sh_gu", "sh_gu"), - (r"l\d+_sh_dn", "sh_dn"), - (r"l\d+_router", "router"), -) - - -def _precompile_mid(mid: dict) -> None: - """Compile every `mid_rows_mm` variant the decode dispatch can reach now, once per weight - geometry: a variant first met by a serving step would stall it for the compile.""" - seen = {} - for name, mw in mid.items(): - kind = next(k for rx, k in _MID_KINDS if re.fullmatch(rx, name)) - lo, hi = _MID_MM_ROWS[kind] - key = (mw.n, mw.k, mw.fp8, None if mw.scale is None else mw.scale.numel()) - if lo <= hi and key not in seen: - seen[key] = mid_rows_precompile(mw, max(lo, 1), hi) - - -#: Decode batches up to this many rows run each trunk hyper-connection site as -#: `qwen4_exp_hc_decode_first_cute` + `qwen4_exp_hc_decode_mix_cute` (or its NVFP4 -#: variant at an MoE site; two GEMV launches); larger ones as `qwen4_exp_hc_inject_norm` -#: and the mixer. A dispatch threshold from the GB300 micro-bench (the GEMVs win at 1 -#: row and lose from 2), at most the entries' MAX_ROWS. -_HC_DECODE_GEMV_MAX_ROWS = 1 # capacity: fixed: dispatch threshold, not a buffer bound - -#: Decode row counts up to which a hyper-connection mixer (so from -#: `_HC_DECODE_GEMV_MAX_ROWS + 1` rows) runs its first half (projection and -#: activations) as the one tcgen05 launch `qwen4_exp_hc_decode_tc_first` instead of -#: cuBLAS split-K (+ its reduce) + `qwen4_exp_hc_lowrank_act`: the entry's MAX_ROWS -#: (two 128-row MMA tiles); larger batches keep cuBLAS. GB300 micro-bench (96 sites, -#: cold weights), us a site for the whole site at 256 rows: 17.8 against 22.4. -_HC_TC_FIRST_MAX_ROWS = ( - 256 # capacity: fixed: dispatch threshold (two MMA tiles), not a buffer bound -) - -#: Prefill steps of at least this many rows run each trunk hyper-connection site on the -#: un-normalized residual: `qwen4_exp_hc_inject_rstd` stores `hc_out` and a `[T, C]` rstd -#: instead of the `[T, C*H]` normed stream, and the mixer -#: (`qwen4_exp_hc_first_rescale_tc` + `qwen4_exp_hc_up_gate_mean_rstd(_fp4)`) forms -#: `normed` on chip. Smaller steps keep `qwen4_exp_hc_inject_norm` and the stored-normed -#: mixer. GB300 micro-bench, us a site (stored normed vs rstd): 2048 rows 57.0 vs 56.7, -#: 3072 rows 82.4 vs 76.1, 4096 rows 104.0 vs 90.8, 8192 rows 186.2 vs 164.9, 16384 rows -#: 360.1 vs 293.7. -_HC_RSTD_MIN_ROWS = 3072 # capacity: fixed: a dispatch threshold, not a buffer size - -#: FP8 block-scale block edge, on both weight-scale axes and the activation row. -_FP8_BLOCK = _weights.FP8_BLOCK -#: Values per e4m3 block scale in an NVFP4 KV pool. -_NVFP4_BLOCK = 16 - -#: False: plain `amax / 448` scale (the runner's recipe). True: power-of-two -#: (UE8M0), a different quantization of the same input. -_MTP_MOE_USE_UE8M0 = False - -#: Set to 1 to run the draft layer's routed experts in NVFP4 on the trunk's trtllm-gen runner: -#: at load the checkpoint's FP8 128x128 block-scaled experts are requantized (per-expert weight -#: global scales, 16-element e4m3 block scales) into the trunk's layout and the FP8 copies are -#: freed. The draft only proposes tokens -- the verify step decides what is emitted -- so this -#: can change the acceptance length, not outputs. 0 (default): the FP8 block-scale runner. -_MTP_MOE_NVFP4_ENV = "TRTLLM_MODELING_V2_MTP_MOE_NVFP4" - -#: The NVFP4 draft MoE's two activation global scales are static, `448 * 6 / A` with `A` this -#: factor times the largest calibrated amax of the trunk's routed experts (its `input_scale`s, -#: FC1 input and FC2 input respectively; the draft layer carries none). A global scale only has -#: to keep the e4m3 block scales inside their range, whose 3-bit precision is the same from -#: 2^-6 to 448: a block whose amax exceeds `A` is clipped, one below `A / 28672` loses scale -#: precision, so the margin costs nothing until blocks are ~1e-3 of `A`. -_MTP_MOE_NVFP4_FC1_MARGIN = 8.0 -_MTP_MOE_NVFP4_FC2_MARGIN = 4.0 - - -def _mtp_moe_nvfp4_enabled() -> bool: - mode = os.environ.get(_MTP_MOE_NVFP4_ENV, "0") - assert mode in ("0", "1"), f"{_MTP_MOE_NVFP4_ENV}={mode!r}: expected 0 or 1" - return mode == "1" - - -#: The history depth of the GDN cached-replay verify path (the manager's -#: `replay_history_size`). The replay entry is certified at 16 only, and a manager that -#: reports another depth is refused at bind. -_GDN_REPLAY_HISTORY = 16 - -#: The most query rows a verify step carries per generation request, `max_draft_len + 1`. -#: It reaches `thop_attention` as `predicted_tokens_per_seq` and the cached replay as its step -#: width `T`; the replay kernel asserts `T <= 8`, and above it the engine would verify with the -#: legacy per-step states, which this target refuses past `_MAX_LEGACY_VERIFY_WIDTH` rows. -_MAX_VERIFY_WIDTH = 8 - -#: The same bound on the legacy verify path (`TRTLLM_USE_GDN_REPLAY=0`): the entries only it -#: calls (`pack_gdn_decode_qkv`, `fused_gdn_gating`, the multi-token -#: `fused_recurrent_gated_delta_rule_update`) are certified at up to 4 rows per request, and its -#: per-step recurrent buffer holds one state per verify row, request and layer (6.75 GiB at 8 -#: rows and 8 requests in fp32). -_MAX_LEGACY_VERIFY_WIDTH = 4 - -#: The recurrent-state pool dtypes this target serves (`kv_cache_config.mamba_ssm_cache_dtype`): -#: the checkpoint's fp32, or bf16, which rounds the state to bf16 once per decode token, once -#: per prefill call and at every cached-replay checkpoint commit. The engine also accepts fp16. -_SSM_STATE_DTYPES = (torch.float32, torch.bfloat16) - - -# ----------------------------------------------- KV cache manager binding --- -# -# The ops are plain functions over raw tensors, and the pool views taken from the -# KV cache manager are `TensorWrapper`s over raw device addresses that do not keep -# memory alive, so a view is only valid for the manager that produced it. `_Bound` -# is everything derived from one manager; the core replaces it as a unit -# (`_bind_cache_manager`) and the hot path pays one identity comparison. The MTP -# draft layer's step can see a second manager (the separate draft KV cache), bound -# the same way into `_DraftBound` (`_bind_draft_cache_manager`). -# -# The `check_*` functions hold the shape/dtype contract. They read only -# `.shape` / `.dtype`, so unit tests run them on meta tensors with no engine. -# -# For a later target with pools that are not in `attn_metadata` (mamba/GDN state, -# side caches): the reusable parts are (1) one immutable bound object per manager, -# replaced atomically and keyed on manager identity, (2) pure contract checks, -# (3) refusing a bind during CUDA graph capture that no eager step validated first -# (a re-bind for the trunk's, a first bind for the draft layer's). Once a second -# target needs it, the next step is a declarative -# `PoolSpec(name, layers, getter, tail_shape, dtype)` list consumed by one shared -# binder under `modeling_v2/` (not `catalog/`), leaving each target only its specs -# and shape constants. Pure-attention targets that read the pool pointers from -# `attn_metadata` on every call (deepseek_v3) do not need this. - - -class _PoolGeometry(NamedTuple): - """The pool shapes this target expects, taken from its config and the metadata.""" - - tokens_per_block: int - heads_kv: int - head_dim: int - index_kv_heads: int - index_head_dim: int - conv_dim: int - conv_kernel: int - num_k_heads: int - num_v_heads: int - head_v_dim: int - head_k_dim: int - stream: int - ple_state_len: int - ngram_context_len: int - dtype: torch.dtype # the activation dtype - kv_dtype: torch.dtype # the paged K/V pool's dtype: bf16, e4m3 (fp8) or uint8 (NVFP4) - kv_width: int # its innermost extent: head_dim, or head_dim // 2 for NVFP4's e2m1 pairs - - -class _KvPlanes(NamedTuple): - """One layer's paged K/V views, and under an NVFP4 pool its e4m3 block-scale views.""" - - k: torch.Tensor - v: torch.Tensor - k_scale: torch.Tensor | None = None - v_scale: torch.Tensor | None = None - - -class _Replay(NamedTuple): - """The GDN cached-replay verify state of one manager: per slot, the recurrent pool - row is a checkpoint and a compact history holds the tokens accepted since it.""" - - old: dict # GDN layer -> (old_x, old_B, old_dt, old_dA_cumsum) history buffers - cache_buf_idx: torch.Tensor # per slot: which history half is pending - prev_num_accepted: torch.Tensor # per slot: pending history entries - history: int - width: int # verify rows per generation request - all_layer_commit: bool # the manager commits partitioned steps over every layer - cute: bool # bf16 pool: the CuTe replay kernel (the Triton kernel's commit contract) - - -class _Bound(NamedTuple): - """Everything derived from one KV cache manager.""" - - cm: object - pool_idx: dict # attention layer -> pool offset (what thop_attention addresses) - index_cache: dict # attention layer -> QSA index plane - position_cache: torch.Tensor # shared QSA position plane - kv: dict # attention layer -> `_KvPlanes` - conv: dict # GDN layer -> conv state pool - ssm: dict # GDN layer -> recurrent state pool - ple_state: tuple # (PLE conv state, PLE n-gram id history) - # MTP only: the verify step's per-step conv buffers, and one recurrent mode. Legacy: - # per-step recurrent buffers (`int_ssm`). Cached replay: `replay`, no `int_ssm`. - int_idx: torch.Tensor | None = None - int_conv: dict | None = None - int_ssm: dict | None = None - replay: _Replay | None = None - - -class _DraftBound(NamedTuple): - """The MTP draft layer's views of the manager its draft step sees: the separate - draft manager when the engine built one, else the trunk's.""" - - cm: object - index_cache: torch.Tensor # the draft layer's QSA index plane - position_cache: torch.Tensor # that manager's QSA position plane - kv: _KvPlanes # the draft layer's paged K/V views - share_slot: int | None # its row in the metadata's shared selection buffers (IndexShare) - - -def _require(cm, names, what: str) -> None: - for name in names: - assert hasattr(cm, name), ( - f"the KV cache manager {type(cm).__name__} does not expose {name}; this target {what}" - ) - - -def _check_pool(t, what, lead, tail, dtype=None) -> None: - """`t` has `lead` leading dims of any size, then exactly `tail`, at `dtype`.""" - assert t is not None, f"{what}: the runtime did not allocate it" - expected = ["*"] * lead + [str(n) for n in tail] - assert t.dim() == lead + len(tail) and tuple(t.shape[lead:]) == tuple(tail), ( - f"{what} is {tuple(t.shape)}, expected [{', '.join(expected)}]" - ) - assert dtype is None or t.dtype == dtype, f"{what} is {t.dtype}, expected {dtype}" - - -def check_position_plane(g: _PoolGeometry, position) -> None: - shape = (g.tokens_per_block, _QSA_POSITION_AXES) - _check_pool(position, "QSA position plane", 1, shape, torch.int32) - - -def check_attention_pools(g: _PoolGeometry, layer: int, index_cache, pool) -> None: - """QSA index plane and paged K/V for one attention layer.""" - index_shape = (g.tokens_per_block, g.index_kv_heads, g.index_head_dim) - _check_pool(index_cache, f"layer {layer}: QSA index plane", 1, index_shape, g.dtype) - # The sparse path needs adjacent HND K/V role planes: [pages, 2, H, block, D]. - kv_shape = (g.heads_kv, g.tokens_per_block, g.kv_width) - _check_pool(pool, f"layer {layer}: paged K/V", 2, kv_shape, g.kv_dtype) - assert pool.shape[1] == 2, f"layer {layer}: paged K/V has {pool.shape[1]} role planes, not 2" - - -def check_kv_scale_planes(g: _PoolGeometry, layer: int, pool, scales) -> None: - """An NVFP4 layer's e4m3 block scales: one uint8 per `_NVFP4_BLOCK` values, laid out like - the data pages (`[slots, 2, H, block, D / 16]`), and holding as many pages per slot as the - data pool does. - - The attention op addresses a layer's data pages and its scale pages with one block-offset - table, so a scale pool whose pages-per-slot differs from the data pool's reads and writes - other slots' scales without any error. Both counts are read off the views' slot strides.""" - scale_shape = (g.heads_kv, g.tokens_per_block, g.head_dim // _NVFP4_BLOCK) - _check_pool(scales, f"layer {layer}: NVFP4 block scales", 2, scale_shape, torch.uint8) - assert scales.shape[1] == 2 and scales.shape[0] == pool.shape[0], ( - f"layer {layer}: the block scales are {tuple(scales.shape)} against paged K/V " - f"{tuple(pool.shape)}; they must hold the same slots, K then V" - ) - pages = [] - for what, t in (("paged K/V", pool), ("block scales", scales)): - page = t[0, 0].numel() - assert t.stride(0) % page == 0, ( - f"layer {layer}: the {what} slot stride {t.stride(0)} is not a whole number of " - f"{page}-element pages" - ) - pages.append(t.stride(0) // page) - assert pages[0] == pages[1], ( - f"layer {layer}: the paged K/V pool holds {pages[0]} pages per slot and the block-scale " - f"pool {pages[1]}; the attention op would address the scales of other slots" - ) - - -def check_recurrent_pools(g: _PoolGeometry, conv, ssm) -> None: - _check_pool(conv, "conv state pool", 1, (g.conv_dim, g.conv_kernel - 1), g.dtype) - # Key dim innermost. - _check_pool(ssm, "recurrent state pool", 1, (g.num_v_heads, g.head_v_dim, g.head_k_dim)) - assert ssm.dtype in _SSM_STATE_DTYPES, ( - f"the recurrent state pool is {ssm.dtype}; this target serves " - f"{' or '.join(str(d) for d in _SSM_STATE_DTYPES)} (kv_cache_config.mamba_ssm_cache_dtype)" - ) - - -def check_ple_pools(g: _PoolGeometry, ple_conv, ple_ngram, conv) -> None: - _check_pool(ple_conv, "PLE conv state", 1, (g.stream, g.ple_state_len), g.dtype) - # The id-history entry writes int64 into it; an int32 pool raises. - _check_pool(ple_ngram, "PLE n-gram context pool", 1, (g.ngram_context_len,), torch.int64) - # The forward addresses all three with one scheduler slot vector. - assert ple_conv.shape[0] == conv.shape[0] == ple_ngram.shape[0], ( - "the PLE and Gated-DeltaNet pools must share one slot space" - ) - - -def check_step_buffers(g: _PoolGeometry, int_conv, conv, int_idx) -> None: - """MTP verify buffers both recurrent modes use: the per-step conv windows - `[slots, steps, ...]` at the pool's dtype (the promotion copies between them without - converting) and the rows `update_mamba_states` reads its sources from.""" - assert int_conv is not None, ( - "the runtime allocated no per-step conv buffer; a partially accepted draft " - "cannot be rolled back" - ) - _check_pool(int_conv, "per-step conv buffer", 2, (g.conv_dim, g.conv_kernel - 1), conv.dtype) - assert int_idx is not None and int_idx.dtype == torch.int32, ( - "the cache manager exposes no int32 intermediate_state_indices" - ) - - -def check_legacy_state(g: _PoolGeometry, int_ssm, ssm) -> None: - """Legacy verify: per-step recurrent buffer `[slots, steps, ...]` at the pool's dtype.""" - assert int_ssm is not None, ( - "the cache manager runs the legacy verify path but allocated no per-step recurrent " - "buffer; a partially accepted draft cannot be rolled back" - ) - ssm_shape = (g.num_v_heads, g.head_v_dim, g.head_k_dim) - _check_pool(int_ssm, "per-step recurrent buffer", 2, ssm_shape, ssm.dtype) - assert int_ssm.shape[1] <= _MAX_LEGACY_VERIFY_WIDTH, ( - f"the legacy verify path runs {int_ssm.shape[1]} rows per generation request; its entries " - f"are certified at up to {_MAX_LEGACY_VERIFY_WIDTH} (max_draft_len " - f"{_MAX_LEGACY_VERIFY_WIDTH - 1}). Leave TRTLLM_USE_GDN_REPLAY unset (cached replay) for " - "a longer draft" - ) - - -def check_replay_state(history, width, prev_accepted, cache_buf_idx, ssm) -> None: - """Cached replay, manager-wide: the geometry the replay entry is certified at, and the - per-slot bookkeeping, which must cover every pool slot.""" - assert ssm.dtype in _SSM_STATE_DTYPES, ( - f"the recurrent pool is {ssm.dtype}; the replay entry is certified at an fp32 or a " - "bf16 checkpoint" - ) - assert history == _GDN_REPLAY_HISTORY, ( - f"replay_history_size {history}; the replay entry is certified at {_GDN_REPLAY_HISTORY}" - ) - assert 2 <= width <= min(_MAX_VERIFY_WIDTH, history), ( - f"replay_step_width {width}; the replay entry is certified at 2..{_MAX_VERIFY_WIDTH} " - f"verify rows and never above the history ({history})" - ) - for name, t in (("prev_num_accepted_tokens", prev_accepted), ("cache_buf_idx", cache_buf_idx)): - assert t.dim() == 1 and t.shape[0] >= ssm.shape[0] and t.dtype == torch.int32, ( - f"{name} is {tuple(t.shape)} {t.dtype}; the kernel reads one int32 per pool slot " - f"({ssm.shape[0]})" - ) - - -def check_replay_history( - g: _PoolGeometry, layer: int, history, pool_slots, old_x, old_B, old_dt -) -> None: - """One GDN layer's replay history: `[slots, 2 halves, ...]`, addressed by `history` and - the head geometry, never by the buffer's own shape.""" - for name, t in (("old_x", old_x), ("old_B", old_B), ("old_dt", old_dt)): - assert t is not None, f"layer {layer}: the manager's layer cache carries no {name}" - assert old_x.shape[0] >= pool_slots, ( - f"layer {layer}: the replay history has {old_x.shape[0]} slots for a {pool_slots}-slot pool" - ) - bf16, fp32 = torch.bfloat16, torch.float32 - x_tail, b_tail = (history, g.num_v_heads, g.head_v_dim), (history, g.num_k_heads, g.head_k_dim) - _check_pool(old_x, f"layer {layer}: old_x", 2, x_tail, bf16) - _check_pool(old_B, f"layer {layer}: old_B", 2, b_tail, bf16) - _check_pool(old_dt, f"layer {layer}: old_dt", 2, (g.num_v_heads, history), fp32) - assert old_x.shape[1] == old_B.shape[1] == old_dt.shape[1] == 2, ( - f"layer {layer}: the replay history is not double-buffered" - ) - - -def check_verify_width(rows: int) -> None: - """The rows per generation request one step carries: `max_draft_len + 1` (`runtime_draft_len - + 1` at run time), a `thop_attention` `predicted_tokens_per_seq` the entry is certified at.""" - assert 1 <= rows <= _MAX_VERIFY_WIDTH, ( - f"predicted_tokens_per_seq {rows} is outside the 1..{_MAX_VERIFY_WIDTH} range " - "attention/thop_attention is certified at" - ) - - -def check_replay_width(replay, rows: int) -> None: - """Under the cached replay a single-row generation step cannot take the ordinary - recurrence: the pool row is a checkpoint whose pending history only the replay kernel - applies. The manager's step width is the one row count per sequence a step may carry.""" - assert replay is None or rows == replay.width, ( - f"the batch carries {rows} query rows per generation sequence and the GDN cached " - f"replay runs at step width {replay.width}; a runtime_draft_len of {rows - 1} (a " - "draft-length schedule, or skipped drafting) cannot run under it. Export " - "TRTLLM_USE_GDN_REPLAY=0 for variable draft lengths" - ) - - -def check_shared_selection(shared, slots, layer: int, block_topk: int) -> None: - """IndexShare: the metadata's shared selection buffers have a row for `layer` and are - `block_topk` wide.""" - assert shared is not None, "IndexShare is on but the QSA metadata allocated no shared buffers" - assert layer in slots and slots[layer] < shared.shape[0], ( - f"the KV cache manager gives layer {layer} no row of the {shared.shape[0]}-row shared " - f"selection buffers (qsa_shared_topk_slots={slots})" - ) - assert shared.shape[2] == block_topk, ( - f"the shared selection rows are {shared.shape[2]} wide and this target selects " - f"{block_topk} blocks per row" - ) - - -def _moe_layer_scalars(w, i: int, backend: str) -> dict[str, torch.Tensor]: - """Layer `i`'s routed-expert scalars for `backend`, built from the checkpoint's - `input_scale` / `weight_scale_2` rows. - - The checkpoint stores modelopt reciprocals (`input_scale = 1 / g_act`, - `weight_scale_2 = 1 / g_w`), so a quantizer's global scale is `1 / input_scale` and - the runner's scalars are products of the two. - """ - isc1, isc1_up = w[f"l{i}_e_isc1"], w[f"l{i}_e_isc1_up"] - assert torch.equal(isc1, isc1_up), ( - f"layer {i}: expert gate/up input_scale differ; the fused FC1 " - "reads one activation quantization for both halves" - ) - assert torch.equal(w[f"l{i}_e_ws2_1"], w[f"l{i}_e_ws2_1_up"]), ( - f"layer {i}: expert gate/up weight_scale_2 differ; the fused " - "FC1 carries one alpha per expert for both halves" - ) - # One activation quantization feeds all 512 experts, so FC1's activation - # global scale must be one number (the shared expert is bf16, no scale). - assert torch.equal(isc1.min(), isc1.max()), ( - f"layer {i}: the routed experts do not share one " - "gate_proj.input_scale, so a single activation quantization " - "cannot feed them all and this assembly is wrong for this " - "checkpoint" - ) - e_isc2, ws2_2 = w[f"l{i}_e_isc2"], w[f"l{i}_e_ws2_2"] - gate1 = (isc1 * w[f"l{i}_e_ws2_1"]).contiguous() - # Routed FC1 activation global scale (one per layer). - scalars = dict(global_scale=(1.0 / isc1[:1]).contiguous()) - if backend == "CUTEDSL": - # The CuteDSL FC1 epilogue requantizes with one global scale for the whole layer, - # where trtllm-gen folds each expert's `down_proj.input_scale` into its FC1 - # scalar. Quantizing at `1 / max(input_scale)` (the smallest per-expert global - # scale, so no block scale saturates) and putting the same factor into FC2's alpha - # makes `global_sf * alpha2[e] == ws2_2[e]` for every expert, whatever the spread. - isc2_max = e_isc2.max() - scalars.update( - # [num_experts] each: activation scale times weight scale. - alpha1=gate1, - alpha2=(isc2_max * ws2_2).contiguous(), - # The FC2 activation's quantization scale, numel 1. - global_sf=(1.0 / isc2_max).reshape(1).contiguous(), - ) - else: - scalars.update( - # [num_experts] each. Swapping the first two is finite and silent. - output1_scale_scalar=(gate1 / e_isc2).contiguous(), - output1_scale_gate_scalar=gate1, - output2_scale_scalar=(e_isc2 * ws2_2).contiguous(), - ) - return scalars - - -class Qwen4ExpTp1ModelingV2(DecoderModel): - def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig): - """`composite_config` is passed in because the text sub-config adopted as - `model_config.pretrained_config` carries neither the vision half nor the - media token ids.""" - super().__init__(model_config) - cfg = model_config.pretrained_config - assert cfg is not None - - # This target IS this geometry — assert, never adapt. - assert torch.cuda.get_device_capability() == _SM, ( - f"target certified on sm_{_SM[0]}{_SM[1]}, running on " - f"sm_{''.join(map(str, torch.cuda.get_device_capability()))}" - ) - assert model_config.mapping.tp_size == 1, "tp1 target" - assert model_config.mapping.pp_size == 1, "tp1 target" - assert not model_config.mapping.enable_attention_dp, "tp1 target" - - # `cfg` is the engine-normalized Qwen4ExpTextConfig (pool geometry is read - # from it). - assert getattr(cfg, "model_type", None) == "qwen4_exp_text", ( - f"the core must be built on the normalized text config; got " - f"model_type {getattr(cfg, 'model_type', None)!r}" - ) - - dt = model_config.torch_dtype - assert dt == torch.bfloat16, f"bf16 target, engine resolved {dt}" - assert cfg.torch_dtype == torch.bfloat16, cfg.torch_dtype - assert not cfg.tie_word_embeddings, "untied lm_head" - assert not cfg.attention_bias, "no q/k/v/o bias anywhere" - assert cfg.hidden_act == "silu", "SwiGLU over a silu gate" - - self.num_layers = cfg.num_hidden_layers - self.hidden = cfg.hidden_size - self.vocab = cfg.vocab_size - self.eps = cfg.rms_norm_eps - - # `layer_types` is authoritative; `full_attention_interval` is the runtime - # helper's fallback, so both must agree. - layer_types = list(cfg.layer_types) - assert len(layer_types) == self.num_layers, (len(layer_types), self.num_layers) - assert set(layer_types) <= {"full_attention", "linear_attention"}, set(layer_types) - self.is_attn = [t == "full_attention" for t in layer_types] - interval = cfg.full_attention_interval - assert self.is_attn == [i % interval == interval - 1 for i in range(self.num_layers)], ( - "layer_types disagrees with full_attention_interval" - ) - self.attn_layers = [i for i, a in enumerate(self.is_attn) if a] - self.gdn_layers = [i for i, a in enumerate(self.is_attn) if not a] - assert self.attn_layers and self.gdn_layers, "this target is the hybrid assembly" - - # The checkpoint picks the speculative mode (`MTP_EAGLE_ONE_MODEL`, one module - # replayed `max_draft_len` times); the module is built only when the engine - # asked for speculation. - spec_config = getattr(model_config, "spec_config", None) - self.mtp_enabled = spec_config is not None - # Rows per request of the widest verify step (`max_draft_len + 1`); 1 without MTP. - self.max_verify_rows = spec_config.max_draft_len + 1 if self.mtp_enabled else 1 - declared = getattr(cfg, "mtp_num_hidden_layers", None) - assert declared in (None, _MTP_NUM_LAYERS), ( - f"this target implements the one-layer MTP module this checkpoint " - f"ships; the text config declares mtp_num_hidden_layers {declared}" - ) - # The module shares the trunk's embedding and the shell's lm_head. - assert not getattr(cfg, "mtp_use_dedicated_embeddings", False), ( - "this checkpoint's MTP module shares the decoder's embedding " - "(`mtp_use_dedicated_embeddings: false`) and ships no " - "`mtp.embed_tokens` tensor; a dedicated-embedding module is a " - "different assembly" - ) - # The nested `mtp` block, when kept: one `full_attention` layer at the trunk's - # rope base. - mtp_block = getattr(cfg, "mtp", None) - self.mtp_theta = None - if mtp_block is not None: - assert isinstance(mtp_block, dict), ( - f"the text config's `mtp` block is {type(mtp_block).__name__}; " - "this target reads it as the published mapping" - ) - assert mtp_block.get("num_hidden_layers", _MTP_NUM_LAYERS) == _MTP_NUM_LAYERS, ( - f"mtp.num_hidden_layers is {mtp_block.get('num_hidden_layers')}" - ) - assert list(mtp_block.get("layer_types", ["full_attention"])) == ["full_attention"], ( - f"the MTP module's layer_types is {mtp_block.get('layer_types')}; this " - "target implements the full-attention form, which is the only " - "one whose weights the checkpoint carries" - ) - # `null` means the trunk's own output. - assert mtp_block.get("mtp_use_hidden_state_from_layer") is None, ( - f"the MTP module asks for the hidden state of layer " - f"{mtp_block.get('mtp_use_hidden_state_from_layer')}; this " - "target hands it the end of the stack, which is what the " - "runtime passes and all this file can reach" - ) - if "rope_theta" in mtp_block: - self.mtp_theta = float(mtp_block["rope_theta"]) - - # ---- hyper-connections ------------------------------------------ - self.hc_count = cfg.hc_count - self.hc_lowrank = cfg.hc_lowrank - assert self.hc_count >= 2, ( - f"the hyper-connection residual carries hc_count streams; this " - f"config declares {self.hc_count}, which is not a multi-stream " - f"residual and is a different assembly" - ) - self.stream = self.hc_count * self.hidden - # Columns of the padded first-projection weight (`_fuse_hc_first_projections`) - # the cuBLAS first projection multiplies: the real `[down; block_inject]` rows - # rounded up to 32 (352 here). GB300, fp32 out, K 10240: 16384 rows 81.5 us - # against 91.5 at the padded 384, 8193 rows 41.6 against 51.2, never slower - # below; the real columns are bitwise the 384-wide product's at every row - # count measured (1-16384; at 320 columns they are not). - self.hc_first_cols = -(-(self.hc_lowrank + self.hc_count) // 32) * 32 - - # ---- full attention --------------------------------------------- - self.heads_q = cfg.num_attention_heads - self.heads_kv = cfg.num_key_value_heads - self.head_dim = cfg.head_dim - assert self.head_dim in (64, 128, 256), "fused_qk_norm_rope head_dim set" - assert self.heads_q % self.heads_kv == 0, "thop_attention GQA rule" - self.q_width = self.heads_q * self.head_dim - self.kv_width = self.heads_kv * self.head_dim - self.qkv_width = self.q_width + 2 * self.kv_width - # The per-head interleaved output gate makes q_proj twice as wide as q. No - # config flag exists, so it is derived from the weight (asserted in weights.py); - # the per-head split below would otherwise address the wrong columns. - assert not hasattr(cfg, "attn_output_gate"), ( - "this config gained an attn_output_gate flag; the gate is derived " - "from q_proj's width here, and the two must be reconciled" - ) - - # ---- RoPE -------------------------------------------------------- - # Partial, neox half-split pairs, no scaling; `rope_theta` is inside - # `rope_parameters` (mirrored as `rope_scaling`). - rope = getattr(cfg, "rope_parameters", None) or getattr(cfg, "rope_scaling", None) - assert isinstance(rope, dict), f"rope parameters must be a dict, got {rope!r}" - self.theta = float(rope["rope_theta"]) - assert self.theta > 0.0, "rope theta" - partial = float(rope.get("partial_rotary_factor", cfg.partial_rotary_factor)) - self.rotary_dim = int(self.head_dim * partial) - assert self.rotary_dim % 2 == 0 and 0 < self.rotary_dim <= self.head_dim, ( - f"rotary_dim {self.rotary_dim} must be even and at most head_dim " - f"{self.head_dim}; the op enforces only the evenness half" - ) - # The mRoPE declaration is what makes the engine hand this forward a - # [3, 1, T] position buffer. - self.mrope_declared = str(rope.get("type", rope.get("rope_type"))) == "mrope" - # `fused_qk_norm_rope` implements the interleaved mRoPE layout only (pair `j` - # reads coordinate 1 when `j % 3 == 1 and j < 3*s1`, 2 when `j % 3 == 2 and - # j < 3*s2`, else 0); the contiguous variant is a different operator. A text - # gate cannot see the split. - assert rope.get("mrope_interleaved") is True, ( - f"this target passes the interleaved mRoPE layout, which is the only " - f"one attention/fused_qk_norm_rope implements; the config declares " - f"mrope_interleaved={rope.get('mrope_interleaved')!r}" - ) - section = list(rope["mrope_section"]) - assert len(section) == _QSA_POSITION_AXES, ( - f"mrope_section must carry one entry per position coordinate; got {section}" - ) - rotary_pairs = self.rotary_dim // 2 - assert sum(section) == rotary_pairs, ( - f"mrope_section {section} sums to {sum(section)} but rotary_dim " - f"{self.rotary_dim} makes {rotary_pairs} rotation pairs; the split " - f"partitions pairs, not dimensions" - ) - self.mrope_section1 = int(section[1]) - self.mrope_section2 = int(section[2]) - # Row-1 section is load-bearing here; row-2 saturates (10, 11, 12 are - # bit-identical at 32 pairs). `use_gemma` with `use_mrope` is certified only - # at the geometry asserted below. - assert ( - self.heads_q, - self.heads_kv, - self.head_dim, - self.rotary_dim, - self.theta, - ) == (24, 2, 256, 64, 1e7), ( - f"the main attention rope runs use_gemma together with use_mrope, " - f"which attention/fused_qk_norm_rope certifies only at " - f"(24, 2, 256, 64, 1e7); this config gives " - f"{(self.heads_q, self.heads_kv, self.head_dim, self.rotary_dim, self.theta)}" - ) - # No YaRN or other scaling. - assert "factor" not in rope, ( - f"this target builds a plain (unscaled) rope; rope carries {rope!r}" - ) - self.max_positions = cfg.max_position_embeddings - - # ---- QSA sparse index -------------------------------------------- - self.index_heads = cfg.indexer_n_heads - self.index_kv_heads = cfg.indexer_kv_heads - self.index_head_dim = cfg.indexer_head_dim - self.token_topk = cfg.indexer_budget - self.compress_ratio = cfg.indexer_compress_ratio - assert self.index_kv_heads == _QSA_INDEX_KV_HEADS, ( - f"QSA requires exactly one index KV head; this config declares {self.index_kv_heads}" - ) - assert self.compress_ratio >= 2 and self.token_topk % self.compress_ratio == 0, ( - f"QSA selection geometry: token_topk {self.token_topk} must divide " - f"by compress_ratio {self.compress_ratio}, which must be at least 2" - ) - # Fused index kernels: both powers of two, index head twice the rotary width. - assert self.index_head_dim == self.rotary_dim * _QSA_INDEX_HEAD_TO_ROTARY_RATIO, ( - f"the fused QSA index kernels require index_head_dim " - f"({self.index_head_dim}) == 2 * rotary_dim ({self.rotary_dim})" - ) - assert self.index_head_dim & (self.index_head_dim - 1) == 0, "index_head_dim power of two" - assert self.compress_ratio & (self.compress_ratio - 1) == 0, "compress_ratio power of two" - self.block_topk = self.token_topk // self.compress_ratio - self.expanded_topk = self.token_topk + self.compress_ratio - 1 - # Dense-attention cutoff: QSA resolves max(token_topk, seq_len_threshold); - # checked against the engine's value below because the CUDA-graph runner - # splits its short and long decode families on it. - self.dense_threshold = self.token_topk - - # ---- the engine's sparse configuration, checked at construction ----- - # The index/position side caches, `qsa_*` metadata and sparse kernels all hang - # off `sparse_attention_config`; checking here fails at cold start rather than - # at the first forward. - sparse_cfg = getattr(model_config, "sparse_attention_config", None) - assert getattr(sparse_cfg, "algorithm", None) == "qsa", ( - f"this target is a QSA sparse-attention assembly and the engine " - f"resolved sparse_attention_config={sparse_cfg!r}. QSA is this " - f"architecture's own model default, so an engine that reaches here " - f"without it dropped the default on the way in rather than being " - f"asked for dense attention" - ) - # The engine's own lowering on this core's config: a mismatch with the side - # cache sizing would otherwise be a silent mis-index. - qsa_params = sparse_cfg.to_sparse_params(pretrained_config=cfg) - assert ( - qsa_params.index_n_heads, - qsa_params.index_kv_heads, - qsa_params.index_head_dim, - qsa_params.token_topk, - qsa_params.compress_ratio, - ) == ( - self.index_heads, - self.index_kv_heads, - self.index_head_dim, - self.token_topk, - self.compress_ratio, - ), ( - f"the engine lowered the QSA index geometry to {qsa_params!r} and " - f"this target derived heads {self.index_heads}, kv heads " - f"{self.index_kv_heads}, head_dim {self.index_head_dim}, token_topk " - f"{self.token_topk}, compress_ratio {self.compress_ratio}; the side " - f"caches are sized from the engine's copy and the kernels are " - f"called with this one" - ) - assert qsa_params.dense_seq_len_threshold == self.dense_threshold, ( - f"the engine's resolved QSA cutoff is " - f"{qsa_params.dense_seq_len_threshold} and this forward branches at " - f"{self.dense_threshold}. The runner captures its short decode " - f"graph family at the engine's number, so a disagreement freezes " - f"the dense body into the graph that long batches replay — which " - f"is finite, plausible output. A serving config that raises " - f"seq_len_threshold above token_topk is the way to get here, and " - f"it is a configuration this target does not implement" - ) - # IndexShare for the MTP draft loop (`sparse_attention_config. - # index_share_for_mtp_iteration`), as the engine resolved it: the same lowering - # sizes the metadata's shared selection buffers. - self.mtp_index_share = bool(qsa_params.mtp_index_share) - self.index_width = (self.index_heads + self.index_kv_heads) * self.index_head_dim - # Packed [q | k | v] scratch for the fused norm+rope call; the spare KV head - # is the v slot the op needs but never reads. - self.index_scratch_width = self.index_width + self.index_head_dim - # The fused input projection's output row (`weights._attn_in_w`): - # [q | k | v | gate | index q, k | v slot]. Its leading columns are the - # packed QKV the main norm+rope rewrites in place, and its tail is the - # index scratch's layout, so the row is `qkv_width + q_width + - # index_scratch_width` wide. The norm+rope addresses a row as - # `(Hq + Hk + Hv) * head_dim` and never touches the v heads, so the - # row must be a whole number of heads: everything after k is "v" to it. - self.attn_in_width = self.qkv_width + self.q_width + self.index_scratch_width - assert self.attn_in_width % self.head_dim == 0, ( - f"the fused projection row ({self.attn_in_width}) is not a whole number of " - f"{self.head_dim}-wide heads, so the main norm+rope cannot address it in place" - ) - self.attn_in_tail_heads = self.attn_in_width // self.head_dim - self.heads_q - self.heads_kv - self.softmax_scale = 1.0 / math.sqrt(self.head_dim) - - # ---- gated deltanet ---------------------------------------------- - self.num_k_heads = cfg.linear_num_key_heads - self.num_v_heads = cfg.linear_num_value_heads - self.head_k_dim = cfg.linear_key_head_dim - self.head_v_dim = cfg.linear_value_head_dim - self.conv_kernel = cfg.linear_conv_kernel_dim - assert self.num_v_heads % self.num_k_heads == 0, "GDN key-head broadcast" - assert self.head_k_dim == 128 and self.head_v_dim == 128, ( - f"the delta-rule entries certify head dims 128/128; this config " - f"declares {self.head_k_dim}/{self.head_v_dim}" - ) - assert self.conv_kernel in (2, 3, 4), "causal_conv1d width domain" - self.key_dim = self.num_k_heads * self.head_k_dim - self.value_dim = self.num_v_heads * self.head_v_dim - self.conv_dim = 2 * self.key_dim + self.value_dim - # Column widths of the one input-projection GEMM's output: [q|k|v | z | a | b]. - self.in_proj_split = (self.conv_dim, self.value_dim, self.num_v_heads, self.num_v_heads) - # Gate activation of the gated RMS norm (`silu` or `sigmoid`). - gate_type = getattr(cfg, "output_gate_type", None) - assert gate_type == "sigmoid", ( - f"gate activation {gate_type!r} is not the sigmoid form this " - f"target passes to rms_norm_gated_token_major" - ) - - # ---- MoE ---------------------------------------------------------- - self.num_experts = cfg.num_experts - self.topk = cfg.num_experts_per_tok - self.moe_inter = cfg.moe_intermediate_size - self.shared_inter = cfg.shared_expert_intermediate_size - assert cfg.norm_topk_prob, "renorm_moe_routing_op always renormalizes" - assert self.num_experts <= 512 and self.topk <= 16, ( - f"renorm_moe_routing_op takes at most 512 experts and top-16; this " - f"config declares {self.num_experts} / {self.topk}" - ) - # Every layer is MoE; the dense `intermediate_size` is dead config. - assert getattr(cfg, "decoder_sparse_step", 1) == 1, "every layer must be MoE" - assert not getattr(cfg, "mlp_only_layers", []), "every layer must be MoE" - # trtllm-gen NVFP4 runner operand-width rules and fp4_quantize's block rule; they - # also cover the CuteDSL FC1 / FC2 rules (K and I in 64-element scale groups, the - # gated N in 128-row pairs of 64 up rows and 64 gate rows). - assert self.hidden % 256 == 0, "fp4_block_scale_moe_runner hidden alignment" - assert (2 * self.moe_inter) % 128 == 0, "fp4_block_scale_moe_runner 2I alignment" - assert (self.hidden // _SF_VEC) % 4 == 0 and (self.moe_inter // _SF_VEC) % 4 == 0, ( - "the 128x4 block-scale swizzle needs both scale widths to be a multiple of 4" - ) - # flashinfer_silu_and_mul's precondition is on the half width. - assert (self.shared_inter * 2) % 16 == 0, "flashinfer_silu_and_mul half alignment" - # Which kernels run the trunk's routed experts; it also fixes the load-time layout - # of the expert stacks (`weights.py`), so it is decided before anything is declared. - requested = model_config.moe_backend_requested.upper() - assert requested in _TRUNK_MOE_BACKENDS, ( - f"moe_config.backend={model_config.moe_backend_requested!r} is not implemented by " - f"this target; the supported requests are {sorted(_TRUNK_MOE_BACKENDS)}" - ) - self.trunk_moe_backend = _TRUNK_MOE_BACKENDS[requested] - self.moe_prefill_copy = _moe_prefill_copy( - self.trunk_moe_backend, model_config.max_num_tokens - ) - self.max_num_tokens = model_config.max_num_tokens - - # ---- PLE ----------------------------------------------------------- - ple_ids = list(getattr(cfg, "ple_layer_ids", None) or []) - assert len(ple_ids) == 1, ( - f"this target implements the single PLE decoder layer this " - f"checkpoint ships; the config declares ple_layer_ids={ple_ids}" - ) - # One-based in the config. - self.ple_layer = ple_ids[0] - 1 - assert 0 <= self.ple_layer < self.num_layers, self.ple_layer - assert not self.is_attn[self.ple_layer], ( - f"the PLE lifecycle state rides the recurrent pools, so its layer " - f"must be a linear-attention layer; layer {self.ple_layer} is not" - ) - self.ngram_size = cfg.ngram_size - self.heads_per_ngram = cfg.heads_per_ngram - self.ngram_heads = self.heads_per_ngram * (self.ngram_size - 1) - self.ple_embed_dim = cfg.ple_embed_dim - self.ple_head_dim = self.ple_embed_dim // self.ngram_heads - assert self.ple_head_dim * self.ngram_heads == self.ple_embed_dim, ( - f"{self.ngram_heads} n-gram heads must tile ple_embed_dim {self.ple_embed_dim}" - ) - self.ple_conv_kernel = cfg.ple_conv_kernel_size - # Length of the runtime's paged conv state. - self.ple_dilation = self.ngram_size - self.ple_state_len = (self.ple_conv_kernel - 1) * self.ple_dilation - self.ngram_context_len = self.ngram_size - 1 - # Shard count and padding rule are config; the table height is not, so the - # table is declared at load time (`weights.py`). - self.ple_shards = cfg.split_ngram_parts - self.ngram_vocab_divisor = cfg.make_ngram_vocab_size_divisible_by - assert self.ple_shards > 0 and self.ngram_vocab_divisor > 0, ( - self.ple_shards, - self.ngram_vocab_divisor, - ) - # Fresh-sequence seed for the n-gram id history and the hash's segment token - # (element 0 when a list, as in the HF reference). - eos = cfg.eos_token_id - self.eos_token_id = int(eos[0] if isinstance(eos, (list, tuple)) else eos) - - # ---- the vision tower ------------------------------------------------ - # A 27-block ViT (no KV cache, deepstack or quantization) feeding a 2x2 patch - # merger. Constants the config omits: `_VISION_ROPE_THETA`, - # `_VISION_NORM_EPS`, and the merger's erf GELU. - vcfg = getattr(composite_config, "vision_config", None) - assert vcfg is not None, ( - "this target implements the composite export's vision tower; the " - "composite config carries no vision_config" - ) - assert getattr(vcfg, "model_type", None) == "qwen4_exp_vision", ( - f"vision_config.model_type is {getattr(vcfg, 'model_type', None)!r}; " - "this target is the Qwen4-Exp vision tower" - ) - self.v_depth = vcfg.depth - self.v_hidden = vcfg.hidden_size - self.v_heads = vcfg.num_heads - self.v_inter = vcfg.intermediate_size - self.v_merge = vcfg.spatial_merge_size - self.v_patch = vcfg.patch_size - self.v_tpatch = vcfg.temporal_patch_size - self.v_channels = vcfg.in_channels - self.v_pos_rows = vcfg.num_position_embeddings - assert self.v_hidden % self.v_heads == 0, (self.v_hidden, self.v_heads) - self.v_head_dim = self.v_hidden // self.v_heads - # vision_rope_2d needs a head divisible by 4 (two axes, one rotation pair - # each). - assert self.v_head_dim in _VISION_ROPE_HEAD_DIMS, ( - f"ops/attention/vision_rope_2d is certified at head widths " - f"{_VISION_ROPE_HEAD_DIMS}; this tower's is {self.v_head_dim}" - ) - assert (self.v_heads, self.v_heads, self.v_head_dim) in _VISION_CERTIFIED_HEADS, ( - f"the tower's attention is the encoder-only (no KV cache) mode of " - f"attention/thop_attention, certified at {sorted(_VISION_CERTIFIED_HEADS)}; " - f"this tower is {(self.v_heads, self.v_heads, self.v_head_dim)}" - ) - # `hidden_act` names the block MLP's activation only; the merger runs the erf - # `nn.GELU()`. Applying it to both is wrong by under one bf16 ulp, so there - # are two entries (`gelu_tanh`, `gelu_erf`). - assert vcfg.hidden_act == "gelu_pytorch_tanh", ( - f"the vision blocks' MLP activation is {vcfg.hidden_act!r}; this " - "target calls ops/activation/gelu_tanh for it" - ) - # No deepstack (`routing.py` refuses it; this is the second lock). - assert not list(getattr(vcfg, "deepstack_visual_indexes", []) or []), ( - f"deepstack_visual_indexes={list(vcfg.deepstack_visual_indexes)} asks " - "for per-layer visual feature injection, which this target does not " - "implement" - ) - assert vcfg.out_hidden_size == self.hidden, ( - f"the merger projects into the language model's hidden width; " - f"out_hidden_size {vcfg.out_hidden_size} against hidden {self.hidden}" - ) - # pos_embed_interpolate asserts a square table; fail before the first image. - side = int(math.isqrt(self.v_pos_rows)) - assert side * side == self.v_pos_rows, ( - f"num_position_embeddings {self.v_pos_rows} is not a perfect square, " - "so it is not a square position grid" - ) - self.v_grid_side = side - # `spatial_merge_size ** 2` patches become one token at the merger. - self.v_merge_dim = self.v_hidden * self.v_merge * self.v_merge - # One flat pixel row per patch, `(in_channels, kt, kh, kw)` channel-slowest. - self.v_pixel_width = self.v_channels * self.v_tpatch * self.v_patch * self.v_patch - # Placeholder token ids live on the composite config's top level. - self.image_token_id = int(composite_config.image_token_id) - self.video_token_id = int(composite_config.video_token_id) - assert self.image_token_id < self.vocab and self.video_token_id < self.vocab, ( - f"the media placeholders ({self.image_token_id}, {self.video_token_id}) " - f"are in-vocabulary against vocab_size {self.vocab}, so the engine " - "cannot find them with its out-of-vocabulary fallback and this target " - "must publish them as `multimodal_token_ids`" - ) - - # ---- KV pool storage ------------------------------------------------ - # `None` keeps the checkpoint's bf16 K/V (`quant_mode` 0); `FP8`, what - # `kv_cache_config.dtype: fp8` resolves to, stores them as e4m3 under - # QM_FP8_KV_CACHE; `NVFP4` as e2m1 pairs with e4m3 block scales under - # QM_NVFP4_KV_CACHE. Anything else is refused: the QSA backend serves its sparse path - # over half, bf16, fp8 and (with its own kernels here) NVFP4 pools, and silently falls - # back to dense otherwise. The index and position planes are kernel contracts, not KV - # storage, and keep their types under all three. - kv_algo = model_config.quant_config.kv_cache_quant_algo - assert kv_algo in (None, QuantAlgo.FP8, QuantAlgo.NVFP4), ( - f"this target stores K/V as the checkpoint's bf16 (kv_cache_quant_algo None), " - f"as fp8-e4m3 (FP8) or as NVFP4; the engine resolved kv_cache_quant_algo {kv_algo!r}" - ) - self.quant_mode = { - None: 0, - QuantAlgo.FP8: QM_FP8_KV_CACHE, - QuantAlgo.NVFP4: QM_NVFP4_KV_CACHE, - }[kv_algo] - # What `check_attention_pools` holds every manager's K/V pool to, the draft layer's - # included: the dtype and the innermost extent (an NVFP4 byte holds two values). - self.kv_pool_dtype = { - None: dt, - QuantAlgo.FP8: torch.float8_e4m3fn, - QuantAlgo.NVFP4: torch.uint8, - }[kv_algo] - self.kv_pool_width = self.head_dim // 2 if kv_algo == QuantAlgo.NVFP4 else self.head_dim - # The weight table reads the tower's geometry off this. - self.vision_config = vcfg - - self.mtp_layer_idx = self.num_layers - if self.mtp_theta is None: - self.mtp_theta = self.theta - assert self.mtp_theta > 0.0, self.mtp_theta - if self.mtp_enabled: - mode = spec_config.spec_dec_mode - assert mode.is_mtp_eagle_one_model(), ( - f"this target implements the one-model MTP-Eagle draft loop, " - f"where one module is replayed; the engine resolved " - f"spec_dec_mode {mode!r}" - ) - resolved = getattr(spec_config, "num_nextn_predict_layers", None) - assert resolved == _MTP_NUM_LAYERS, ( - f"the engine resolved num_nextn_predict_layers={resolved} from " - f"the checkpoint; this target carries the weights of " - f"{_MTP_NUM_LAYERS} MTP module" - ) - # `max_draft_len + 1` becomes `predicted_tokens_per_seq` on every attention call - # and the cached replay's step width, certified at up to `_MAX_VERIFY_WIDTH`. - max_draft = spec_config.max_draft_len - assert 1 <= max_draft <= _MAX_VERIFY_WIDTH - 1, ( - f"max_draft_len={max_draft} needs {max_draft + 1} verify rows per request, and " - f"attention/thop_attention and the GDN cached replay are certified at up to " - f"{_MAX_VERIFY_WIDTH} (max_draft_len {_MAX_VERIFY_WIDTH - 1}); 8 is also the " - "replay kernel's ceiling" - ) - assert self.trunk_moe_backend != "CUTEDSL" or max_draft <= _CUTEDSL_MAX_DRAFT_LEN, ( - f"max_draft_len={max_draft} with moe_config.backend=CUTEDSL: the CuteDSL expert " - f"pipeline is certified at max_draft_len up to {_CUTEDSL_MAX_DRAFT_LEN}" - ) - # A linear tree keeps the Blackwell spec-dec mask surface inert. - assert not getattr(spec_config, "use_dynamic_tree", False), ( - "dynamic-tree drafting is not implemented by this target" - ) - assert getattr(spec_config, "eagle_choices", None) is None, ( - "static-tree drafting is not implemented by this target" - ) - assert spec_config.max_total_draft_tokens == max_draft, ( - f"a linear tree mirrors max_total_draft_tokens from " - f"max_draft_len; the engine resolved " - f"{spec_config.max_total_draft_tokens} against {max_draft}" - ) - - # The recurrent half of the verify step has two forms, and the cache manager - # picks one when it is built (`TRTLLM_USE_GDN_REPLAY`, default on): the cached - # replay or the legacy per-step recurrent buffers. Both are implemented; - # `_bind_step_buffers` reads which one a manager allocated. - # The FP8 block-scale runner's own shape rules. - assert self.hidden % _FP8_BLOCK == 0 and self.moe_inter % _FP8_BLOCK == 0, ( - f"fp8_block_scale_moe_runner needs hidden ({self.hidden}) and " - f"intermediate ({self.moe_inter}) to be multiples of {_FP8_BLOCK}" - ) - assert self.topk < self.num_experts and self.topk <= 32, ( - f"the no-groups renormalize routing kernel takes 0 < top_k <= 32 " - f"and top_k < num_experts; this config declares {self.topk} of " - f"{self.num_experts}" - ) - - # The dtype the table is declared at in weights.py. - self.ple_table_dtype = torch.float8_e4m3fn - # Resolved at construction; read by `weights._declare_ngram_table`. - self.ple_host_offload = _ple_host_offload_enabled() - # CPU allocation behind the offloaded table. The `ple_table` parameter is a - # device-address view that owns none of those pages; dropping this reference - # frees memory a kernel still points into. None when resident. - self._ple_table_host = None - - # Every weight but the n-gram table; meta-init intercepts torch.empty. - self.w = _weights.MODEL_WEIGHTS.declare(_weights.MODEL_WEIGHTS.dims(self)) - - # Derived once after load (`post_load_weights`). - self._moe_scalars: list | None = None - self._moe_scalars_cd: list | None = None - self._gdn_gating: dict | None = None - self._cos_sin: torch.Tensor | None = None - # The attention op's KV scales: fp32 [1] ones over an fp8 pool (never None there: the - # generation kernels read `kvCacheScale[0]`), fp32 [3] ones (q, K, V) over an NVFP4 - # one, None over bf16. Unit, because the fp8 QSA store and sparse GQA take no scale - # while a request's history is written by the dense and the sparse path alike, so the - # two agree on what a slot holds only at 1.0 (the NVFP4 entries take the [3] tensor - # and agree at any value, but the checkpoint has no KV scales to calibrate with). - # Lossy: fp8 K/V are stored by a plain cast, so a magnitude past the e4m3 maximum (448) - # saturates (`_KvRangeProbe` reports it under TRTLLM_MODELING_V2_VALIDATE). - self.kv_scale_orig_quant: torch.Tensor | None = None - self.kv_scale_quant_orig: torch.Tensor | None = None - # MTP module operands; None with speculation off. - self._mtp: _MtpWeights | None = None - # The draft experts' NVFP4 copy (`_mtp_moe_nvfp4`), built once: it frees the FP8 ones. - self._mtp_moe_nvfp4: dict | None = None - self._mtp_cos_sin: torch.Tensor | None = None - # The verify step's side-cache rollback, registered with the spec worker; None - # with speculation off. - self.spec_rollback: SpeculativeRollback | None = ( - SpeculativeRollback() if self.mtp_enabled else None - ) - self._prefill: "PrefillTarget | None" = None - self._decode: "DecodeTarget | None" = None - self._mtp_fused: "_MtpFused | None" = None - # The draft head's FP8 copy of lm_head: (e4m3 [vocab, hidden], fp32 [1, vocab] row - # scales). - self._mtp_head_fp8: "tuple[torch.Tensor, torch.Tensor] | None" = None - # A packed block-scaled e4m3 copy of lm_head for decode rows (`fp8_gemv`): the draft - # head's, and the target head's under the dense FP8 mode. - self._lm_head_g = None - # The draft head's NVFP4 copy (`_MTP_HEAD_NVFP4_ENV`), read by `nvfp4_gemv` at decode - # rows in place of the packed FP8 copy. - self._mtp_head_nv = None - # The shared expert's stream; see `_moe_block`. Made in `post_load_weights`: - # a stream created inside a graph capture is not capturable. - self._side_stream: torch.cuda.Stream | None = None - self._moe_memset_stream: torch.cuda.Stream | None = None - self.vision_tower: VisionTower | None = None - # Everything derived from the KV cache manager (`_bind_cache_manager`); - # None until the first forward. - self._bound: _Bound | None = None - self._binds = 0 - # Set by `begin_step`, cleared by `forward`: a forward without a prepared step is a bug. - self._step_begun = False - # The draft layer's own bind, keyed on the manager its draft step sees. - self._draft_bound: _DraftBound | None = None - - # ---------------------------------------------------------------- load --- - - def _norm_slices(self, key: str) -> tuple: - """Split a grouped-norm weight into its per-group slices (taken once).""" - return tuple(torch.split(self.w[key], self.hidden, dim=0)) - - def _fuse_hc_first_projections(self) -> None: - """Concatenate each trunk hyper-connection's `down` and `block_inject` - weights into one `[rows, stream]` tensor, zero-padded to a multiple of - 128 rows, so a site's first projection is one GEMM - (`_hc_first[(layer, tag)]`; `_hc_first["fhc"]` is the stack-final - mixer's `down`, padded the same way). cuBLAS was measured 1.8x (decode) - to 4.5x (8192 rows) faster at N = 384 than at the unpadded 324; its calls - multiply the first `hc_first_cols` (352) columns, faster still. - - The checkpoint's parameters are rebound as row views of the padded - tensors, so this costs only the padding rows (~1.2 MB per site), and - anything still reading them separately reads the same values.""" - w = self.w - - def fuse(*params): - rows = sum(t.shape[0] for t in params) - padded = -(-rows // 128) * 128 - pad = torch.zeros( - [padded - rows, self.stream], dtype=params[0].dtype, device=params[0].device - ) - both = torch.cat([t.data for t in params] + [pad], 0) - start = 0 - for t in params: - n = t.shape[0] - t.data = both[start : start + n] - start += n - return both - - self._hc_first = { - (i, tag): fuse(w[f"l{i}_{tag}_down"], w[f"l{i}_{tag}_inject"]) - for i in range(self.num_layers) - for tag in ("ahc", "mhc") - } - self._hc_first["fhc"] = fuse(w["fhc_down"]) - if self.mtp_enabled: - p = f"l{self.mtp_layer_idx}_mtp_" - self._hc_first["mtp_ahc"] = fuse(w[f"{p}ahc_down"], w[f"{p}ahc_inject"]) - self._hc_first["mtp_mhc"] = fuse(w[f"{p}mhc_down"], w[f"{p}mhc_inject"]) - self._hc_first["mtp_fhc"] = fuse(w[f"{p}fhc_down"]) - - def post_load_weights(self) -> None: - """Derive what the targets read, then build them, now that the weights are real. - - Not an override (DecoderModel has no such hook); named for the shell hook.""" - w = self.w - device = w["embed"].device - moe_scalars = [ - _moe_layer_scalars(w, i, self.trunk_moe_backend) for i in range(self.num_layers) - ] - # With the second expert copy (`moe_prefill_copy`), its CuteDSL scalars, which - # prefill-sized launches run on (`_CUTE_DSL_MOE_MIN_ROWS`). - moe_scalars_cd = ( - [_moe_layer_scalars(w, i, "CUTEDSL") for i in range(self.num_layers)] - if self.moe_prefill_copy - else None - ) - self._moe_scalars = moe_scalars - self._moe_scalars_cd = moe_scalars_cd - # Both delta-rule entries index A_log and dt_bias as `ptr + head` and certify - # fp32; the checkpoint is bf16. - self._gdn_gating = { - i: ( - w[f"l{i}_A_log"].float().contiguous(), - w[f"l{i}_dt_bias"].float().contiguous(), - ) - for i in self.gdn_layers - } - # Reconcile the table height against the checkpoint's id-space tensors: an - # index past the end is a device-side assert that poisons the process. - table_rows = w["ple_table"].shape[0] - used_rows = int(w["ple_offs"][-1]) + int(w["ple_vocab"][-1]) - assert used_rows <= table_rows, ( - f"the n-gram table has {table_rows} rows but its per-head offsets " - f"and vocabulary sizes address {used_rows}" - ) - assert table_rows % self.ngram_vocab_divisor == 0, ( - f"the n-gram table's {table_rows} rows are not a multiple of " - f"make_ngram_vocab_size_divisible_by={self.ngram_vocab_divisor}" - ) - assert table_rows - used_rows < self.ngram_vocab_divisor, ( - f"the n-gram table carries {table_rows - used_rows} padding rows, " - f"more than one divisor of {self.ngram_vocab_divisor}: the shards " - f"do not concatenate to the padded id space this checkpoint " - f"declares" - ) - self._cos_sin = self._index_rope_table(self.theta, device) - if self.mtp_enabled: - # The draft layer's index plane rotates at the module's own rope base, like its - # attention rope. - self._mtp_cos_sin = ( - self._cos_sin - if self.mtp_theta == self.theta - else self._index_rope_table(self.mtp_theta, device) - ) - if self.quant_mode in (QM_FP8_KV_CACHE, QM_NVFP4_KV_CACHE): - # Built here, not at construction: that runs under meta init. - scales = 3 if self.quant_mode == QM_NVFP4_KV_CACHE else 1 - self.kv_scale_quant_orig = torch.ones(scales, dtype=torch.float32, device=device) - self.kv_scale_orig_quant = 1.0 / self.kv_scale_quant_orig - self._report_ple_residency() - if self.mtp_enabled: - self._mtp = _mtp_views(self) - if _mtp_moe_nvfp4_enabled() and device.type == "cuda": - if self._mtp_moe_nvfp4 is None: - self._mtp_moe_nvfp4 = _mtp_moe_nvfp4(self) - self._mtp = self._mtp._replace(**self._mtp_moe_nvfp4) - - self._fuse_hc_first_projections() - if device.type == "cuda": - # Every first-half tile `DecodeTarget` reaches up to `_HC_TC_FIRST_MAX_ROWS` rows. - hc_tc_first_precompile(self.stream, self.hc_lowrank, self.hc_count, device) - if self.mtp_enabled: - self._mtp_fused = _mtp_fused_views(self) - # FP8 copies of the trunk's dense projections (`_DENSE_FP8_ENV`), beside the bf16 - # weights the 1-row GEMV / fused paths and the draft layer keep reading. - self._dense8 = {} - # ... and packed FP8 copies with 128x128 block scales for the decode-row `fp8_gemv` - # (the shared expert's gate|up without its gate-logit row, which stays bf16). - self._gemv8 = {} - # ... and `mid_rows_mm`'s view of the same weights for decode rows in `_MID_MM_ROWS`: - # the per-tensor e4m3 codes above, the shared expert's and the router's bf16 weights - # (no new copies). - self._mid = {} - if _dense_fp8_enabled(): - for name in w: - if re.fullmatch(r"l\d+_(attn_in_w|o_w|in_proj_w|out_w)", name): - codes, scale = fp8_per_tensor_weight(w[name]) - self._dense8[name] = (codes.t(), scale) - self._gemv8[name] = pack_from_bf16(w[name]) - self._mid[name] = MidWeight(codes, scale) - elif re.fullmatch(r"l\d+_sh_gu", name): - self._gemv8[name] = pack_from_bf16(w[name][: 2 * self.shared_inter]) - self._mid[name] = MidWeight(w[name][: 2 * self.shared_inter]) - elif re.fullmatch(r"l\d+_sh_dn", name): - self._gemv8[name] = pack_from_bf16(w[name]) - self._mid[name] = MidWeight(w[name]) - elif re.fullmatch(r"l\d+_router", name): - self._mid[name] = MidWeight(w[name]) - if device.type == "cuda": - _precompile_mid(self._mid) - # NVFP4 copies of the GDN input projections (`_GDN_INPROJ_NVFP4_ENV`) for decode rows, - # with every variant a decode step can reach compiled now. - self._gemv_nv = {} - if _gdn_inproj_nvfp4_enabled() and device.type == "cuda": - for name in w: - if re.fullmatch(r"l\d+_in_proj_w", name): - self._gemv_nv[name] = nvfp4_pack_split_from_bf16(w[name]) - if self._gemv_nv: - precompile_split( - next(iter(self._gemv_nv.values())), - _INPROJ_NVFP4_MAX_ROWS, - self.conv_dim, - _PAD_SLOT_ID, - tuple(range(2, self.max_verify_rows + 1)), - ) - - # The CPU-only load checks build the targets over host weights; they launch nothing. - self._side_stream = torch.cuda.Stream(device=device) if device.type == "cuda" else None - # The prefill CuteDSL MoE's output memset, beside its FC1 (`fp4_moe_renorm_cute_dsl`): - # its own stream, so the join before FC2 does not also wait for the shared expert. - self._moe_memset_stream = ( - torch.cuda.Stream(device=device) if device.type == "cuda" else None - ) - self._prefill = PrefillTarget(self) - self._decode = DecodeTarget(self) - self.vision_tower = VisionTower(self) - - def _index_rope_table(self, theta: float, device) -> torch.Tensor: - """The index branch's rope table at `theta` in the QSA kernels' layout: - `[max_positions, 2 * rotary_pairs]` fp32, cos block then sin block.""" - pairs = self.rotary_dim // 2 - exponent = torch.arange(0, self.rotary_dim, 2, dtype=torch.float64, device=device) - inv_freq = (1.0 / (theta ** (exponent / self.rotary_dim))).to(torch.float32) - positions = torch.arange(self.max_positions, dtype=torch.float32, device=device) - angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0) - table = torch.cat([torch.cos(angles), torch.sin(angles)], dim=1).contiguous() - assert table.shape == (self.max_positions, 2 * pairs), table.shape - return table - - def _report_ple_residency(self) -> None: - """Check where the n-gram table ended up, and report it once. - - Offload is output-identical, so a stray `.to()` or a driver that declined - the mapping would pass every smoke test; the pointer's memory kind read - back from the driver is the signal that cannot be faked. The report is - printed here, before the KV pool is sized from what is left.""" - table = self.w["ple_table"] - kind = _weights.memory_kind(table.data_ptr()) - expected = ( - _weights.CUDA_MEMORY_TYPE_HOST - if self.ple_host_offload - else _weights.CUDA_MEMORY_TYPE_DEVICE - ) - assert kind == expected, ( - f"the n-gram table was declared with {_PLE_HOST_OFFLOAD_ENV}=" - f"{'1' if self.ple_host_offload else '0'}, so its pages must be " - f"{'host' if self.ple_host_offload else 'device'} memory, but the " - f"driver reports cudaMemoryType {kind} for the pointer the forward " - f"gathers from. The offload is output-identical, so nothing " - f"downstream of this would notice" - ) - mib = 1024.0 * 1024.0 - table_mib = table.numel() * table.element_size() / mib - free, total = torch.cuda.mem_get_info() - print( - f"ple n-gram table: {tuple(table.shape)} {table.dtype}, " - f"{table_mib:.1f} MiB, " - f"{'HOST-resident (device-addressable)' if self.ple_host_offload else 'DEVICE-resident'}" - f" [{_PLE_HOST_OFFLOAD_ENV}=" - f"{os.environ.get(_PLE_HOST_OFFLOAD_ENV, '1 (default)')}]; " - f"after load: allocated {torch.cuda.memory_allocated() / mib:.1f} MiB, " - f"reserved {torch.cuda.memory_reserved() / mib:.1f} MiB, " - f"driver-used {(total - free) / mib:.1f} MiB of {total / mib:.1f} MiB", - flush=True, - ) - - # ------------------------------------------------------------ fail-fast --- - - def _pool_geometry(self, tokens_per_block: int) -> _PoolGeometry: - return _PoolGeometry( - tokens_per_block=tokens_per_block, - heads_kv=self.heads_kv, - head_dim=self.head_dim, - index_kv_heads=self.index_kv_heads, - index_head_dim=self.index_head_dim, - conv_dim=self.conv_dim, - conv_kernel=self.conv_kernel, - num_k_heads=self.num_k_heads, - num_v_heads=self.num_v_heads, - head_v_dim=self.head_v_dim, - head_k_dim=self.head_k_dim, - stream=self.stream, - ple_state_len=self.ple_state_len, - ngram_context_len=self.ngram_context_len, - dtype=self.w["embed"].dtype, - kv_dtype=self.kv_pool_dtype, - kv_width=self.kv_pool_width, - ) - - def _kv_planes(self, cm, g: _PoolGeometry, layer: int, index_cache) -> _KvPlanes: - """One layer's paged K/V views from `cm`, checked against the contract: the data - planes, and under NVFP4 the block-scale planes that go with them.""" - pool = cm.get_buffers(layer, kv_layout="HND") - check_attention_pools(g, layer, index_cache, pool) - if self.quant_mode != QM_NVFP4_KV_CACHE: - return _KvPlanes(pool[:, 0], pool[:, 1]) - _require(cm, ("get_block_scale_buffers",), "reads an NVFP4 pool's e4m3 block-scale planes") - scales = cm.get_block_scale_buffers(layer, kv_layout="HND") - check_kv_scale_planes(g, layer, pool, scales) - return _KvPlanes(pool[:, 0], pool[:, 1], scales[:, 0], scales[:, 1]) - - def _bind_cache_manager(self, md, position_ids) -> None: - """Bind this module to one KV cache manager: fail-fast on the contract, then - take every pool view from it. - - Called once per manager, not per model: with capacity estimation on, the - engine builds a throwaway manager, warms up, shuts it down and builds the - real one on the same module, and a view held across the swap dangles - silently. The result is built in locals and installed as one `_Bound`, so a - failed bind leaves the previous binding untouched. The contract checks - re-run on every bind.""" - # Checked first: some fields below are properties that assert on a - # manager-less metadata instead of raising AttributeError. - cm = md.kv_cache_manager - assert cm is not None, "a paged KV cache is required" - # A re-bind during capture would bake two managers' pointers into one graph; - # the engine releases graphs before its own swap, so this is an ordering - # violation. - assert not (self._binds and torch.cuda.is_current_stream_capturing()), ( - f"the KV cache manager changed to {type(cm).__name__} during CUDA graph capture" - ) - self._binds += 1 - - missing = [name for name in _QSA_FIELDS if not hasattr(md, name)] - assert not missing, ( - f"metadata fields missing: {missing}; the engine did not select QSA " - "metadata for this model" - ) - - # `position_ids` is the int32 [3, 1, T] buffer the engine builds on an mRoPE - # config even for text-only batches; row 0 is the text position. - assert position_ids.dtype == torch.int32, position_ids.dtype - assert self.mrope_declared, "this target takes the mRoPE position layout" - assert position_ids.dim() == 3 and position_ids.shape[0] == _QSA_POSITION_AXES, ( - f"expected the [{_QSA_POSITION_AXES}, 1, T] mRoPE position buffer, got " - f"{tuple(position_ids.shape)}" - ) - - # thop_attention addresses the pool by the layer's pool offset, not its - # ordinal among attention layers; counting them aliases layers with finite, - # plausible output. - pool_idx = {i: cm.layer_offsets[i] for i in self.attn_layers} - pools = {row[0] for row in md.host_kv_cache_pool_mapping.tolist()} - assert pools <= {0, 1}, ( - f"KV addressing over {len(pools)} pools is not certified " - f"(thop_attention's cells cover one and two); layer->pool ids {sorted(pools)}" - ) - - # No sliding window: both targets pass `max_seq_len` as the window. The vector - # is indexed by pool offset and recurrent layers carry a large negative - # sentinel, so only attention offsets are read. - windows = cm.max_attention_window_vec - attn_windows = {windows[off] for off in pool_idx.values()} - assert attn_windows <= {None, md.max_seq_len}, ( - f"the attention layers' pool windows are {sorted(map(repr, attn_windows))} " - f"but this target passes max_seq_len {md.max_seq_len} as the window; " - "a sliding window is not implemented here" - ) - - # The NVFP4 attention kernels exist for the paged-context path only; without it the - # dense op would not read the pool the sparse path wrote. - assert self.quant_mode != QM_NVFP4_KV_CACHE or md.use_paged_context_fmha, ( - "an NVFP4 KV cache is read by the paged-context attention kernels only, which the " - "engine enables for chunked prefill, KV cache block reuse or speculative decoding; " - "none of the three is on" - ) - g = self._pool_geometry(md.tokens_per_block) - _require( - cm, - ("get_index_k_buffer", "get_qsa_position_buffer", "get_buffers"), - "reads the QSA index plane, the shared position plane and the paged K/V", - ) - # `qsa_topk_select` checks only that the buffer is at least `index_topk` wide, - # then addresses it as flat `index_topk`-wide rows ignoring the real stride, so - # a width mismatch silently mis-writes every row but the first. - assert md.qsa_topk_indices.shape[1] == self.block_topk, ( - f"the engine sized the QSA selection buffer {md.qsa_topk_indices.shape[1]} " - f"wide and this target selects {self.block_topk} blocks per row" - ) - position_cache = cm.get_qsa_position_buffer() - check_position_plane(g, position_cache) - index_cache, kv = {}, {} - for i in self.attn_layers: - index_cache[i] = cm.get_index_k_buffer(i) - kv[i] = self._kv_planes(cm, g, i, index_cache[i]) - - # A missed hybrid predicate would still build and pass a smoke run with 36 - # layers carrying no state. - mm = md.mamba_metadata - assert mm is not None and mm is not False, ( - "attn_metadata.mamba_metadata is absent: the runtime allocated no " - f"recurrent-state pool, so the {len(self.gdn_layers)} linear-attention " - "layers would run with no state" - ) - mm_missing = [name for name in _MAMBA_FIELDS if not hasattr(mm, name)] - assert not mm_missing, f"mamba metadata fields missing: {mm_missing}" - # `ple_layer_cache` exists only on the V2 hybrid manager; under V1 the PLE - # state silently does not exist. - _require( - cm, - ("get_conv_states", "get_ssm_states", "ple_layer_cache"), - "reads the recurrent, conv and PLE pools", - ) - conv = {i: cm.get_conv_states(i) for i in self.gdn_layers} - ssm = {i: cm.get_ssm_states(i) for i in self.gdn_layers} - conv0, ssm0 = conv[self.gdn_layers[0]], ssm[self.gdn_layers[0]] - check_recurrent_pools(g, conv0, ssm0) - - ple_cache = cm.ple_layer_cache(self.ple_layer) - assert ple_cache is not None, ( - f"no PLE lifecycle state for decoder layer {self.ple_layer}: the runtime " - "allocated neither the dilated conv state nor the n-gram id history" - ) - ple_conv, ple_ngram = ple_cache - check_ple_pools(g, ple_conv, ple_ngram, conv0) - - assert md.effective_beam_width == 1, "beam search is not implemented" - assert md.cache_indirection is None, "beam search is not implemented" - assert not md.is_cross, "cross attention is not implemented" - # With speculation on, the Blackwell trtllm-gen spec-dec mask machinery is off - # for a linear tree at sm_103: every `spec_decoding_*` tensor stays None and - # `predicted_tokens_per_seq` alone produces the within-block mask. - assert not md.is_spec_decoding_enabled, ( - "the attention metadata reports is_spec_decoding_enabled; on sm_103 a " - "linear-tree MTP forces that machinery off" - ) - bound = _Bound( - cm, pool_idx, index_cache, position_cache, kv, conv, ssm, (ple_conv, ple_ngram) - ) - if self.spec_rollback is not None: - self.spec_rollback.set_index_planes([index_cache[i] for i in self.attn_layers]) - if self.mtp_enabled: - bound = bound._replace(**self._bind_step_buffers(cm, g, mm, conv0, ssm0)) - self._bound = bound - - def _bind_step_buffers(self, cm, g: _PoolGeometry, mm, conv0, ssm0) -> dict: - """Per-manager fail-fast for the MTP increment: the verify step needs - runtime-owned buffers the base path never touches. Every row's conv window goes - to a per-step buffer in both recurrent modes; the recurrent state is either the - cached replay (a checkpoint per slot plus a compact history of accepted tokens, - no per-step recurrent buffer) or, with `TRTLLM_USE_GDN_REPLAY=0`, per-step - recurrent buffers. Returns the `_Bound` fields it adds.""" - assert cm.is_speculative(), ( - f"the KV cache manager {type(cm).__name__} reports is_speculative() False, " - "so it allocated none of the per-step state buffers the verify path writes" - ) - _require( - cm, - ("mamba_layer_cache", "get_replay_state_update_metadata"), - "binds its per-step state buffers and replay metadata", - ) - # Bind the per-layer views once for subsequent forwards. - layer_caches = {i: cm.mamba_layer_cache(i) for i in self.gdn_layers} - int_conv = {i: cache.intermediate_conv_window for i, cache in layer_caches.items()} - int_idx = getattr(cm, "intermediate_state_indices", None) - g0 = self.gdn_layers[0] - check_step_buffers(g, int_conv[g0], conv0, int_idx) - int_ssm0 = layer_caches[g0].intermediate_ssm - if cm.get_replay_state_update_metadata() is None: - check_legacy_state(g, int_ssm0, ssm0) - int_ssm = {i: cache.intermediate_ssm for i, cache in layer_caches.items()} - return dict(int_idx=int_idx, int_conv=int_conv, int_ssm=int_ssm) - assert int_ssm0 is None, ( - "the cache manager runs the cached-replay verify path and also allocated a " - "per-step recurrent buffer; the runtime allocates one or the other" - ) - return dict(int_idx=int_idx, int_conv=int_conv, replay=self._bind_replay(cm, g, mm, ssm0)) - - def _bind_replay(self, cm, g: _PoolGeometry, mm, ssm0) -> _Replay: - """The cached-replay history buffers and bookkeeping of one manager.""" - meta = cm.get_replay_state_update_metadata() - assert meta is not None, ( - "the cached-replay binding requires replay metadata " - "(prev_num_accepted_tokens / cache_buf_idx / widths)" - ) - history, width = meta.replay_history_size, meta.replay_step_width - check_replay_state(history, width, meta.prev_num_accepted_tokens, meta.cache_buf_idx, ssm0) - old = {} - for i in self.gdn_layers: - layer_cache = cm.mamba_layer_cache(i) - old[i] = ( - layer_cache.old_x, - layer_cache.old_B, - layer_cache.old_dt, - layer_cache.old_dA_cumsum, - ) - check_replay_history(g, i, history, ssm0.shape[0], *old[i][:3]) - # The partitioned replay is the manager's to commit: it commits exactly when this - # flag holds and a step carries `CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE` or more - # generation requests, and the forward hands the kernel the work items under that - # same condition. - all_layer = bool(getattr(cm, "use_gdn_cached_replay_all_layer_commit", False)) - # A bf16 pool replays on the CuTe kernel, an fp32 pool on the Triton one; both commit a - # wrapping checkpoint in-kernel when given no work items and leave it to the manager's - # all-layer commit when given them. - cute = ssm0.dtype == torch.bfloat16 - if all_layer: - missing = [n for n in ("replay_work_items", "replay_n_writes") if not hasattr(mm, n)] - assert not missing, ( - f"the manager's all-layer replay commit is on but the mamba metadata lacks " - f"{missing}; the partitioned verify call needs both" - ) - return _Replay( - old, meta.cache_buf_idx, meta.prev_num_accepted_tokens, history, width, all_layer, cute - ) - - def _bind_draft_cache_manager(self, md) -> None: - """Bind the draft layer to the manager its draft step sees: its own index plane, - that manager's position plane and its paged K/V. - - With a separate draft KV cache that is the draft manager the worker swaps in for - the draft loop, whose pools hold the draft layer's pages and none of the trunk's; - otherwise the trunk's manager, which maps the draft layer as well. Keyed on the - manager like `_bind_cache_manager`. Never first bound under capture: the runner - runs every captured batch shape eagerly first, so a bind under capture would bake - views no eager step validated into the graph.""" - cm = md.kv_cache_manager - assert not torch.cuda.is_current_stream_capturing(), ( - f"the draft layer is not bound to KV cache manager {type(cm).__name__} (first use, " - "or the manager changed since its last bind) and CUDA graph capture is active; " - "an eager step must take and check its views first" - ) - _require( - cm, - ("get_index_k_buffer", "get_qsa_position_buffer", "get_buffers"), - "reads the draft layer's QSA index plane, position plane and paged K/V", - ) - g = self._pool_geometry(md.tokens_per_block) - layer = self.mtp_layer_idx - index_cache = cm.get_index_k_buffer(layer) - position_cache = cm.get_qsa_position_buffer() - check_position_plane(g, position_cache) - kv = self._kv_planes(cm, g, layer, index_cache) - share_slot = None - if self.mtp_index_share: - # The metadata sizes its shared selection buffers from the same resolved knob, - # one row per sparse layer of whichever manager is bound. - assert md.sparse_metadata_params.mtp_index_share, ( - "this target resolved IndexShare on but the engine's QSA metadata did not" - ) - slots = cm.qsa_shared_topk_slots - check_shared_selection(md.qsa_shared_topk_indices, slots, layer, self.block_topk) - share_slot = slots[layer] - self._draft_bound = _DraftBound(cm, index_cache, position_cache, kv, share_slot) - - # -------------------------------------------------------------- forward --- - - def _select_target(self, attn_metadata) -> "PrefillTarget | DecodeTarget": - """Which target runs this step: decode if `num_contexts == 0`, else prefill. - - A draft loop rewrites `num_contexts` in place between draft steps, so the - draft step routes through here per call.""" - return self._decode if attn_metadata.num_contexts == 0 else self._prefill - - def begin_step( - self, - attn_metadata: AttentionMetadata, - input_ids: torch.IntTensor | None, - position_ids: torch.IntTensor | None, - inputs_embeds: torch.FloatTensor | None = None, - lora_params: dict | None = None, - spec_metadata=None, - vision: tuple | None = None, - ) -> None: - """Check what holds for every step, bind the cache manager, choose the target and - let the prefill target prepare the step. - - The shell calls it before `forward`, on every step. It is separate because a - breakable prefill CUDA graph captures `forward` and replays it without calling it: - whatever a step reads from the host (metadata, host buffers, the manager) is read - here, where it runs eagerly on a replayed step too. `vision` is `(rows, - mm_token_indices)` when the step carries media: the tower's `[R, hidden]` rows and - the `[R]` packed-batch rows they replace; the prefill target splices them inside - its forward.""" - assert self._prefill is not None and self._decode is not None, ( - "load_weights must run before forward" - ) - assert position_ids is not None - assert isinstance(attn_metadata, TrtllmAttentionMetadata) - # Unimplemented inputs must fail loudly, not be dropped. - assert lora_params is None, "LoRA is not implemented by this target" - # `input_ids` is required on every step: the PLE block hashes token ids whether or - # not vision rows replace some embeddings, and text gates would not catch a mis-hash. - assert input_ids is not None, ( - "this target's PLE layer hashes token ids, so input_ids is required " - "on every step, including one that carries vision rows" - ) - # The token embedding is this target's own and vision rows are spliced into it - # inside the forward; an embedding handed in from outside would be dropped on every - # replay of a breakable prefill graph. - assert inputs_embeds is None, ( - "this target computes its token embedding itself and splices vision rows into " - "it; an engine-supplied inputs_embeds is not implemented" - ) - # `spec_metadata` is present exactly when drafting. On MTP_EAGLE_ONE_MODEL it - # sets `layers_to_capture = ()`, so the trunk owes no capture hook. - assert (spec_metadata is not None) == self.mtp_enabled, ( - f"spec_metadata {'arrived' if spec_metadata is not None else 'is absent'} " - f"but the model was built with mtp_enabled={self.mtp_enabled}; the " - "draft path and the weights it loads are decided at construction" - ) - # Keyed on the manager, not a first-forward flag: capacity estimation swaps the - # manager under the same module (see `_bind_cache_manager`). - if self._bound is None or attn_metadata.kv_cache_manager is not self._bound.cm: - self._bind_cache_manager(attn_metadata, position_ids) - target = self._select_target(attn_metadata) - if target is self._prefill: - target.prepare(attn_metadata, input_ids, position_ids, vision, spec_metadata) - else: - assert vision is None, ( - "a generation-only step carries vision rows; the executor attaches media " - "to context requests only" - ) - self._step_begun = True - - def forward( - self, - attn_metadata: AttentionMetadata, - input_ids: torch.IntTensor | None = None, - position_ids: torch.IntTensor | None = None, - inputs_embeds: torch.FloatTensor | None = None, - lora_params: dict | None = None, - **kwargs, - ): - """Run the step `begin_step` prepared on the target it chose. - - Under a breakable prefill CUDA graph this is the body the engine captures for a - step with context rows; a replayed step never enters it, which is why nothing here - may read the step from the host. Returns the logits input, or under speculation the - decode target's `(final, hc)` pair and the prefill target's one packed tensor - (`split_spec_output`).""" - assert self._step_begun, ( - "the shell calls begin_step before every forward; this forward has no prepared step" - ) - self._step_begun = False - return self._select_target(attn_metadata).forward( - attn_metadata, input_ids=input_ids, position_ids=position_ids, **kwargs - ) - - def split_spec_output(self, output) -> tuple[torch.Tensor, torch.Tensor]: - """`(final, hc)` out of what `forward` returned under speculation: `final` is the - `[T, hidden]` mix `lm_head` reads and `hc` the `[T, hc_count * hidden]` stream the - MTP module reads. The decode target returns the pair; the prefill target returns one - `[T, hidden + stream]` tensor, of which the two are column views.""" - if isinstance(output, tuple): - return output - final, hc = torch.split(output, [self.hidden, self.stream], -1) - return final, hc - - -class _LayerWeights(NamedTuple): - """One trunk layer's weights as the views the ops read (GEMM operands are `.t()` - views). Attention fields are None on a Gated-DeltaNet layer and vice versa.""" - - ahc_norm: torch.Tensor - ahc_first: torch.Tensor - ahc_up: torch.Tensor - mhc_norm: torch.Tensor - mhc_first: torch.Tensor - mhc_up: torch.Tensor - attn_in: torch.Tensor | None - o: torch.Tensor | None - q_norm: torch.Tensor | None - k_norm: torch.Tensor | None - idx_qn: torch.Tensor | None - idx_kn: torch.Tensor | None - in_proj: torch.Tensor | None - out: torch.Tensor | None - conv_w: torch.Tensor | None - a_log: torch.Tensor | None - dt_bias: torch.Tensor | None - gnorm: torch.Tensor | None - router: torch.Tensor - moe_gscale: torch.Tensor - fc1_w: torch.Tensor - fc1_s: torch.Tensor - fc2_w: torch.Tensor - fc2_s: torch.Tensor - # trtllm-gen scalars; None on CuteDSL. - out1_scale: torch.Tensor | None - out1_gate_scale: torch.Tensor | None - out2_scale: torch.Tensor | None - # CuteDSL scalars; None on trtllm-gen. - alpha1: torch.Tensor | None - alpha2: torch.Tensor | None - global_sf: torch.Tensor | None - # Under trtllm-gen, the second expert copy in the CuteDSL layout and its scalars, - # which prefill-sized launches run on (`_CUTE_DSL_MOE_MIN_ROWS`); None on CuteDSL and - # without the copy (`moe_prefill_copy`). - fc1_w_cd: torch.Tensor | None - fc1_s_cd: torch.Tensor | None - fc2_w_cd: torch.Tensor | None - fc2_s_cd: torch.Tensor | None - alpha1_cd: torch.Tensor | None - alpha2_cd: torch.Tensor | None - global_sf_cd: torch.Tensor | None - sh_gu: torch.Tensor - sh_dn: torch.Tensor - sh_gate: torch.Tensor - # `[gate | up | gate logit | pad]` (`weights._cat_rows`): `sh_gu` and `sh_gate` are its - # column blocks. - sh_gug: torch.Tensor - # FP8 copies `(codes_t, scale)` of the dense projections (`_DENSE_FP8_ENV`); None when off - # and on the other layer kind. - attn_in8: tuple | None = None - o8: tuple | None = None - in_proj8: tuple | None = None - out8: tuple | None = None - # Packed FP8 decode copies (`fp8_gemv`, 128x128 block scales) under the same mode. - attn_in_g: object | None = None - o_g: object | None = None - in_proj_g: object | None = None - out_g: object | None = None - sh_gu_g: object | None = None - sh_dn_g: object | None = None - # The GDN in_proj's NVFP4 decode copy (`_GDN_INPROJ_NVFP4_ENV`, `nvfp4_gemv.SplitNvfp4Weight`). - in_proj_nv: object | None = None - # `mid_rows_mm` views (`MidWeight`) under the same mode, for `_MID_MM_ROWS` decode rows. - attn_in_m: object | None = None - o_m: object | None = None - in_proj_m: object | None = None - out_m: object | None = None - sh_gu_m: object | None = None - sh_dn_m: object | None = None - router_m: object | None = None - # `1 + ahc_norm` / `1 + mhc_norm` in fp32, the rstd mixers' operand - # (`qwen4_exp_hc_up_gate_mean_rstd`); set by `PrefillTarget` only. - ahc_norm_w1: torch.Tensor | None = None - mhc_norm_w1: torch.Tensor | None = None - - -def _layer_views(core: Qwen4ExpTp1ModelingV2) -> list[_LayerWeights]: - """Every trunk layer's `_LayerWeights`, read off the realized weights.""" - w = core.w - # The expert stacks are uint8 storage; each backend reads the packed nibbles and the - # block scales as its own dtypes. - if core.trunk_moe_backend == "CUTEDSL": - w_dtype, s_dtype = torch.float4_e2m1fn_x2, torch.uint8 - else: - w_dtype, s_dtype = torch.uint8, torch.float8_e4m3fn - layers = [] - for i in range(core.num_layers): - attn: dict = dict( - attn_in=None, - o=None, - q_norm=None, - k_norm=None, - idx_qn=None, - idx_kn=None, - ) - gdn: dict = dict( - in_proj=None, - out=None, - conv_w=None, - a_log=None, - dt_bias=None, - gnorm=None, - ) - if core.is_attn[i]: - attn = dict( - attn_in=w[f"l{i}_attn_in_w"].t(), - o=w[f"l{i}_o_w"].t(), - q_norm=w[f"l{i}_q_norm"], - k_norm=w[f"l{i}_k_norm"], - idx_qn=w[f"l{i}_idx_qn"], - idx_kn=w[f"l{i}_idx_kn"], - ) - else: - a_log, dt_bias = core._gdn_gating[i] - gdn = dict( - in_proj=w[f"l{i}_in_proj_w"].t(), - out=w[f"l{i}_out_w"].t(), - # [conv_dim, 1, width] -> [conv_dim, width]: a view (unit width stride). - conv_w=w[f"l{i}_conv_w"].reshape(core.conv_dim, core.conv_kernel), - a_log=a_log, - dt_bias=dt_bias, - gnorm=w[f"l{i}_gnorm"], - ) - scalars = core._moe_scalars[i] - cute_dsl_copy: dict = dict( - fc1_w_cd=None, - fc1_s_cd=None, - fc2_w_cd=None, - fc2_s_cd=None, - alpha1_cd=None, - alpha2_cd=None, - global_sf_cd=None, - ) - if core._moe_scalars_cd is not None: - scalars_cd = core._moe_scalars_cd[i] - cute_dsl_copy = dict( - fc1_w_cd=w[f"l{i}_fc1_w_cd"], - fc1_s_cd=w[f"l{i}_fc1_s_cd"], - fc2_w_cd=w[f"l{i}_fc2_w_cd"], - fc2_s_cd=w[f"l{i}_fc2_s_cd"], - alpha1_cd=scalars_cd["alpha1"], - alpha2_cd=scalars_cd["alpha2"], - global_sf_cd=scalars_cd["global_sf"], - ) - layers.append( - _LayerWeights( - ahc_norm=w[f"l{i}_ahc_norm"], - ahc_first=core._hc_first[(i, "ahc")].t(), - ahc_up=w[f"l{i}_ahc_up"].t(), - mhc_norm=w[f"l{i}_mhc_norm"], - mhc_first=core._hc_first[(i, "mhc")].t(), - mhc_up=w[f"l{i}_mhc_up"].t(), - **attn, - **gdn, - router=w[f"l{i}_router"].t(), - moe_gscale=scalars["global_scale"], - fc1_w=w[f"l{i}_fc1_w"].view(w_dtype), - fc1_s=w[f"l{i}_fc1_s"].view(s_dtype), - fc2_w=w[f"l{i}_fc2_w"].view(w_dtype), - fc2_s=w[f"l{i}_fc2_s"].view(s_dtype), - out1_scale=scalars.get("output1_scale_scalar"), - out1_gate_scale=scalars.get("output1_scale_gate_scalar"), - out2_scale=scalars.get("output2_scale_scalar"), - alpha1=scalars.get("alpha1"), - alpha2=scalars.get("alpha2"), - global_sf=scalars.get("global_sf"), - **cute_dsl_copy, - sh_gu=w[f"l{i}_sh_gu"][: 2 * core.shared_inter].t(), - sh_dn=w[f"l{i}_sh_dn"].t(), - sh_gate=w[f"l{i}_sh_gu"][2 * core.shared_inter : 2 * core.shared_inter + 1].t(), - sh_gug=w[f"l{i}_sh_gu"].t(), - attn_in8=getattr(core, "_dense8", {}).get(f"l{i}_attn_in_w"), - o8=getattr(core, "_dense8", {}).get(f"l{i}_o_w"), - in_proj8=getattr(core, "_dense8", {}).get(f"l{i}_in_proj_w"), - out8=getattr(core, "_dense8", {}).get(f"l{i}_out_w"), - attn_in_g=getattr(core, "_gemv8", {}).get(f"l{i}_attn_in_w"), - o_g=getattr(core, "_gemv8", {}).get(f"l{i}_o_w"), - in_proj_g=getattr(core, "_gemv8", {}).get(f"l{i}_in_proj_w"), - out_g=getattr(core, "_gemv8", {}).get(f"l{i}_out_w"), - sh_gu_g=getattr(core, "_gemv8", {}).get(f"l{i}_sh_gu"), - sh_dn_g=getattr(core, "_gemv8", {}).get(f"l{i}_sh_dn"), - in_proj_nv=getattr(core, "_gemv_nv", {}).get(f"l{i}_in_proj_w"), - attn_in_m=getattr(core, "_mid", {}).get(f"l{i}_attn_in_w"), - o_m=getattr(core, "_mid", {}).get(f"l{i}_o_w"), - in_proj_m=getattr(core, "_mid", {}).get(f"l{i}_in_proj_w"), - out_m=getattr(core, "_mid", {}).get(f"l{i}_out_w"), - sh_gu_m=getattr(core, "_mid", {}).get(f"l{i}_sh_gu"), - sh_dn_m=getattr(core, "_mid", {}).get(f"l{i}_sh_dn"), - router_m=getattr(core, "_mid", {}).get(f"l{i}_router"), - ) - ) - return layers - - -class _StackWeights(NamedTuple): - """The weights outside the layer loop: the stack-final mixer (this checkpoint's - final norm) and the PLE side path.""" - - fhc_norm: torch.Tensor - fhc_first: torch.Tensor - fhc_up: torch.Tensor - ple_key: torch.Tensor - ple_val: torch.Tensor - ple_nk_w: torch.Tensor - ple_nq_w: torch.Tensor - ple_nc_w: torch.Tensor - ple_conv_w: torch.Tensor - ple_mult: torch.Tensor - ple_vocab: torch.Tensor - ple_offs: torch.Tensor - ple_tscale: torch.Tensor - # `1 + fhc_norm` in fp32 (see `_LayerWeights.ahc_norm_w1`); set by `PrefillTarget` only. - fhc_norm_w1: torch.Tensor | None = None - - -def _norm_w1(norm: torch.Tensor) -> torch.Tensor: - """A grouped norm's `(1 + w)` in fp32, contiguous: what - `qwen4_exp_hc_up_gate_mean_rstd` multiplies the scaled residual by (the `1 + w` - `qwen4_exp_hc_inject_norm` forms per element, bitwise).""" - return (1.0 + norm.float()).contiguous() - - -def _stack_views(core: Qwen4ExpTp1ModelingV2) -> _StackWeights: - w = core.w - return _StackWeights( - fhc_norm=w["fhc_norm"], - fhc_first=core._hc_first["fhc"].t(), - fhc_up=w["fhc_up"].t(), - ple_key=w["ple_key_w"].t(), - ple_val=w["ple_val_w"].t(), - # The three PLE norm weights whole: the fused tail normalizes each stream inside - # one program. - ple_nk_w=w["ple_nk"], - ple_nq_w=w["ple_nq"], - ple_nc_w=w["ple_nc"], - ple_conv_w=w["ple_conv_w"], - ple_mult=w["ple_mult"], - ple_vocab=w["ple_vocab"], - ple_offs=w["ple_offs"], - # `dequantize_e4m3_per_tensor` needs the scale at the input's rank (3), all 1s. - ple_tscale=w["ple_tscale"].reshape(1, 1, 1), - ) - - -class _MtpWeights(NamedTuple): - """The MTP module's weights (keyed at its layer index, `num_hidden_layers`).""" - - pre_e_norm: torch.Tensor - pre_h_norm: tuple - fc_e: torch.Tensor - fc_h: torch.Tensor - ahc_norm: tuple - ahc_down: torch.Tensor - ahc_up: torch.Tensor - ahc_inject: torch.Tensor - mhc_norm: tuple - mhc_down: torch.Tensor - mhc_up: torch.Tensor - mhc_inject: torch.Tensor - q: torch.Tensor - k: torch.Tensor - v: torch.Tensor - o: torch.Tensor - q_norm: torch.Tensor - k_norm: torch.Tensor - idx: torch.Tensor - idx_qn: torch.Tensor - idx_kn: torch.Tensor - router: torch.Tensor - fc1_w: torch.Tensor - fc1_s: torch.Tensor - fc2_w: torch.Tensor - fc2_s: torch.Tensor - sh_gu: torch.Tensor - sh_dn: torch.Tensor - sh_gate: torch.Tensor - fhc_norm: tuple - fhc_down: torch.Tensor - fhc_up: torch.Tensor - # With `_MTP_MOE_NVFP4_ENV`, `fc1_w` .. `fc2_s` are the NVFP4 copy in the trunk's - # trtllm-gen layout and these are its static activation global scale `[1]` and the - # runner's three per-expert scalars; None on the FP8 block-scale path. - moe_gscale: torch.Tensor | None = None - out1_scale: torch.Tensor | None = None - out1_gate_scale: torch.Tensor | None = None - out2_scale: torch.Tensor | None = None - - -def _mtp_views(core: Qwen4ExpTp1ModelingV2) -> _MtpWeights: - w = core.w - p = f"l{core.mtp_layer_idx}_mtp_" - return _MtpWeights( - pre_e_norm=w[f"{p}pre_e_norm"], - pre_h_norm=core._norm_slices(f"{p}pre_h_norm"), - fc_e=w[f"{p}fc_e"].t(), - fc_h=w[f"{p}fc_h"].t(), - ahc_norm=core._norm_slices(f"{p}ahc_norm"), - ahc_down=w[f"{p}ahc_down"], - ahc_up=w[f"{p}ahc_up"], - ahc_inject=w[f"{p}ahc_inject"], - mhc_norm=core._norm_slices(f"{p}mhc_norm"), - mhc_down=w[f"{p}mhc_down"], - mhc_up=w[f"{p}mhc_up"], - mhc_inject=w[f"{p}mhc_inject"], - q=w[f"{p}q_w"].t(), - k=w[f"{p}k_w"].t(), - v=w[f"{p}v_w"].t(), - o=w[f"{p}o_w"].t(), - q_norm=w[f"{p}q_norm"], - k_norm=w[f"{p}k_norm"], - idx=w[f"{p}idx_w"].t(), - idx_qn=w[f"{p}idx_qn"], - idx_kn=w[f"{p}idx_kn"], - router=w[f"{p}router"].t(), - fc1_w=w[f"{p}fc1_w"], - fc1_s=w[f"{p}fc1_s"], - fc2_w=w[f"{p}fc2_w"], - fc2_s=w[f"{p}fc2_s"], - sh_gu=w[f"{p}sh_gu"][: 2 * core.shared_inter].t(), - sh_dn=w[f"{p}sh_dn"].t(), - sh_gate=w[f"{p}sh_gu"][2 * core.shared_inter : 2 * core.shared_inter + 1].t(), - fhc_norm=core._norm_slices(f"{p}fhc_norm"), - fhc_down=w[f"{p}fhc_down"], - fhc_up=w[f"{p}fhc_up"], - ) - - -def _mtp_moe_nvfp4(core: Qwen4ExpTp1ModelingV2) -> dict: - """The draft layer's routed experts requantized to NVFP4 (`_MTP_MOE_NVFP4_ENV`), as the - `_MtpWeights` fields `_mtp_routed_nvfp4` reads; frees the FP8 copies. - - The runner's scalars follow `_moe_layer_scalars` (trtllm-gen) with `1 / g_in` and - `1 / g_mid` standing for the checkpoint's `input_scale`s and `1 / g_w` for its - `weight_scale_2`: `out1_gate = 1 / (g_in * g_w1[e])`, `out1 = out1_gate * g_mid`, - `out2 = 1 / (g_mid * g_w2[e])`. The activation scales are static - (`_MTP_MOE_NVFP4_FC1_MARGIN` / `_FC2_MARGIN`).""" - w = core.w - p = f"l{core.mtp_layer_idx}_mtp_" - keys = [f"{p}{k}" for k in ("fc1_w", "fc1_s", "fc2_w", "fc2_s")] - fc1_w, fc1_s, fc2_w, fc2_s, g_w1, g_w2 = _weights.mtp_experts_to_nvfp4(*(w[k] for k in keys)) - # Nothing else reads the FP8 experts: `_mtp_views` handed them to `_MtpWeights` only, and - # the caller replaces those fields with the NVFP4 copy. - for k in keys: - w[k].data = w[k].data.new_empty(0) - torch.cuda.empty_cache() - trunk = range(core.num_layers) - in_amax = max(float(w[f"l{i}_e_isc1"].max()) for i in trunk) * _weights.NVFP4_AMAX_TO_GLOBAL - mid_amax = max(float(w[f"l{i}_e_isc2"].max()) for i in trunk) * _weights.NVFP4_AMAX_TO_GLOBAL - # `1 / g` of each activation quantization. - inv_g_in = _MTP_MOE_NVFP4_FC1_MARGIN * in_amax / _weights.NVFP4_AMAX_TO_GLOBAL - inv_g_mid = _MTP_MOE_NVFP4_FC2_MARGIN * mid_amax / _weights.NVFP4_AMAX_TO_GLOBAL - gate1 = (inv_g_in / g_w1).contiguous() - print( - f"draft MoE: {len(g_w1)} experts requantized FP8 -> NVFP4 ({_MTP_MOE_NVFP4_ENV}=1), " - f"static activation amax FC1 {inv_g_in * _weights.NVFP4_AMAX_TO_GLOBAL:.1f}, " - f"FC2 {inv_g_mid * _weights.NVFP4_AMAX_TO_GLOBAL:.1f}", - flush=True, - ) - return dict( - fc1_w=fc1_w, - fc1_s=fc1_s.view(torch.float8_e4m3fn), - fc2_w=fc2_w, - fc2_s=fc2_s.view(torch.float8_e4m3fn), - moe_gscale=torch.full((1,), 1.0 / inv_g_in, dtype=torch.float32, device=fc1_w.device), - out1_scale=(gate1 / inv_g_mid).contiguous(), - out1_gate_scale=gate1, - out2_scale=(inv_g_mid / g_w2).contiguous(), - ) - - -def _mtp_routed_nvfp4(core: Qwen4ExpTp1ModelingV2, M: _MtpWeights, x: torch.Tensor): - """The draft layer's routed experts on their NVFP4 copy: the activation quantized at the - static global scale (linear block scales, as the trunk's MoE sites write them), then the - trunk's logits-routed trtllm-gen call (its Renormalize routing is the FP8 path's - `renorm_moe_routing_op`). Returns the combined `[T, H]` bf16 rows.""" - logits = cublas_mm(x, M.router) - xq, xsf = fp4_quantize(x.contiguous(), M.moe_gscale, _SF_VEC, False, False) - return fp4_block_scale_moe_renorm_routed( - logits, - xq, - xsf.view(torch.float8_e4m3fn), - M.fc1_w, - M.fc1_s, - M.fc2_w, - M.fc2_s, - M.out1_scale, - M.out1_gate_scale, - M.out2_scale, - core.num_experts, - core.topk, - core.moe_inter, - ) - - -def _mtp_moe_nvfp4_block(core: Qwen4ExpTp1ModelingV2, M: _MtpWeights, x: torch.Tensor): - """The draft layer's MoE on the NVFP4 experts (`_mtp_routed_nvfp4`) plus its bf16 shared - expert, combined as on the FP8 path.""" - routed = _mtp_routed_nvfp4(core, M, x) - shared = cublas_mm(flashinfer_silu_and_mul(cublas_mm(x, M.sh_gu)), M.sh_dn) - gate_logits = cublas_mm(x, M.sh_gate) - return fused_sigmoid_gate_mul_add(routed, gate_logits, shared) - - -class _MtpFused(NamedTuple): - """The draft layer's operands in the trunk's fused layouts, built after load: the - `[q | k | v | gate | index | pad]` input projection (`weights._attn_in_w`), each - hyper-connection's `[down; block_inject; 0]` first projection (`_fuse_hc_first_projections`) - with its up projection and whole grouped-norm weight, and a unit injection weight (the - front end's per-stream `x_h + x_e` as the first site's injection). The embedding branch - enters all four streams; nothing in the checkpoint or config says so, and entering stream - 0 alone would leave the text bit-correct with every draft rejected, which only the - acceptance gate detects.""" - - attn_in: torch.Tensor - ahc_norm: torch.Tensor - ahc_first: torch.Tensor - ahc_up: torch.Tensor - mhc_norm: torch.Tensor - mhc_first: torch.Tensor - mhc_up: torch.Tensor - fhc_norm: torch.Tensor - fhc_first: torch.Tensor - fhc_up: torch.Tensor - unit_inject: torch.Tensor - - -def _mtp_fused_views(core: Qwen4ExpTp1ModelingV2) -> _MtpFused: - w = core.w - p = f"l{core.mtp_layer_idx}_mtp_" - attn_in = _weights._attn_in_w((w[f"{p}q_w"], w[f"{p}k_w"], w[f"{p}v_w"], w[f"{p}idx_w"]), core) - assert attn_in.shape == (core.attn_in_width, core.hidden), attn_in.shape - return _MtpFused( - attn_in=attn_in.t(), - ahc_norm=w[f"{p}ahc_norm"], - ahc_first=core._hc_first["mtp_ahc"].t(), - ahc_up=w[f"{p}ahc_up"].t(), - mhc_norm=w[f"{p}mhc_norm"], - mhc_first=core._hc_first["mtp_mhc"].t(), - mhc_up=w[f"{p}mhc_up"].t(), - fhc_norm=w[f"{p}fhc_norm"], - fhc_first=core._hc_first["mtp_fhc"].t(), - fhc_up=w[f"{p}fhc_up"].t(), - unit_inject=torch.ones( - [core.max_num_tokens, core.hc_count], dtype=attn_in.dtype, device=attn_in.device - ), - ) - - -def _select(core, scores, row_starts, row_ends, topk_indices, live_cols=None) -> None: - """Select each row's top-`block_topk` groups into `topk_indices` with `qsa_topk_select`, - then rewrite the row canonically, in place. - - The radix engine compacts its picks through a shared-memory atomic, so the order is - launch-dependent, and where the tie class at the budget's last score straddles the - budget -- every zero-score group of a row with fewer positive scores than the budget, - on this checkpoint's relu scores -- which tied groups are kept can be too. The - expansion keeps the row's order and the sparse GQA reduces in it, so a raw row past the - dense threshold is not reproducible run to run. `qsa_canonical_selection` rewrites it - into the window's canonical selection -- the groups strictly above the budget's last - score, then the highest-index (most recent) ties, ascending, `-1` last -- a function of - the scores alone, and the engine's own set in ascending order wherever the budget's - last score is untied. A canonical row is prefix-packed, which the expansion and the - IndexShare capture and reuse require. - - `index_topk` is the block budget: the entry's default (2048) is 4x the buffer width and - the op ignores the row stride, so it mis-writes silently. - - The rewrite reads every column it is given, and `scores` is as wide as the block table, - not the live window. A caller that knows a host-side bound on the rows' windows passes - it as `live_cols`: the rewrite then reads only those columns and returns the same rows.""" - qsa_topk_select(scores, row_starts, row_ends, topk_indices, index_topk=core.block_topk) - if live_cols is not None: - scores = scores[:, :live_cols] - qsa_canonical_selection(scores, row_starts, row_ends, topk_indices) - - -def _store_kv(core, kv, k, v, request_indices, logical_positions, block_table, tokens_per_block): - """Append this step's K/V rows to the layer's paged pool: a plain move over a bf16 or fp8 - pool (an fp8 pool takes e4m3 unscaled, as the dense op does at the unit KV scales), the - quantizing store over an NVFP4 one, at the same second-level scales as the dense op.""" - if kv.k_scale is None: - qsa_paged_kv_store( - k, v, kv.k, kv.v, request_indices, logical_positions, block_table, tokens_per_block - ) - return - qsa_paged_kv_store_nvfp4( - k, - v, - kv.k, - kv.v, - kv.k_scale, - kv.v_scale, - request_indices, - logical_positions, - block_table, - tokens_per_block, - core.kv_scale_orig_quant, - ) - - -def _sparse_gqa( - core, - kv, - q, - block_table, - selected_tokens, - request_indices, - tokens_per_block, - query_positions, - compress_ratio, - output_gate=None, - rows_per_request=1, -): - """Sparse GQA over the layer's paged pool (dequantizing an NVFP4 one). The prefill pair - `query_positions` / `compress_ratio` is a bound on a context row's scan, or None, None. - `output_gate` (NVFP4 pool, no prefill pair only) returns the gated output. - `rows_per_request` (NVFP4 pool, no prefill pair only): the rows come that many consecutive - rows a request, a verify step's generation rows; 1 for any other batch.""" - if kv.k_scale is None: - assert output_gate is None, "the output gate is fused on the NVFP4 decode-row path only" - return qsa_paged_sparse_gqa( - q=q, - k_cache=kv.k, - v_cache=kv.v, - block_table=block_table, - selected_tokens=selected_tokens, - request_indices=request_indices, - tokens_per_block=tokens_per_block, - softmax_scale=core.softmax_scale, - query_positions=query_positions, - compress_ratio=compress_ratio, - ) - if query_positions is None: - # Rows without the prefill pair (decode, draft, a mixed step's generation rows): the - # CuTe split-K + merge, with the output gate fused when the caller hands it over. - if rows_per_request > 1: - # A verify step's rows, `rows_per_request` a request: the request-grouped entry - # loads each selected K/V token once for a group of a request's rows (from eight - # groups up; below that it is the per-row kernel). - return qsa_paged_sparse_gqa_nvfp4_verify( - q, - kv.k, - kv.v, - kv.k_scale, - kv.v_scale, - core.kv_scale_quant_orig, - block_table, - selected_tokens, - request_indices, - tokens_per_block, - core.softmax_scale, - rows_per_request, - output_gate=output_gate, - ) - return qsa_paged_sparse_gqa_nvfp4_decode( - q, - kv.k, - kv.v, - kv.k_scale, - kv.v_scale, - core.kv_scale_quant_orig, - block_table, - selected_tokens, - request_indices, - tokens_per_block, - core.softmax_scale, - output_gate=output_gate, - ) - assert output_gate is None, "the output gate is fused on the decode-row path only" - return qsa_paged_sparse_gqa_nvfp4( - q=q, - k_cache=kv.k, - v_cache=kv.v, - k_scale_cache=kv.k_scale, - v_scale_cache=kv.v_scale, - kv_scale_quant_orig=core.kv_scale_quant_orig, - block_table=block_table, - selected_tokens=selected_tokens, - request_indices=request_indices, - tokens_per_block=tokens_per_block, - softmax_scale=core.softmax_scale, - query_positions=query_positions, - compress_ratio=compress_ratio, - ) - - -def _attn_in_heads(core, proj): - """The main q, k and v heads of the fused attention projection's output `proj` - `[T, attn_in_width]` (`weights._attn_in_w`), as `[T, heads, head_dim]` views.""" - q, k, v, _ = torch.split( - torch.reshape(proj, [proj.shape[0], core.attn_in_width // core.head_dim, core.head_dim]), - [core.heads_q, core.heads_kv, core.heads_kv, core.attn_in_tail_heads - core.heads_kv], - 1, - ) - return q, k, v - - -def _cute_dsl_routed(core, L, logits, xq, xsf, routed_out) -> torch.Tensor: - """The trunk's routed experts on the CuteDSL pipeline: Renormalize routing, the sort - into per-expert row tiles, FC1 gathering its rows from the unpermuted NVFP4 activation - `xq` / `xsf` (linear scales) with SwiGLU and the requantize to the layer's one FC2 input - scale fused, FC2, and the routing-weighted unpermute in a fixed slot order, which - writes every row of `routed_out` `[T, hidden]` and returns it.""" - # fp32 weights: the dtype the unpermute scales by. - topk_ids, topk_w = renorm_moe_routing_op(logits, topk=core.topk, output_dtype=torch.float32) - # The sort pads each expert's (token, expert) pairs to whole tiles and emits the tile - # tables both GEMMs walk. - tile_expert, tile_limit, expanded_to_permuted, permuted_to_expanded, _, live_tiles = moe_sort( - token_selected_experts=topk_ids, - token_final_scales=topk_w, - num_experts=core.num_experts, - top_k=core.topk, - local_expert_offset=0, - local_num_experts=core.num_experts, - tile_tokens_dim=_CUTEDSL_TILE_SIZE, - ) - h, h_sf = cute_dsl_nvfp4_gather_grouped_gemm_act_fusion( - input=xq.view(torch.float4_e2m1fn_x2), - weight=L.fc1_w, - input_scale=xsf.view(torch.uint8).view(xq.shape[0], core.hidden // _SF_VEC), - weight_scale=L.fc1_s, - alpha=L.alpha1, - tile_idx_to_group_idx=tile_expert, - tile_idx_to_mn_limit=tile_limit, - permuted_idx_to_expanded_idx=permuted_to_expanded, - num_non_exiting_tiles=live_tiles, - global_sf=L.global_sf, - num_experts=core.num_experts, - top_k=core.topk, - num_local_experts=core.num_experts, - local_expert_offset=0, - tile_size=_CUTEDSL_TILE_SIZE, - scaling_vector_size=_SF_VEC, - activation_type=int(ActivationType.Swiglu), - ) - # FC2 writes one bf16 row per permuted row. - y = cute_dsl_nvfp4_grouped_gemm( - input=h, - weight=L.fc2_w, - input_scale=h_sf, - weight_scale=L.fc2_s, - alpha=L.alpha2, - tile_idx_to_group_idx=tile_expert, - num_non_exiting_tiles=live_tiles, - num_experts=core.num_experts, - top_k=core.topk, - num_local_experts=core.num_experts, - local_expert_offset=0, - tile_size=_CUTEDSL_TILE_SIZE, - output_dtype=torch.bfloat16, - scaling_vector_size=_SF_VEC, - ) - return moe_unpermute_inplace( - permuted_input=y, - output=routed_out, - expanded_idx_to_permuted_idx=expanded_to_permuted, - topk_scales=topk_w, - ) - - -def _nvfp4_context_groups(pages: list, rows: list, max_pages: int) -> list: - """Consecutive context requests in groups of at most `max_pages` pages, a request with - more pages alone: `(first request, end request, first row, end row)` a group, where - request `r` has `pages[r]` pages and the next `rows[r]` rows.""" - groups = [] - first = first_row = row = listed = 0 - for r, (n_pages, n_rows) in enumerate(zip(pages, rows)): - if r > first and listed + n_pages > max_pages: - groups.append((first, r, first_row, row)) - first, first_row, listed = r, row, 0 - listed += n_pages - row += n_rows - if pages: - groups.append((first, len(pages), first_row, row)) - return groups - - -def _nvfp4_context_plan( - core, md, req_idx, block_table, tpb: int, one_row_per_request: bool = False -) -> tuple: - """The pages a step's context requests read, for the bf16 row-tile pass over an NVFP4 - pool: request `r` lists its first `ceil(kv_len / tpb)` table pages. The requests come in - groups that fit `_NVFP4_CONTEXT_SCRATCH_BYTES` (`_nvfp4_context_groups`); within a group - request `r` takes scratch pages `offsets[r]` on. Returns `(groups, table, k_scratch, - v_scratch)`: per group `(first row, end row, slots, counts, offsets)`, the operands - `qsa_nvfp4_pages_to_bf16` fills per layer and group, the rows those of its requests (one - a request with `one_row_per_request`, else their step rows); `table` (all -1 here) gets - the requests' slot rows, the scratch pools `[pages, Hkv, tpb, D]` bf16 the largest group.""" - num_contexts = md.num_contexts - device = block_table.device - # The host's and the device's KV lengths are the same per-step values: the host's - # group the requests and size the scratch, the device's place the pages. - pages = ((md.kv_lens_runtime[:num_contexts] + (tpb - 1)) // tpb).tolist() - rows = [1] * num_contexts if one_row_per_request else md.seq_lens[:num_contexts].tolist() - page_bytes = 2 * core.heads_kv * tpb * core.head_dim * 2 # K and V, bf16 - groups = _nvfp4_context_groups(pages, rows, max(1, _NVFP4_CONTEXT_SCRATCH_BYTES // page_bytes)) - lens = md.seq_lens_cuda[:num_contexts].long() - slots = req_idx[torch.cumsum(lens, 0) - lens] - counts = torch.div( - md.kv_lens_cuda_runtime[:num_contexts].long() + (tpb - 1), tpb, rounding_mode="floor" - ) - planned = [] - for first, end, first_row, end_row in groups: - group_counts = counts[first:end] - planned.append( - ( - first_row, - end_row, - slots[first:end], - group_counts, - torch.cumsum(group_counts, 0) - group_counts, - ) - ) - listed = max(sum(pages[first:end]) for first, end, _, _ in groups) - table = torch.full_like(block_table, -1) - shape = (listed, core.heads_kv, tpb, core.head_dim) - k_scratch = torch.empty(shape, dtype=torch.bfloat16, device=device) - v_scratch = torch.empty(shape, dtype=torch.bfloat16, device=device) - return planned, table, k_scratch, v_scratch - - -def _sparse_gqa_nvfp4_context( - core, - kv, - plan, - q, - block_table, - selected, - req_idx, - logical, - tpb: int, - num_ctx: int, - out, - rows_per_request: int = 1, -) -> torch.Tensor: - """Sparse GQA of a step with context rows over an NVFP4 pool, into `out` `[rows, - q_width]`. With a plan (`_nvfp4_context_plan`) the first `num_ctx` rows, the context - rows, take the row-tile pass over the plan's scratch, filled here from this layer's - pool a group of requests at a time; the other rows, and every row without a plan, go - split-K over the pool itself (with the prefill pair the entry would take its per-row - prefill kernel). The generation - rows after the context rows come `rows_per_request` a request: more than one (a verify - step's) takes the request-grouped entry for them on their own.""" - rows = q.shape[0] - head = 0 - if plan is not None: - groups, table, k_scratch, v_scratch = plan - assert groups[-1][1] == num_ctx, ( - f"the plan's groups end at row {groups[-1][1]}, not at the {num_ctx} context rows" - ) - for first_row, end_row, slots, counts, offsets in groups: - qsa_nvfp4_pages_to_bf16( - kv.k, - kv.v, - kv.k_scale, - kv.v_scale, - core.kv_scale_quant_orig, - block_table, - slots, - counts, - offsets, - table, - k_scratch, - v_scratch, - ) - context = qsa_sparse_gqa_rowtile( - q=q[first_row:end_row], - k_cache=k_scratch, - v_cache=v_scratch, - block_table=table, - selected_tokens=selected[first_row:end_row], - request_indices=req_idx[first_row:end_row], - tokens_per_block=tpb, - softmax_scale=core.softmax_scale, - query_positions=logical[first_row:end_row], - compress_ratio=core.compress_ratio, - ) - out[first_row:end_row].copy_( - torch.reshape(context, [end_row - first_row, core.q_width]) - ) - head = num_ctx - elif ( - rows_per_request > 1 - and 0 < num_ctx < rows - and _verify_takes_grouped( - q[num_ctx:], kv.k, block_table, selected[num_ctx:], tpb, rows_per_request - ) - ): - # Without a plan the context rows would share the generation rows' split-K launch; - # when those are enough for the request-grouped kernel the context rows go on their own. - context = _sparse_gqa( - core, - kv, - q=q[:num_ctx], - block_table=block_table, - selected_tokens=selected[:num_ctx], - request_indices=req_idx[:num_ctx], - tokens_per_block=tpb, - query_positions=None, - compress_ratio=None, - ) - out[:num_ctx].copy_(torch.reshape(context, [num_ctx, core.q_width])) - head = num_ctx - if head < rows: - rest = _sparse_gqa( - core, - kv, - q=q[head:], - block_table=block_table, - selected_tokens=selected[head:], - request_indices=req_idx[head:], - tokens_per_block=tpb, - query_positions=None, - compress_ratio=None, - # the generation rows alone, or every row of the step - rows_per_request=rows_per_request if head == num_ctx else 1, - ) - out[head:rows].copy_(torch.reshape(rest, [rows - head, core.q_width])) - return out - - -def _live_score_columns(core, md, batch: int) -> int: - """A host-side bound on every row's selection window in a step with context rows: the - longest sequence's compressed length, with the verify width as slack for the draft - rows that sit past the host's KV lengths.""" - longest = int(md.kv_lens_runtime[:batch].max()) if batch > 0 else 0 - return (longest + _MAX_VERIFY_WIDTH) // core.compress_ratio + 1 - - -def _replay_verify( - core, mq_out, a_d, b_d, ssm_states, i, L, si_d, n, p, replay, pdl=False -) -> torch.Tensor: - """One GDN layer's verify recurrence under the cached replay: `n` generation sequences - of `p` rows over the conv output `mq_out` (`[n * p, conv_dim]`, token-major) and the - raw `a`/`b` projections, with the fp32 gating parameters fused into the kernel. - - `replay` is `(work_items, n_writes, partitioned)`: the work items are given exactly when - the manager runs its all-layer commit after the step (without that commit every history - wrap goes uncommitted, with it a wrap committed here is committed twice); without them - the kernel commits every wrap itself. An fp32 pool runs the Triton kernel, a bf16 pool - the CuTe one (`gdn_cached_replay_cute`) under the same contract. `mq_out` - is the conv's fresh, contiguous output, so the kernel reads q/k/v from it directly - (`packed_qkv`) and never from the views; the CuTe kernel reads `a`/`b` through their - row stride (column views of the input projection). `pdl`: the kernel right before is - the in_proj GEMV with the conv as its epilogue, which signals its dependents at its - start -- the CuTe kernel then reads the checkpoint and the history (which nothing in - the step writes ahead of it) before `griddepcontrol.wait`. The Triton conv does not - signal dependents, so after it the kernel launches without PDL.""" - work_items, n_writes, partitioned = replay - state = core._bound.replay - old_x, old_B, old_dt, old_dA = state.old[i] - qv, kv, vv = torch.split(mq_out, [core.key_dim, core.key_dim, core.value_dim], 1) - entry = ( - gdn_cached_replay_cute - if state.cute - else fused_recurrent_gated_delta_rule_cached_replay_update - ) - out = entry( - q=torch.reshape(qv, [n, p, core.num_k_heads, core.head_k_dim]), - k=torch.reshape(kv, [n, p, core.num_k_heads, core.head_k_dim]), - v=torch.reshape(vv, [n, p, core.num_v_heads, core.head_v_dim]), - g=torch.reshape(a_d, [n, p, core.num_v_heads]), - beta=torch.reshape(b_d, [n, p, core.num_v_heads]), - ssm_states=ssm_states, - state_indices=si_d, - old_u=old_x, - old_k=old_B, - old_G=old_dt, - old_beta=old_dA, - cache_buf_idx=state.cache_buf_idx, - prev_num_accepted_tokens=state.prev_num_accepted, - history_size=state.history, - scale=None, - use_qk_l2norm_in_kernel=True, - A_log=L.a_log, - dt_bias=L.dt_bias, - launch_with_pdl=bool(pdl and state.cute), - replay_work_items=work_items, - n_writes=n_writes, - block_v=None, - num_warps=None, - use_l2_state_cache=None, - use_l2_streaming_inputs=None, - use_l2_shared_inputs=None, - main_num_stages=None, - packed_qkv=mq_out, - use_all_layer_commit=partitioned, - output=None, - ) - return torch.reshape(out, [1, n * p, core.num_v_heads, core.head_v_dim]) - - -def _replay_work(core, mm, num_decodes: int) -> tuple: - """`_replay_verify`'s `replay` argument for a verify step of `num_decodes` generation - requests: the partitioned replay exactly when `update_mamba_states` runs the manager's - all-layer commit after the step.""" - partitioned = ( - core._bound.replay.all_layer_commit - and num_decodes >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE - ) - return ( - mm.replay_work_items[:num_decodes] if partitioned else None, - mm.replay_n_writes if partitioned else None, - partitioned, - ) - - -class _KvRangeProbe: - """`TRTLLM_MODELING_V2_VALIDATE` over an fp8 KV pool: reports, once, the largest K/V - magnitude the first eager forward's attention layers store. - - The pool takes K/V at the unit scale, so anything past the e4m3 maximum saturates - silently (a bf16 inf included). One host sync per attention layer of that one forward; - a forward under CUDA graph capture is passed over without a read.""" - - def __init__(self, last_layer: int) -> None: - self._last = last_layer - self._worst = 0.0 - - def observe(self, layer: int, k: torch.Tensor, v: torch.Tensor) -> bool: - """Fold one attention layer's K/V in; False once the forward's last layer has - reported.""" - if torch.cuda.is_current_stream_capturing(): - return True - peak = float(torch.maximum(k.abs().amax(), v.abs().amax())) - # A NaN stays: `max` would drop it. - self._worst = math.nan if math.isnan(peak + self._worst) else max(self._worst, peak) - if layer != self._last: - return True - limit = torch.finfo(torch.float8_e4m3fn).max - # Written so that NaN fails too. - if not self._worst <= limit: - logger.warning( - f"FP8 KV cache: the first forward stores a K/V magnitude of {self._worst:g}, " - f"past the e4m3 maximum {limit:g}; the pool is written at the unit scale, " - "so those values saturate" - ) - return False - - -class _PrefillStep(NamedTuple): - """One step's host-side state, built by `PrefillTarget.prepare` and read by the forward's - breaks. - - `num_tokens` is the batch's row count `n` and `physical_tokens` the `T` the forward runs - over (`n`, or the prefill CUDA graph bucket it was padded to); `input_ids` is the batch's - `n` ids. `attn` is thop_attention's per-step arguments; `gdn`, `ple` and `qsa` are the - recurrent, PLE and QSA bookkeeping tuples their breaks unpack, each cut to the batch's - rows; `stage` says the speculative rollback stages this verify step; `vision` is the - shell's `(rows, mm_token_indices)` or None; `nvfp4_context` is `_nvfp4_context_plan`'s - scratch for the context rows over an NVFP4 pool, or None. A new one replaces it every - step, so no break reads a previous step's.""" - - num_tokens: int - physical_tokens: int - input_ids: torch.Tensor - attn: dict - gdn: tuple - ple: tuple - qsa: tuple - stage: bool - vision: tuple | None - nvfp4_context: tuple | None = None - - -class PrefillTarget: - """Any batch with context rows (`num_contexts > 0`), possibly with generation rows - beside them; the recurrent layers' decode-only branches never run here.""" - - def __init__(self, core: Qwen4ExpTp1ModelingV2) -> None: - """Take the per-layer weight views once; MTP operands are read off `core._mtp`.""" - self.core = core - # Off unless TRTLLM_MODELING_V2_VALIDATE asks; each target owns its flag (the - # two `step_args` projections differ). - self._contract_pending = step_contract_enabled() - # Dropped once it has reported. - self._kv_probe = ( - _KvRangeProbe(core.attn_layers[-1]) - if self._contract_pending and core.quant_mode == QM_FP8_KV_CACHE - else None - ) - # The rstd sites' `(1 + w)` operands (`_HC_RSTD_MIN_ROWS`), formed once here. - self._layers = [ - L._replace(ahc_norm_w1=_norm_w1(L.ahc_norm), mhc_norm_w1=_norm_w1(L.mhc_norm)) - for L in _layer_views(core) - ] - stack = _stack_views(core) - self._stack = stack._replace(fhc_norm_w1=_norm_w1(stack.fhc_norm)) - # The draft sites' `(1 + w)` operands, by norm weight, formed on first use. - self._mtp_norm_w1: dict = {} - # The step `prepare` built, read by the breaks. - self._step: _PrefillStep | None = None - - def _check_step_contract(self, md) -> None: - """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): calling - `step_args` surfaces an upstream metadata rename here, not mid-forward.""" - if not self._contract_pending: - return - try: - self.step_args(md) - except AttributeError as exc: - raise AssertionError(f"attention metadata surface drifted: {exc}") from exc - self._contract_pending = False - - def step_args(self, md: TrtllmAttentionMetadata) -> dict: - """Project the prepared metadata onto thop_attention's batch state, once per - forward. Tensors are engine-owned persistent buffers refreshed in place; - Python ints are per-capture constants.""" - return dict( - sequence_length=md.kv_lens_cuda_runtime, - host_past_key_value_lengths=md.kv_lens_runtime, - host_total_kv_lens=md.host_total_kv_lens, - context_lengths=md.prompt_lens_cuda_runtime, - host_context_lengths=md.prompt_lens_cpu_runtime, - host_request_types=md.host_request_types_runtime, - kv_cache_block_offsets=md.kv_cache_block_offsets, - host_kv_cache_pool_pointers=md.host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=md.host_kv_cache_pool_mapping, - workspace=md.effective_workspace, - tokens_per_block=md.tokens_per_block, - max_num_requests=md.max_num_requests, - max_context_length=md.max_context_length, - max_seq_len=md.max_seq_len, - # No sliding window: the full admitted length on every attention layer. - attention_window_size=md.max_seq_len, - num_contexts=md.num_contexts, - num_ctx_tokens=md.num_ctx_tokens, - trtllm_gen_jit_warmup=md.trtllm_gen_jit_warmup, - use_paged_context_fmha=md.use_paged_context_fmha, - beam_width=md.effective_beam_width, - cache_indirection=md.cache_indirection, - block_ids_per_seq=md.block_ids_per_seq, - max_context_q_len_override=md.max_context_q_len_override, - is_cross=md.is_cross, - is_spec_decoding_enabled=md.is_spec_decoding_enabled, - use_spec_decoding=md.use_spec_decoding, - is_spec_dec_tree=md.is_spec_dec_tree, - spec_decoding_generation_lengths=md.spec_decoding_generation_lengths, - spec_decoding_position_offsets=get_spec_decoding_position_offsets(md), - spec_decoding_packed_mask=md.spec_decoding_packed_mask, - spec_decoding_bl_tree_mask_offset=md.spec_decoding_bl_tree_mask_offset, - spec_decoding_bl_tree_mask=md.spec_decoding_bl_tree_mask, - spec_decoding_target_max_draft_tokens=md.max_total_draft_tokens, - spec_bl_tree_first_sparse_mask_offset_kv=md.spec_bl_tree_first_sparse_mask_offset_kv, - num_sparse_topk=md.num_sparse_topk, - flash_mla_tile_scheduler_metadata=md.flash_mla_tile_scheduler_metadata, - flash_mla_num_splits=md.flash_mla_num_splits, - max_num_sequences=md.max_num_sequences, - force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask, - ) - - # --------------------------------------------------------------- blocks --- - - def _hc_mixer(self, normed, first, up, with_inject): - """One hyper-connection mixer from the site's normed stream: the first - projection over `[down; block_inject; 0]` (fp32 logits), the SiLU and - injection activations, then the up projection with the sigmoid gate and - the stream mean in its epilogue (the gate logits never leave the GEMM). - Returns `(mixed, inject)`; `inject` is None without `with_inject` (the - stack-final mixer).""" - core = self.core - low, inject = qwen4_exp_hc_lowrank_act( - cublas_mm(normed, first[:, : core.hc_first_cols], None, torch.float32), - core.hc_lowrank, - core.hc_count, - with_inject, - normed.dtype, - ) - return qwen4_exp_hc_up_gate_mean(low, up, normed, core.hc_count), inject - - def _hc_moe_mixer(self, normed, first, up, gscale): - """`_hc_mixer` at an MoE site (always with injection) whose fused up - projection also writes the routed experts' NVFP4 input (linear scales, - the layer's global scale `gscale`). Returns `(mixed, inject, xq, xsf)`.""" - core = self.core - low, inject = qwen4_exp_hc_lowrank_act( - cublas_mm(normed, first[:, : core.hc_first_cols], None, torch.float32), - core.hc_lowrank, - core.hc_count, - True, - normed.dtype, - ) - mixed, xq, xsf = qwen4_exp_hc_up_gate_mean_fp4(low, up, normed, core.hc_count, gscale) - return mixed, inject, xq, xsf - - def _hc_mixer_rstd(self, hc, rstd, norm, norm_w1, first, up, with_inject): - """`_hc_mixer` from the site's residual `hc` and its `[T, C]` rstd - (`qwen4_exp_hc_inject_rstd`) instead of a stored normed stream: the first - projection forms `bf16(hc * rstd * (1 + norm))` on its A tiles, the up - projection's epilogue `hc * rstd * norm_w1` in fp32.""" - core = self.core - low, inject = qwen4_exp_hc_first_rescale_tc( - hc, rstd, norm, first, core.hc_count, core.hc_lowrank, with_inject - ) - mixed = qwen4_exp_hc_up_gate_mean_rstd(low, up, hc, rstd, norm_w1, core.hc_count) - return mixed, inject - - def _hc_moe_mixer_rstd(self, hc, rstd, norm, norm_w1, first, up, gscale): - """`_hc_moe_mixer` on the rstd form (see `_hc_mixer_rstd`). Returns - `(mixed, inject, xq, xsf)`.""" - core = self.core - low, inject = qwen4_exp_hc_first_rescale_tc( - hc, rstd, norm, first, core.hc_count, core.hc_lowrank, True - ) - mixed, xq, xsf = qwen4_exp_hc_up_gate_mean_rstd_fp4( - low, up, hc, rstd, norm_w1, core.hc_count, gscale - ) - return mixed, inject, xq, xsf - - def _grouped_norm(self, x: torch.Tensor, weights: tuple) -> torch.Tensor: - """`Qwen4ExpTextRMSNorm(group_size=hidden)`: an independent RMS per `hidden`-wide - group, then `(1 + w)` over the full width; one gemma-RMS-norm per column - slice (row stride certified by that entry) plus a concat.""" - core = self.core - parts = torch.split(x, [core.hidden] * core.hc_count, 1) - return torch.cat( - [flashinfer_gemma_rmsnorm(part, w, core.eps) for w, part in zip(weights, parts)], 1 - ) - - def _moe_block(self, x, xq, xsf, L: _LayerWeights, routed_out) -> tuple: - """One MoE layer: NVFP4 routed experts plus a bf16 shared expert, left - uncombined. Returns `(routed, (gate_logits, shared))`: the next - `qwen4_exp_hc_inject_norm` forms `routed + sigmoid(gate_logits) * shared` - inside its kernel. - - The shared expert and its per-token gate logit run on the core's side - stream, concurrently with the routed path, and join it before this - returns, so every consumer on the main stream is ordered after both. - `x` is the bf16 hidden states (router and shared expert); `xq` / `xsf` - are the same states in NVFP4 with **linear** block scales, written by the - MoE-site mixer (`_hc_moe_mixer`). On trtllm-gen the routed path hands the - runner the router logits (its Renormalize routing is bitwise the - `renorm_moe_routing_op` + pre-routed composition), and a launch of at - least `_CUTE_DSL_MOE_MIN_ROWS` rows runs the CuteDSL MoE over the second - expert copy instead when the core keeps it (`moe_prefill_copy`); on CuteDSL - every launch is `_cute_dsl_routed`, which writes `routed_out`.""" - core = self.core - main = torch.cuda.current_stream() - side = core._side_stream - # Fork: the side stream reads `x` only after everything queued on `main` - # so far. What it allocates is consumed on `main` after the join below, - # and the next layer's fork orders any reuse of those blocks after that. - side.wait_stream(main) - with torch.cuda.stream(side): - # One GEMM forms gate|up and the gate logit (`sh_gug`'s extra rows): a separate - # gate GEMV at prefill sizes keeps the SMs from the routed path's persistent - # FC1 until it drains. - gug = cublas_mm(x, L.sh_gug) - shared = cublas_mm(silu_and_mul_rows(gug, core.shared_inter), L.sh_dn) - gate_logits = gug[:, 2 * core.shared_inter : 2 * core.shared_inter + 1] - # The routing reads its logits as a dense row-major buffer from data_ptr() - # and ignores strides; cublas_mm allocates them contiguous. - logits = cublas_mm(x, L.router) - if core.trunk_moe_backend == "CUTEDSL": - routed = _cute_dsl_routed(core, L, logits, xq, xsf, routed_out) - elif core.moe_prefill_copy and x.shape[0] >= _CUTE_DSL_MOE_MIN_ROWS: - routed = fp4_moe_renorm_cute_dsl( - logits, - xq, - xsf, - L.fc1_w_cd, - L.fc1_s_cd, - L.alpha1_cd, - L.fc2_w_cd, - L.fc2_s_cd, - L.alpha2_cd, - L.global_sf_cd, - core.num_experts, - core.topk, - memset_stream=core._moe_memset_stream, - ) - else: - routed = fp4_block_scale_moe_renorm_routed( - logits, - xq, - xsf.view(torch.float8_e4m3fn), - L.fc1_w, - L.fc1_s, - L.fc2_w, - L.fc2_s, - L.out1_scale, - L.out1_gate_scale, - L.out2_scale, - core.num_experts, - core.topk, - core.moe_inter, - ) - main.wait_stream(side) - return routed, (gate_logits, shared) - - def _attention_block(self, x, i, L, proj_qkv, index_scratch, attn_out, pos3): - """One gated GQA layer plus its QSA index branch. - - The four input projections are one GEMM against the load-time fused - operand (`weights._attn_in_w`), whose output row is - `[q | k | v | gate | index q, k | v slot]`. One norm+rope launch rewrites - the main q/k heads in place on that output and writes the normalized, - rotated index query into the scratch; v, the gate and the raw index key - are not touched. Both are row-wise and run over all `T` rows here; the - cache writes, the selection and the attention core read the step's batch - and run in the `_attention_core` break, which leaves the attention output - in `attn_out`. The gate is read through its row stride.""" - core = self.core - num_tokens = x.shape[0] - - proj = _dense_mm(x, L.attn_in, L.attn_in8, cublas_mm, "in") - gate = torch.split(proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1)[1] - # The index query lands in the q slots of the scratch; its k and v slots - # are layout only and nothing reads them. - q_index = torch.split( - torch.reshape( - index_scratch, - [num_tokens, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim], - ), - [core.index_heads, core.index_kv_heads, core.index_kv_heads], - 1, - )[0] - # Both planes' Gemma norm + rope in one launch, over the target's fp32 - # cos|sin table: the main q/k heads in place (the real interleaved mRoPE; - # three agreeing rows reduce exactly to the plain rope on row 0), the index - # query from the projection's index columns into `q_index` (coordinate 0 - # alone, matching `_QSA_MROPE_SECTION` on the key side). - index_section = (_QSA_MROPE_SECTION or (0, 0, 0)) if _QSA_INDEX_USE_MROPE else (0, 0, 0) - qwen4_exp_qk_norm_rope( - proj, - q_index, - pos3, - core._cos_sin, - L.q_norm, - L.k_norm, - L.idx_qn, - num_heads_q=core.heads_q, - num_heads_k=core.heads_kv, - head_dim=core.head_dim, - index_heads=core.index_heads, - index_head_dim=core.index_head_dim, - index_q_col=core.qkv_width + core.q_width, - rotary_dim=core.rotary_dim, - eps=core.eps, - mrope_section1=core.mrope_section1, - mrope_section2=core.mrope_section2, - index_mrope_section1=index_section[1], - index_mrope_section2=index_section[2], - ) - if self._kv_probe is not None: - _, k_dst, v_dst = _attn_in_heads(core, proj) - if not self._kv_probe.observe(i, k_dst, v_dst): - self._kv_probe = None - eager_on_graph(proj, proj_qkv, index_scratch, attn_out, i, inner=self._attention_core) - fused_sigmoid_mul_inplace(attn_out, gate) - return _dense_mm(attn_out, L.o, L.o8, cublas_mm, "out") - - def _attention_core(self, proj, proj_qkv, index_scratch, attn_out, i) -> None: - """Break: the index plane's store and compression, the dense/sparse decision and the - attention core over the step's `n` rows. Writes rows `[:n]` of `attn_out` - `[T, q_width]` in place and zeroes rows `[n:]`. `proj` is the layer's fused - projection output after the norm+rope, `proj_qkv` the dense op's packed QKV - scratch. - - The index plane advances on every step, including below the threshold, so a - request growing past it finds its prefix compressed; above it the main K/V append - becomes this layer's own call. The raw index key is stored before any - norm/rotation: the compression step applies both itself, so a pre-normalized key - is double-normalized (invisible below the sparse threshold). - - On a verify step (`stage`) the generation rows' index slots (and, on the first - attention layer, their position slots) are snapshotted for `SpeculativeRollback` - before the store and compression overwrite them.""" - core = self.core - L = self._layers[i] - step = self._step - ( - req_idx, - logical, - seq_lens_q, - visible_blocks, - topk_indices, - topk_row_starts, - block_table, - tokens_per_block, - coords, - sparse, - live_cols, - ) = step.qsa - num_tokens = step.num_tokens - proj = proj[:num_tokens] - qkv_view, _, index_view = torch.split( - proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1 - ) - q_dst, k_dst, v_dst = _attn_in_heads(core, proj) - - index_cache = core._bound.index_cache[i] - token_k = torch.split( - torch.reshape( - index_view, - [num_tokens, core.index_scratch_width // core.index_head_dim, core.index_head_dim], - ), - [core.index_heads, core.index_kv_heads, 1], - 1, - )[1] - if step.stage: - core.spec_rollback.stage_index(index_cache, core._bound.position_cache) - qsa_paged_index_store( - token_k, - coords, - index_cache, - core._bound.position_cache, - req_idx, - logical, - block_table, - tokens_per_block, - ) - q_index = torch.split( - torch.reshape( - index_scratch[:num_tokens], - [num_tokens, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim], - ), - [core.index_heads, core.index_kv_heads, core.index_kv_heads], - 1, - )[0] - qsa_prefill_compress( - logical_positions=logical, - request_indices=req_idx, - block_table=block_table, - index_cache=index_cache, - position_cache=core._bound.position_cache, - k_norm_weight=L.idx_kn, - cos_sin=core._cos_sin, - eps=core.eps, - tokens_per_block=tokens_per_block, - compress_ratio=core.compress_ratio, - mrope_section=_QSA_MROPE_SECTION, - ) - - if sparse: - kv = core._bound.kv[i] - # Append the fused-norm+rope'd K/V here: nothing else writes the pool on - # this path, and selection reads the current token from it. Over an fp8 pool the - # store and the sparse GQA below take e4m3 unscaled, as the dense op does at the - # unit KV scales; over an NVFP4 pool they quantize and dequantize at the dense op's - # second-level scales. - _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) - # The step's context rows come first, consecutive rows of a request, and score - # as query tiles; its generation rows, one a request, score row by row (a query - # tile of them would walk every request's keys once per row). Each part selects - # into its own rows of the top-k. - num_ctx = step.attn["num_ctx_tokens"] - parts = [(0, num_ctx, True)] - if num_ctx < num_tokens: - parts.append((num_ctx, num_tokens, False)) - # Each part's canonical block top-k and its expansion are one launch writing its - # rows of the step's selection (QSA's row starts are zero); a score table that - # kernel does not serve (wider than its 4096 columns) takes the composition. - selected = torch.empty( - [num_tokens, core.token_topk + core.compress_ratio - 1], - dtype=torch.int32, - device=proj.device, - ) - for lo, hi, context_rows in parts: - scores = qsa_paged_index_scores( - q=q_index[lo:hi], - index_cache=index_cache, - block_table=block_table, - query_positions=logical[lo:hi], - request_indices=req_idx[lo:hi], - tokens_per_block=tokens_per_block, - compress_ratio=core.compress_ratio, - only_visible_blocks=True, - context_rows=context_rows, - ) - if _topk_expand_supports(scores): - qsa_topk_expand_decode( - scores, - visible_blocks[lo:hi], - logical[lo:hi], - seq_lens_q[lo:hi], - topk_indices[lo:hi], - core.compress_ratio, - core.token_topk, - out=selected[lo:hi], - ) - else: - _select( - core, - scores, - topk_row_starts[lo:hi], - visible_blocks[lo:hi], - topk_indices[lo:hi], - live_cols, - ) - selected[lo:hi].copy_( - qsa_expand_block_indices( - topk_indices[lo:hi], - logical[lo:hi], - seq_lens_q[lo:hi], - compress_ratio=core.compress_ratio, - token_topk=core.token_topk, - ) - ) - if kv.k_scale is not None: - _sparse_gqa_nvfp4_context( - core, - kv, - step.nvfp4_context, - q_dst, - block_table, - selected, - req_idx, - logical, - tokens_per_block, - num_ctx, - attn_out[:num_tokens], - rows_per_request=step.attn["predicted_tokens_per_seq"], - ) - elif num_tokens >= _ROWTILE_MIN_ROWS and kv.k.dtype == q_dst.dtype: - # A context launch of thousands of rows over a bf16 pool takes the row-tile - # pass: consecutive rows of one request share their K/V tiles; rows of - # groups that span requests, the generation rows of a mixed step among them, - # are served row by row. Below the threshold the per-row kernel is faster. - out = qsa_sparse_gqa_rowtile( - q=q_dst, - k_cache=kv.k, - v_cache=kv.v, - block_table=block_table, - selected_tokens=selected, - request_indices=req_idx, - tokens_per_block=tokens_per_block, - softmax_scale=core.softmax_scale, - query_positions=logical, - compress_ratio=core.compress_ratio, - ) - attn_out[:num_tokens].copy_(torch.reshape(out, [num_tokens, core.q_width])) - else: - out = _sparse_gqa( - core, - kv, - q=q_dst, - block_table=block_table, - selected_tokens=selected, - request_indices=req_idx, - tokens_per_block=tokens_per_block, - # A bound, not a mask: lets the kernel stop scanning a row early - # (context-batch branch only). - query_positions=logical, - compress_ratio=core.compress_ratio, - ) - attn_out[:num_tokens].copy_(torch.reshape(out, [num_tokens, core.q_width])) - else: - # The dense op reads its packed QKV flat: one compact copy. - qkv = proj_qkv[:num_tokens] - qkv.copy_(qkv_view) - thop_attention( - q=qkv, - output=attn_out[:num_tokens], - local_layer_idx=core._bound.pool_idx[i], - num_heads=core.heads_q, - num_kv_heads=core.heads_kv, - head_size=core.head_dim, - quant_mode=core.quant_mode, - kv_scale_orig_quant=core.kv_scale_orig_quant, - kv_scale_quant_orig=core.kv_scale_quant_orig, - **step.attn, - **_CALL_CONSTANTS, - ) - attn_out[num_tokens:].zero_() - - def _gdn_verify( - self, mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, n, p, replay - ) -> torch.Tensor: - """The recurrent half of a speculative verify step: `n` generation sequences of - `p = runtime_draft_len + 1` tokens, written so that the accepted prefix, unknown - until the sampler has run, can be kept and the rest dropped. - - The conv writes each row's window to the runtime's per-step buffer and leaves the - conv pool alone; `update_mamba_states` promotes window `k - 1`. - `causal_conv1d_triton_update` is the only conv with per-step snapshots (its result - for a pad-sentinel slot is uninitialized, safe because every row is a real - sequence). The recurrence has two forms, bound per manager - (`_bind_step_buffers`): - - * cached replay (`replay` given): the pool row is a checkpoint; one call rebuilds - the current state from it and the slot's pending history, runs the `p` rows and - appends them to the history. Rejected rows cost nothing to undo, because the - manager advances the pending count by the accepted tokens only. A full history - advances the checkpoint inside the kernel below - `CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE` requests and in the manager's - all-layer commit from there on, hence `replay` carries the work items exactly - under the manager's own condition. - * legacy (`replay` None): every row's recurrent state goes to the per-step buffer - instead of the pool, and the promotion copies step `k - 1` into it. Advancing - the pool directly would leave rejected tokens' state in every recurrent layer, - invisible to the acceptance rate. `pack_gdn_decode_qkv` does no Q/K L2 norm, - hence `use_qk_l2norm_in_kernel=True` below.""" - core = self.core - # [n*p, C] -> [n, C, p]: the multi-token conv reads channel-major; its cells - # certify this transposed view directly. - x3 = torch.transpose(torch.reshape(mq_d, [n, p, core.conv_dim]), 1, 2) - conv_out = causal_conv1d_triton_update( - x=x3, - conv_state=core._bound.conv[i], - weight=L.conv_w, - bias=None, - activation="silu", - conv_state_indices=si_d, - pad_slot_id=_PAD_SLOT_ID, - # Snapshots, not the accepted-token roll. - num_accepted_tokens=None, - intermediate_conv_window=core._bound.int_conv[i], - intermediate_state_indices=int_idx, - launch_dependent_kernels=False, - ) - # Back to token-major (a view). - mq_out = torch.reshape(torch.transpose(conv_out, 1, 2), [n * p, core.conv_dim]) - if replay is not None: - return _replay_verify(core, mq_out, a_d, b_d, ssm_states, i, L, si_d, n, p, replay) - q, k, v = pack_gdn_decode_qkv( - mq_out, core.num_k_heads, core.head_k_dim, core.num_v_heads, core.head_v_dim - ) - g = torch.reshape( - fused_gdn_gating( - A_log=L.a_log, - a=a_d, - dt_bias=L.dt_bias, - beta=_SOFTPLUS_BETA, - threshold=_SOFTPLUS_THRESHOLD, - ), - [n, p, core.num_v_heads], - ) - beta = torch.reshape(torch.sigmoid(b_d), [n, p, core.num_v_heads]) - # Initial states are gathered from the live pool into a batch-scoped copy. The - # per-step buffer is indexed by batch row (`intermediate_state_indices[:n]`), - # not slot id: slot ids would write rows nothing promotes, or past the end. - src = torch.reshape( - nn.functional.embedding(si_d, torch.reshape(ssm_states, [ssm_states.shape[0], -1])), - [n, core.num_v_heads, core.head_v_dim, core.head_k_dim], - ) - out = fused_recurrent_gated_delta_rule_update( - torch.reshape(q, [n, p, core.num_k_heads, core.head_k_dim]), - torch.reshape(k, [n, p, core.num_k_heads, core.head_k_dim]), - torch.reshape(v, [n, p, core.num_v_heads, core.head_v_dim]), - g, - beta, - initial_state_source=src, - initial_state_indices=int_idx, - scale=None, - use_qk_l2norm_in_kernel=True, - cu_seqlens=None, - disable_state_update=True, - disable_output_calculation=False, - intermediate_states_buffer=core._bound.int_ssm[i], - cache_steps=p, - output=None, - ) - return torch.reshape(out, [1, n * p, core.num_v_heads, core.head_v_dim]) - - def _gdn_block(self, x, i, L, core_out) -> torch.Tensor: - """One Gated-DeltaNet layer over a batch with a context prefix: the one input - projection GEMM (q|k|v, z, a and b are column views of its output) and the - gated-norm epilogue run over all `T` rows here, the recurrence in the `_gdn_core` - break, which leaves its output in `core_out` `[1, T, num_v_heads, head_v_dim]`.""" - core = self.core - proj = _dense_mm(x, L.in_proj, L.in_proj8, cublas_mm, "in") - z = torch.split(proj, core.in_proj_split, 1)[1] - eager_on_graph(proj, core_out, i, inner=self._gdn_core) - return self._gdn_epilogue(core_out, z, L, x.shape[0]) - - def _gdn_core(self, proj, out, i) -> None: - """Break: the recurrence over the step's `n` rows. Writes rows `[:n]` of `out` - `[1, T, num_v_heads, head_v_dim]` in place and zeroes rows `[n:]`. Every branch - is on a host int of the step, never a tensor value. - - Outside a verify step, over either pool dtype, the context rows take one fused - conv + post-conv pass (`_gdn_prefill_prep`) and FlashInfer's chunked recurrence, and - a mixed step's one-token generation rows the decode step's recurrence - (`gdn_decode_rows`); a verify step takes `_gdn_core_verify`. A context sequence - carrying nothing starts from zeros without its pool slots being cleared: the conv - pass reads no carried window for it and publishes the whole window, and the - recurrence (`has_initial_state`) gathers zeros for it and scatters its final state.""" - core = self.core - L = self._layers[i] - step = self._step - ( - num_prefill_tokens, - num_decode_tokens, - si_p, - si_all, - si_d, - hi_p, - hi_all, - qsl_p, - qsl_all, - qsl_long, - is_verify, - num_decodes, - gen_p, - int_idx, - qsl_long_p, - replay, - ) = step.gdn - num_tokens = step.num_tokens - if is_verify: - self._gdn_core_verify(proj, out, i) - return - ssm_states = core._bound.ssm[i] - mixed_qkv, _, a, b = torch.split(proj[:num_tokens], core.in_proj_split, 1) - conv_states = core._bound.conv[i] - - if num_decode_tokens > 0: - mq_p, mq_d = torch.split(mixed_qkv, [num_prefill_tokens, num_decode_tokens], 0) - a_p, b_p = a[:num_prefill_tokens], b[:num_prefill_tokens] - q, k, v, g, beta = self._gdn_prefill_prep(mq_p, a_p, b_p, L, i, qsl_p, si_p, hi_p) - causal_conv1d_update( - x=mq_d, - conv_state=conv_states, - weight=L.conv_w, - bias=None, - silu_activation=True, - cache_seqlens=None, - conv_state_indices=si_d, - pad_slot_id=_PAD_SLOT_ID, - ) - # The two halves run separate recurrences over disjoint slots and write - # disjoint row ranges of the output. - out_p = out[:, :num_prefill_tokens] - flashinfer_chunk_gated_delta_rule( - q, - k, - v, - g, - beta, - scale=None, - initial_state=ssm_states, - initial_state_indices=si_p, - inplace_indexed_state_update=True, - output_final_state=False, - cu_seqlens=qsl_long_p, - head_first=False, - use_qk_l2norm_in_kernel=False, - output=out_p, - state_workspace=None, - has_initial_state=hi_p, - linear_gate=True, - ) - # The generation rows: the decode step's recurrence (its kernel without the - # gated norm, which runs with the context rows' below) over their rows of the - # projection, whose q/k/v columns the conv step rewrote; it computes both gates - # and the q/k norm itself and reads the rows' pool slots before its PDL wait - # (the context rows' recurrence touches only theirs). - gdn_decode_rows( - proj[num_prefill_tokens:num_tokens], - ssm_states, - si_d, - L.a_log, - L.dt_bias, - core.num_k_heads, - core.head_k_dim, - out[0, num_prefill_tokens:num_tokens].view(num_decode_tokens, core.value_dim), - ) - else: - q, k, v, g, beta = self._gdn_prefill_prep( - mixed_qkv, a, b, L, i, qsl_all, si_all, hi_all - ) - # q and k arrive L2-normalized; the recurrence must not repeat it. - flashinfer_chunk_gated_delta_rule( - q, - k, - v, - g, - beta, - scale=None, - initial_state=ssm_states, - initial_state_indices=si_all, - inplace_indexed_state_update=True, - output_final_state=False, - cu_seqlens=qsl_long, - head_first=False, - use_qk_l2norm_in_kernel=False, - output=out[:, :num_tokens], - state_workspace=None, - has_initial_state=hi_all, - linear_gate=True, - ) - out[:, num_tokens:].zero_() - - def _gdn_core_verify(self, proj, out, i) -> None: - """`_gdn_core` on a verify step: the context rows' fused conv + post-conv pass and - FlashInfer's chunked delta rule (as without verify rows), and the verify path for the - generation rows. A context row - advances its slot as on the base path; a generation row is a draft chain of unknown - accepted length, so it goes to the rollback buffers. The q|k|v columns are read in - place, as the column view of the projection every consumer here is certified on; - the a and b columns are copied out first, as the verify path reads them contiguous. - The chunked delta rule writes its rows of `out` directly.""" - core = self.core - L = self._layers[i] - step = self._step - ( - num_prefill_tokens, - num_decode_tokens, - si_p, - si_all, - si_d, - hi_p, - hi_all, - qsl_p, - qsl_all, - qsl_long, - is_verify, - num_decodes, - gen_p, - int_idx, - qsl_long_p, - replay, - ) = step.gdn - num_tokens = step.num_tokens - mixed_qkv, _, a, b = torch.split(proj[:num_tokens], core.in_proj_split, 1) - a, b = a.contiguous(), b.contiguous() - - ssm_states = core._bound.ssm[i] - - # A context sequence carrying nothing starts from zeros (`has_initial_state` - # below; its slots are not cleared first, see `_gdn_core`). - mq_p, mq_d = torch.split(mixed_qkv, [num_prefill_tokens, num_decode_tokens], 0) - a_p, a_d = torch.split(a, [num_prefill_tokens, num_decode_tokens], 0) - b_p, b_d = torch.split(b, [num_prefill_tokens, num_decode_tokens], 0) - # The context rows as on a step without verify rows: one fused conv + post-conv - # pass and FlashInfer's chunked recurrence (q and k arrive L2-normalized). - q, k, v, g, beta = self._gdn_prefill_prep(mq_p, a_p, b_p, L, i, qsl_p, si_p, hi_p) - flashinfer_chunk_gated_delta_rule( - q, - k, - v, - g, - beta, - scale=None, - initial_state=ssm_states, - initial_state_indices=si_p, - inplace_indexed_state_update=True, - output_final_state=False, - cu_seqlens=qsl_long_p, - head_first=False, - use_qk_l2norm_in_kernel=False, - output=out[:, :num_prefill_tokens], - state_workspace=None, - has_initial_state=hi_p, - linear_gate=True, - ) - out_d = self._gdn_verify( - mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, num_decodes, gen_p, replay - ) - out[:, num_prefill_tokens:num_tokens].copy_(out_d) - out[:, num_tokens:].zero_() - - def _gdn_prefill_prep(self, mq, a, b, L, i, qsl, si, hi): - """The context rows' recurrence inputs: the causal conv against each sequence's - carried window (published back to the pool), SiLU, the q|k|v split with q/k - L2-normalized, and the decay/beta gates (the decay as the factor `exp(g)`, which - the chunked recurrence takes with `linear_gate=True`) -- one pass over `mq`, the - token-major q|k|v column view of the projection.""" - core = self.core - return gdn_prefill_conv_post( - x=mq, - weight=L.conv_w, - conv_states=core._bound.conv[i], - query_start_loc=qsl, - cache_indices=si, - has_initial_state=hi, - a=a, - b=b, - A_log=L.a_log, - dt_bias=L.dt_bias, - num_k_heads=core.num_k_heads, - head_k_dim=core.head_k_dim, - num_v_heads=core.num_v_heads, - head_v_dim=core.head_v_dim, - pad_slot_id=_PAD_SLOT_ID, - l2_norm_eps=_L2_NORM_EPS, - softplus_threshold=_SOFTPLUS_THRESHOLD, - linear_gate=True, - ) - - def _gdn_epilogue(self, out, z, L, num_tokens) -> torch.Tensor: - """Gated RMS norm and output projection shared by every recurrence.""" - core = self.core - # Sigmoid gate and a plain (non-gemma) weight (`output_gate_type`). - o = rms_norm_gated_token_major( - torch.reshape(out, [-1, core.head_v_dim]), - torch.reshape(z, [num_tokens, core.num_v_heads, core.head_v_dim]), - weight=L.gnorm, - eps=core.eps, - fp8_scale=None, - gate_activation="sigmoid", - ) - o = torch.reshape(o, [num_tokens, core.value_dim]) - return _dense_mm(o, L.out, L.out8, cublas_mm, "out") - - def _ple_block(self, hc, hc_out, normed_out) -> None: - """Break: the n-gram side path on one decoder layer, over the step's `n` rows, - added into the stream and normalized for the layer's attention-side site. Writes - rows `[:n]` of `hc_out` (the stream) and `normed_out` (its grouped norm), both - `[T, stream]`, in place and zeroes rows `[n:]`. - - Per-request stateful: ids become one `ngram_size`-wide window per token - (predecessors from the paged id history); each window hashes to one row per - head; rows gather fp8 vectors from one table and dequantize by one per-tensor - scale; key/value projections and the gate scale one value into each stream; - a dilated depthwise short conv over a paged state is added back. A sequence - carrying nothing is seeded with `eos_token_id`, not 0: id 0 is a real token - the hash cannot distinguish. The id history and the conv state publish the - state after the whole chunk. - - The window, hash, gather and dequantize are `ple_ngram_embed`, and the gate, the - conv, the residual add and the next site's grouped norm `ple_gate_conv_inject_norm`, - which writes both outputs in place. On a verify step (`stage`) the generation rows' - pre-step pool rows and ids are staged before the first launch publishes, and the - conv input the fused launch writes out after it, for `SpeculativeRollback` to - rewrite both pools to the accepted prefix once the sampler has decided it.""" - core = self.core - S = self._stack - step = self._step - conv_states, ngram_pool, si_all, carry_all, qsl = step.ple - num_tokens = step.num_tokens - next_norm = self._layers[core.ple_layer].ahc_norm - hc = hc[:num_tokens] - conv_input = None - if step.stage: - core.spec_rollback.stage_ple_verify(ngram_pool, conv_states, step.input_ids) - conv_input = torch.empty_like(hc) - ngram = ple_ngram_embed( - step.input_ids, - ngram_pool, - si_all, - carry_all, - qsl, - core.w["ple_table"], - S.ple_tscale, - S.ple_mult, - S.ple_vocab, - S.ple_offs, - core.eos_token_id, - num_tokens, - ) - ple_gate_conv_inject_norm( - hc, - cublas_mm(ngram, S.ple_key), - cublas_mm(ngram, S.ple_val), - S.ple_nk_w, - S.ple_nq_w, - S.ple_nc_w, - S.ple_conv_w, - conv_states, - si_all, - qsl, - next_norm, - core.eps, - core.hc_count, - _PLE_CLAMP_MIN, - core.ple_dilation, - _PAD_SLOT_ID, - num_tokens, - carry_all, - conv_input=conv_input, - out=(hc_out[:num_tokens], normed_out[:num_tokens]), - ) - if step.stage: - core.spec_rollback.stage_ple_conv_input(conv_input) - hc_out[num_tokens:].zero_() - normed_out[num_tokens:].zero_() - - def _splice(self, h) -> None: - """Break: write the step's vision rows over their placeholder rows of the token - embedding `h` `[T, hidden]` in place, and zero rows `[n:]`. - - `mm_token_indices` names every placeholder row of the packed batch in the order the - per-request rows concatenate in. An indexed write, not a slice copy: two images in - one prompt are separated by framing tokens, so the placeholder rows are several - disjoint runs. A text step writes nothing but the padding. The rows are the step's: - a replayed graph never sees the previous step's.""" - step = self._step - if step.vision is not None: - rows, mm_token_indices = step.vision - h.index_copy_(0, mm_token_indices, rows) - h[step.num_tokens :].zero_() - - # -------------------------------------------------------------- forward --- - - def prepare( - self, - attn_metadata: AttentionMetadata, - input_ids: torch.IntTensor, - position_ids: torch.IntTensor, - vision: tuple | None, - spec_metadata, - ) -> None: - """Everything this step reads from the host, kept for the forward's breaks as one - `_PrefillStep`: the contract check, the attention op's per-step arguments and the - batch's bookkeeping, built once and shared by every layer. - - Called by the core's `begin_step` on every step with context rows, before the - forward. On a step a breakable prefill graph replays, this and the breaks are all of - this target that runs.""" - core = self.core - md = attn_metadata - self._check_step_contract(md) - attn_args = self.step_args(md) - mm = md.mamba_metadata - - # **The activation rows are the batch's rows, or a prefill CUDA graph bucket's.** A - # breakable prefill graph pads the step's rows up to `padded_num_tokens` (zero ids, - # zero positions), and so does a padded step that runs eagerly; the forward runs over - # all of them and every break slices the batch's own rows back out. Anything else is - # an activation buffer longer than the batch, which would push padding rows through - # the recurrences into the state pools. - physical_tokens = input_ids.shape[0] - assert physical_tokens in (md.num_tokens, md.padded_num_tokens), ( - f"the engine handed {physical_tokens} activation rows for a batch its " - f"metadata describes as {md.num_tokens} tokens " - f"(padded_num_tokens={md.padded_num_tokens})" - ) - - # The batch's bookkeeping, built once and shared by every layer. - num_tokens = md.num_tokens - batch = md.num_seqs - num_prefills = md.num_contexts - num_decodes = batch - num_prefills - num_prefill_tokens = md.num_ctx_tokens - num_decode_tokens = num_tokens - num_prefill_tokens - # Rows per generation sequence: 1 without speculation, else - # `runtime_draft_len + 1`; read from the metadata because the worker rewrites - # it between draft steps. - gen_p = num_decode_tokens // num_decodes if num_decodes else 1 - assert gen_p * num_decodes == num_decode_tokens, ( - f"{num_decode_tokens} generation rows do not divide over " - f"{num_decodes} generation sequences; a ragged per-sequence draft " - "length can be expressed neither to the attention call nor to the " - "multi-token recurrence" - ) - assert core.mtp_enabled or gen_p == 1, ( - f"{num_decode_tokens} generation rows over {num_decodes} generation " - "sequences with speculation off; the base path implements one query " - "token per generation sequence, and both the attention call's " - "within-block mask and the recurrent path's cumulative-length " - "vector assume it" - ) - # `gen_p > 1` is load-bearing: verify parks a generation row in the rollback - # buffers, which `update_mamba_states` promotes only if the worker drafts. - # With `runtime_draft_len == 0` (a `draft_len_schedule` can do this) it skips - # drafting before the promotion, leaving every recurrent layer one token - # stale; that row is the always-accepted golden token, so the ordinary decode - # path is right. - if spec_metadata is not None and num_decodes > 0: - declared_p = spec_metadata.runtime_draft_len + 1 - assert gen_p == declared_p, ( - f"the batch carries {gen_p} query rows per generation sequence " - f"but the worker's runtime_draft_len says {declared_p}; the " - "per-step conv snapshots and the recurrent rollback buffer are " - "sized from the second and written from the first" - ) - # Per-iteration counterpart of the core's `max_draft_len` bound. - check_verify_width(gen_p) - check_replay_width(core._bound.replay, gen_p) - is_verify = spec_metadata is not None and num_decodes > 0 and gen_p > 1 - attn_args["predicted_tokens_per_seq"] = gen_p - - si_all = mm.state_indices[:batch] - hi_all = mm.has_initial_states[:batch] - si_p, si_d = torch.split(si_all, [num_prefills, num_decodes], 0) - # The single-token decode kernel takes its slot vector through a 16-byte-aligned - # pointer, and the generation half of the int32 slot vector starts `num_prefills` - # rows in, aligned only when that count is a multiple of four: one copy per step - # makes it aligned for every layer. - si_d = torch.clone(si_d) - hi_p = hi_all[:num_prefills] - # `query_start_loc` covers the context prefix; `_long` the whole batch, and its - # prefix the context rows' chunked recurrence when generation rows run beside them. - qsl_all = mm.query_start_loc - qsl_p = None if qsl_all is None else qsl_all[: num_prefills + 1] - qsl_long = mm.query_start_loc_long - qsl_long_p = qsl_long[: num_prefills + 1] - int_idx = None - replay = None - if is_verify: - # The runtime's arange over the slot pool, from the cache manager (what - # `update_mamba_states` reads). - assert core._bound.int_idx is not None, "the cache-manager bind has not run" - int_idx = core._bound.int_idx[:num_decodes] - if core._bound.replay is not None: - replay = _replay_work(core, mm, num_decodes) - gdn_state = ( - num_prefill_tokens, - num_decode_tokens, - si_p, - si_all, - si_d, - hi_p, - hi_all, - qsl_p, - qsl_all, - qsl_long, - is_verify, - num_decodes, - gen_p, - int_idx, - qsl_long_p, - replay, - ) - - # PLE carry flags must cover the whole batch, but `has_initial_states` is - # filled for the context prefix only (generation entries are stale); every - # generation sequence carries, so pad with True (host-int sized, graph-safe). - ple_carry = nn.functional.pad(hi_p, [0, num_decodes], mode="constant", value=1.0) - # Under speculation `Mamba2Metadata.prepare`'s `query_start_loc_long` has one - # row per generation sequence, so with `gen_p > 1` it is short and both PLE - # stateful calls are silently mis-segmented; rebuild it from `seq_lens_cuda`, - # which still holds the true counts during the trunk forward. - if gen_p == 1: - ple_qsl = qsl_long - else: - ple_qsl = cu_seqlens_from_seq_lens(md.seq_lens_cuda[:batch], torch.int64) - ple_state = (core._bound.ple_state[0], core._bound.ple_state[1], si_all, ple_carry, ple_qsl) - - # Dense/sparse decision once per forward, from a host-side metadata read (as - # QSA's backend does): a per-capture constant, matching the engine's separate - # short/long decode graph families. - sparse = batch > 0 and int(md.kv_lens_runtime[:batch].max()) > core.dense_threshold - - # The `[T, 3]` transpose of the `[3, 1, T]` position buffer, which the QSA position - # plane stores. - coords = torch.transpose(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]), 0, 1) - qsa = ( - md.qsa_req_idx_per_token[:num_tokens], - md.qsa_logical_positions[:num_tokens], - md.qsa_sequence_lengths[:num_tokens], - md.qsa_visible_blocks[:num_tokens], - md.qsa_topk_indices[:num_tokens], - md.qsa_topk_row_starts[:num_tokens], - md.qsa_block_table, - md.tokens_per_block, - coords[:num_tokens], - sparse, - _live_score_columns(core, md, batch), - ) - - # A verify step's generation rows write the PLE pools and the QSA side caches for - # every draft row, accepted or not; `SpeculativeRollback` keeps what its commit - # needs to undo the rejected ones. Reopened on every forward, so a step whose - # commit never ran (a failed forward) leaves nothing for the next one. Context - # rows are never staged. - stage = is_verify - rollback = core.spec_rollback - if rollback is not None: - rollback.discard() - if stage: - rollback.begin( - num_contexts=num_prefills, - prefix_tokens=num_prefill_tokens, - num_gens=num_decodes, - rows_per_gen=gen_p, - slots=si_d, - request_indices=md.qsa_req_idx_per_token[num_prefill_tokens:num_tokens], - logical_positions=md.qsa_logical_positions[num_prefill_tokens:num_tokens], - block_table=md.qsa_block_table, - tokens_per_block=md.tokens_per_block, - ) - - nvfp4_context = None - if ( - sparse - and core.quant_mode == QM_NVFP4_KV_CACHE - and num_prefill_tokens > _NVFP4_ROWTILE_MIN_ROWS - ): - nvfp4_context = _nvfp4_context_plan( - core, md, qsa[0], md.qsa_block_table, md.tokens_per_block - ) - - self._step = _PrefillStep( - num_tokens=num_tokens, - physical_tokens=physical_tokens, - input_ids=input_ids[:num_tokens], - attn=attn_args, - gdn=gdn_state, - ple=ple_state, - qsa=qsa, - stage=stage, - vision=vision, - nvfp4_context=nvfp4_context, - ) - - def forward( - self, - attn_metadata: AttentionMetadata, - input_ids: torch.IntTensor, - position_ids: torch.IntTensor, - **kwargs, - ): - """The step `prepare` built, over the `T` activation rows the engine handed in: the - batch's `n` rows, or the prefill CUDA graph bucket they were padded to. - - **This is the body a breakable prefill CUDA graph captures**, once per bucket, and - replays for every step it pads to that bucket without running this Python. So - nothing here reads the step: every op is row-wise and sized by `T` alone, its inputs - are the engine's persistent `input_ids` / `position_ids` buffers, and every buffer is - allocated here. What depends on the step's batch runs in one of four breaks that - run eagerly between the captured segments and read `self._step`: `_splice` (once), - `_ple_block` (once), `_attention_core` (every attention layer) and `_gdn_core` (every - Gated-DeltaNet layer). Each writes rows `[:n]` of a fixed output allocated here and - zeroes rows `[n:]`, so the padding rows stay finite through the row-wise ops and - reach neither a state pool nor a real row. With the backend off every step is - eager and unpadded, and the breaks are plain calls. - - Every hyper-connection site is `qwen4_exp_hc_inject_norm` (the previous block's - injection fused with this site's grouped norm; after a MoE block it also forms that - block's output) and then the mixer (`_hc_mixer`, `_hc_moe_mixer`). From - `_HC_RSTD_MIN_ROWS` rows the site is `qwen4_exp_hc_inject_rstd` (the same - injection, the norm's `[T, C]` rstd instead of the normed stream) and the rstd - mixer (`_hc_mixer_rstd`, `_hc_moe_mixer_rstd`); the PLE layer's attention-side - site keeps the stored-normed mixer either way, on the normed stream the PLE - kernel writes.""" - core = self.core - num_tokens = input_ids.shape[0] - assert num_tokens == self._step.physical_tokens, ( - f"the forward got {num_tokens} activation rows and prepare saw " - f"{self._step.physical_tokens}" - ) - - # `pos3`: contiguous `[3, T]` for `use_mrope=True` (the engine hands a slice of a - # larger buffer, whose row stride the ops reject, hence the copy). - pos3 = torch.empty( - [_QSA_POSITION_AXES, position_ids.shape[-1]], - dtype=torch.int32, - device=position_ids.device, - ) - pos3.copy_(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1])) - - # The token embedding with this step's vision rows spliced over their placeholders. - h = nn.functional.embedding(input_ids, core.w["embed"]) - eager_on_graph(h, inner=self._splice) - dt = h.dtype - dev = h.device - - # Scratch shared by every layer that uses it; each break rewrites its whole output. - proj_qkv = torch.empty([num_tokens, core.qkv_width], dtype=dt, device=dev) - index_scratch = torch.empty([num_tokens, core.index_scratch_width], dtype=dt, device=dev) - attn_out = torch.empty([num_tokens, core.q_width], dtype=dt, device=dev) - gdn_out = torch.empty( - [1, num_tokens, core.num_v_heads, core.head_v_dim], dtype=dt, device=dev - ) - ple_hc = torch.empty([num_tokens, core.stream], dtype=dt, device=dev) - ple_normed = torch.empty([num_tokens, core.stream], dtype=dt, device=dev) - # The CuteDSL routed output, shared by all layers: each layer's unpermute rewrites - # every row, and the next site's injection consumes it before the next layer's MoE. - routed_out = ( - torch.empty([num_tokens, core.hidden], dtype=dt, device=dev) - if core.trunk_moe_backend == "CUTEDSL" - else None - ) - - # The embedding repeated across the hc_count residual streams. - hc_count = core.hc_count - hc = torch.cat([h] * hc_count, -1) - layers, S = self._layers, self._stack - # A host decision on `T` alone (each captured bucket takes one branch). - by_rstd = num_tokens >= _HC_RSTD_MIN_ROWS - normed = rstd = None - if by_rstd: - _, rstd = qwen4_exp_hc_inject_rstd(hc, None, None, core.eps, hc_count) - else: - _, normed = qwen4_exp_hc_inject_norm( - hc, None, None, layers[0].ahc_norm, core.eps, hc_count - ) - for i, L in enumerate(layers): - if i == core.ple_layer: - eager_on_graph(hc, ple_hc, ple_normed, inner=self._ple_block) - hc, normed = ple_hc, ple_normed - # The PLE kernel stores this site's normed stream: the stored-normed - # mixer reads it (cheaper than a second pass for its rstd). - mixed, inject = self._hc_mixer(normed, L.ahc_first, L.ahc_up, True) - elif by_rstd: - mixed, inject = self._hc_mixer_rstd( - hc, rstd, L.ahc_norm, L.ahc_norm_w1, L.ahc_first, L.ahc_up, True - ) - else: - mixed, inject = self._hc_mixer(normed, L.ahc_first, L.ahc_up, True) - if core.is_attn[i]: - block_out = self._attention_block( - mixed, i, L, proj_qkv, index_scratch, attn_out, pos3 - ) - else: - block_out = self._gdn_block(mixed, i, L, gdn_out) - if by_rstd: - hc, rstd = qwen4_exp_hc_inject_rstd(hc, block_out, inject, core.eps, hc_count) - mixed, inject, xq, xsf = self._hc_moe_mixer_rstd( - hc, rstd, L.mhc_norm, L.mhc_norm_w1, L.mhc_first, L.mhc_up, L.moe_gscale - ) - else: - hc, normed = qwen4_exp_hc_inject_norm( - hc, block_out, inject, L.mhc_norm, core.eps, hc_count - ) - mixed, inject, xq, xsf = self._hc_moe_mixer( - normed, L.mhc_first, L.mhc_up, L.moe_gscale - ) - next_norm = layers[i + 1].ahc_norm if i + 1 < len(layers) else S.fhc_norm - # The MoE block output is formed inside the injection: routed + - # sigmoid(gate) * shared, rounded once. Both [T, hidden] operands are - # released here, not at the next layer's MoE block, whose own peak they - # would otherwise add to. - routed, (gate_logits, shared) = self._moe_block(mixed, xq, xsf, L, routed_out) - if by_rstd: - hc, rstd = qwen4_exp_hc_inject_rstd( - hc, routed, inject, core.eps, hc_count, shared, gate_logits - ) - else: - hc, normed = qwen4_exp_hc_inject_norm( - hc, routed, inject, next_norm, core.eps, hc_count, shared, gate_logits - ) - del routed, gate_logits, shared - # The stack-final mixer with no injection weight is the final norm. - if by_rstd: - final, _ = self._hc_mixer_rstd( - hc, rstd, S.fhc_norm, S.fhc_norm_w1, S.fhc_first, S.fhc_up, False - ) - else: - final, _ = self._hc_mixer(normed, S.fhc_first, S.fhc_up, False) - if not core.mtp_enabled: - return final - # `lm_head` reads the `[T, hidden]` mix; the MTP module reads the - # `[T, hc_count * hidden]` stream it was collapsed from. They leave packed in one - # tensor, because a breakable prefill graph's body returns one tensor; the shell - # takes the two column views (`split_spec_output`). - return torch.cat([final, hc], -1) - - # ------------------------------------------------------------- the draft --- - - def _mtp_site(self, hc, block_out, inject, norm, first, up) -> tuple: - """One of the draft layer's hyper-connection sites on the trunk's prefill kernels: - `block_out` injected into the stream `hc`, then the mixer (with injection) from the - stored normed stream, or from its rstd from `_HC_RSTD_MIN_ROWS` rows. Returns - `(hc, mixed, inject)`.""" - core = self.core - if hc.shape[0] < _HC_RSTD_MIN_ROWS: - hc, normed = qwen4_exp_hc_inject_norm( - hc, block_out, inject, norm, core.eps, core.hc_count - ) - return (hc, *self._hc_mixer(normed, first, up, True)) - w1 = self._mtp_norm_w1.get(norm.data_ptr()) - if w1 is None: - w1 = self._mtp_norm_w1[norm.data_ptr()] = _norm_w1(norm) - hc, rstd = qwen4_exp_hc_inject_rstd(hc, block_out, inject, core.eps, core.hc_count) - return (hc, *self._hc_mixer_rstd(hc, rstd, norm, w1, first, up, True)) - - def _mtp_moe_block(self, x, dt) -> torch.Tensor: - """The draft layer's MoE: FP8 block-scale routed experts plus a bf16 shared expert. - - The scale buffer comes back transposed `[hidden/128, tokens]` and is passed - straight through; never reshape it. At `tokens == hidden/128` (20, an - ordinary decode batch) the shapes coincide, so nothing would catch a - re-materialization. Pre-routed; `routing_bias` must be `None` (the entry - refuses one). With `_MTP_MOE_NVFP4_ENV` the routed experts run NVFP4 - (`_mtp_moe_nvfp4_block`).""" - core = self.core - M = core._mtp - if M.moe_gscale is not None: - return _mtp_moe_nvfp4_block(core, M, x) - logits = cublas_mm(x, M.router) - topk_ids, topk_w = renorm_moe_routing_op(logits, topk=core.topk, output_dtype=dt) - xq, xsf = fp8_quantize_1x128(x, use_ue8m0=_MTP_MOE_USE_UE8M0) - routed = fp8_block_scale_moe_runner( - routing_logits=None, - routing_bias=None, - hidden_states=xq, - hidden_states_scale=xsf, - gemm1_weights=M.fc1_w, - gemm1_weights_scale=M.fc1_s, - gemm2_weights=M.fc2_w, - gemm2_weights_scale=M.fc2_s, - num_experts=core.num_experts, - top_k=core.topk, - num_fused_shared_experts=None, - n_group=None, - topk_group=None, - intermediate_size=core.moe_inter, - local_expert_offset=0, - local_num_experts=core.num_experts, - routed_scaling_factor=None, - routing_method_type=_ROUTING_METHOD_RENORM, - topk_weights=topk_w, - topk_ids=topk_ids, - act_type=_ACT_TYPE_SWIGLU, - gemm1_clamp_limit=None, - ) - shared = cublas_mm(flashinfer_silu_and_mul(cublas_mm(x, M.sh_gu)), M.sh_dn) - gate_logits = cublas_mm(x, M.sh_gate) - return fused_sigmoid_gate_mul_add(routed, gate_logits, shared) - - def _capture_selection(self, md, topk_indices, visible_blocks, compact=False) -> None: - """IndexShare: stash, per request, the selection of the row the draft loop continues - from, with that row's complete-group count, in the draft layer's row of the shared - buffers. - - That row is the one the worker gathers for the next draft step: a context - request's last row of this chunk, and a generation request's last accepted row - (`num_accepted - 1` past its first). Capturing here rather than at the first draft - decode step keeps the indexer at one call per draft loop for a batch with context - rows too. A `compact` selection holds each context request's last row, then every - generation row (the draft step's `keep`).""" - batch = md.num_seqs - assert md.qsa_mtp_num_accepted is not None, ( - "IndexShare needs the worker's accepted counts to pick each request's capture " - "row; the draft loop did not supply them" - ) - seq_lens = md.seq_lens_cuda[:batch].to(torch.int64) - starts = torch.cumsum(seq_lens, 0) - seq_lens - accepted = md.qsa_mtp_num_accepted[:batch].to(torch.int64) - offsets = torch.minimum(torch.clamp(accepted - 1, min=0), seq_lens - 1) - offsets[: md.num_contexts] = seq_lens[: md.num_contexts] - 1 - rows = starts + offsets - if compact: - num_contexts = md.num_contexts - rows[:num_contexts] = torch.arange(num_contexts, device=rows.device) - rows[num_contexts:] -= md.num_ctx_tokens - num_contexts - slot = self.core._draft_bound.share_slot - md.qsa_shared_topk_indices[slot, :batch].copy_(torch.index_select(topk_indices, 0, rows)) - md.qsa_shared_topk_visible_blocks[slot, :batch].copy_( - torch.index_select(visible_blocks, 0, rows) - ) - # Host-side readiness, recorded in program order: under capture this runs while the - # graph is recorded, so a graph holds the reuse only after the capture it reads. - md.qsa_shared_topk_captured_slots.add(slot) - - def mtp_forward( - self, - embed_tokens: torch.Tensor, - all_rank_num_tokens, - input_ids: torch.Tensor, - position_ids: torch.Tensor, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - **kwargs, - ) -> torch.Tensor: - """One draft step whose metadata still carries context rows. - - `spec_metadata` and `all_rank_num_tokens` are accepted and unused.""" - core = self.core - assert isinstance(attn_metadata, TrtllmAttentionMetadata) - assert core._bound is not None, ( - "the trunk's cache-manager bind has not run; the shell calls the " - "core before the worker, so this cannot be reached first" - ) - md = attn_metadata - rows = hidden_states.shape[0] - assert input_ids.shape[0] == rows, ( - f"the draft step's {input_ids.shape[0]} token ids and {rows} " - "hidden-state rows must describe the same tokens" - ) - assert hidden_states.shape[1] == core.stream, ( - f"the MTP module reads the trunk's {core.stream}-wide four-stream " - f"residual; the runtime handed it {hidden_states.shape[1]} columns" - ) - dt = hidden_states.dtype - dev = hidden_states.device - - # Projected on every call: the worker rewrites `attn_metadata` in place - # between draft step 0 and 1+ (`_seq_lens` to 1, `num_contexts` and - # `num_ctx_tokens` to 0), so a first-call value is wrong on the rest. - step = self.step_args(md) - M = core._mtp - num_ctx = md.num_ctx_tokens - gen_rows = rows - num_ctx - gen_seqs = md.num_seqs - md.num_contexts - p = gen_rows // gen_seqs if gen_seqs else 1 - assert p * gen_seqs == gen_rows, ( - f"{gen_rows} generation rows do not divide over {gen_seqs} " - "generation sequences; a ragged per-sequence draft length cannot be " - "expressed to the attention call" - ) - - # The KV cache written here is not the trunk's: with a separate draft KV - # manager the worker swaps `attn_metadata.kv_cache_manager` and - # `kv_cache_block_offsets` for the draft loop. The layer index must be taken - # from the current manager every call; caching the trunk's addresses another - # layer's pages with finite, plausible drafts. - cm = md.kv_cache_manager - assert cm is not None, "a paged KV cache is required" - offsets = cm.layer_offsets - assert core.mtp_layer_idx in offsets, ( - f"the MTP module's attention is global layer index " - f"{core.mtp_layer_idx}, which the cache manager " - f"{type(cm).__name__} does not map (it maps {len(offsets)} layers, " - f"is_draft={getattr(cm, 'is_draft', None)}); the engine allocated " - "no KV layer for the draft module" - ) - layer_idx = offsets[core.mtp_layer_idx] - window = cm.max_attention_window_vec[layer_idx] - assert window is None or window == md.max_seq_len, ( - f"the MTP layer's pool window is {window!r} but this call passes " - f"max_seq_len {md.max_seq_len} as the attention window; a sliding " - "window is not implemented here" - ) - - # The draft layer's QSA planes and the page table that addresses them. With a - # separate draft manager the engine's QSA metadata swaps `qsa_block_table` for the - # table it decodes from the draft manager's offsets (`prepare_for_draft_forward`); - # a table still decoding the trunk's offsets would read and write another pool's - # pages. - draft = core._draft_bound - if draft is None or cm is not draft.cm: - core._bind_draft_cache_manager(md) - draft = core._draft_bound - assert cm is core._bound.cm or md.qsa_block_table is getattr( - md, "qsa_draft_block_table", None - ), ( - "the draft step is bound to a separate draft KV cache manager but " - "qsa_block_table is not that manager's table: the engine's QSA metadata did " - "not swap its page table for the draft forward" - ) - # The `qsa_*` row mapping is sized from the metadata's token count, which the - # worker rewrites together with it between draft steps. - assert rows == md.num_tokens, ( - f"the draft step carries {rows} rows for a batch its metadata describes as " - f"{md.num_tokens} tokens; the QSA row mapping covers the second" - ) - # The trunk's dense/sparse decision from the same host buffer: the worker advances - # only the device KV lengths through the draft loop, so every draft step of one - # engine forward takes the branch its trunk took, in the CUDA-graph family the - # runner chose for that trunk. - batch = md.num_seqs - sparse = batch > 0 and int(md.kv_lens_runtime[:batch].max()) > core.dense_threshold - # IndexShare: the only draft step with context rows is the loop's first, its - # draft-extend step. It always runs the indexer and, with sharing on, captures the - # selection the draft decode steps reuse. - share = core.mtp_index_share and md.qsa_in_mtp_draft_loop - assert not (share and md.qsa_indexer_skip_topk), ( - "the worker asked a draft step with context rows to reuse a selection; only " - "the draft decode steps after it reuse" - ) - - # Copy into a fresh contiguous `[3, N]`: a sliced view keeps the source row - # stride and the op rejects it under `use_mrope`. - assert position_ids.shape[0] == _QSA_POSITION_AXES, ( - f"expected the mRoPE position layout with {_QSA_POSITION_AXES} " - f"coordinate rows, got {tuple(position_ids.shape)}" - ) - pos3 = torch.empty( - [_QSA_POSITION_AXES, position_ids.numel() // _QSA_POSITION_AXES], - dtype=torch.int32, - device=position_ids.device, - ) - pos3.copy_(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1])) - # The trunk's two other views, from the copy: `pos` (coordinate 0 alone, which the - # index rope reads) and `coords` (the `[N, 3]` transpose the position plane stores). - pos = torch.reshape(torch.split(pos3, [1, _QSA_POSITION_AXES - 1], 0)[0], [-1]) - coords = torch.transpose(pos3, 0, 1) - - # ---- the front end -------------------------------------------------- - e = nn.functional.embedding(input_ids, embed_tokens) - x_e = cublas_mm(flashinfer_gemma_rmsnorm(e, M.pre_e_norm, core.eps), M.fc_e) - # `fc_hidden` is one `[hidden, hidden]` matrix applied to each of the four - # streams (the checkpoint's norm statistics say they are distinct streams). - h_n = self._grouped_norm(hidden_states, M.pre_h_norm) - x_h = torch.reshape( - cublas_mm(torch.reshape(h_n, [rows * core.hc_count, core.hidden]), M.fc_h), - [rows, core.stream], - ) - # Both hyper-connection sites run the trunk's site and mixers, as - # `DecodeTarget._mtp_layer` does: the front end's per-stream `x_h + x_e` is the first - # site's injection at unit weight. - F = core._mtp_fused - x, mixed, inject = self._mtp_site( - x_h, x_e, F.unit_inject[:rows], F.ahc_norm, F.ahc_first, F.ahc_up - ) - - # ---- the decoder layer ---------------------------------------------- - # The draft loop reads one row per request from this step: a context request's last - # row and one row of each generation run (the worker's `gather_ids`). Every row writes - # the layer's KV and index planes; past those writes only `keep` -- each context's - # last row, then every generation row -- goes on, the query side included on the - # sparse path (the dense call's metadata describes every row). The other returned - # rows keep the layer's input stream. - keep = None - if num_ctx > md.num_contexts: - ctx_lens = md.seq_lens_cuda[: md.num_contexts].to(torch.long) - keep = torch.cat( - [torch.cumsum(ctx_lens, 0) - 1, torch.arange(num_ctx, rows, device=dev)] - ) - q_keep = keep if sparse else None - q_rows = rows if q_keep is None else q_keep.shape[0] - qkv = torch.empty([rows, core.qkv_width], dtype=dt, device=dev) - attn_out = torch.empty([q_rows, core.q_width], dtype=dt, device=dev) - qg = cublas_mm(mixed if q_keep is None else mixed[q_keep], M.q) - k = cublas_mm(mixed, M.k) - v = cublas_mm(mixed, M.v) - qkv_h = torch.reshape(qkv, [rows, core.heads_q + 2 * core.heads_kv, core.head_dim]) - q_dst, k_dst, v_dst = torch.split(qkv_h, [core.heads_q, core.heads_kv, core.heads_kv], 1) - q_src, gate = torch.split( - torch.reshape(qg, [q_rows, core.heads_q, 2 * core.head_dim]), - [core.head_dim, core.head_dim], - 2, - ) - if q_keep is None: - q_dst.copy_(q_src) - else: - # The other rows' query heads stay unwritten: only the rope reads them, row-local. - q_dst.index_copy_(0, q_keep, q_src) - k_dst.copy_(torch.reshape(k, [rows, core.heads_kv, core.head_dim])) - v_dst.copy_(torch.reshape(v, [rows, core.heads_kv, core.head_dim])) - # The trunk rope's certified cell at the module's own base (`mtp.rope_theta`). - fused_qk_norm_rope( - qkv, - num_heads_q=core.heads_q, - num_heads_k=core.heads_kv, - num_heads_v=core.heads_kv, - head_dim=core.head_dim, - rotary_dim=core.rotary_dim, - eps=core.eps, - q_weight=M.q_norm, - k_weight=M.k_norm, - base=core.mtp_theta, - is_neox=True, - position_ids=pos3, - is_qk_norm=True, - use_gemma=True, - use_mrope=True, - mrope_section1=core.mrope_section1, - mrope_section2=core.mrope_section2, - ) - # The QSA index branch, as in the trunk's `_attention_block`: the raw key stored - # before anything normalizes or rotates it, the query's own norm and rotation, - # then the compression of every group this step completes. Maintained on every - # draft step at both ends of the threshold, so the draft plane is complete when a - # request crosses it. - req_idx = md.qsa_req_idx_per_token[:rows] - logical = md.qsa_logical_positions[:rows] - block_table = md.qsa_block_table - tokens_per_block = md.tokens_per_block - index_cache = draft.index_cache - qk_index = cublas_mm(mixed, M.idx) - token_k = torch.split( - torch.reshape(qk_index, [rows, core.index_heads + core.index_kv_heads, -1]), - [core.index_heads, core.index_kv_heads], - 1, - )[1] - qsa_paged_index_store( - token_k, - coords, - index_cache, - draft.position_cache, - req_idx, - logical, - block_table, - tokens_per_block, - ) - index_scratch = torch.empty([q_rows, core.index_scratch_width], dtype=dt, device=dev) - index_scratch[:, : core.index_width].copy_(qk_index if q_keep is None else qk_index[q_keep]) - q_pos3, q_pos, q_req, q_logical = pos3, pos, req_idx, logical - if q_keep is not None: - q_pos3 = torch.index_select(pos3, 1, q_keep) - q_pos, q_req, q_logical = pos[q_keep], req_idx[q_keep], logical[q_keep] - fused_qk_norm_rope( - index_scratch, - num_heads_q=core.index_heads, - num_heads_k=core.index_kv_heads, - num_heads_v=core.index_kv_heads, - head_dim=core.index_head_dim, - rotary_dim=core.rotary_dim, - eps=core.eps, - q_weight=M.idx_qn, - k_weight=M.idx_kn, - base=core.mtp_theta, - is_neox=True, - position_ids=q_pos3 if _QSA_INDEX_USE_MROPE else q_pos, - is_qk_norm=True, - use_gemma=True, - use_mrope=_QSA_INDEX_USE_MROPE, - mrope_section1=(_QSA_MROPE_SECTION or (0, 0, 0))[1], - mrope_section2=(_QSA_MROPE_SECTION or (0, 0, 0))[2], - ) - q_index = torch.split( - torch.reshape( - index_scratch, - [q_rows, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim], - ), - [core.index_heads, core.index_kv_heads, core.index_kv_heads], - 1, - )[0] - qsa_prefill_compress( - logical_positions=logical, - request_indices=req_idx, - block_table=block_table, - index_cache=index_cache, - position_cache=draft.position_cache, - k_norm_weight=M.idx_kn, - cos_sin=core._mtp_cos_sin, - eps=core.eps, - tokens_per_block=tokens_per_block, - compress_ratio=core.compress_ratio, - mrope_section=_QSA_MROPE_SECTION, - ) - - if sparse: - kv = draft.kv - # The append the dense op would have owned, then selection over the draft - # layer's own compressed keys and the sparse GQA over its pool. - _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) - row_starts = md.qsa_topk_row_starts[:rows] - visible_blocks = md.qsa_visible_blocks[:rows] - seq_lengths = md.qsa_sequence_lengths[:rows] - if q_keep is not None: - row_starts, visible_blocks = row_starts[q_keep], visible_blocks[q_keep] - seq_lengths = seq_lengths[q_keep] - scores = qsa_paged_index_scores( - q=q_index, - index_cache=index_cache, - block_table=block_table, - query_positions=q_logical, - request_indices=q_req, - tokens_per_block=tokens_per_block, - compress_ratio=core.compress_ratio, - only_visible_blocks=True, - # A context's run of rows shares one key gather; `keep` has none. - context_rows=q_keep is None, - ) - topk_indices = md.qsa_topk_indices[:q_rows] - _select( - core, - scores, - row_starts, - visible_blocks, - topk_indices, - _live_score_columns(core, md, batch), - ) - if share: - self._capture_selection(md, topk_indices, visible_blocks, q_keep is not None) - selected = qsa_expand_block_indices( - topk_indices, - q_logical, - seq_lengths, - compress_ratio=core.compress_ratio, - token_topk=core.token_topk, - ) - q_att = q_dst if q_keep is None else q_dst[q_keep] - if kv.k_scale is not None: - # `keep` holds one row per context request. - q_ctx = num_ctx if q_keep is None else md.num_contexts - plan = None - if q_ctx > _NVFP4_ROWTILE_MIN_ROWS: - plan = _nvfp4_context_plan( - core, - md, - req_idx, - block_table, - tokens_per_block, - one_row_per_request=q_keep is not None, - ) - _sparse_gqa_nvfp4_context( - core, - kv, - plan, - q_att, - block_table, - selected, - q_req, - q_logical, - tokens_per_block, - q_ctx, - attn_out, - rows_per_request=p, - ) - else: - attn_out = torch.reshape( - _sparse_gqa( - core, - kv, - q=q_att, - block_table=block_table, - selected_tokens=selected, - request_indices=q_req, - tokens_per_block=tokens_per_block, - query_positions=q_logical, - compress_ratio=core.compress_ratio, - ), - [q_rows, core.q_width], - ) - else: - attn_out = thop_attention( - q=qkv, - output=attn_out, - local_layer_idx=layer_idx, - predicted_tokens_per_seq=p, - num_heads=core.heads_q, - num_kv_heads=core.heads_kv, - head_size=core.head_dim, - quant_mode=core.quant_mode, - kv_scale_orig_quant=core.kv_scale_orig_quant, - kv_scale_quant_orig=core.kv_scale_quant_orig, - **step, - **_CALL_CONSTANTS, - ) - fused_sigmoid_mul_inplace(attn_out, gate) - if keep is not None: - if q_keep is None: - attn_out = attn_out[keep] - out, inject = x, inject[keep] - x = x[keep] - x, mixed, inject = self._mtp_site( - x, cublas_mm(attn_out, M.o), inject, F.mhc_norm, F.mhc_first, F.mhc_up - ) - x = qwen4_exp_post_inject(x, self._mtp_moe_block(mixed, dt), inject) - return x if keep is None else out.index_copy_(0, keep, x) - - -def _position_key(position_ids: torch.Tensor) -> tuple: - """What identifies a step's positions buffer: its storage, dtype and `[3, N]` layout.""" - view = torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]) - return (view.data_ptr(), view.dtype, tuple(view.shape), view.stride()) - - -def _copy_positions(position_ids: torch.Tensor) -> torch.Tensor: - """A fresh contiguous int32 `[3, N]` copy of the step's mRoPE positions.""" - pos3 = torch.empty( - [_QSA_POSITION_AXES, position_ids.numel() // _QSA_POSITION_AXES], - dtype=torch.int32, - device=position_ids.device, - ) - pos3.copy_(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1])) - return pos3 - - -class DecodeTarget: - """Generation-only batches (no context rows): drops the context-row recurrent - ops and adds the fused single-token recurrence.""" - - def __init__(self, core: Qwen4ExpTp1ModelingV2) -> None: - """Take the per-layer weight views once; MTP operands are read off `core._mtp`.""" - self.core = core - # See PrefillTarget.__init__. - self._contract_pending = step_contract_enabled() - self._kv_probe = ( - _KvRangeProbe(core.attn_layers[-1]) - if self._contract_pending and core.quant_mode == QM_FP8_KV_CACHE - else None - ) - self._layers = _layer_views(core) - self._stack = _stack_views(core) - # `(position key, pos3)` of the last trunk step: the first draft step of the same - # forward reads the trunk's positions buffer and reuses its int32 copy. - self._trunk_pos3: tuple | None = None - - def _check_step_contract(self, md) -> None: - """Opt-in first-forward fail-fast; see PrefillTarget._check_step_contract.""" - if not self._contract_pending: - return - try: - self.step_args(md) - except AttributeError as exc: - raise AssertionError(f"attention metadata surface drifted: {exc}") from exc - self._contract_pending = False - - def step_args(self, md: TrtllmAttentionMetadata) -> dict: - """See PrefillTarget.step_args. Context counts are stated, not read: routing - guarantees them, and a decode-only CUDA graph would freeze whatever it read.""" - return dict( - sequence_length=md.kv_lens_cuda_runtime, - host_past_key_value_lengths=md.kv_lens_runtime, - host_total_kv_lens=md.host_total_kv_lens, - context_lengths=md.prompt_lens_cuda_runtime, - host_context_lengths=md.prompt_lens_cpu_runtime, - host_request_types=md.host_request_types_runtime, - kv_cache_block_offsets=md.kv_cache_block_offsets, - host_kv_cache_pool_pointers=md.host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=md.host_kv_cache_pool_mapping, - workspace=md.effective_workspace, - tokens_per_block=md.tokens_per_block, - max_num_requests=md.max_num_requests, - max_context_length=md.max_context_length, - max_seq_len=md.max_seq_len, - # No sliding window: the full admitted length on every attention layer. - attention_window_size=md.max_seq_len, - num_contexts=0, - num_ctx_tokens=0, - trtllm_gen_jit_warmup=md.trtllm_gen_jit_warmup, - use_paged_context_fmha=md.use_paged_context_fmha, - beam_width=md.effective_beam_width, - cache_indirection=md.cache_indirection, - block_ids_per_seq=md.block_ids_per_seq, - max_context_q_len_override=md.max_context_q_len_override, - is_cross=md.is_cross, - is_spec_decoding_enabled=md.is_spec_decoding_enabled, - use_spec_decoding=md.use_spec_decoding, - is_spec_dec_tree=md.is_spec_dec_tree, - spec_decoding_generation_lengths=md.spec_decoding_generation_lengths, - spec_decoding_position_offsets=get_spec_decoding_position_offsets(md), - spec_decoding_packed_mask=md.spec_decoding_packed_mask, - spec_decoding_bl_tree_mask_offset=md.spec_decoding_bl_tree_mask_offset, - spec_decoding_bl_tree_mask=md.spec_decoding_bl_tree_mask, - spec_decoding_target_max_draft_tokens=md.max_total_draft_tokens, - spec_bl_tree_first_sparse_mask_offset_kv=md.spec_bl_tree_first_sparse_mask_offset_kv, - num_sparse_topk=md.num_sparse_topk, - flash_mla_tile_scheduler_metadata=md.flash_mla_tile_scheduler_metadata, - flash_mla_num_splits=md.flash_mla_num_splits, - max_num_sequences=md.max_num_sequences, - force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask, - ) - - # --------------------------------------------------------------- blocks --- - - def _hc_first(self, normed, first, with_inject): - """The mixer's first half from the site's normed stream: the projection - over `[down; block_inject; 0]` and the SiLU / injection activations. - Returns `(low, inject)`. Up to `_HC_TC_FIRST_MAX_ROWS` rows (a per-capture - constant) one tcgen05 launch, `qwen4_exp_hc_decode_tc_first`; larger batches - cuBLAS's fp32 logits and `qwen4_exp_hc_lowrank_act`. Only sites past - `_HC_DECODE_GEMV_MAX_ROWS` rows reach a mixer.""" - core = self.core - if normed.shape[0] <= _HC_TC_FIRST_MAX_ROWS: - return qwen4_exp_hc_decode_tc_first( - normed, first, core.hc_lowrank, core.hc_count, with_inject - ) - return qwen4_exp_hc_lowrank_act( - cublas_mm(normed, first[:, : core.hc_first_cols], None, torch.float32), - core.hc_lowrank, - core.hc_count, - with_inject, - normed.dtype, - ) - - def _hc_mixer(self, normed, first, up, with_inject): - """One hyper-connection mixer from the site's normed stream: the first - projection over `[down; block_inject; 0]` (fp32 logits), the SiLU and - injection activations, then the up projection with the sigmoid gate and the - stream mean in its epilogue. Returns `(mixed, inject)`; `inject` is None - without `with_inject` (the stack-final mixer). - - Neither projection takes the step's `gemm` (`low_m_bf16_mm` under CUDA - graphs): neither FlashInfer low-M kernel writes the first projection's - fp32 logits, and the up projection's K = 320 tiles neither of them.""" - core = self.core - low, inject = self._hc_first(normed, first, with_inject) - return qwen4_exp_hc_up_gate_mean(low, up, normed, core.hc_count), inject - - def _hc_site(self, hc, block_out, inject, norm, first, up, with_inject, normed=None, moe=None): - """One trunk hyper-connection site: the previous block's injection (none at - the stack's first site and after the PLE add), this site's grouped norm and - its mixer. Returns `(hc, mixed, inject)`. The row count is a per-capture - constant: at most `_HC_DECODE_GEMV_MAX_ROWS` rows run two GEMV launches, - larger batches `qwen4_exp_hc_inject_norm` and `_hc_mixer`. `normed`, when - given, is this site's grouped norm of `hc` already computed (the PLE tail's, - with no injection pending); the cuBLAS path starts from it. - - After a MoE block, `block_out` is its routed output and `moe` the - `(gate_logits, shared, topk_weights, topk_rows)` `_moe_block` left - uncombined: the site's first kernel forms the block output itself. With - the top-k pair `block_out` is the runner's permuted FC2 rows and that - kernel (on every route) also forms their finalize. A site handed - `normed` has no pending block, so it never carries `moe`.""" - core = self.core - gate_logits, shared, topk_w, topk_rows = moe if moe is not None else (None,) * 4 - if _HC_DECODE_GEMV_MAX_ROWS < hc.shape[0] <= _HC_DECODE_ROWS_MAX: - # Two tensor-core launches; with the top-k pair the first also forms the finalize. - n_cols = core.hc_lowrank + (core.hc_count if with_inject else 0) - hc, low, inject, rstd = qwen4_exp_hc_decode_rows_first( - hc, - block_out, - inject, - norm, - first, - core.hc_count, - n_cols, - core.eps, - shared, - gate_logits, - topk_w, - topk_rows, - ) - mixed, _, _ = qwen4_exp_hc_decode_rows_mix(low, rstd, up, hc, norm, core.hc_count) - return hc, mixed, inject - if hc.shape[0] <= _HC_DECODE_GEMV_MAX_ROWS: - n_cols = core.hc_lowrank + (core.hc_count if with_inject else 0) - hc, part, ss = qwen4_exp_hc_decode_first_cute( - hc, - block_out, - inject, - norm, - first, - core.hc_count, - n_cols, - shared, - gate_logits, - topk_w, - topk_rows, - ) - mixed, inject = qwen4_exp_hc_decode_mix_cute( - part, ss, up, hc, norm, core.eps, core.hc_count, with_inject - ) - return hc, mixed, inject - if normed is None: - hc, normed = qwen4_exp_hc_inject_norm( - hc, - block_out, - inject, - norm, - core.eps, - core.hc_count, - shared, - gate_logits, - topk_w, - topk_rows, - ) - mixed, inject = self._hc_mixer(normed, first, up, with_inject) - return hc, mixed, inject - - def _hc_moe_site(self, hc, block_out, inject, norm, first, up, gscale): - """`_hc_site` at an MoE site (always with injection), returning the - routed experts' NVFP4 input (linear scales, the layer's global scale - `gscale`) beside the mix, written by the mixer's own last launch: on the - GEMV route `qwen4_exp_hc_decode_mix_fp4_cute`, otherwise the fused up - projection `qwen4_exp_hc_up_gate_mean_fp4`. Returns - `(hc, mixed, inject, xq, xsf)`.""" - core = self.core - if _HC_DECODE_GEMV_MAX_ROWS < hc.shape[0] <= _HC_DECODE_ROWS_MAX: - hc, low, inject, rstd = qwen4_exp_hc_decode_rows_first( - hc, - block_out, - inject, - norm, - first, - core.hc_count, - core.hc_lowrank + core.hc_count, - core.eps, - ) - mixed, xq, xsf = qwen4_exp_hc_decode_rows_mix( - low, rstd, up, hc, norm, core.hc_count, gscale - ) - return hc, mixed, inject, xq, xsf - if hc.shape[0] <= _HC_DECODE_GEMV_MAX_ROWS: - hc, part, ss = qwen4_exp_hc_decode_first_cute( - hc, block_out, inject, norm, first, core.hc_count, core.hc_lowrank + core.hc_count - ) - mixed, inject, xq, xsf = qwen4_exp_hc_decode_mix_fp4_cute( - part, ss, up, hc, norm, core.eps, core.hc_count, True, gscale - ) - return hc, mixed, inject, xq, xsf - hc, normed = qwen4_exp_hc_inject_norm(hc, block_out, inject, norm, core.eps, core.hc_count) - low, inject = self._hc_first(normed, first, True) - mixed, xq, xsf = qwen4_exp_hc_up_gate_mean_fp4(low, up, normed, core.hc_count, gscale) - return hc, mixed, inject, xq, xsf - - def _grouped_norm(self, x: torch.Tensor, weights: tuple) -> torch.Tensor: - """See `PrefillTarget._grouped_norm`.""" - core = self.core - parts = torch.split(x, [core.hidden] * core.hc_count, 1) - return torch.cat( - [flashinfer_gemma_rmsnorm(part, w, core.eps) for w, part in zip(weights, parts)], 1 - ) - - def _moe_block( - self, x, xq, xsf, L: _LayerWeights, routed_out, gemm=cublas_mm, fold=False - ) -> tuple: - """One MoE layer: NVFP4 routed experts plus a bf16 shared expert, left - uncombined. Returns `(routed, (gate_logits, shared, topk_weights, - topk_rows))`: the next hyper-connection site forms `routed + - sigmoid(gate_logits) * shared` inside its first kernel (`_hc_site`'s - `moe`). With `fold` (trtllm-gen, and the next site takes the combine; a - per-capture constant) the runner also skips its finalize: `routed` is then - its permuted FC2 rows and `topk_weights` / `topk_rows` its routing weights - and index map, which that site's first kernel - (`qwen4_exp_hc_decode_first_cute`, `qwen4_exp_hc_decode_rows_first` or - `qwen4_exp_hc_inject_norm`) gathers; - without it both are None and `routed` is the finalized `[T, H]` rows. - - The shared expert and its per-token gate logit run on the core's side - stream, concurrently with the routed path, and join it before this - returns; from `_ROUTER_FIRST_MIN_ROWS` rows they launch after the router. - `x` is the bf16 hidden states (router and shared expert); `xq` / - `xsf` are the same states in NVFP4 with **linear** block scales, from - `_hc_moe_site`. On trtllm-gen the routed path hands the runner the router - logits (its Renormalize routing is bitwise the `renorm_moe_routing_op` + - pre-routed composition); on CuteDSL it is `_cute_dsl_routed`, which writes - `routed_out`. `gemm` is the step's bf16 projection entry (see `forward`); it - launches on the current stream, so the shared expert's GEMMs run on the side - stream.""" - core = self.core - main = torch.cuda.current_stream() - side = core._side_stream - # Fork: the side stream reads `x` only after everything queued on `main` - # so far. What it allocates is consumed on `main` after the join below, - # and the next layer's fork orders any reuse of those blocks after that. - side.wait_stream(main) - # The routing reads its logits as a dense row-major buffer from data_ptr() and - # ignores strides; every projection entry allocates them contiguous. From - # `_ROUTER_FIRST_MIN_ROWS` rows the router launches before the shared expert (the - # fork above precedes it, so the side stream does not wait for it). - router_first = x.shape[0] >= _ROUTER_FIRST_MIN_ROWS - if router_first: - logits = _dense_mm(x, L.router, None, gemm, "router", mid=L.router_m) - with torch.cuda.stream(side): - gu = _dense_mm(x, L.sh_gu, None, gemm, "sh_gu", L.sh_gu_g, mid=L.sh_gu_m) - shared = _dense_mm( - flashinfer_silu_and_mul(gu), L.sh_dn, None, gemm, "sh_dn", L.sh_dn_g, mid=L.sh_dn_m - ) - gate_logits = gemm(x, L.sh_gate) - if not router_first: - logits = _dense_mm(x, L.router, None, gemm, "router", mid=L.router_m) - if core.trunk_moe_backend == "CUTEDSL": - routed = _cute_dsl_routed(core, L, logits, xq, xsf, routed_out) - main.wait_stream(side) - return routed, (gate_logits, shared, None, None) - if fold: - moe_op = fp4_block_scale_moe_renorm_permuted - else: - moe_op = fp4_block_scale_moe_renorm_routed - out = moe_op( - logits, - xq, - xsf.view(torch.float8_e4m3fn), - L.fc1_w, - L.fc1_s, - L.fc2_w, - L.fc2_s, - L.out1_scale, - L.out1_gate_scale, - L.out2_scale, - core.num_experts, - core.topk, - core.moe_inter, - ) - main.wait_stream(side) - if fold: - routed, topk_w, topk_rows = out - return routed, (gate_logits, shared, topk_w, topk_rows) - return out, (gate_logits, shared, None, None) - - def _attention_prologue( - self, proj, q_index, index_view, index_cache, index_section, L, qsa, side, stage - ) -> None: - """The norm + rope of both planes on this stream and the QSA index branch (the - verify step's staging, the raw-key store, then compression) on the side stream, - which reads only the projection's index-key columns and the index pools; the - attention core joins it before reading either.""" - core = self.core - num_tokens = proj.shape[0] - req_idx, logical, _, _, _, _, block_table, tokens_per_block, pos3, coords, _ = qsa - side.wait_stream(torch.cuda.current_stream()) - qwen4_exp_qk_norm_rope( - proj, - q_index, - pos3, - core._cos_sin, - L.q_norm, - L.k_norm, - L.idx_qn, - num_heads_q=core.heads_q, - num_heads_k=core.heads_kv, - head_dim=core.head_dim, - index_heads=core.index_heads, - index_head_dim=core.index_head_dim, - index_q_col=core.qkv_width + core.q_width, - rotary_dim=core.rotary_dim, - eps=core.eps, - mrope_section1=core.mrope_section1, - mrope_section2=core.mrope_section2, - index_mrope_section1=index_section[1], - index_mrope_section2=index_section[2], - ) - - # ---- the QSA index branch ------------------------------------------ - # The raw key is stored **before** anything normalizes or rotates it, because - # the compression step pools raw keys and applies the norm and the rotation - # itself at the group's first member's coordinates; a key normalized on the way - # in double-normalizes every anchor that reads it, which is finite, silent, and - # invisible below the sparse threshold. The norm + rope above never writes the - # index key columns. - with torch.cuda.stream(side): - token_k = torch.split( - torch.reshape( - index_view, - [ - num_tokens, - core.index_scratch_width // core.index_head_dim, - core.index_head_dim, - ], - ), - [core.index_heads, core.index_kv_heads, 1], - 1, - )[1] - if stage: - core.spec_rollback.stage_index(index_cache, core._bound.position_cache) - qsa_paged_index_store( - token_k, - coords, - index_cache, - core._bound.position_cache, - req_idx, - logical, - block_table, - tokens_per_block, - ) - qsa_prefill_compress( - logical_positions=logical, - request_indices=req_idx, - block_table=block_table, - index_cache=index_cache, - position_cache=core._bound.position_cache, - k_norm_weight=L.idx_kn, - cos_sin=core._cos_sin, - eps=core.eps, - tokens_per_block=tokens_per_block, - compress_ratio=core.compress_ratio, - mrope_section=_QSA_MROPE_SECTION, - ) - - def _attention_block( - self, x, i, L, qkv, index_scratch, attn_out, qsa, step, stage, gemm=cublas_mm - ): - """One gated GQA layer plus its QSA index branch for a generation-only batch. - - The four input projections are one GEMM against the load-time fused - operand (`weights._attn_in_w`), whose output row is - `[q | k | v | gate | index q, k | v slot]`. One norm+rope launch rewrites - the main q/k heads in place on that output and writes the normalized, - rotated index query into the scratch; the sparse path reads q/k/v as views - of it, the dense path takes one compact copy of its leading `[q | k | v]` - (the dense op reads its QKV flat), the gate is read through its row stride, - and the index branch reads its raw key from it. - - **The index plane is maintained on every step, at both ends of the - threshold.** Below it the dense core runs and the index cache is still - advanced, so a request that grows past the threshold finds its prefix - already compressed; above it the dense core is replaced outright and the - main K/V append becomes this layer's own call, because the op that would - have done it is not the one running. Above the threshold over a bf16 or NVFP4 pool, - everything up to the index scores is one launch, `qwen4_exp_attn_decode_prologue` - (over an NVFP4 pool it quantizes the K/V append). On a verify step the rollback - stages the index slots first, and at a group's last token the launch takes the - raw keys of the group's earlier rows in the same step from their projection - rows, as the index plane holds them only after the launch.""" - core = self.core - ( - req_idx, - logical, - seq_lens_q, - visible_blocks, - topk_indices, - topk_row_starts, - block_table, - tokens_per_block, - pos3, - coords, - sparse, - ) = qsa - num_tokens = x.shape[0] - - main = torch.cuda.current_stream() - side = core._side_stream - index_cache = core._bound.index_cache[i] - kv = core._bound.kv[i] - # A bf16 or an NVFP4 pool (an fp8 one stays on the composed prologue). - fused_prologue = sparse and (kv.k_scale is not None or kv.k.dtype == x.dtype) - if fused_prologue and stage: - # The rollback's copy of the slots the prologue launch overwrites, taken on - # the side stream under the input projection. - side.wait_stream(main) - with torch.cuda.stream(side): - core.spec_rollback.stage_index(index_cache, core._bound.position_cache) - proj = _dense_mm(x, L.attn_in, L.attn_in8, gemm, "in", L.attn_in_g, mid=L.attn_in_m) - qkv_view, gate, index_view = torch.split( - proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1 - ) - q_dst, k_dst, v_dst = _attn_in_heads(core, proj) - # The index query lands in the q slots of the scratch; its k and v slots - # are layout only and nothing reads them. - q_index = torch.split( - torch.reshape( - index_scratch, - [num_tokens, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim], - ), - [core.index_heads, core.index_kv_heads, core.index_kv_heads], - 1, - )[0] - index_section = (_QSA_MROPE_SECTION or (0, 0, 0)) if _QSA_INDEX_USE_MROPE else (0, 0, 0) - # Up to SCORE_ROWS decode rows the index scores ride the prologue launch too. - scores = None - if fused_prologue and not stage and num_tokens <= _PROLOGUE_SCORE_ROWS: - scores = torch.empty( - (num_tokens, block_table.shape[1] * tokens_per_block // core.compress_ratio), - dtype=torch.float32, - device=proj.device, - ) - if fused_prologue: - if stage: - main.wait_stream(side) - # The norm + rope of both planes, the K/V append the dense op would have - # owned (quantizing over an NVFP4 pool), and the index branch's raw-key and - # coordinate store and group compression in one launch on this stream. - qwen4_exp_attn_decode_prologue( - proj, - q_index, - pos3, - coords, - core._cos_sin, - L.q_norm, - L.k_norm, - L.idx_qn, - L.idx_kn, - kv.k, - kv.v, - index_cache, - core._bound.position_cache, - req_idx, - logical, - block_table, - tokens_per_block, - core.compress_ratio, - num_heads_q=core.heads_q, - num_heads_k=core.heads_kv, - head_dim=core.head_dim, - index_heads=core.index_heads, - index_head_dim=core.index_head_dim, - index_q_col=core.qkv_width + core.q_width, - rotary_dim=core.rotary_dim, - eps=core.eps, - mrope_section1=core.mrope_section1, - mrope_section2=core.mrope_section2, - index_mrope_section1=index_section[1], - index_mrope_section2=index_section[2], - compress_mrope_section=_QSA_MROPE_SECTION, - k_scale_cache=kv.k_scale, - v_scale_cache=kv.v_scale, - kv_scale_orig_quant=None if kv.k_scale is None else core.kv_scale_orig_quant, - rows_per_request=step["predicted_tokens_per_seq"], - index_scores=scores, - ) - else: - self._attention_prologue( - proj, q_index, index_view, index_cache, index_section, L, qsa, side, stage - ) - if self._kv_probe is not None and not self._kv_probe.observe(i, k_dst, v_dst): - self._kv_probe = None - - gated = False - if sparse: - if not fused_prologue: - # Selection reads the index plane the side stream wrote. - main.wait_stream(side) - # Append the fused-norm+rope'd K/V; the selection reads the current - # token from the pool. Over an fp8 or NVFP4 pool as in the prefill target. - _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) - if scores is None: - scores = qsa_paged_index_scores( - q=q_index, - index_cache=index_cache, - block_table=block_table, - query_positions=logical, - request_indices=req_idx, - tokens_per_block=tokens_per_block, - compress_ratio=core.compress_ratio, - only_visible_blocks=True, - context_rows=False, - # Its predecessor is the prologue launch and its consumer the top-k, - # which waits before reading the scores. With one row per request the - # prologue writes no compressed key but the one the entry re-reads - # after its wait, so at decode-sized batches the pages and keys are - # read before it; a verify step's later rows also see the key its - # group-end row compresses, so they read every key after the wait. - use_pdl=fused_prologue, - prefetch=fused_prologue and not stage and num_tokens <= _SPARSE_DECODE_MAX_ROWS, - ) - if _topk_expand_supports(scores): - # The canonical block top-k and its expansion to tokens in one launch - # (QSA's row starts are zero). - selected = qsa_topk_expand_decode( - scores, - visible_blocks, - logical, - seq_lens_q, - topk_indices, - core.compress_ratio, - core.token_topk, - ) - else: - _select(core, scores, topk_row_starts, visible_blocks, topk_indices) - selected = qsa_expand_block_indices( - topk_indices, - logical, - seq_lens_q, - compress_ratio=core.compress_ratio, - token_topk=core.token_topk, - ) - if kv.k_scale is not None: - # The output gate rides the CuTe kernel's store (one split) or merge. A verify - # step's rows come `predicted_tokens_per_seq` a request. - attn = _sparse_gqa( - core, - kv, - q=q_dst, - block_table=block_table, - selected_tokens=selected, - request_indices=req_idx, - tokens_per_block=tokens_per_block, - # The prefill pair bounds a context row's scan; none here. - query_positions=None, - compress_ratio=None, - output_gate=gate, - rows_per_request=step["predicted_tokens_per_seq"], - ) - gated = True - else: - # The output gate rides the split-K merge whenever the launch has - # one (more than one split); otherwise it is its own kernel below. - gated = qsa_paged_sparse_gqa.fuses_output_gate(q_dst, kv.k, selected, gate) - if ( - gated - and num_tokens <= _SPARSE_DECODE_MAX_ROWS - and _sparse_decode_supports(q_dst, kv.k, selected) - ): - # The split-K pass, then the merge and the gate as one launch spread - # over the SMs. - attn = qsa_paged_sparse_gqa_decode( - q_dst, - kv.k, - kv.v, - block_table, - selected, - req_idx, - tokens_per_block, - core.softmax_scale, - gate, - ) - else: - attn = qsa_paged_sparse_gqa( - q=q_dst, - k_cache=kv.k, - v_cache=kv.v, - block_table=block_table, - selected_tokens=selected, - request_indices=req_idx, - tokens_per_block=tokens_per_block, - softmax_scale=core.softmax_scale, - query_positions=None, - compress_ratio=None, - output_gate=gate if gated else None, - ) - out = torch.reshape(attn, [num_tokens, core.q_width]) - else: - main.wait_stream(side) - # The dense op reads its packed QKV flat: one compact copy. - qkv.copy_(qkv_view) - out = thop_attention( - q=qkv, - output=attn_out, - local_layer_idx=core._bound.pool_idx[i], - num_heads=core.heads_q, - num_kv_heads=core.heads_kv, - head_size=core.head_dim, - quant_mode=core.quant_mode, - kv_scale_orig_quant=core.kv_scale_orig_quant, - kv_scale_quant_orig=core.kv_scale_quant_orig, - **step, - **_CALL_CONSTANTS, - ) - if not gated: - fused_sigmoid_mul_inplace(out, gate) - return _dense_mm(out, L.o, L.o8, gemm, "out", L.o_g, mid=L.o_m) - - def _gdn_verify( - self, mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, n, p, replay, mq_out=None - ) -> torch.Tensor: - """Recurrent half of a speculative verify step; see - `PrefillTarget._gdn_verify`. `mq_out` is the conv output when the in_proj ran the - conv step as its epilogue (`fp8_gemv_conv1d_verify`, then `mq_d` is None); else the - conv runs here over the projection's q|k|v columns `mq_d`.""" - core = self.core - # The in_proj GEMV ran the conv as its epilogue and signals its dependents early. - after_gemv = mq_out is not None - if mq_out is None: - # [n*p, C] -> [n, C, p]: the multi-token conv reads channel-major. - x3 = torch.transpose(torch.reshape(mq_d, [n, p, core.conv_dim]), 1, 2) - conv_out = causal_conv1d_triton_update( - x=x3, - conv_state=core._bound.conv[i], - weight=L.conv_w, - bias=None, - activation="silu", - conv_state_indices=si_d, - pad_slot_id=_PAD_SLOT_ID, - # Snapshots, not the accepted-token roll. - num_accepted_tokens=None, - intermediate_conv_window=core._bound.int_conv[i], - intermediate_state_indices=int_idx, - launch_dependent_kernels=False, - ) - # Back to token-major (a view). - mq_out = torch.reshape(torch.transpose(conv_out, 1, 2), [n * p, core.conv_dim]) - if replay is not None: - return _replay_verify( - core, mq_out, a_d, b_d, ssm_states, i, L, si_d, n, p, replay, pdl=after_gemv - ) - q, k, v = pack_gdn_decode_qkv( - mq_out, core.num_k_heads, core.head_k_dim, core.num_v_heads, core.head_v_dim - ) - g = torch.reshape( - fused_gdn_gating( - A_log=L.a_log, - a=a_d, - dt_bias=L.dt_bias, - beta=_SOFTPLUS_BETA, - threshold=_SOFTPLUS_THRESHOLD, - ), - [n, p, core.num_v_heads], - ) - beta = torch.reshape(torch.sigmoid(b_d), [n, p, core.num_v_heads]) - # Initial states are gathered from the live pool; the per-step buffer is - # indexed by batch row (`intermediate_state_indices[:n]`), not slot id. - src = torch.reshape( - nn.functional.embedding(si_d, torch.reshape(ssm_states, [ssm_states.shape[0], -1])), - [n, core.num_v_heads, core.head_v_dim, core.head_k_dim], - ) - out = fused_recurrent_gated_delta_rule_update( - torch.reshape(q, [n, p, core.num_k_heads, core.head_k_dim]), - torch.reshape(k, [n, p, core.num_k_heads, core.head_k_dim]), - torch.reshape(v, [n, p, core.num_v_heads, core.head_v_dim]), - g, - beta, - initial_state_source=src, - initial_state_indices=int_idx, - scale=None, - use_qk_l2norm_in_kernel=True, - cu_seqlens=None, - disable_state_update=True, - disable_output_calculation=False, - intermediate_states_buffer=core._bound.int_ssm[i], - cache_steps=p, - output=None, - ) - return torch.reshape(out, [1, n * p, core.num_v_heads, core.head_v_dim]) - - def _gdn_block(self, x, i, L, gdn_state, gemm=cublas_mm, fused_conv=False) -> torch.Tensor: - """One Gated-DeltaNet layer over a generation-only batch: one token per - sequence, or under speculation the whole draft chain through the verify - path. Every branch is on a host int the engine fixes for a capture. - - One GEMM for all four input projections: q|k|v, z, a and b are column views of - its output. A graph-captured single-token step at up to `_GDN_CONV_FUSED_MAX_ROWS` - rows runs the conv step as that GEMM's epilogue (`gdn_in_proj_conv1d_pdl`), and - over an fp32 recurrent pool at up to `_GDN_DECODE_FUSED_MAX_ROWS` rows the - recurrence and the gated norm as one launch (`gdn_decode_gated_norm_pdl`); past - them, and at every batch size over a bf16 pool, after the conv step, as one - streaming launch (`gdn_decode_rows_gated_norm`). That kernel computes each (row, - v-head) on its own, so over a bf16 pool a row's output and state do not depend on - its batch. - - A graph-captured verify step whose in_proj takes `fp8_gemv` (dense FP8 mode, up to - `_VERIFY_CONV_FUSED_MAX_ROWS` rows) runs the multi-token conv step -- window snapshots and - the pool advance included -- as that GEMM's epilogue too (`fp8_gemv_conv1d_verify`, - bitwise the projection + `causal_conv1d_triton_update` it replaces); the conv output - is its own contiguous buffer, and z, a and b are column views of the rest. - - With the in_proj's NVFP4 copy (`_GDN_INPROJ_NVFP4_ENV`) every step of up to - `_INPROJ_NVFP4_MAX_ROWS` rows reads it on the same kernel instead, a captured - single-token step at all of those rows with the conv step as the epilogue.""" - core = self.core - si_all, is_verify, num_decodes, gen_p, int_idx, replay = gdn_state - num_tokens = x.shape[0] - ssm_states = core._bound.ssm[i] - - conv_fused = fused_conv and not is_verify and num_tokens <= _GDN_CONV_FUSED_MAX_ROWS - # The NVFP4 copy takes every decode row count within its range, the conv step fused - # at each of them in a captured step. - nv = ( - L.in_proj_nv - if L.in_proj_nv is not None and num_tokens <= _INPROJ_NVFP4_MAX_ROWS and _gemv_ready(x) - else None - ) - nv_conv = nv is not None and fused_conv and not is_verify - fp8_conv = ( - not nv_conv - and fused_conv - and not is_verify - and L.in_proj_g is not None - and _gemv_ready(x) - and num_tokens <= _FP8_CONV_FUSED_MAX_ROWS - ) - # The verify conv as the in_proj's epilogue reads the pool, taps and slots before its - # PDL wait, like the single-token fusion: captured steps only. - verify_fused = ( - fused_conv - and is_verify - and (L.in_proj_g is not None or nv is not None) - and _gemv_ready(x) - and num_tokens <= min(_FP8_GEMV_ROWS["in"][1], _VERIFY_CONV_FUSED_MAX_ROWS) - ) - mq_out = None - if verify_fused: - verify_op = ( - fp8_gemv_conv1d_verify_op if nv is None else nvfp4_split_gemv_conv1d_verify_op - ) - mq_out, zab = verify_op( - x, - L.in_proj_g if nv is None else nv, - core._bound.conv[i], - L.conv_w, - si_all, - core._bound.int_conv[i], - int_idx, - core.conv_dim, - gen_p, - _PAD_SLOT_ID, - ) - elif nv_conv: - proj = nvfp4_split_gemv_conv1d_op( - x, nv, core._bound.conv[i], L.conv_w, si_all, core.conv_dim, _PAD_SLOT_ID - ) - elif fp8_conv: - proj = fp8_gemv_conv1d_op( - x, - L.in_proj_g, - core._bound.conv[i], - L.conv_w, - si_all, - core.conv_dim, - _PAD_SLOT_ID, - ) - elif conv_fused: - proj = gdn_in_proj_conv1d_pdl( - x, - L.in_proj, - core._bound.conv[i], - L.conv_w, - si_all, - core.conv_dim, - _PAD_SLOT_ID, - True, - ) - elif nv is not None: - proj = nvfp4_split_gemv_op(x, nv) - else: - proj = _dense_mm(x, L.in_proj, L.in_proj8, gemm, "in", L.in_proj_g, mid=L.in_proj_m) - if verify_fused: - mixed_qkv = None - z, a, b = torch.split(zab, core.in_proj_split[1:], 1) - else: - mixed_qkv, z, a, b = torch.split(proj, core.in_proj_split, 1) - - if ( - ssm_states.dtype == torch.float32 - and conv_fused - and num_tokens <= _GDN_DECODE_FUSED_MAX_ROWS - ): - # The recurrence and the gated norm as one launch that reads this layer's - # pool rows before its PDL wait: in a captured single-token step nothing - # ahead of it writes them, and the slot indices are filled before the - # graph is launched. - o = gdn_decode_gated_norm_pdl( - proj, - ssm_states, - si_all, - L.a_log, - L.dt_bias, - L.gnorm, - core.eps, - core.num_k_heads, - core.head_k_dim, - ) - return _dense_mm(o, L.out, L.out8, gemm, "out", L.out_g, mid=L.out_m) - if is_verify: - # The CuTe replay kernel reads `a`/`b` through their row stride: no copies. - strided_ab = replay is not None and core._bound.replay.cute - out = self._gdn_verify( - mixed_qkv, - a if strided_ab else a.contiguous(), - b if strided_ab else b.contiguous(), - ssm_states, - i, - L, - si_all, - int_idx, - num_decodes, - gen_p, - replay, - mq_out, - ) - return self._gdn_epilogue(out, z, L, num_tokens, gemm) - if not (conv_fused or fp8_conv or nv_conv): - causal_conv1d_update( - x=mixed_qkv, - conv_state=core._bound.conv[i], - weight=L.conv_w, - bias=None, - silu_activation=True, - cache_seqlens=None, - conv_state_indices=si_all, - pad_slot_id=_PAD_SLOT_ID, - ) - # The recurrence and the gated norm as one launch over `proj`, whose q/k/v - # columns the conv step rewrote in place; it computes both gates and the q/k L2 - # norm itself, at the softplus pair (1, 20) and epsilon 1e-6 this target's - # constants name, and reads this layer's pool rows before its PDL wait (nothing - # in a single-token step writes them ahead of it). - o = gdn_decode_rows_gated_norm( - proj, - ssm_states, - si_all, - L.a_log, - L.dt_bias, - L.gnorm, - core.eps, - core.num_k_heads, - core.head_k_dim, - ) - return _dense_mm(o, L.out, L.out8, gemm, "out", L.out_g, mid=L.out_m) - - def _gdn_epilogue(self, out, z, L, num_tokens, gemm=cublas_mm) -> torch.Tensor: - """Gated RMS norm (sigmoid gate, plain non-gemma weight) and output projection - shared by every recurrence.""" - core = self.core - o = rms_norm_gated_token_major( - torch.reshape(out, [-1, core.head_v_dim]), - torch.reshape(z, [num_tokens, core.num_v_heads, core.head_v_dim]), - weight=L.gnorm, - eps=core.eps, - fp8_scale=None, - gate_activation="sigmoid", - ) - o = torch.reshape(o, [num_tokens, core.value_dim]) - return _dense_mm(o, L.out, L.out8, gemm, "out", L.out_g, mid=L.out_m) - - def _ple_block(self, hc, input_ids, ple_state, next_norm, stage, gemm=cublas_mm) -> tuple: - """The n-gram side path on one decoder layer, added into the stream and - normalized for the next site. Returns `(hc, normed)`. - - The same steps as the prefill target's `_ple_block` -- id history, hash, table - gather and dequant, the gated injection, the dilated short conv -- in three - launches plus the two projections, because a decode batch carries at most - `gen_p` rows per sequence: - - 1. `ple_ngram_embed`: window, hash, gather, dequantize, and the id - history published back, one program per sequence; - 2. the key and value projections, on the step's bf16 projection entry - `gemm` (see `forward`); - 3. `ple_gate_conv_inject_norm`: the key / query grouped norms, the gate, - the conv-input grouped norm, the conv against the paged conv state - (published back), the residual add into the stream and this - layer's attention-side grouped norm, one program per - (sequence, stream). - - Every generation sequence carries its history by construction, so no carry - flags are passed. A verify step (`stage`) stages what `SpeculativeRollback` - rewrites the two pools from after sampling: both pools' pre-step rows before - either launch publishes, and the conv input the third launch writes out.""" - core = self.core - S = self._stack - conv_states, ngram_pool, si_all, carry_all, qsl, gen_p = ple_state - if stage: - core.spec_rollback.stage_ple_verify(ngram_pool, conv_states, input_ids) - ngram = ple_ngram_embed( - input_ids, - ngram_pool, - si_all, - None, - qsl, - core.w["ple_table"], - S.ple_tscale, - S.ple_mult, - S.ple_vocab, - S.ple_offs, - core.eos_token_id, - gen_p, - ) - conv_input = torch.empty_like(hc) if stage else None - out = ple_gate_conv_inject_norm( - hc, - gemm(ngram, S.ple_key), - gemm(ngram, S.ple_val), - S.ple_nk_w, - S.ple_nq_w, - S.ple_nc_w, - S.ple_conv_w, - conv_states, - si_all, - qsl, - next_norm, - core.eps, - core.hc_count, - _PLE_CLAMP_MIN, - core.ple_dilation, - _PAD_SLOT_ID, - gen_p, - conv_input=conv_input, - ) - if stage: - core.spec_rollback.stage_ple_conv_input(conv_input) - return out - - # -------------------------------------------------------------- forward --- - - def forward( - self, - attn_metadata: AttentionMetadata, - input_ids: torch.IntTensor | None = None, - position_ids: torch.IntTensor | None = None, - **kwargs, - ): - core = self.core - md = attn_metadata - self._check_step_contract(md) - step = self.step_args(md) - spec_metadata = kwargs.get("spec_metadata") - mm = md.mamba_metadata - # The bf16 block projections of a graph-captured step run as PDL launches - # (`_decode_projection`); an eager step keeps cuBLAS. - gemm = _decode_projection if md.is_cuda_graph else cublas_mm - - # A captured verify step builds `pos3`, the PLE segmentation vector and the - # rollback's staging in one launch (`verify_step_setup`, below). - rollback = core.spec_rollback - fused_setup = ( - md.is_cuda_graph - and spec_metadata is not None - and 0 < md.num_seqs < _VERIFY_SETUP_MAX_SEQS - and position_ids.shape[-1] > md.num_seqs - and position_ids.stride(-1) == 1 - and rollback is not None - and rollback.index_planes is not None - ) - # Two views of one `[3, 1, T]` buffer (see PrefillTarget.forward): contiguous - # `pos3` and `coords`. - coords = torch.transpose(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]), 0, 1) - pos3 = None if fused_setup else _copy_positions(position_ids) - - h = nn.functional.embedding(input_ids, core.w["embed"]) - num_tokens = h.shape[0] - dt = h.dtype - dev = h.device - - # The batch's bookkeeping, built once; every sequence is a generation sequence. - batch = md.num_seqs - # Host-int check (graph-safe); see PrefillTarget.prepare. - assert num_tokens == md.num_tokens, ( - f"the engine handed {num_tokens} activation rows for a batch its " - f"metadata describes as {md.num_tokens} tokens " - f"(padded_num_tokens={getattr(md, 'padded_num_tokens', None)}). A " - "prefill/piecewise CUDA-graph backend pads the activation buffer " - "past the batch and expects the model to slice it back; this " - "target does not implement that configuration. Decode CUDA graphs " - "pad with real dummy requests instead, so they are in the metadata " - "and this assert is silent for them." - ) - # Rows per generation sequence: 1 without speculation, else - # `runtime_draft_len + 1`. - gen_p = num_tokens // batch if batch else 1 - assert gen_p * batch == num_tokens, ( - f"{num_tokens} generation rows do not divide over {batch} generation " - "sequences; a ragged per-sequence draft length can be expressed " - "neither to the attention call nor to the multi-token recurrence" - ) - assert core.mtp_enabled or gen_p == 1, ( - f"{num_tokens} generation rows over {batch} generation sequences with " - "speculation off; the base path implements one query token per " - "generation sequence" - ) - # `gen_p > 1` is load-bearing for verify (see PrefillTarget.prepare). - if spec_metadata is not None and batch > 0: - declared_p = spec_metadata.runtime_draft_len + 1 - assert gen_p == declared_p, ( - f"the batch carries {gen_p} query rows per generation sequence " - f"but the worker's runtime_draft_len says {declared_p}; the " - "per-step conv snapshots and the recurrent rollback buffer are " - "sized from the second and written from the first" - ) - check_verify_width(gen_p) - check_replay_width(core._bound.replay, gen_p) - is_verify = spec_metadata is not None and batch > 0 and gen_p > 1 - step["predicted_tokens_per_seq"] = gen_p - - si_all = mm.state_indices[:batch] - qsl_long = mm.query_start_loc_long - int_idx = None - replay = None - if is_verify: - # The runtime's arange over the slot pool (what `update_mamba_states` reads). - assert core._bound.int_idx is not None, "the cache-manager bind has not run" - int_idx = core._bound.int_idx[:batch] - if core._bound.replay is not None: - replay = _replay_work(core, mm, batch) - gdn_state = (si_all, is_verify, batch, gen_p, int_idx, replay) - - # PLE segmentation vector: rebuilt from `seq_lens_cuda` when `gen_p > 1` (see - # PrefillTarget.prepare). Every generation sequence carries its n-gram history, so - # the PLE block takes no carry flags. - fused_setup = fused_setup and is_verify - if pos3 is None and not fused_setup: - pos3 = _copy_positions(position_ids) - if fused_setup: - planes = rollback.index_planes - (pos3, ple_qsl, staged_slots, staged_pages, staged_within, staged_index) = ( - verify_step_setup( - position_ids, - md.seq_lens_cuda[:batch], - si_all, - md.qsa_req_idx_per_token[:num_tokens], - md.qsa_logical_positions[:num_tokens], - md.qsa_block_table, - md.tokens_per_block, - *planes, - ) - ) - elif gen_p == 1: - ple_qsl = qsl_long - else: - ple_qsl = cu_seqlens_from_seq_lens(md.seq_lens_cuda[:batch], torch.int64) - ple_state = ( - core._bound.ple_state[0], - core._bound.ple_state[1], - si_all, - None, - ple_qsl, - gen_p, - ) - self._trunk_pos3 = (_position_key(position_ids), pos3) - - # Dense/sparse decision once per forward, from a host-side metadata read: a - # per-capture constant matching the engine's short/long decode graph families. - sparse = batch > 0 and int(md.kv_lens_runtime[:batch].max()) > core.dense_threshold - - qsa = ( - md.qsa_req_idx_per_token[:num_tokens], - md.qsa_logical_positions[:num_tokens], - md.qsa_sequence_lengths[:num_tokens], - md.qsa_visible_blocks[:num_tokens], - md.qsa_topk_indices[:num_tokens], - md.qsa_topk_row_starts[:num_tokens], - md.qsa_block_table, - md.tokens_per_block, - pos3, - coords, - sparse, - ) - - # Staging for the speculative rollback, as on the prefill target (see - # PrefillTarget.prepare); every row here is a generation row. - stage = is_verify - if rollback is not None: - rollback.discard() - if fused_setup: - rollback.begin_staged( - num_gens=batch, - rows_per_gen=gen_p, - slots=staged_slots, - pages=staged_pages, - within=staged_within, - index_all=staged_index, - ) - elif stage: - rollback.begin( - num_contexts=0, - prefix_tokens=0, - num_gens=batch, - rows_per_gen=gen_p, - slots=si_all, - request_indices=md.qsa_req_idx_per_token[:num_tokens], - logical_positions=md.qsa_logical_positions[:num_tokens], - block_table=md.qsa_block_table, - tokens_per_block=md.tokens_per_block, - fused=md.is_cuda_graph, - ) - - # Scratch shared by all attention layers: thop_attention overwrites rows - # [:num_tokens] of `output` in place, and the packed buffers are fully - # rewritten every layer (the QKV one on the dense path only). - qkv = torch.empty([num_tokens, core.qkv_width], dtype=dt, device=dev) - index_scratch = torch.empty([num_tokens, core.index_scratch_width], dtype=dt, device=dev) - attn_out = torch.empty([num_tokens, core.q_width], dtype=dt, device=dev) - # The CuteDSL routed output, shared by all layers: each layer's unpermute rewrites - # every row, and the next site's injection consumes it before the next layer's MoE. - routed_out = ( - torch.empty([num_tokens, core.hidden], dtype=dt, device=dev) - if core.trunk_moe_backend == "CUTEDSL" - else None - ) - - # The embedding repeated across the hc_count residual streams. - hc_count = core.hc_count - hc = torch.cat([h] * hc_count, -1) - layers, S = self._layers, self._stack - # Every hyper-connection site (`_hc_site`): the previous block's pending - # injection, this site's grouped norm and its mixer. The stack's first - # site has no pending injection; at the PLE layer it is applied alone - # (the PLE block reads the injected stream), and the PLE tail returns the - # stream with its output added plus this site's grouped norm of it, which - # the site's cuBLAS mixer starts from (a GEMV site normalizes inside its - # first launch and reads the stream alone). A MoE block's output is - # pending uncombined (`moe`): the next site's first kernel forms it, - # except before the PLE add, where the combine runs alone. - block_out = inject = moe = None - for i, L in enumerate(layers): - if i == core.ple_layer: - if block_out is not None: - if moe is not None: - # `_moe_block` never folds the finalize before this layer. - block_out = fused_sigmoid_gate_mul_add(block_out, *moe[:2]) - hc = qwen4_exp_post_inject(hc, block_out, inject) - block_out = inject = moe = None - hc, normed = self._ple_block(hc, input_ids, ple_state, L.ahc_norm, stage, gemm) - else: - normed = None - hc, mixed, inject = self._hc_site( - hc, - block_out, - inject, - L.ahc_norm, - L.ahc_first, - L.ahc_up, - True, - normed=normed, - moe=moe, - ) - if core.is_attn[i]: - block_out = self._attention_block( - mixed, i, L, qkv, index_scratch, attn_out, qsa, step, stage, gemm - ) - else: - block_out = self._gdn_block(mixed, i, L, gdn_state, gemm, md.is_cuda_graph) - hc, mixed, inject, xq, xsf = self._hc_moe_site( - hc, block_out, inject, L.mhc_norm, L.mhc_first, L.mhc_up, L.moe_gscale - ) - # The trtllm-gen runner's finalize folds into the next site's first kernel - # when that site takes the combine (every one but the PLE layer's). - fold = core.trunk_moe_backend == "TRTLLM" and i + 1 != core.ple_layer - block_out, moe = self._moe_block(mixed, xq, xsf, L, routed_out, gemm, fold=fold) - # The stack-final mixer with no injection weight: it returns the mix - # alone, and that mix *is* this checkpoint's final norm. - hc, final, _ = self._hc_site( - hc, block_out, inject, S.fhc_norm, S.fhc_first, S.fhc_up, False, moe=moe - ) - if spec_metadata is None: - return final - # `lm_head` reads the `[T, hidden]` mix; the MTP module reads the - # `[T, hc_count * hidden]` stream it was collapsed from. - return final, hc - - # ------------------------------------------------------------- the draft --- - - def _mtp_moe_block(self, x, dt) -> torch.Tensor: - """FP8 block-scale (or, with `_MTP_MOE_NVFP4_ENV`, NVFP4) draft MoE; see - `PrefillTarget._mtp_moe_block`.""" - core = self.core - M = core._mtp - if M.moe_gscale is not None: - return _mtp_moe_nvfp4_block(core, M, x) - logits = cublas_mm(x, M.router) - topk_ids, topk_w = renorm_moe_routing_op(logits, topk=core.topk, output_dtype=dt) - xq, xsf = fp8_quantize_1x128(x, use_ue8m0=_MTP_MOE_USE_UE8M0) - routed = fp8_block_scale_moe_runner( - routing_logits=None, - routing_bias=None, - hidden_states=xq, - hidden_states_scale=xsf, - gemm1_weights=M.fc1_w, - gemm1_weights_scale=M.fc1_s, - gemm2_weights=M.fc2_w, - gemm2_weights_scale=M.fc2_s, - num_experts=core.num_experts, - top_k=core.topk, - num_fused_shared_experts=None, - n_group=None, - topk_group=None, - intermediate_size=core.moe_inter, - local_expert_offset=0, - local_num_experts=core.num_experts, - routed_scaling_factor=None, - routing_method_type=_ROUTING_METHOD_RENORM, - topk_weights=topk_w, - topk_ids=topk_ids, - act_type=_ACT_TYPE_SWIGLU, - gemm1_clamp_limit=None, - ) - shared = cublas_mm(flashinfer_silu_and_mul(cublas_mm(x, M.sh_gu)), M.sh_dn) - gate_logits = cublas_mm(x, M.sh_gate) - return fused_sigmoid_gate_mul_add(routed, gate_logits, shared) - - def _mtp_layer( - self, - embed_tokens, - input_ids, - hidden_states, - md, - step, - draft, - layer_idx, - rows, - gen_seqs, - p, - sparse, - share, - reuse, - pos3, - pos, - coords, - ) -> torch.Tensor: - """The draft layer on the trunk's decode machinery: the front end's two projections, - both hyper-connection sites as `_hc_site` (the front end's per-stream `x_h + x_e` is - the first site's injection at unit weight, the attention output the second's), the - attention as `_mtp_attention`, the FP8 MoE as `_mtp_moe_block`. Returns the stream - with the MoE output injected, as the composed step does. The hyper-connection - mixers compute the composed `qwen4_exp_gated_residual`'s formulas with fp32 logits - and an fp32 bottleneck (the trunk sites' numerics), so the draft's tokens may differ - from the composed step's in the last place; the verify step decides acceptance.""" - core = self.core - M, F = core._mtp, core._mtp_fused - gemm = _decode_projection if md.is_cuda_graph else cublas_mm - e = nn.functional.embedding(input_ids, embed_tokens) - x_e = gemm(flashinfer_gemma_rmsnorm(e, M.pre_e_norm, core.eps), M.fc_e) - h_n = self._grouped_norm(hidden_states, M.pre_h_norm) - x_h = torch.reshape( - gemm(torch.reshape(h_n, [rows * core.hc_count, core.hidden]), M.fc_h), - [rows, core.stream], - ) - x, mixed, inject = self._hc_site( - x_h, x_e, F.unit_inject[:rows], F.ahc_norm, F.ahc_first, F.ahc_up, True - ) - attn = self._mtp_attention( - mixed, - md, - step, - draft, - layer_idx, - rows, - gen_seqs, - p, - sparse, - share, - reuse, - pos3, - pos, - coords, - gemm, - ) - x, mixed, inject = self._hc_site(x, attn, inject, F.mhc_norm, F.mhc_first, F.mhc_up, True) - return qwen4_exp_post_inject(x, self._mtp_moe_block(mixed, x.dtype), inject) - - def _mtp_attention( - self, - x, - md, - step, - draft, - layer_idx, - rows, - gen_seqs, - p, - sparse, - share, - reuse, - pos3, - pos, - coords, - gemm, - ) -> torch.Tensor: - """The draft layer's gated GQA and QSA index branch, as `_attention_block` runs a - trunk layer over the draft manager's planes (`draft`): one input projection - against `_MtpFused.attn_in`, the composed prologue (norm + rope of both planes at - the module's rope base, the index store and compression on the side stream), the - K/V append, the index scores and the one-launch top-k + expansion, sparse GQA, the - output gate and projection. Under IndexShare a draft decode step (`reuse`) skips the - index branch and rebuilds its selection from the draft-extend step's capture; the - draft-extend step (`share`) captures it.""" - core = self.core - M, F = core._mtp, core._mtp_fused - dt, dev = x.dtype, x.device - proj = gemm(x, F.attn_in) - qkv_view, gate, index_view = torch.split( - proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1 - ) - q_dst, k_dst, v_dst = _attn_in_heads(core, proj) - index_scratch = torch.empty([rows, core.index_scratch_width], dtype=dt, device=dev) - q_index = torch.split( - torch.reshape( - index_scratch, - [rows, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim], - ), - [core.index_heads, core.index_kv_heads, core.index_kv_heads], - 1, - )[0] - index_section = (_QSA_MROPE_SECTION or (0, 0, 0)) if _QSA_INDEX_USE_MROPE else (0, 0, 0) - req_idx = md.qsa_req_idx_per_token[:rows] - logical = md.qsa_logical_positions[:rows] - block_table = md.qsa_block_table - tokens_per_block = md.tokens_per_block - main = torch.cuda.current_stream() - side = core._side_stream - side.wait_stream(main) - qwen4_exp_qk_norm_rope( - proj, - q_index, - pos3, - core._mtp_cos_sin, - M.q_norm, - M.k_norm, - M.idx_qn, - num_heads_q=core.heads_q, - num_heads_k=core.heads_kv, - head_dim=core.head_dim, - index_heads=core.index_heads, - index_head_dim=core.index_head_dim, - index_q_col=core.qkv_width + core.q_width, - rotary_dim=core.rotary_dim, - eps=core.eps, - mrope_section1=core.mrope_section1, - mrope_section2=core.mrope_section2, - index_mrope_section1=index_section[1], - index_mrope_section2=index_section[2], - ) - if not reuse: - # Raw key stored before anything normalizes it (see `_attention_prologue`). - with torch.cuda.stream(side): - token_k = torch.split( - torch.reshape( - index_view, - [ - rows, - core.index_scratch_width // core.index_head_dim, - core.index_head_dim, - ], - ), - [core.index_heads, core.index_kv_heads, 1], - 1, - )[1] - qsa_paged_index_store( - token_k, - coords, - draft.index_cache, - draft.position_cache, - req_idx, - logical, - block_table, - tokens_per_block, - ) - qsa_prefill_compress( - logical_positions=logical, - request_indices=req_idx, - block_table=block_table, - index_cache=draft.index_cache, - position_cache=draft.position_cache, - k_norm_weight=M.idx_kn, - cos_sin=core._mtp_cos_sin, - eps=core.eps, - tokens_per_block=tokens_per_block, - compress_ratio=core.compress_ratio, - mrope_section=_QSA_MROPE_SECTION, - ) - main.wait_stream(side) - gated = False - if sparse: - kv = draft.kv - _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) - topk_indices = md.qsa_topk_indices[:rows] - visible_blocks = md.qsa_visible_blocks[:rows] - seq_lens_q = md.qsa_sequence_lengths[:rows] - if reuse: - slot = draft.share_slot - assert slot in md.qsa_shared_topk_captured_slots, ( - "a draft decode step would reuse a selection that no draft-extend step " - "of this loop captured" - ) - qsa_mtp_reuse_topk( - shared_indices=md.qsa_shared_topk_indices[slot], - shared_visible_blocks=md.qsa_shared_topk_visible_blocks[slot], - request_indices=req_idx, - visible_blocks=visible_blocks, - output=topk_indices, - ) - selected = qsa_expand_block_indices( - topk_indices, - logical, - seq_lens_q, - compress_ratio=core.compress_ratio, - token_topk=core.token_topk, - ) - else: - scores = qsa_paged_index_scores( - q=q_index, - index_cache=draft.index_cache, - block_table=block_table, - query_positions=logical, - request_indices=req_idx, - tokens_per_block=tokens_per_block, - compress_ratio=core.compress_ratio, - only_visible_blocks=True, - context_rows=False, - ) - if _topk_expand_supports(scores): - selected = qsa_topk_expand_decode( - scores, - visible_blocks, - logical, - seq_lens_q, - topk_indices, - core.compress_ratio, - core.token_topk, - ) - else: - _select( - core, scores, md.qsa_topk_row_starts[:rows], visible_blocks, topk_indices - ) - selected = qsa_expand_block_indices( - topk_indices, - logical, - seq_lens_q, - compress_ratio=core.compress_ratio, - token_topk=core.token_topk, - ) - if share: - self._capture_selection(md, topk_indices, visible_blocks, gen_seqs, p) - if kv.k_scale is not None: - attn = _sparse_gqa( - core, - kv, - q=q_dst, - block_table=block_table, - selected_tokens=selected, - request_indices=req_idx, - tokens_per_block=tokens_per_block, - query_positions=None, - compress_ratio=None, - output_gate=gate, - ) - gated = True - else: - gated = qsa_paged_sparse_gqa.fuses_output_gate(q_dst, kv.k, selected, gate) - attn = qsa_paged_sparse_gqa( - q=q_dst, - k_cache=kv.k, - v_cache=kv.v, - block_table=block_table, - selected_tokens=selected, - request_indices=req_idx, - tokens_per_block=tokens_per_block, - softmax_scale=core.softmax_scale, - query_positions=None, - compress_ratio=None, - output_gate=gate if gated else None, - ) - out = torch.reshape(attn, [rows, core.q_width]) - else: - qkv = torch.empty([rows, core.qkv_width], dtype=dt, device=dev) - qkv.copy_(qkv_view) - attn_out = torch.empty([rows, core.q_width], dtype=dt, device=dev) - out = thop_attention( - q=qkv, - output=attn_out, - local_layer_idx=layer_idx, - predicted_tokens_per_seq=p, - num_heads=core.heads_q, - num_kv_heads=core.heads_kv, - head_size=core.head_dim, - quant_mode=core.quant_mode, - kv_scale_orig_quant=core.kv_scale_orig_quant, - kv_scale_quant_orig=core.kv_scale_quant_orig, - **step, - **_CALL_CONSTANTS, - ) - if not gated: - fused_sigmoid_mul_inplace(out, gate) - return gemm(out, M.o) - - def _capture_selection(self, md, topk_indices, visible_blocks, batch, p) -> None: - """IndexShare: stash, per request, the selection of its last accepted draft-extend - row with that row's complete-group count, in the draft layer's row of the shared - buffers. Every request carries `p` rows here, and the worker continues the draft - loop from row `num_accepted - 1` of each (the golden token counts as accepted).""" - assert md.qsa_mtp_num_accepted is not None, ( - "IndexShare needs the worker's accepted counts to pick each request's capture " - "row; the draft loop did not supply them" - ) - accepted = md.qsa_mtp_num_accepted[:batch].to(torch.int64) - rows = torch.arange(batch, dtype=torch.int64, device=topk_indices.device) * p - rows = rows + torch.clamp(accepted - 1, min=0, max=p - 1) - slot = self.core._draft_bound.share_slot - md.qsa_shared_topk_indices[slot, :batch].copy_(torch.index_select(topk_indices, 0, rows)) - md.qsa_shared_topk_visible_blocks[slot, :batch].copy_( - torch.index_select(visible_blocks, 0, rows) - ) - # Host-side readiness; see PrefillTarget._capture_selection. - md.qsa_shared_topk_captured_slots.add(slot) - - def mtp_forward( - self, - embed_tokens: torch.Tensor, - all_rank_num_tokens, - input_ids: torch.Tensor, - position_ids: torch.Tensor, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - **kwargs, - ) -> torch.Tensor: - """One draft step over generation rows only (every step past the first in an - engine forward, and the first when the batch had no context). - - `spec_metadata` and `all_rank_num_tokens` are accepted and unused.""" - core = self.core - assert isinstance(attn_metadata, TrtllmAttentionMetadata) - assert core._bound is not None, ( - "the trunk's cache-manager bind has not run; the shell calls the " - "core before the worker, so this cannot be reached first" - ) - md = attn_metadata - rows = hidden_states.shape[0] - assert input_ids.shape[0] == rows, ( - f"the draft step's {input_ids.shape[0]} token ids and {rows} " - "hidden-state rows must describe the same tokens" - ) - assert hidden_states.shape[1] == core.stream, ( - f"the MTP module reads the trunk's {core.stream}-wide four-stream " - f"residual; the runtime handed it {hidden_states.shape[1]} columns" - ) - - # Projected on every call: the worker rewrites `attn_metadata` in place - # between draft steps (see PrefillTarget.mtp_forward). - step = self.step_args(md) - gen_seqs = md.num_seqs - p = rows // gen_seqs if gen_seqs else 1 - assert p * gen_seqs == rows, ( - f"{rows} generation rows do not divide over {gen_seqs} " - "generation sequences; a ragged per-sequence draft length cannot be " - "expressed to the attention call" - ) - - # The KV cache written here is the draft manager's, swapped in by the worker - # per call (see PrefillTarget.mtp_forward); take the layer index from the - # current manager every time. - cm = md.kv_cache_manager - assert cm is not None, "a paged KV cache is required" - offsets = cm.layer_offsets - assert core.mtp_layer_idx in offsets, ( - f"the MTP module's attention is global layer index " - f"{core.mtp_layer_idx}, which the cache manager " - f"{type(cm).__name__} does not map (it maps {len(offsets)} layers, " - f"is_draft={getattr(cm, 'is_draft', None)}); the engine allocated " - "no KV layer for the draft module" - ) - layer_idx = offsets[core.mtp_layer_idx] - window = cm.max_attention_window_vec[layer_idx] - assert window is None or window == md.max_seq_len, ( - f"the MTP layer's pool window is {window!r} but this call passes " - f"max_seq_len {md.max_seq_len} as the attention window; a sliding " - "window is not implemented here" - ) - - # The draft layer's QSA planes and the page table that addresses them. With a - # separate draft manager the engine's QSA metadata swaps `qsa_block_table` for the - # table it decodes from the draft manager's offsets (`prepare_for_draft_forward`); - # a table still decoding the trunk's offsets would read and write another pool's - # pages. - draft = core._draft_bound - if draft is None or cm is not draft.cm: - core._bind_draft_cache_manager(md) - draft = core._draft_bound - assert cm is core._bound.cm or md.qsa_block_table is getattr( - md, "qsa_draft_block_table", None - ), ( - "the draft step is bound to a separate draft KV cache manager but " - "qsa_block_table is not that manager's table: the engine's QSA metadata did " - "not swap its page table for the draft forward" - ) - # The `qsa_*` row mapping is sized from the metadata's token count, which the - # worker rewrites together with it between draft steps. - assert rows == md.num_tokens, ( - f"the draft step carries {rows} rows for a batch its metadata describes as " - f"{md.num_tokens} tokens; the QSA row mapping covers the second" - ) - # The trunk's dense/sparse decision from the same host buffer: the worker advances - # only the device KV lengths through the draft loop, so every draft step of one - # engine forward takes the branch its trunk took, in the CUDA-graph family the - # runner chose for that trunk. - sparse = gen_seqs > 0 and int(md.kv_lens_runtime[:gen_seqs].max()) > core.dense_threshold - # IndexShare: the loop's draft-extend step (`p` rows per request) runs the indexer - # and, with sharing on, captures each request's selection; every draft decode step - # after it (one row per request) reuses that capture instead of running the indexer. - share = core.mtp_index_share and md.qsa_in_mtp_draft_loop - reuse = share and md.qsa_indexer_skip_topk - assert not share or reuse == (p == 1), ( - f"the worker's reuse flag ({md.qsa_indexer_skip_topk}) disagrees with the " - f"step's shape ({p} rows per request): only the draft decode steps after the " - "draft-extend step reuse its selection" - ) - - # Copy into a fresh contiguous `[3, N]`: a sliced view keeps the source row - # stride and the op rejects it under `use_mrope`. - assert position_ids.shape[0] == _QSA_POSITION_AXES, ( - f"expected the mRoPE position layout with {_QSA_POSITION_AXES} " - f"coordinate rows, got {tuple(position_ids.shape)}" - ) - # The first draft step reads the trunk's positions buffer, whose int32 copy the trunk - # made (taken once); a contiguous int32 `position_ids` (every later draft step) is - # read in place. The step's ops only read positions. - trunk, self._trunk_pos3 = self._trunk_pos3, None - if trunk is not None and trunk[0] == _position_key(position_ids): - pos3 = trunk[1] - elif position_ids.dtype == torch.int32 and position_ids.is_contiguous(): - pos3 = torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]) - else: - pos3 = _copy_positions(position_ids) - # The trunk's two other views, from `pos3`: `pos` (coordinate 0 alone, which the - # index rope reads) and `coords` (the `[N, 3]` transpose the position plane stores). - pos = torch.reshape(torch.split(pos3, [1, _QSA_POSITION_AXES - 1], 0)[0], [-1]) - coords = torch.transpose(pos3, 0, 1) - - return self._mtp_layer( - embed_tokens, - input_ids, - hidden_states, - md, - step, - draft, - layer_idx, - rows, - gen_seqs, - p, - sparse, - share, - reuse, - pos3, - pos, - coords, - ) - - -class _VisionBlockWeights(NamedTuple): - """One vision block's weights (`.t()` GEMM views). LayerNorm has a bias, unlike - every norm in the language tower.""" - - n1_w: torch.Tensor - n1_b: torch.Tensor - qkv: torch.Tensor - qkv_b: torch.Tensor - proj: torch.Tensor - proj_b: torch.Tensor - n2_w: torch.Tensor - n2_b: torch.Tensor - fc1: torch.Tensor - fc1_b: torch.Tensor - fc2: torch.Tensor - fc2_b: torch.Tensor - - -class _VisionWeights(NamedTuple): - """The tower's weights: patch embedding, position table, blocks and patch merger.""" - - patch_w: torch.Tensor - patch_b: torch.Tensor - pos: torch.Tensor - blocks: list - m_n_w: torch.Tensor - m_n_b: torch.Tensor - m_fc1: torch.Tensor - m_fc1_b: torch.Tensor - m_fc2: torch.Tensor - m_fc2_b: torch.Tensor - - -def _vision_views(core: Qwen4ExpTp1ModelingV2) -> _VisionWeights: - w = core.w - return _VisionWeights( - # Not transposed: the entry takes the checkpoint's `[out, in, kt, kh, kw]`. - patch_w=w["vis_patch_w"], - patch_b=w["vis_patch_b"], - pos=w["vis_pos"], - blocks=[ - _VisionBlockWeights( - n1_w=w[f"l{b}_vis_n1_w"], - n1_b=w[f"l{b}_vis_n1_b"], - # Row layout (q heads, k, v) is already thop_attention's packed form. - qkv=w[f"l{b}_vis_qkv_w"].t(), - qkv_b=w[f"l{b}_vis_qkv_b"], - proj=w[f"l{b}_vis_proj_w"].t(), - proj_b=w[f"l{b}_vis_proj_b"], - n2_w=w[f"l{b}_vis_n2_w"], - n2_b=w[f"l{b}_vis_n2_b"], - fc1=w[f"l{b}_vis_fc1_w"].t(), - fc1_b=w[f"l{b}_vis_fc1_b"], - fc2=w[f"l{b}_vis_fc2_w"].t(), - fc2_b=w[f"l{b}_vis_fc2_b"], - ) - for b in range(core.v_depth) - ], - m_n_w=w["vis_m_n_w"], - m_n_b=w["vis_m_n_b"], - m_fc1=w["vis_m_fc1_w"].t(), - m_fc1_b=w["vis_m_fc1_b"], - m_fc2=w["vis_m_fc2_w"].t(), - m_fc2_b=w["vis_m_fc2_b"], - ) - - -def _frame_calls( - grid: list[tuple[int, int, int]], max_rows: int -) -> list[list[tuple[int, int, int]]]: - """Cut `grid` (one `(frames, height, width)` row per image or video) into tower calls - of whole frames, in order, each at most `max_rows` patch rows. A call is its grid - rows: the frames of one item that fall in it are one `(frames, height, width)` row.""" - calls: list[list[tuple[int, int, int]]] = [[]] - room = max_rows - for frames, height, width in grid: - rows = height * width - assert rows <= max_rows, ( - f"one frame is {height}x{width} = {rows} patch rows; the vision ops are " - f"certified up to {max_rows} rows in one call" - ) - while frames: - take = min(frames, room // rows) - if take == 0: - calls.append([]) - room = max_rows - continue - calls[-1].append((take, height, width)) - room -= take * rows - frames -= take - return calls - - -class VisionTower: - """One request's images or videos through the 27-block ViT and its patch merger. - - Not a target: runs only for a context request carrying media, from the shell's - splice. Attention is bidirectional per frame (a video's frame group is a frame), - expressed to `thop_attention` as one context "sequence" per frame; running the - rows as one sequence is 315x-582x outside tolerance. The position-embedding - interpolation and 2-D rope must both emit in spatial-merge order (what the - merger's 4-to-1 reshape consumes); a raster-order emitter gives a correctly - shaped, silently wrong tower. - - The frames run in calls of at most `_VISION_CALL_MAX_PATCHES` rows, cut only between - frames: attention is per frame, the position embedding and rope are per grid row - and repeat per frame, the norms and MLPs are row-wise, and a frame's `h * w` is a - multiple of `merge**2`, so a cut keeps every merger block whole. One frame past the - bound cannot be served.""" - - def __init__(self, core: Qwen4ExpTp1ModelingV2) -> None: - self.core = core - self._w = _vision_views(core) - # Persistent scratch: `thop_attention` grows `workspace` in place; never - # the engine's own workspace. - self._ws: torch.Tensor | None = None - - def __call__(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: - """`pixel_values` is `[sum(t*h*w), in_channels*kt*kh*kw]`, one flat patch per - row, channel-slowest. `grid_thw` is one `(frames, height, width)` row per - image or video and stays on the host (no synchronization). Returns - `[sum(t*h*w) / merge**2, out_hidden_size]`, one row per placeholder, in grid - row order.""" - core = self.core - dt = pixel_values.dtype - assert dt == core.w["embed"].dtype, ( - f"the tower runs at the activation dtype {core.w['embed'].dtype}; the " - f"runtime handed pixel_values as {dt}. The input processor's `dtype` " - "property is what decides this, and no cast happens here" - ) - assert pixel_values.dim() == 2 and pixel_values.shape[1] == core.v_pixel_width, ( - f"pixel_values must be [patches, {core.v_pixel_width}] flat patch rows; " - f"got {tuple(pixel_values.shape)}" - ) - - grid: list[tuple[int, int, int]] = [] - for row in grid_thw.tolist(): - assert len(row) == 3, f"grid_thw rows are (t, h, w); got {row}" - frames, height, width = (int(v) for v in row) - assert frames >= 1 and height >= 1 and width >= 1, row - assert height % core.v_merge == 0 and width % core.v_merge == 0, ( - f"grid sides must divide by spatial_merge_size {core.v_merge}; got " - f"height={height}, width={width}" - ) - grid.append((frames, height, width)) - assert grid, "grid_thw describes no image or video" - described = sum(t * h * w for t, h, w in grid) - assert described == pixel_values.shape[0], ( - f"grid_thw describes {described} patch rows and pixel_values carries " - f"{pixel_values.shape[0]}" - ) - - outputs = [] - start = 0 - for call in _frame_calls(grid, _VISION_CALL_MAX_PATCHES): - rows = sum(t * h * w for t, h, w in call) - outputs.append( - self._run(pixel_values[start : start + rows], call, grid_thw.new_tensor(call)) - ) - start += rows - return outputs[0] if len(outputs) == 1 else torch.cat(outputs, 0) - - def _run( - self, pixel_values: torch.Tensor, grid: list[tuple[int, int, int]], grid_thw: torch.Tensor - ) -> torch.Tensor: - """One tower call over whole frames: `grid` and `grid_thw` (the same rows, as a - host list and a host tensor) describe exactly the rows of `pixel_values`.""" - core = self.core - dt, dev = pixel_values.dtype, pixel_values.device - assert pixel_values.is_cuda, "pixel_values must reach the tower on the device" - - # ---- the frame segmentation, host-side ------------------------------- - # One attention segment per frame: a video's frames do not attend to each other. - seg_lens: list[int] = [] - for frames, height, width in grid: - seg_lens.extend([height * width] * frames) - num_patches = pixel_values.shape[0] - num_segments = len(seg_lens) - max_seg = max(seg_lens) - - # Batch-state tensors for the encoder-only call; they describe the frames. - host_lens = torch.tensor(seg_lens, dtype=torch.int32) - host_types = torch.zeros([num_segments], dtype=torch.int32) - host_total = torch.tensor([num_patches, 0], dtype=torch.int32) - dev_lens = host_lens.to(dev) - if self._ws is None: - self._ws = torch.empty([0], dtype=torch.int8, device=dev) - - # ---- patch embedding + absolute position ----------------------------- - V = self._w - h = torch.add( - patch_embed_conv3d(pixel_values, V.patch_w, V.patch_b), - pos_embed_interpolate(V.pos, grid, core.v_merge), - ) - attn_out = torch.empty([num_patches, core.v_hidden], dtype=dt, device=dev) - - for B in V.blocks: - qkv = cublas_mm(layer_norm(h, B.n1_w, B.n1_b, _VISION_NORM_EPS), B.qkv, B.qkv_b) - qkv_h = torch.reshape(qkv, [num_patches, 3 * core.v_heads, core.v_head_dim]) - q_view, k_view, _ = torch.split(qkv_h, [core.v_heads] * 3, 1) - q_rot, k_rot = vision_rope_2d( - q_view, k_view, grid_thw, core.v_merge, theta=_VISION_ROPE_THETA - ) - q_view.copy_(q_rot) - k_view.copy_(k_rot) - attn_out = thop_attention( - q=qkv, - output=attn_out, - workspace=self._ws, - sequence_length=dev_lens, - host_past_key_value_lengths=host_lens, - host_total_kv_lens=host_total, - context_lengths=dev_lens, - host_context_lengths=host_lens, - host_request_types=host_types, - max_num_requests=num_segments, - max_context_length=max_seg, - max_seq_len=max_seg, - attention_window_size=max_seg, - num_contexts=num_segments, - num_ctx_tokens=num_patches, - num_heads=core.v_heads, - num_kv_heads=core.v_heads, - head_size=core.v_head_dim, - **_VISION_CALL_CONSTANTS, - ) - h = torch.add(h, cublas_mm(attn_out, B.proj, B.proj_b)) - mlp = cublas_mm(layer_norm(h, B.n2_w, B.n2_b, _VISION_NORM_EPS), B.fc1, B.fc1_b) - # Non-gated MLP with the tanh GELU (not the merger's erf). - h = torch.add(h, cublas_mm(gelu_tanh(mlp), B.fc2, B.fc2_b)) - - # ---- the patch merger ------------------------------------------------ - # The norm runs over `hidden` before the 4-to-1 shuffle; the reshape is legal - # only because both order-bearing steps emitted in spatial-merge order. - merged = torch.reshape( - layer_norm(h, V.m_n_w, V.m_n_b, _VISION_NORM_EPS), [-1, core.v_merge_dim] - ) - # The erf GELU (`nn.GELU()`); within one bf16 ulp of tanh, hence a separate entry. - return cublas_mm(gelu_erf(cublas_mm(merged, V.m_fc1, V.m_fc1_b)), V.m_fc2, V.m_fc2_b) - - -class MTPLayer: - """The checkpoint's MTP module, replayed once per draft step: one more decoder - layer (two hyper-connections, gated GQA, 512-expert MoE) behind a two-branch - front end (`fc_embedding` over the NEXT token's embedding, `fc_hidden` over the - hidden state per stream). It returns `x` un-collapsed; `shared_head` owns the - module's mixer. The step runs on the targets (`mtp_forward`) over `core._mtp`. - - The MoE is FP8 128x128 weight blocks (the trunk's is NVFP4); re-quantizing it - would change the checkpoint's meaning, visible only as lower acceptance. - - Attention is the trunk's QSA layer over the module's own indexer. Every draft step - maintains the layer's index plane, then runs the dense core while the trunk's host KV - lengths stay within `dense_seq_len_threshold`, and selection plus sparse GQA above it. - With a separate draft KV manager the layer's K/V, index and position pages live in - that manager's pools: the worker swaps `attn_metadata.kv_cache_manager` and the block - offsets for the draft loop, `QSAAttentionMetadata.prepare_for_draft_forward` swaps - `qsa_block_table` to the table it decodes from the draft manager's offsets, and the - core binds the draft manager's planes once per manager (`_DraftBound`). The planes need - no rollback: each engine step's first draft call rewrites every position from the - newly committed length on. - - IndexShare (`sparse_attention_config.index_share_for_mtp_iteration`) runs the draft - indexer once per MTP iteration instead of once per draft step. The draft-extend step, - the loop's first, runs it and captures per request the selection of the row the loop - continues from, into the metadata's shared buffers; every draft decode step after it - skips the index branch and rebuilds its selection from that capture with - `qsa_mtp_reuse_topk`, which appends the groups completed since; a full row gives up - its oldest ids to them. The expansion adds the incomplete tail, so the tokens drafted - since the capture stay visible. Every selection is canonical first (ascending, - padding last), so the reuse keeps the order. - - Not an `nn.Module`: the weights live in the core's `ParameterDict`, and a second - registration would break the weight load's both-directions coverage assert.""" - - def __init__(self, core: Qwen4ExpTp1ModelingV2, logits_processor) -> None: - self.core = core - self.logits_processor = logits_processor - - def __call__(self, *args, **kwargs) -> torch.Tensor: - return self.forward(*args, **kwargs) - - def shared_head( - self, hidden_states, lm_head, attn_metadata, return_context_logits - ) -> torch.Tensor: - """The module's own head: its stack-final hyper-connection mixer (a full gated - mixer with its own parameters, not a plain RMSNorm), then the trunk's - `lm_head` through the shell's `logits_processor`.""" - core = self.core - F = core._mtp_fused - # The decode target's norm-only site: the GEMV pair at one row, the fused norm and - # mixer above (fp32 logits, as the trunk's stack-final mixer). - _, final, _ = core._decode._hc_site( - hidden_states, None, None, F.fhc_norm, F.fhc_first, F.fhc_up, False - ) - # The draft head reads its FP8 copy of `lm_head` (`post_load_weights`). - lm_head = _Fp8DraftHead(*core._mtp_head_fp8, core._lm_head_g, core._mtp_head_nv) - return self.logits_processor.forward( - final, - lm_head, - attn_metadata, - return_context_logits, - widen=not getattr(self.logits_processor, "draft_argmax_only", False), - ) - - def forward( - self, - embed_tokens: torch.Tensor, - all_rank_num_tokens, - input_ids: torch.Tensor, - position_ids: torch.Tensor, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - **kwargs, - ) -> torch.Tensor: - """Route the draft step by the phase this call's metadata shows, per call: the - worker rewrites `attn_metadata` between draft steps, so one engine step's - draft calls can cross from prefill to decode.""" - return self.core._select_target(attn_metadata).mtp_forward( - embed_tokens, - all_rank_num_tokens, - input_ids, - position_ids, - hidden_states, - attn_metadata, - **kwargs, - ) - - -class DraftModel: - """The container the spec worker reaches the drafter through. The runtime - inspects only `mtp_layers` (`[0]`: MTP-Eagle replays one layer), - `embed_tokens`, `lm_head`, and `model.d2t` (correctly absent: shared - vocabulary). - - `embed_tokens` resolves through the core on every read: this container is - built while parameters are meta tensors, and materialization replaces the - tensor objects, so a captured reference stays on meta.""" - - def __init__(self, core: Qwen4ExpTp1ModelingV2, lm_head, logits_processor) -> None: - self.core = core - self.mtp_layers = [MTPLayer(core, logits_processor)] - self.lm_head = lm_head - - @property - def embed_tokens(self) -> torch.Tensor: - return self.core.w["embed"] - - -class SpeculativeRollback: - """The target's auxiliary speculative-state handler: it undoes what a verify step's - rejected draft rows wrote into the four side-cache planes the runtime's own state - promotion does not cover. - - A verify step runs `p = runtime_draft_len + 1` rows per generation request through the - trunk, and four ops write per-request state for every row without knowing how many the - sampler keeps: the PLE n-gram id history (`ops/ple/ngram_token_history`), the PLE - dilated-conv state (`ops/ssm/dilated_short_conv_paged`), each attention layer's QSA - index-K plane (`qsa_paged_index_store`, `qsa_prefill_compress`) and the shared QSA - position plane. The Gated-DeltaNet state is not here: the runtime promotes it. - - For a generation request that accepted `a` rows (golden token included, so - `1 <= a <= p`) the commit rewrites: - - * the n-gram pool row to the id window of row `a - 1` minus its oldest id; - * the conv pool row to columns `[a, a + state_len)` of `[pre-step state | conv input - of the request's p rows]`; - * every index-K and position slot written by a row of ordinal `>= a` to its - pre-verify contents. - - The targets stage during the forward (`begin`, then `stage_ple_verify` before the - fused PLE launches and `stage_ple_conv_input` after them in the PLE block, and - `stage_index` before each attention layer's store) and the spec worker calls - `commit_speculative_states` after its draft loop, inside the same model forward. - Context rows and non-verify steps stage nothing, and a commit with nothing staged - returns. Staging and commit use metadata-sized - allocations, gathers, `torch.where` and indexed writes only, and everything staged is - allocated by the forward and consumed by that forward's commit, so a captured graph - holds both ends and replay re-reads the engine's buffers in place. - - A step opened with `fused` (the decode target's graph-captured steps) stages the PLE - pools with one `spec_rollback_stage` launch and, with the index planes bound, commits - with one `spec_rollback_commit` launch instead of the gathers and indexed writes.""" - - def __init__(self) -> None: - self._pending: dict | None = None - # (base addresses, page strides, plane count) of every attention layer's index-K - # plane, once the cache manager is bound: `begin` then gathers all of them at once and - # the commit restores them in one launch (`ops/attention/qsa_rollback_planes`). - self._planes: tuple | None = None - - @property - def index_planes(self) -> tuple | None: - """`(base addresses, page strides, plane count)` of the bound index-K planes.""" - return self._planes - - def begin_staged( - self, - *, - num_gens: int, - rows_per_gen: int, - slots: torch.Tensor, - pages: torch.Tensor, - within: torch.Tensor, - index_all: torch.Tensor, - ) -> None: - """Open a generation-only verify step's staging from what `verify_step_setup` - resolved: the requests' int64 recurrent slots, the rows' int64 QSA pages and offsets, - and every bound index-K plane at them. The PLE staging and the commit are the - one-launch ones (`begin(..., fused=True)`).""" - n, p = num_gens, rows_per_gen - assert n > 0 and p > 1, ( - f"a verify step stages generation requests carrying drafts; got {n} " - f"requests of {p} rows" - ) - assert pages.shape[0] == n * p and index_all.shape[1] == n * p, ( - f"{n} generation requests of {p} rows are {n * p} rows; got {pages.shape[0]} slots" - ) - self._pending = dict( - num_contexts=0, - prefix=0, - n=n, - p=p, - slots=slots, - pages=pages, - within=within, - index=[], - position=None, - index_all=index_all, - fused=True, - ngram_stream=None, - ) - - def set_index_planes(self, planes: list) -> None: - """Bind the attention layers' index-K planes (static for the manager's lifetime).""" - ptrs, strides = _plane_table(planes) - self._planes = (ptrs, strides, len(planes)) - - # ------------------------------------------------------------- staging --- - - def discard(self) -> None: - """Drop whatever an earlier forward staged.""" - self._pending = None - - def begin( - self, - *, - num_contexts: int, - prefix_tokens: int, - num_gens: int, - rows_per_gen: int, - slots: torch.Tensor, - request_indices: torch.Tensor, - logical_positions: torch.Tensor, - block_table: torch.Tensor, - tokens_per_block: int, - fused: bool = False, - ) -> None: - """Open a verify step's staging. The generation rows follow the `prefix_tokens` - context rows as `num_gens` runs of `rows_per_gen`, so row `s * p + t` is ordinal `t` - of generation request `s`. `slots` are those requests' recurrent slots; - `request_indices` and `logical_positions` are the generation rows' own. Each row's - QSA slot is resolved here, as `qsa_paged_index_store` resolves it. `fused` selects - the one-launch PLE staging and commit.""" - n, p = num_gens, rows_per_gen - assert n > 0 and p > 1, ( - f"a verify step stages generation requests carrying drafts; got {n} " - f"requests of {p} rows" - ) - assert logical_positions.shape[0] == n * p, ( - f"{n} generation requests of {p} rows are {n * p} rows; got " - f"{logical_positions.shape[0]} positions" - ) - positions = logical_positions.to(torch.long) - pages = block_table[request_indices.to(torch.long), positions // tokens_per_block] - self._pending = dict( - num_contexts=num_contexts, - prefix=prefix_tokens, - n=n, - p=p, - slots=slots.to(torch.long), - pages=pages.to(torch.long), - within=positions % tokens_per_block, - index=[], - position=None, - index_all=None, - fused=fused, - ngram_stream=None, - ) - if self._planes is not None: - # Every plane's slots before the trunk writes any of them: what each layer's - # `stage_index` would capture before its own store, in one launch. - ptrs, strides, count = self._planes - pending = self._pending - pending["index_all"] = qsa_rollback_gather( - ptrs, strides, pending["pages"], pending["within"], count - ) - - def stage_ple_pools(self, ngram_pool: torch.Tensor, conv_pool: torch.Tensor) -> None: - """Before the PLE block's first write: both pools, and the generation slots' - pre-step conv rows.""" - pending = self._pending - slots = pending["slots"] - pending["ngram_pool"] = ngram_pool - pending["conv_pool"] = conv_pool - pending["conv_prev"] = conv_pool[slots] - - def stage_ple_verify( - self, ngram_pool: torch.Tensor, conv_pool: torch.Tensor, input_ids: torch.Tensor - ) -> None: - """A verify step's PLE staging, before either pool is written: - `stage_ple_pools`, plus the n-gram candidates rebuilt from the slots' pre-step id - history and the rows' ids. Row `t`'s history after the step is the last - `context` ids of `[history | ids of rows 0..t]` -- what the id window of row `t` - minus its oldest id is, every generation sequence carrying its history.""" - pending = self._pending - prefix, n, p = pending["prefix"], pending["n"], pending["p"] - if pending["fused"]: - pending["ngram_pool"] = ngram_pool - pending["conv_pool"] = conv_pool - pending["ngram_stream"], pending["conv_prev"] = spec_rollback_stage( - ngram_pool, conv_pool, pending["slots"], input_ids[prefix : prefix + n * p], p - ) - return - self.stage_ple_pools(ngram_pool, conv_pool) - context = ngram_pool.shape[1] - ids = torch.reshape(input_ids[prefix : prefix + n * p], [n, p]).to(ngram_pool.dtype) - stream = torch.cat([ngram_pool[pending["slots"]], ids], 1) - # Window k of the unfold is the history after k rows. - pending["ngram_candidates"] = stream.unfold(1, context, 1)[:, 1:] - - def stage_ple_conv_input(self, conv_input: torch.Tensor) -> None: - """The generation rows' conv input (`[T, C]`, token-major), as the fused PLE - launch wrote it into a buffer the forward allocated for this staging alone. A - mixed batch's rows are copied out, so its context prefix is not kept alive - through the draft loop.""" - pending = self._pending - prefix, n, p = pending["prefix"], pending["n"], pending["p"] - conv_x = torch.reshape(conv_input[prefix:], [n, p, -1]) - pending["conv_x"] = conv_x.clone() if prefix else conv_x - - def stage_index(self, index_cache: torch.Tensor, position_cache: torch.Tensor) -> None: - """Before an attention layer's index store: the generation rows' slots of its - index-K plane, and, on the step's first attention layer (the first to write the - shared plane), of the position plane.""" - pending = self._pending - pages, within = pending["pages"], pending["within"] - if pending["position"] is None: - pending["position"] = (position_cache, position_cache[pages, within]) - if pending["index_all"] is not None: - pending["index"].append(None) - return - pending["index"].append((index_cache, index_cache[pages, within])) - - # -------------------------------------------------------------- commit --- - - def commit_speculative_states( - self, - num_accepted_tokens: torch.Tensor, - state_indices: torch.Tensor, - num_contexts: int, - ) -> None: - """Rewrite the staged planes to the accepted prefix. Called by the spec worker.""" - pending, self._pending = self._pending, None - if pending is None: - return - n, p = pending["n"], pending["p"] - assert num_contexts == pending["num_contexts"], ( - f"the worker commits with {num_contexts} context requests and the " - f"forward staged {pending['num_contexts']}" - ) - assert state_indices.shape[0] == num_contexts + n, ( - f"the worker hands {state_indices.shape[0]} slots for a batch the " - f"forward staged as {num_contexts} + {n} requests" - ) - assert pending.get("conv_x") is not None and len(pending["index"]) > 0, ( - "a verify step staged neither its PLE rows nor any attention layer; " - "the forward did not run the staging it opened" - ) - if pending["ngram_stream"] is not None and pending["index_all"] is not None: - ptrs, strides, _ = self._planes - cache, old = pending["position"] - spec_rollback_commit( - num_accepted_tokens[num_contexts : num_contexts + n], - state_indices[num_contexts:], - pending["ngram_pool"], - pending["ngram_stream"], - pending["conv_pool"], - pending["conv_prev"], - pending["conv_x"], - cache, - old, - pending["pages"], - pending["within"], - ptrs, - strides, - pending["index_all"], - ) - return - accepted = num_accepted_tokens[num_contexts : num_contexts + n].to(torch.long) - slots = state_indices[num_contexts:].to(torch.long) - rows = torch.arange(n, device=accepted.device) - ngram_pool, conv_pool = pending["ngram_pool"], pending["conv_pool"] - state_len = conv_pool.shape[-1] - ngram_pool[slots] = pending["ngram_candidates"][rows, accepted - 1] - # [n, C, state_len + p]; window k of the unfold is the state after k tokens. - history = torch.cat([pending["conv_prev"], torch.transpose(pending["conv_x"], 1, 2)], 2) - conv_pool[slots] = history.unfold(2, state_len, 1)[rows, :, accepted] - - # [n * p]: True at every ordinal the sampler rejected. - rejected = torch.reshape( - torch.arange(p, device=accepted.device).unsqueeze(0) >= accepted.unsqueeze(1), [-1] - ) - pages, within = pending["pages"], pending["within"] - if pending["index_all"] is not None: - ptrs, strides, _ = self._planes - qsa_rollback_restore(ptrs, strides, pages, within, rejected, pending["index_all"]) - else: - for cache, old in pending["index"]: - cache[pages, within] = torch.where( - rejected[:, None, None], old, cache[pages, within] - ) - cache, old = pending["position"] - cache[pages, within] = torch.where(rejected[:, None], old, cache[pages, within]) - - -#: Per modality, the keys under `multimodal_data[modality]`: the flat patch rows and the -#: host-side `(t, h, w)` grid. -_MEDIA_KEYS = { - "image": ("pixel_values", "image_grid_thw"), - "video": ("pixel_values_videos", "video_grid_thw"), -} - - -class _WidenedHead: - """`lm_head` whose bf16 logits come back as fp32 from `bf16_to_fp32` (one vectorized - kernel; torch's `.float()` is an unvectorized copy at under half the bandwidth), so the - logits processor's own `.float()` returns them as they are. The same values bit for - bit.""" - - def __init__(self, lm_head) -> None: - self.lm_head = lm_head - - def __call__(self, hidden_states: torch.Tensor) -> torch.Tensor: - logits = self.lm_head(hidden_states) - if logits.dtype == torch.bfloat16 and logits.is_contiguous(): - return bf16_to_fp32(logits) - return logits - - -class _Fp8DraftHead: - """The draft head's projection over an FP8 copy of `lm_head`: up to `fp8_gemv`'s rows on - the packed block-scaled copy (bf16 activations, one launch at HBM rate) -- or, given the - NVFP4 copy `nv` (`_MTP_HEAD_NVFP4_ENV`), `nvfp4_gemv` over it --, more rows as per-token - e4m3 activations (`quantize_e4m3_activation`) times the per-row e4m3 weight, scaled in - cuBLASLt's epilogue; bf16 out (`_WidenedHead` widens it as it does the bf16 head's).""" - - def __init__( - self, weight_fp8: torch.Tensor, weight_scale: torch.Tensor, packed=None, nv=None - ) -> None: - self.weight_fp8, self.weight_scale, self.packed = weight_fp8, weight_scale, packed - self.nv = nv - # The same per-row copy as `mid_rows_mm` reads it (views). - self.mid = _head_mid(weight_fp8, weight_scale) - - def __call__(self, hidden_states: torch.Tensor) -> torch.Tensor: - rows = hidden_states.shape[0] - if rows <= _FP8_GEMV_HEAD_ROWS and _gemv_ready(hidden_states): - if self.nv is not None: - return nvfp4_gemv_op(hidden_states, self.nv) - if self.packed is not None: - return fp8_gemv_op(hidden_states, self.packed) - if rows <= _MID_HEAD_ROWS and _gemv_ready(hidden_states): - return mid_rows_mm_op(hidden_states, self.mid) - x8, xs = torch.ops.tensorrt_llm.quantize_e4m3_activation(hidden_states) - return torch._scaled_mm( - x8, - self.weight_fp8.t(), - scale_a=xs.float().reshape(-1, 1), - scale_b=self.weight_scale, - out_dtype=hidden_states.dtype, - ) - - -def _head_mid(weight_fp8: torch.Tensor, weight_scale: torch.Tensor) -> MidWeight: - """`mid_rows_mm`'s view of a per-row-scaled e4m3 head copy (`_fp8_rowwise`).""" - return MidWeight(weight_fp8, weight_scale.reshape(-1)) - - -def _fp8_rowwise(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """`weight` [rows, cols] bf16 -> (e4m3 [rows, cols], fp32 [1, rows]): one scale a row, - amax / 448, so each row's largest magnitude maps to e4m3's largest finite value.""" - scale = torch.empty(weight.shape[0], dtype=torch.float32, device=weight.device) - w8 = torch.empty_like(weight, dtype=torch.float8_e4m3fn) - for lo in range(0, weight.shape[0], 16384): - hi = min(lo + 16384, weight.shape[0]) - chunk = weight[lo:hi].float() - scale[lo:hi] = (chunk.abs().amax(dim=1) / 448.0).clamp(min=1e-12) - w8[lo:hi] = (chunk / scale[lo:hi, None]).to(torch.float8_e4m3fn) - return w8, scale.reshape(1, -1).contiguous() - - -class _Fp8VerifyHead: - """The target's `lm_head` under the dense FP8 mode: `_Fp8DraftHead` over the same copies -- - decode rows up to `fp8_gemv`'s on the packed block-scaled copy, more rows (a large verify - step) on the per-row-scaled copy through cuBLASLt with per-token activation scales.""" - - def __init__(self, rowwise: tuple, packed) -> None: - self.head = _Fp8DraftHead(*rowwise, packed) - - def __call__(self, hidden_states: torch.Tensor) -> torch.Tensor: - return self.head(hidden_states) - - -class _WideningLogitsProcessor(LogitsProcessor): - """The shell's logits processor (its row gather, its fp32 logits) over `_WidenedHead`; - with `fp8_head` set (the dense FP8 mode) the target's own head runs `_Fp8VerifyHead`. - - `widen=False` (with `return_context_logits`, the speculative paths) returns the head's - own bf16 logits: for a consumer that only takes their row argmax, which is the same - token over bf16 as over their exact fp32 widening (same order, ties, NaN and +-0).""" - - fp8_head = None # (rowwise copy, packed copy) under the dense FP8 mode - #: Set by the shell's forward each engine step: the step's draft logits feed only the - #: greedy draft sampler's argmax (see `_spec_argmax_only`). - draft_argmax_only = False - - def forward( - self, - hidden_states, - lm_head, - attn_metadata, - return_context_logits: bool = False, - widen: bool = True, - ) -> torch.Tensor: - if self.fp8_head is not None and not isinstance(lm_head, _Fp8DraftHead): - lm_head = _Fp8VerifyHead(*self.fp8_head) - if not widen and return_context_logits: - return lm_head(hidden_states) - return super().forward( - hidden_states, _WidenedHead(lm_head), attn_metadata, return_context_logits - ) - - -def _spec_argmax_only(spec_worker, spec_metadata) -> tuple[bool, bool]: - """`(verify, draft)`: whether this step's verify / draft logits are consumed only by a - row argmax, so they can stay bf16. - - Both need no guided decoder (it edits the logits). The draft logits then feed only - the draft sampler, an argmax unless rejection sampling runs on a non-greedy batch. The - verify logits are also the step's `logits` output and the occurrence penalties' and - relaxed acceptance's input; the one-model speculative sampler admits no request - asking for logits or log probabilities, so they stay bf16 only for an all-greedy batch - with the penalties off (a deploy-time flag) and strict acceptance. Every input is a - deploy-time constant or `is_all_greedy_sample`, part of the CUDA graph key.""" - if getattr(spec_worker, "guided_decoder", None) is not None: - return False, False - draft = not spec_metadata.wants_advanced_draft_sampling - verify = ( - spec_metadata.is_all_greedy_sample - and not getattr(spec_metadata, "enable_penalty", False) - and not getattr( - getattr(spec_worker, "spec_config", None), "use_relaxed_acceptance_for_thinking", False - ) - ) - return verify, draft - - -class ModelingV2Qwen38FlashNextNvfp4Sm103Tp1( - DecoderModelForCausalLM[Qwen4ExpTp1ModelingV2, PretrainedConfig] -): - # Whether the engine takes the one-engine speculative fast paths for this model; set in - # `__init__` when it builds its speculative branch. - spec_decode_fast_paths: bool = False - # The step-loop hooks this model opts into (`runtime_hooks_of`); None: the defaults. - runtime_hooks: RuntimeHooks | None = None - - @classmethod - def get_preferred_kv_cache_manager_version(cls, pretrained_config=None) -> str: - """Prefer KV cache manager V2: PLE state (`ple_layer_cache`) and QSA's index and - position planes exist only there. Under V1 the model builds and passes a - smoke run with the PLE layer stateless. The engine resolves this from the - built-in class; the declaration serves a direct instantiation.""" - return "V2" - - def __init__(self, model_config: ModelConfig): - cfg = model_config.pretrained_config - assert cfg is not None - - # Adopt the engine-normalized text tower (`Qwen4ExpTextConfig`) as the model's - # config: everything that sizes the pools reads - # `model_engine.model.model_config.pretrained_config`, and the composite - # parent mirrors text geometry through a pass that silently skips keys the - # parent already defaults. The swap is in place because the engine keeps its - # own reference to this `ModelConfig`. - text = getattr(cfg, "text_config", None) - assert text is not None, ( - "this target is the language tower of the composite Qwen4-Exp " - "export; pretrained_config carries no text_config" - ) - assert getattr(text, "architectures", None) == ["Qwen4ExpForCausalLM"], ( - f"the text tower must be the normalized Qwen4-Exp text config; got " - f"architectures {getattr(text, 'architectures', None)!r}" - ) - assert text.torch_dtype is not None, "the text tower must declare a dtype" - # Captured before the swap: the core needs the vision half and media token ids. - composite = cfg - model_config.pretrained_config = text - - super().__init__( - Qwen4ExpTp1ModelingV2(model_config, composite), - config=model_config, - hidden_size=text.hidden_size, - vocab_size=text.vocab_size, - ) - # The logits' bf16 -> fp32 cast as one vectorized kernel (the draft head shares it). - self.logits_processor = _WideningLogitsProcessor() - # Media placeholder ids for the executor; built lazily (`__init__` runs under - # meta init). - self._mm_token_ids_value = (self.model.image_token_id, self.model.video_token_id) - self._mm_token_ids_cache: torch.Tensor | None = None - # `lm_head` is outside the quantization scheme: a dense bf16 `LMHead`; a - # quantized head would register a different parameter set. - assert self.lm_head.weight.dtype == torch.bfloat16, ( - f"expected a bf16 lm_head; the shell built {self.lm_head.weight.dtype}, " - "which means the engine resolved a quantized output head for a " - "checkpoint that stores it unquantized" - ) - - # The speculative branch, built only when `speculative_config` asked for it. - # Written out rather than inherited: the in-tree one-engine shell builds its - # drafter via a module-level function dispatching on `model_type` (still the - # upstream family name), which would construct trtllm's own MTP layer. - self.spec_config = getattr(model_config, "spec_config", None) - self.draft_model = None - self.spec_worker = None - # No separate draft checkpoint; the loader reads this, so None is stated. - self.draft_config = None - if self.spec_config is not None: - assert self.model.mtp_enabled, "core built without the MTP module" - assert hasattr(self, "logits_processor"), ( - "the inherited shell no longer exposes logits_processor; the " - "draft head and the shell's own gather both project through it" - ) - self.draft_model = DraftModel(self.model, self.lm_head, self.logits_processor) - # Must not be left at `get_spec_worker`'s default False: the resource - # manager decides separately (`_should_create_separate_draft_kv_cache`) and - # builds the target manager from target layer masks alone, so left False - # the MTP layer is handed a manager with no `layer_offsets` entry for it. - self.use_separate_draft_kv_cache = should_use_separate_draft_kv_cache(self.spec_config) - # The one-engine speculative fast paths, as `SpecDecOneEngineForCausalLM` takes them - # where validated: the draft loop's step-0 epilogue, the recorded steady decode-step - # prepare and the overlap loop on the execution stream. At TP1 no logits are - # vocabulary-sharded, so the paths for those do not apply. - self.spec_decode_fast_paths = True - self.runtime_hooks = RuntimeHooks(loop_on_execution_stream=True) - self.spec_worker = get_spec_worker( - self.spec_config, - model_config, - model_config.mapping, - use_separate_draft_kv_cache=self.use_separate_draft_kv_cache, - spec_decode_fast_paths=self.spec_decode_fast_paths, - ) - assert self.spec_worker is not None, ( - f"the engine resolved no spec worker for spec_dec_mode " - f"{self.spec_config.spec_dec_mode!r}" - ) - # Caches the draft->target vocab offset map (correctly None here). - self.spec_worker.set_draft_model(self.draft_model) - # The PLE pools and the QSA side caches are written per verify row and nothing - # the runtime owns rolls them back; the worker calls this handler after its - # draft loop, inside the same forward. - self.spec_worker.register_auxiliary_state_handler(self.model.spec_rollback) - - # ----------------------------------------------------------- multimodal --- - # Names the executor reads off a model class by `getattr`; each is silently - # optional and its absence costs correctness rather than raising. - - @property - def multimodal_data_device_paths(self) -> list[str]: - """Tensors under ``multimodal_data`` the executor moves to CUDA. Absent, it moves - everything, putting `image_grid_thw` and `video_grid_thw` on the device where - `pos_embed_interpolate` and `vision_rope_2d` would each pay a D2H sync.""" - return [ - "image.pixel_values", - "video.pixel_values_videos", - "mrope_config.mrope_position_ids", - "mrope_config.mrope_position_deltas", - ] - - @property - def multimodal_token_ids(self) -> torch.Tensor: - """Placeholder ids whose rows the vision embeddings replace. Absent, the - executor uses `input_ids >= vocab_size`, but this checkpoint's placeholders - (248056, 248057) are in vocabulary, so nothing is selected and the engine - generates fluent text while blind. Built on first read (meta init); not a - parameter or buffer.""" - if self._mm_token_ids_cache is None: - self._mm_token_ids_cache = torch.tensor(self._mm_token_ids_value, dtype=torch.int32) - return self._mm_token_ids_cache - - def validate_request(self, request) -> None: - """Admission check, run by the executor on every new request. A media request - the forward would assert on fails here as that request's own error; an assert - inside the step ends the executor under the overlap scheduler. Reads only the - request's host-side payload, and a request without media returns at once. - - Refused: a frame past `_VISION_CALL_MAX_PATCHES` patches (the one bound the - tower keeps), a malformed grid or pixel payload, media whose rows disagree with - the prompt's placeholder count, missing or mis-sized `mrope_position_ids`, and a - request with both images and videos whose item metadata is missing or does not - list its items.""" - data = request.py_multimodal_data - if not data: - return - grids = {} - for modality, (pixel_key, grid_key) in _MEDIA_KEYS.items(): - media = data.get(modality) - if media is not None: - grids[modality] = self._media_grid(modality, media, pixel_key, grid_key) - if not grids: - return - - mrope_ids = (data.get("mrope_config") or {}).get("mrope_position_ids") - if ( - mrope_ids is None - or mrope_ids.dim() != 3 - or mrope_ids.shape[0] != _QSA_POSITION_AXES - or mrope_ids.shape[2] != request.prompt_len - ): - raise ValueError( - f"a media request needs the [{_QSA_POSITION_AXES}, 1, {request.prompt_len}] " - "mrope_config.mrope_position_ids from the input processor, one position per " - "prompt token; got " - f"{None if mrope_ids is None else tuple(mrope_ids.shape)}" - ) - - if len(grids) > 1: - metadata = get_multimodal_item_metadata(data) - self._check_item_order(grids, metadata.item_refs if metadata is not None else None) - - cumsum = data.get("multimodal_embed_mask_cumsum") - if cumsum is not None: - merge_unit = self.model.v_merge * self.model.v_merge - rows = sum(t * h * w for grid in grids.values() for t, h, w in grid) // merge_unit - placeholders = int(cumsum[-1]) - if placeholders != rows: - raise ValueError( - f"the prompt holds {placeholders} image/video placeholder tokens and its " - f"{' and '.join(sorted(grids))} payload produces {rows} embedding rows" - ) - - def _media_grid(self, modality, media, pixel_key, grid_key): - """One modality's `(frames, height, width)` rows, checked on the host.""" - core = self.model - pixel_values, grid_thw = media.get(pixel_key), media.get(grid_key) - if pixel_values is None or grid_thw is None: - raise ValueError( - f"the {modality} payload carries {sorted(media)}; this target reads " - f"{pixel_key!r} and {grid_key!r} from it" - ) - if grid_thw.dim() != 2 or grid_thw.shape[0] == 0 or grid_thw.shape[1] != 3: - raise ValueError( - f"{grid_key} must hold one (t, h, w) row per {modality}; got shape " - f"{tuple(grid_thw.shape)}" - ) - grid = [tuple(int(v) for v in row) for row in grid_thw.tolist()] - for item, (frames, height, width) in enumerate(grid): - if min(frames, height, width) < 1 or height % core.v_merge or width % core.v_merge: - raise ValueError( - f"{modality} {item} has grid {(frames, height, width)}; every entry must " - f"be positive and the sides divisible by spatial_merge_size {core.v_merge}" - ) - if height * width > _VISION_CALL_MAX_PATCHES: - raise ValueError( - f"{modality} {item} has frames of {height}x{width} = {height * width} " - f"patches ({height * core.v_patch}x{width * core.v_patch} pixels); the " - f"vision tower serves frames of at most {_VISION_CALL_MAX_PATCHES} " - "patches, so send it at a lower resolution" - ) - patches = sum(t * h * w for t, h, w in grid) - if tuple(pixel_values.shape) != (patches, core.v_pixel_width): - raise ValueError( - f"{pixel_key} is {tuple(pixel_values.shape)}; its grids describe " - f"{patches} patch rows of {core.v_pixel_width}" - ) - return grid - - @staticmethod - def _check_item_order(grids, order) -> None: - """A request with several modalities lists its items in prompt order.""" - if not order: - raise ValueError( - f"a request carrying {sorted(grids)} needs item metadata to splice them in " - "prompt order; the payload holds each modality's rows apart" - ) - listed = {modality: [] for modality in grids} - for entry in order: - try: - modality, index = entry - listed[modality].append(int(index)) - except (KeyError, TypeError, ValueError) as error: - raise ValueError( - f"item_refs entry {entry!r} is not an item of this request, which " - f"carries {sorted(grids)}; entries are (modality, index) pairs" - ) from error - for modality, grid in grids.items(): - if sorted(listed[modality]) != list(range(len(grid))): - raise ValueError( - f"item_refs lists {modality} items {listed[modality]}; the payload " - f"carries {len(grid)}" - ) - - def _multimodal_embeds(self, kwargs): - """Run the vision tower for this batch; returns `(rows, mm_token_indices)` (the - tower's `[R, hidden]` rows in packed-batch order and the `[R]` placeholder rows they - replace, which the prefill target splices into its embedding), or None when the step - carries no media. - - Lives in the shell because the inherited `DecoderModelForCausalLM.forward` - drops `multimodal_params` from `**kwargs` on the no-speculation path (a - booting, fluent, blind engine). Host-side bookkeeping: `multimodal_params` - is a subset of the context requests, so never zip it positionally; - `multimodal_runtime` (chunked prefill / KV reuse only) gives the cached and - in-chunk row counts over the request's images and videos together, in prompt - order; `mm_token_indices` names every placeholder row. - - The tower runs once per modality a request carries, and `_in_prompt_order` - merges the two outputs. Positions are not derived here: every coordinate of - a vision span, a video's temporal one included, is in the payload's - `mrope_config.mrope_position_ids`, which the engine slices into - `position_ids`.""" - mm_params = kwargs.get("multimodal_params") or [] - if not mm_params: - return None - rows = [] - for params in mm_params: - data = getattr(params, "multimodal_data", None) or {} - image, video = data.get("image"), data.get("video") - if image is None and video is None: - continue - mrope_ids = (data.get("mrope_config") or {}).get("mrope_position_ids") - assert ( - mrope_ids is not None - and mrope_ids.dim() == 3 - and mrope_ids.shape[0] == _QSA_POSITION_AXES - ), ( - f"a media request needs mrope_config.mrope_position_ids of shape " - f"[{_QSA_POSITION_AXES}, 1, L]; got " - f"{None if mrope_ids is None else tuple(mrope_ids.shape)}" - ) - runtime = getattr(params, "multimodal_runtime", None) - if runtime is not None and int(runtime.num_mm_tokens_in_chunk) == 0: - # Every placeholder row of this request lies in another chunk. - continue - towers = {} - if image is not None: - towers["image"] = self._run_tower(image, *_MEDIA_KEYS["image"]) - if video is not None: - towers["video"] = self._run_tower(video, *_MEDIA_KEYS["video"]) - embeds = self._in_prompt_order(params, towers) - if runtime is not None: - # Chunked prefill or reused prefix: splice only this chunk's slice. - start = int(runtime.num_cached_mm_tokens) - count = int(runtime.num_mm_tokens_in_chunk) - total = embeds.shape[0] - assert int(runtime.total_embeds_in_request) == total, ( - f"the executor counted {runtime.total_embeds_in_request} " - f"placeholder rows for this request and the tower produced " - f"{total} over its {sorted(towers)}; the two are derived from the " - "same grids and a disagreement means the merger's row count is wrong" - ) - assert 0 <= start <= start + count <= total, (start, count, total) - embeds = embeds[start : start + count] - rows.append(embeds) - if not rows: - return None - vision = rows[0] if len(rows) == 1 else torch.cat(rows, 0) - mm_token_indices = kwargs.get("mm_token_indices") - assert mm_token_indices is not None, ( - "the executor shipped multimodal params but no mm_token_indices; " - "that pairing is what `multimodal_token_ids` on this class buys, and " - "without it there is no sync-free way to find the placeholder rows" - ) - assert mm_token_indices.shape[0] == vision.shape[0], ( - f"the batch has {mm_token_indices.shape[0]} placeholder rows and the " - f"tower produced {vision.shape[0]} embedding rows" - ) - return vision, mm_token_indices - - def _run_tower(self, media, pixel_key, grid_key): - """The tower over one modality's payload: `(rows, grid_thw)`, the rows one per - placeholder, item by item in the order the processor emitted them.""" - pixel_values, grid_thw = media.get(pixel_key), media.get(grid_key) - assert pixel_values is not None and grid_thw is not None, ( - f"the multimodal payload carries {sorted(media)}; this target " - f"reads {pixel_key!r} and {grid_key!r} from it" - ) - assert not grid_thw.is_cuda, ( - f"{grid_key} must stay on the host; see multimodal_data_device_paths" - ) - return self.model.vision_tower(pixel_values, grid_thw), grid_thw - - def _in_prompt_order(self, params, towers): - """One request's tower rows in the order its placeholders appear. - - `towers` maps each modality the request carries to `(rows, grid_thw)`; each - grid row is one item of `t*h*w / merge**2` rows. One modality is already in - prompt order. Two are reassembled item by item from `item_metadata.item_refs` - (`(modality, index)` per item in prompt order, `index` counting that - modality's items); without it there is no host-side way to tell which - placeholder run is an image and which a video.""" - if len(towers) == 1: - return next(iter(towers.values()))[0] - metadata = params.item_metadata - order = metadata.item_refs if metadata is not None else None - assert order, ( - f"a request carrying {sorted(towers)} needs item metadata to splice them " - "in prompt order; the payload holds each modality's rows apart" - ) - merge_unit = self.model.v_merge * self.model.v_merge - items = {} - for modality, (embeds, grid_thw) in towers.items(): - counts = [t * h * w // merge_unit for t, h, w in grid_thw.tolist()] - assert sum(counts) == embeds.shape[0], ( - f"the {modality} grid describes {sum(counts)} rows and the tower " - f"produced {embeds.shape[0]}" - ) - items[modality] = torch.split(embeds, counts, 0) - listed = {modality: [] for modality in items} - for modality, index in order: - assert modality in items, ( - f"item_refs lists a {modality!r} item; this request carries {sorted(items)}" - ) - listed[modality].append(int(index)) - for modality, parts in items.items(): - assert sorted(listed[modality]) == list(range(len(parts))), ( - f"item_refs lists {modality} items {listed[modality]}; the " - f"payload carries {len(parts)}" - ) - return torch.cat([items[modality][int(index)] for modality, index in order], 0) - - def forward( - self, - attn_metadata: AttentionMetadata, - input_ids: torch.IntTensor | None = None, - position_ids: torch.IntTensor | None = None, - inputs_embeds: torch.FloatTensor | None = None, - return_context_logits: bool = False, - spec_metadata=None, - lora_params: dict | None = None, - **kwargs, - ): - """One engine step. - - Without a spec worker this delegates to the base forward with every argument - it was handed (`resource_manager` stays in `**kwargs` deliberately: naming - it would drop it from what the base receives). With one, the shell owns the - logits gather (one-model modes set `without_logits`) and the worker needs - ungathered hidden states; `position_ids` is passed in the engine's layout - (flattening gives silently wrong draft positions). The vision tower runs here - because the base forward would not pass `multimodal_params` through; its rows reach - the core through `begin_step`, and the prefill target splices them.""" - self.model.begin_step( - attn_metadata, - input_ids, - position_ids, - inputs_embeds=inputs_embeds, - lora_params=lora_params, - spec_metadata=spec_metadata, - vision=self._multimodal_embeds(kwargs), - ) - if self.spec_worker is None: - assert spec_metadata is None, ( - "spec_metadata arrived without a spec worker; this engine was " - "built without a speculative_config" - ) - return super().forward( - attn_metadata, - input_ids, - position_ids, - inputs_embeds, - return_context_logits, - spec_metadata, - lora_params, - **kwargs, - ) - assert spec_metadata is not None, ( - "the spec worker is built but the engine passed no spec_metadata" - ) - # `final` is the `[T, hidden]` mix for `lm_head`, `hc` the stream the MTP reads. - final, hc = self.model.split_spec_output( - self.model( - attn_metadata=attn_metadata, - input_ids=input_ids, - position_ids=position_ids, - inputs_embeds=inputs_embeds, - lora_params=lora_params, - spec_metadata=spec_metadata, - **kwargs, - ) - ) - # `gather_ids`: one row per context request, `runtime_draft_len + 1` per - # generation request. Logits only an argmax reads stay bf16. - verify_argmax_only, self.logits_processor.draft_argmax_only = _spec_argmax_only( - self.spec_worker, spec_metadata - ) - logits = self.logits_processor.forward( - nn.functional.embedding(spec_metadata.gather_ids, final), - self.lm_head, - attn_metadata, - True, - widen=not verify_argmax_only, - ) - # A prefill CUDA graph pads the step's rows; the worker reshapes generation rows per - # request, so it takes the batch's own. - if attn_metadata.padded_num_tokens is not None: - n = attn_metadata.num_tokens - input_ids, hc = input_ids[:n], hc[:n] - position_ids = position_ids[..., :n] - return self.spec_worker( - input_ids=input_ids, - position_ids=position_ids, - hidden_states=hc, - logits=logits, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - draft_model=self.draft_model, - resource_manager=kwargs.get("resource_manager"), - ) - - def load_weights(self, weights, *args, **kwargs): - _weights.MODEL_WEIGHTS.load(self, weights) - - def post_load_weights(self): - super().post_load_weights() - self.model.post_load_weights() - # The draft head reads a per-vocabulary-row e4m3 copy of lm_head (cuBLASLt with - # per-token activation scales). It only proposes tokens -- the verify step's bf16 head - # decides what is emitted -- so the copy can change acceptance, not outputs (measured - # unchanged: gsm8k and pass-key acceptance). - draft_fp8 = self.model.mtp_enabled - rowwise = ( - _fp8_rowwise(self.lm_head.weight.data) if draft_fp8 or _dense_fp8_enabled() else None - ) - if draft_fp8: - self.model._mtp_head_fp8 = rowwise - # The draft head's decode rows read an NVFP4 copy instead of the packed FP8 one. - draft_nv = draft_fp8 and _mtp_head_nvfp4_enabled() - # One packed block-scaled copy serves the draft head's decode rows and, under the dense - # FP8 mode, the target's own head's. - if (draft_fp8 and not draft_nv) or _dense_fp8_enabled(): - self.model._lm_head_g = pack_from_bf16(self.lm_head.weight.data) - if _dense_fp8_enabled(): - self.logits_processor.fp8_head = (rowwise, self.model._lm_head_g) - if rowwise is not None and self.lm_head.weight.is_cuda: - # Every head variant decode rows past the GEMV reach, compiled now. - mid_rows_precompile(_head_mid(*rowwise), _FP8_GEMV_HEAD_ROWS + 1, _MID_HEAD_ROWS) - if draft_nv: - self._build_draft_head_nvfp4(self.lm_head.weight.data) - - def _build_draft_head_nvfp4(self, weight: torch.Tensor) -> None: - """The draft head's NVFP4 copy of `weight` (`_MTP_HEAD_NVFP4_ENV`), with both of - `nvfp4_gemv`'s row-octet kernels compiled now rather than at a first serving call.""" - nv = nvfp4_pack_from_bf16(weight) - nvfp4_gemv_precompile(nv) - self.model._mtp_head_nv = nv - print( - f"draft head NVFP4 ({_MTP_HEAD_NVFP4_ENV}=1): lm_head {tuple(weight.shape)} -> " - f"{nv.data.numel() / 2**20:.1f} MiB (e2m1 + 16-element e4m3 scales, global scale " - f"{nv.global_scale:.3e}), decode rows <= {_FP8_GEMV_HEAD_ROWS} on nvfp4_gemv", - flush=True, - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py deleted file mode 100644 index 7987c692ddfd..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py +++ /dev/null @@ -1,100 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""In-place sigmoid output gating: ``attention_output *= sigmoid(gate)``.""" - -from __future__ import annotations - -import torch - -# Registered by the module that defines it: after `import -# tensorrt_llm._torch.custom_ops` alone, torch.ops.trtllm has no -# `fused_sigmoid_mul_inplace` (measured). -import tensorrt_llm._torch.attention.kernels.fused_qk_norm_rope_gate # noqa: F401 -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -#: The checkpoint's gated attention: 24 q heads of 256, q projection [T, 24, 512]. -_CKPT = dict(heads=24, head_dim=256, dtype=torch.bfloat16, gate="interleaved") - - -class FusedSigmoidMulInplace(OpWrapper): - """`attention_output[t, i] *= sigmoid(gate[t, i // D, i % D])` in place, `D = gate.shape[-1]`. - - Returns `attention_output`, the tensor it wrote. Both operands are widened to - fp32 and the product is rounded once to `attention_output.dtype`. `gate` is - read through its own token and head strides, so the gate half of an - interleaved `[T, heads, 2 * D]` q projection is passed as a strided view with - no repack (bitwise equal to its compacted copy, measured). A 2-D `gate` of - `attention_output`'s shape is one head spanning the row. - - Fusion boundary: sigmoid and multiply only. The caller owns producing both - operands and everything downstream (the o projection). `gate` is read-only. - """ - - ARCHS = frozenset({Arch.SM_103}) - - CELLS: tuple[Cell, ...] = ( - Cell( - why="decode 128 tokens: 2-D gate off the fused projection row (stride 14080)", - spec=dict(tokens=128, heads=24, head_dim=256, dtype=torch.bfloat16, gate="fused_row"), - ), - Cell( - why="prefill 8192 tokens: 2-D gate off the fused projection row, 2048-wide blocks", - spec=dict(tokens=8192, heads=24, head_dim=256, dtype=torch.bfloat16, gate="fused_row"), - ), - Cell( - why="MTP decode: 1 token, 24x256 gate as the strided half of q_proj", - spec=dict(tokens=1, **_CKPT), - ), - ) - - note = """ - Metadata consumed: none. One Triton kernel, JIT-compiled on the first call, - no autotuner: grid `(T, cdiv(hidden, block))`, `block = min(next_pow2(hidden), - 1024 if T < 1024 else 2048)`, 4 warps. - - Preconditions, every one enforced loudly by the op's own bare `assert` - (an AssertionError with no message): `attention_output` 2-D with - `stride(-1) == 1` (row stride free); `gate` 2-D of `attention_output`'s - shape, or 3-D with `shape[0] == T` and `shape[1] * shape[2] == hidden`, - `stride(-1) == 1`. A CPU tensor raises ValueError from Triton (measured). - Hence no `is_valid`. Cross-device operands were not probed. - - The pre-call contents of `attention_output` are destroyed; the gate is - left bitwise unchanged. - - Numerics: bf16 bitwise equal to the fp32 reference except ~1e-6 of - elements one ulp off (`tl.sigmoid`'s fast exp); default tolerance holds, - nothing loosened. Feeding the q half instead of the gate half moves the - result by ~2.15 max abs at [17, 6144] (measured), so the per-head indexing - is under test. - - Certified: bf16 operands at the checkpoint's hidden 6144, a compact - `attention_output`; the 2-D window of the fused projection row at 128 - tokens (decode, 1024-wide blocks) and 8192 tokens (prefill, 2048-wide - blocks), and the strided q_proj half (MTP) at 1 token. - - Hard limits - T * row stride token offset computed in int64 boundary 2^63 not probed - (largest run 8192 rows) - - Not certified: 0 tokens, the strided q_proj half above 1 token, a strided - `attention_output`, a gate dtype other than `attention_output.dtype`, - fp16/fp32 operands, other hidden sizes, CUDA-graph capture, torch.compile. - """ - - def __call__(self, attention_output: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: - torch.ops.trtllm.fused_sigmoid_mul_inplace(attention_output, gate) - return attention_output - - def reference(self, attention_output: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: - """What `attention_output` holds after the call; neither input is mutated. - - `gate.reshape(T, hidden)` is the per-head `(t, i // D, i % D)` read, for - the 2-D and 3-D forms alike. - """ - num_tokens, hidden = attention_output.shape - g = gate.reshape(num_tokens, hidden).float() - return (attention_output.float() * torch.sigmoid(g)).to(attention_output.dtype) - - -fused_sigmoid_mul_inplace = FusedSigmoidMulInplace() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_canonical_selection.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_canonical_selection.py deleted file mode 100644 index b2736198b176..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_canonical_selection.py +++ /dev/null @@ -1,214 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""QSA block selection, made canonical: one answer per row, whichever tied columns the -selection engine happened to keep. - -Composed from native torch ops: no registered op rewrites a selection, and none of the -top-k engines (`indexer_topk_prefill`, the CuTe DSL kernels, FlashInfer's top-k) takes a -tie-break rule or fixes the order of its output. -""" - -from __future__ import annotations - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -#: This checkpoint's block budget: `indexer_budget` 2048 tokens over -#: `indexer_compress_ratio` 4. -_BLOCK_TOPK = 512 -#: Score-row width the target hands over: a 260-page block table at 32 tokens -#: per block, compressed 4:1. -_SCORE_COLUMNS = 260 * 32 // 4 -#: The most `[rows, cols]` elements one rewrite pass holds at once; a taller matrix is -#: rewritten in row tiles so the fp32 / int32 temporaries stay a few hundred MiB. -_TILE_ELEMENTS = 1 << 24 - - -def _cell(why: str, rows, windows, logits: str, tail: str = "nan", **extra) -> Cell: - return Cell( - why=why, - spec=dict( - rows=rows, - cols=extra.pop("cols", _SCORE_COLUMNS), - index_topk=_BLOCK_TOPK, - windows=windows, - logits=logits, - tail=tail, - selection=extra.pop("selection", "random"), - **extra, - ), - ) - - -def _rewrite_tile(logits, row_starts, row_ends, indices) -> None: - num_rows, num_cols = logits.shape - index_topk = indices.shape[1] - starts = row_starts[:, None] - columns = torch.arange(num_cols, dtype=torch.int32, device=logits.device)[None, :] - # The window's scores, NaN elsewhere: NaN compares false, so nothing outside the - # window is ever counted or kept. - windowed = torch.where( - (columns >= starts) & (columns < row_ends[:, None]), logits, float("nan") - ) - # The k-th value: the smallest score the engine kept. - picked = torch.gather(logits, 1, (indices + starts).clamp_(0, num_cols - 1).long()) - kth = picked.masked_fill_(indices < 0, float("inf")).amin(1, keepdim=True) - at_least = windowed >= kth - tied = windowed == kth - # Keep every column at or above the k-th value except the lowest-index ties past - # k: a tied column is dropped while no more than `surplus` ties lie at or before - # it. A window of at most k has no surplus. - surplus = at_least.sum(1, keepdim=True, dtype=torch.int32) - index_topk - dropped = tied & (torch.cumsum(tied, 1, dtype=torch.int32) <= surplus) - kept_seen = torch.cumsum(at_least ^ dropped, 1, dtype=torch.int32) - # The j-th kept column is the first whose running count reaches j. - slots = torch.arange(1, index_topk + 1, dtype=torch.int32, device=logits.device) - slots = slots.repeat(num_rows, 1) - torch.searchsorted(kept_seen, slots, out_int32=True, out=indices) - indices.sub_(starts) - indices.masked_fill_(slots > kept_seen[:, -1:], -1) - - -class QsaCanonicalSelection(OpWrapper): - """Rewrite an exact top-k selection, in place, into the window's canonical selection. - - For row `r` with window `[s, e) = [row_starts[r], row_ends[r])` of `logits` and - `k = indices.shape[1]`, `indices[r]` holds on entry what a top-k engine wrote: - window-local offsets of an exact top-k of the window, in any order, `-1` past - `min(k, e - s)`. Let `v` be the smallest selected score. On return the row is - - * every window column scoring strictly above `v`, then the highest-index window - columns scoring exactly `v` until there are `k` -- the k largest - `(score, column)` pairs; or the whole window when `e - s <= k`; - * as window-local offsets in ascending order, `-1` in every remaining slot. - - So the row is a function of `(logits, row_starts, row_ends, k)` alone: two engines, - or two launches of one engine, that broke a tie at the k-th value differently produce - the same row. On a row with no tie at the k-th value this is the engine's own set, - sorted. Ties are float equality: `-0.0` and `+0.0` are one tie class. Returns - `indices`. - - Fusion boundary: the rewrite only. Selecting is the engine's call before it; the - engine must be exact -- `v` is read off its picks, and a pick below the true k-th - value would admit more than `k` columns above it. - """ - - ARCHS = frozenset({Arch.SM_103}) - - # The target calls this right after `qsa_topk_select`, at `k = 512` over the - # `[rows, 2080]` score matrix whose columns past each `row_end` are allocator - # garbage (the cells poison them). `selection` names what produced the input row: - # `random` is an exact top-k that keeps a random subset of the tie class at the k-th - # value, in a random order (the adversary); `radix` is the engine the target calls. - CELLS: tuple[Cell, ...] = ( - _cell( - "decode: 1, 128 and 256 rows (MTP3 and MTP7 at batch 32)", - (1, 128, 256), - ("positions_random", 2051, 8191), - "relu", - ), - _cell( - "prefill: 8192 rows from position 0, rewritten in row tiles", - 8192, - ("chunk", 0), - "relu", - ), - _cell( - "input from the radix engine (qsa_topk_select), zero ties", - 64, - ("positions_random", 2051, 8191), - "relu_sparse", - selection="radix", - ), - ) - - note = """ - Metadata consumed: none. Stateless; every input is an argument, and `indices` is - both the engine's selection read and the canonical row written. - - Preconditions (none checked): - - `indices` holds an EXACT top-k of each window: the smallest selected score is - the window's k-th largest. `qsa_topk_select` is exact; an approximate selection - would be rewritten into a row with more than `k` columns at or above its k-th - pick, truncated to the lowest `k` column indices -- wrong, silently. - - `0 <= row_starts[r] <= row_ends[r] <= logits.shape[1]`. Anything outside the - window is never read for a decision, so garbage, NaN and +inf there are fine - (the cells poison it). - - Scores inside the window are not NaN: a NaN column is never kept, and a NaN pick - empties the row. (The engine returns the identity for a window of at most k - without reading it; this rewrite reads it.) - - Cost: about 30 kernels -- elementwise passes, two int32 prefix sums over - `[rows, cols]` and one sorted search over `[rows, k]` -- with no host read and no - data-dependent shape, so it can be captured into a CUDA graph (not certified: no - cell replays one). It reads the whole `cols`-wide score row, not the live window: a - caller that knows a bound on the windows passes `logits[:, :bound]`, which gives the - same rows. A matrix of more than 2^24 elements is rewritten in row tiles. - """ - - def __call__( - self, - logits: torch.Tensor, - row_starts: torch.Tensor, - row_ends: torch.Tensor, - indices: torch.Tensor, - ) -> torch.Tensor: - num_rows, num_cols = logits.shape - tile = max(1, _TILE_ELEMENTS // max(num_cols, 1)) - for lo in range(0, num_rows, tile): - hi = min(lo + tile, num_rows) - _rewrite_tile(logits[lo:hi], row_starts[lo:hi], row_ends[lo:hi], indices[lo:hi]) - return indices - - def reference( - self, - logits: torch.Tensor, - row_starts: torch.Tensor, - row_ends: torch.Tensor, - indices: torch.Tensor, - ) -> torch.Tensor: - """The k largest `(score, column)` pairs of each window, ascending by column, - `-1`-padded: a top-k over one int64 key per column whose high word orders the - scores as fp32 does and whose low word is the column. `indices` supplies only the - width `k`; its contents are ignored.""" - num_rows, num_cols = logits.shape - index_topk = indices.shape[1] - columns = torch.arange(num_cols, device=logits.device)[None, :] - starts, ends = row_starts.long()[:, None], row_ends.long()[:, None] - in_window = (columns >= starts) & (columns < ends) - # `+ 0.0` folds -0.0 into +0.0, so the two are one tie class. - bits = (logits.float() + 0.0).view(torch.int32).long() - order = torch.where(bits < 0, bits ^ 0x7FFFFFFF, bits) - key = torch.where(in_window, (order << 32) | columns, torch.iinfo(torch.int64).min) - width = min(index_topk, num_cols) - top = torch.topk(key, width, dim=1).indices - taken = torch.gather(in_window, 1, top) - pad = torch.iinfo(torch.int64).max - ordered = torch.sort(torch.where(taken, top, pad), dim=1).values - out = torch.where(ordered == pad, -1, ordered - starts).to(torch.int32) - if width < index_topk: - out = torch.nn.functional.pad(out, (0, index_topk - width), value=-1) - return out - - def is_valid( - self, - logits: torch.Tensor, - row_starts: torch.Tensor, - row_ends: torch.Tensor, - indices: torch.Tensor, - ) -> None: - # A one-element `row_starts` / `row_ends`, or a one-row `indices`, broadcasts over - # every row of `logits` without an error: a [1] `row_starts` against 4 rows - # returns every row framed by row 0's start. - num_rows = logits.shape[0] - assert row_starts.shape == (num_rows,) and row_ends.shape == (num_rows,), ( - f"row_starts and row_ends need one entry per logits row ({num_rows}); got " - f"{tuple(row_starts.shape)} and {tuple(row_ends.shape)}" - ) - assert indices.dim() == 2 and indices.shape[0] == num_rows, ( - f"indices must be [{num_rows}, k]; got {tuple(indices.shape)}" - ) - - -qsa_canonical_selection = QsaCanonicalSelection() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py deleted file mode 100644 index fff162545bc7..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py +++ /dev/null @@ -1,245 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""QSA compressed-block index scoring, read straight out of the paged index-K side cache.""" - -from __future__ import annotations - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -# Approved exception: QSA's selection path is Triton called as Python, with no -# torch.ops boundary. The engine's own sparse hook (`select_qsa_paged_tokens` -# in `_torch/attention/backends/sparse/qsa/module.py`) calls this same callable. -from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import ( - triton_qsa_paged_index_scores as _triton_qsa_paged_index_scores, -) - -_CHECKPOINT = dict( - heads=4, - head_dim=128, - tokens_per_block=32, - compress_ratio=4, - dtype=torch.bfloat16, - # `qsa_logical_positions` is int64, `qsa_req_idx_per_token` int32. - positions_dtype=torch.int64, - request_dtype=torch.int32, - # The target's q is a head slice of its packed [T, 4 + 1 + 1, 128] index - # scratch buffer, so its row stride is 768 rather than 512. - q_layout="scratch_slice", -) - - -def _spec(**overrides) -> dict: - """A cell spec: this checkpoint's geometry unless the cell says otherwise.""" - return {**_CHECKPOINT, **overrides} - - -#: The serving runtime dumps: an [8, 260] table on the warm-up capture path and -#: a [32, 260] one under the mmmu gate's mixed steps. -_WIDTH = 260 - - -#: The same plane under an NVFP4 KV pool: the QSA position page joins each physical -#: page after the index buffers (`page_pad` elements), so the trunk's view is -#: `[pages, 32, 1, 128]` at page stride 12 x 4096 + 192 = 49344. -_TRUNK_POOL_NVFP4 = dict(pool_buffers=12, page_pad=192) - - -class QsaPagedIndexScores(OpWrapper): - """Score every causally visible compressed block of each query row. - - Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels. - triton_qsa_paged_index_scores` (one Triton launch after allocating the - output). For row `r` and compressed block `b` of request - `request_indices[r]`:: - - anchor = b * compress_ratio + compress_ratio - 1 - key_b = index_cache[block_table[req, anchor // tpb], anchor % tpb, 0] - out[r, b] = sum_h relu(q[r, h] . key_b) * index_head_dim ** -0.5 - - The block's key is whatever the compression step left at the group's LAST - position; nothing is pooled here. Row `r` sees `(query_positions[r] + 1) // - compress_ratio` blocks (complete groups only). Columns at or past that are - `-inf` when `only_visible_blocks` is False and **not written** -- the fresh - allocation's garbage -- when it is True. The column count is the block - table's width, `block_table.shape[1] * tpb // compress_ratio`. - - `context_rows` is a layout declaration ("consecutive rows mostly share a - request"), not a phase one: with more than 256 rows and an index-head count - that divides 128 into more than one row (1, 2, 4, ..., 64) it selects a - query-tiled kernel sharing one key gather across `128 // heads` rows; - otherwise the row-wise kernel runs. Each cell names the kernel it was measured on. The - two are not bitwise equal. - - Fusion boundary: the table walk, the gather, the dots, relu, head sum, - scale and causal bound are inside. Producing `q` (norm, RoPE), writing and - compressing the index cache, and everything after the scores are not. - Returns a fresh fp32 `[rows, columns]`; nothing is written in place. - """ - - ARCHS = frozenset({Arch.SM_103}) - - CELLS: tuple[Cell, ...] = ( - Cell( - why="NVFP4 trunk plane, page stride 49344: decode MTP 4 requests x 2 rows", - spec=_spec( - layout="decode", - gen_positions=(7000, 5000, 4200, 2100), - rows_per_request=2, - slots=32, - pages=_WIDTH, - only_visible_blocks=True, - context_rows=False, - branch="row_wise", - **_TRUNK_POOL_NVFP4, - ), - ), - Cell( - why="decode MTP7: 32 slots x 8 rows = 256 rows, the last decode-sized count", - spec=_spec( - layout="decode", - gen_positions=tuple(100 + (499 * g) % 8000 for g in range(32)), - rows_per_request=8, - slots=32, - pages=_WIDTH, - only_visible_blocks=True, - context_rows=False, - branch="row_wise", - ), - ), - Cell( - why="capture: 8191 rows over 8 ragged requests, visible-only", - spec=_spec( - layout="capture", - lengths=(1024, 512, 2048, 256, 1500, 700, 1151, 1000), - slots=8, - pages=_WIDTH, - only_visible_blocks=True, - context_rows=True, - branch="query_tile", - ), - ), - ) - - note = """ - Metadata consumed: none -- every input is an argument. What the target - passes: `block_table` is `QSAAttentionMetadata.qsa_block_table` (raw - pool-page ids, int32, -1 past each sequence); `query_positions` / - `request_indices` are `qsa_logical_positions` / `qsa_req_idx_per_token`; - `index_cache` is the layer's bf16 index-K plane (`get_index_k_buffer`). - - Preconditions (not checked): every block-table entry the causal bound - reaches is a valid page, and `request_indices[r]` is a row of the table -- - the loads are masked by the bound, not by the pool size, so a violation is - an out-of-bounds read (not probed). The compression step must already have - run: an uncompressed anchor scores as a raw index key. `rows == 0` returns - an empty `[0, columns]` tensor without launching. - - The op validates its own metadata loudly (`ValueError`s; the test pins the - one-entry-per-row metadata refusal), so `is_valid` adds nothing. - - `only_visible_blocks=True` really leaves garbage (measured, not pinned by a - test): with a poisoned allocation freed just before the call, the whole - unscored tail came back poisoned. At 8191 rows from position 0 over 2080 columns that is a full - 32-column band (2048..2079) for every row, 8.65M cells. A consumer that - scans `[row_start, row_end)` -- `qsa_topk_select` -- is unaffected. - - fp32 operands are accepted but `tl.dot` rounds them to a tensor-core - format: 8.1e-3 max abs against an fp32 reference on scores to 7.8 (measured), - outside the test's band. No cell covers fp32; the side cache is bf16. - - CUDA graphs: the test captures the decode launch (`use_pdl`, `prefetch`) at - 4 rows and replays it after `q` and the anchor keys were rewritten in - place; every replay matches an eager call bitwise. - - `prefetch` (with `use_pdl` and `only_visible_blocks`, at generation-sized - batches) reads the pages and compressed keys before the PDL wait and q after - it. Sound only behind the decode step's own prologue - (`qwen4_exp_attn_decode_prologue`), which writes no compressed key but the last - visible block's, at a row whose position ends a group; that key is read again - after the wait. - - Every cell runs `only_visible_blocks=True` (the target never passes False), - 4 index heads x 128, compress ratio 4, 32-token pages, bf16 operands, int64 - positions and int32 request indices. The NVFP4-pool page stride is - certified on the row-wise kernel only. - - Not certified: more than 8191 rows, tables wider than 260 pages, more than - 32 table rows, any other head geometry, compress ratio, page size or - operand dtype, a context step's mixed request rows, and the `-inf` padded - mode. Hard limits: not probed; the largest certified extent is 8191 x 2080 - output elements. - """ - - def __call__( - self, - q: torch.Tensor, - index_cache: torch.Tensor, - block_table: torch.Tensor, - query_positions: torch.Tensor, - request_indices: torch.Tensor, - tokens_per_block: int, - compress_ratio: int, - only_visible_blocks: bool, - context_rows: bool, - use_pdl: bool = False, - prefetch: bool = False, - ) -> torch.Tensor: - return _triton_qsa_paged_index_scores( - q=q, - index_cache=index_cache, - block_table=block_table, - query_positions=query_positions, - request_indices=request_indices, - tokens_per_block=tokens_per_block, - compress_ratio=compress_ratio, - only_visible_blocks=only_visible_blocks, - context_rows=context_rows, - use_pdl=use_pdl, - prefetch=prefetch, - ) - - def reference( - self, - q: torch.Tensor, - index_cache: torch.Tensor, - block_table: torch.Tensor, - query_positions: torch.Tensor, - request_indices: torch.Tensor, - tokens_per_block: int, - compress_ratio: int, - only_visible_blocks: bool, - context_rows: bool, - use_pdl: bool = False, - prefetch: bool = False, - ) -> torch.Tensor: - """fp32 scores, `-inf` past each row's bound -- or NaN there, meaning - "unspecified", when `only_visible_blocks` is True. `context_rows` is - ignored: it picks a kernel, not a result. - - Walked in row chunks only to bound the gathered keys, which reach - 8.7 GiB in fp32 at 8191 rows x 2080 columns. - """ - rows, _, head_dim = q.shape - num_columns = block_table.shape[1] * tokens_per_block // compress_ratio - blocks = torch.arange(num_columns, device=q.device) - anchors = blocks * compress_ratio + compress_ratio - 1 - logical_pages, in_page = anchors // tokens_per_block, anchors % tokens_per_block - table = block_table.long() - requests = request_indices.long() - scores = torch.empty(rows, num_columns, dtype=torch.float32, device=q.device) - chunk = max(1, 2**26 // max(num_columns * head_dim, 1)) - for lo in range(0, rows, chunk): - hi = min(lo + chunk, rows) - # Pages past a sequence are -1; their columns are masked below. - physical = table[requests[lo:hi]][:, logical_pages].clamp_min(0) - keys = index_cache[physical, in_page.expand(hi - lo, num_columns)][:, :, 0, :] - per_head = torch.einsum("rhd,rbd->rbh", q[lo:hi].float(), keys.float()) - scores[lo:hi] = torch.relu(per_head).sum(-1) * (head_dim**-0.5) - visible = blocks[None, :] < ((query_positions.long() + 1) // compress_ratio)[:, None] - tail = float("nan") if only_visible_blocks else float("-inf") - return scores.masked_fill(~visible, tail) - - -qsa_paged_index_scores = QsaPagedIndexScores() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py deleted file mode 100644 index 21609fd37846..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa.py +++ /dev/null @@ -1,283 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""QSA sparse GQA: attention over an explicitly selected set of paged K/V tokens.""" - -from __future__ import annotations - -from typing import Optional - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -# Approved exception: QSA's sparse path never reaches the thop attention op's -# sparse argument group, so there is no torch.ops surface for this computation. -from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import ( - can_fuse_qsa_splitk_output_gate as _can_fuse_qsa_splitk_output_gate, -) -from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import ( - triton_qsa_paged_sparse_gqa as _triton_qsa_paged_sparse_gqa, -) - -_CHECKPOINT = dict( - q_heads=24, - kv_heads=2, - head_dim=256, - tokens_per_block=32, - compress_ratio=4, - token_topk=2048, - cache_dtype=torch.bfloat16, - q_dtype=torch.bfloat16, - # q is the first 24 heads of the packed [T, 24 + 2 + 2, 256] QKV buffer; - # "fused_row" is the target's fused input projection row, 55 heads of 256. - q_layout="qkv_slice", - request_dtype=torch.int32, - # The engine's table: 260 pages per slot, -1 past each sequence. - pages=260, -) - -_SPLIT_K = ("splitk", "merge") -_PREFILL = ("prefill",) -_FP8 = dict(cache_dtype=torch.float8_e4m3fn) - - -def _spec(**overrides) -> dict: - """A cell spec: this checkpoint's geometry unless the cell says otherwise. - - `segments` is `(first position, rows)` per request, in row order; - `prefill_pair` says whether `query_positions` / `compress_ratio` are - passed, which the target does exactly when the step has a context request. - """ - return {**_CHECKPOINT, **overrides} - - -def _decode(segments, launches=_SPLIT_K, **overrides) -> dict: - base = dict(segments=segments, slots=32, prefill_pair=False, launches=launches) - return _spec(**{**base, **overrides}) - - -class QsaPagedSparseGqa(OpWrapper): - """Sparse GQA over the selected tokens of an HND paged K/V pool. - - Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels. - triton_qsa_paged_sparse_gqa`. For row `r`, query head `h`, KV head - `h // (Hq // Hkv)` (contiguous groups, not interleaved), over the columns - `j` of `selected_tokens[r]` that name a resident token:: - - t = selected_tokens[r, j] # logical position - page = block_table[request_indices[r], t // tokens_per_block] - s_j = (q[r, h] . k_cache[page, kv, t % tpb]) * softmax_scale - out[r, h] = sum_j softmax(s)_j * v_cache[page, kv, t % tpb] - - A negative index, a page past the table's width, a negative or - out-of-pool table entry, and a request index outside the table's rows all - contribute nothing; a row left with nothing returns zeros. Indices are not - deduplicated: a token named twice is attended twice. Returns a fresh, - contiguous `[rows, Hq, D]` tensor of `q.dtype`, even for a strided `q`. - - Two kernels, chosen inside the call: with `query_positions` and - `compress_ratio` given and `rows * Hkv > 512`, a causal-prefix prefill - kernel; otherwise split-K, plus a merge kernel when it used more than one - split. Each cell names the launches it was measured on. - - **The prefill pair is a bound, not a mask.** The prefill kernel reads only - `min(visible // cr, block_topk) * cr + visible % cr` leading columns of a - row (`visible = query_positions[r] + 1`, rounded up to a 64-wide tile) and - nothing after them. That is exactly what QSA's own top-k + expansion emit, - packed at the front; a selection with valid entries past the bound is - silently truncated (measured: a planted index moving the right answer by - 0.77 moved the op's by nothing). The split-K kernel ignores the pair. - - **The output gate rides the merge when `output_gate` is passed.** Each - merged head row is rounded to `q.dtype`, multiplied in fp32 by - `sigmoid(output_gate)` and rounded again: the arithmetic of the ungated - call followed by `fused_sigmoid_mul_inplace`. `output_gate` is `[rows, - Hq * D]` or `[rows, Hq, D]` in `q.dtype`, unit stride in its last dim, any - row stride (the gate columns of the fused projection row). Only a split-K - launch with more than one split has a merge to carry it; `fuses_output_gate` - says whether a call does, and the op raises `ValueError` if handed a gate - it cannot carry. - - Fusion boundary: the block-table gather, scale, online softmax, P.V, the - split-K merge and, when passed, the output gate are inside. Deciding what - to attend to, appending K/V to the pool, and q's projection, norm and RoPE - are not. - """ - - ARCHS = frozenset({Arch.SM_103}) - - CELLS: tuple[Cell, ...] = ( - # ---- generation: no prefill pair, split-K --------------------------- - Cell( - why="decode MTP7: 32 slots x 8 rows = 256 rows, 512 programs, 4-way split-K, " - "q and the gate off the fused row, the gate fused into the merge", - spec=_decode( - tuple((100 + (499 * g) % 8000, 8) for g in range(32)), - q_layout="fused_row", - gate="fused_row", - ), - ), - # ---- context: the prefill pair is passed ---------------------------- - Cell( - why="fp8 context 272 rows from 2032: causal-prefix prefill kernel", - spec=_decode(((2032, 272),), launches=_PREFILL, prefill_pair=True, **_FP8), - ), - Cell( - why="q off the fused projection row (55 x 256), 8191-row prefill", - spec=_decode(((0, 8191),), launches=_PREFILL, prefill_pair=True, q_layout="fused_row"), - ), - ) - - note = """ - Metadata consumed: none -- every input is an argument. What the target - passes: `block_table` is `qsa_block_table`, physical page ids into dim 0 - of the caches (the *unscaled* table; a reference assuming logical == - physical lands 306x outside tolerance); `selected_tokens` is - `qsa_expand_block_indices`' output, logical positions, `-1`-padded; - `k_cache` / `v_cache` are the layer pool's role planes. - - Preconditions: - - Whenever the prefill kernel can run, each selection row is packed at - the front with exactly the bound's count of valid entries (see the - class docstring). QSA's producer -- `qsa_topk_select` with zero - `row_starts` and `row_ends = (pos + 1) // cr`, then the expansion -- - was measured to satisfy it at 300, 8191 and 8192 rows, including with - NaN/inf planted in the scores. Checking it needs the contents, so it - is not an `is_valid` guard. - - The prefill pair is a property of the launch: passed because some - rows are context, it puts the generation rows on the prefill kernel - too, and the bound applies to them (measured with a planted index; not - pinned by a test). - - There is no KV scale argument: an fp8 pool is read at face value (each code - widens exactly to bf16 under a bf16 q), so it must hold K/V written at unit scale. - - `D` must be a power of two (192 is refused); `Hq % Hkv == 0`. - - The op validates its own inputs loudly (`ValueError`s; the test pins the - refusal of positions passed without a compress ratio; a selection width - the compress ratio cannot produce is refused on the prefill kernel only), - so `is_valid` adds nothing. `rows == 0` returns an empty `[0, Hq, D]` - tensor (not certified). - - PDL: `TRTLLM_ENABLE_PDL` (default on) enables programmatic dependent - launch for multi-split launches of 4 rows or fewer; no cell runs one (the - target sends gated decode batches that small to - `qsa_paged_sparse_gqa_decode`). Calls are bitwise deterministic (pinned - on a 32-split launch and its merge). - - The decode site carries up to 8 rows per request (max_draft_len 7): the - decode cell runs 32 requests of 8 rows, 512 `rows * Hkv` programs and a - 4-way split-K launch with the output gate in its merge. - - Not certified: capture in a CUDA graph (the split-K path allocates its - partials inside the call); PDL launches; head geometries other than 24/2 - x 256; `tokens_per_block` other than 32; more than 8191 rows; the split-K - kernel above 256 rows, with one split (no merge) or without the gate; an - fp8 pool on the split-K kernel; rows whose selection resolves no token - (all padding, a foreign request, unmapped pages); an fp16 pool or q; - standalone contiguous K/V buffers; int64 request indices. - Hard limits: not probed; the largest certified extent is 8191 rows x 2051 - selected tokens through a [32, 260] table. - """ - - def __call__( - self, - *, - q: torch.Tensor, - k_cache: torch.Tensor, - v_cache: torch.Tensor, - block_table: torch.Tensor, - selected_tokens: torch.Tensor, - request_indices: torch.Tensor, - tokens_per_block: int, - softmax_scale: float, - query_positions: Optional[torch.Tensor] = None, - compress_ratio: Optional[int] = None, - output_gate: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - return _triton_qsa_paged_sparse_gqa( - q=q, - k_cache=k_cache, - v_cache=v_cache, - block_table=block_table, - selected_tokens=selected_tokens, - request_indices=request_indices, - tokens_per_block=tokens_per_block, - softmax_scale=softmax_scale, - query_positions=query_positions, - compress_ratio=compress_ratio, - output_gate=output_gate, - ) - - @staticmethod - def fuses_output_gate( - q: torch.Tensor, - k_cache: torch.Tensor, - selected_tokens: torch.Tensor, - output_gate: torch.Tensor, - ) -> bool: - """Whether a call with these operands and no prefill pair carries `output_gate`.""" - return _can_fuse_qsa_splitk_output_gate(q, k_cache, selected_tokens, output_gate) - - def reference( - self, - *, - q: torch.Tensor, - k_cache: torch.Tensor, - v_cache: torch.Tensor, - block_table: torch.Tensor, - selected_tokens: torch.Tensor, - request_indices: torch.Tensor, - tokens_per_block: int, - softmax_scale: float, - query_positions: Optional[torch.Tensor] = None, - compress_ratio: Optional[int] = None, - output_gate: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """fp32 sparse GQA over every resident selected token, rounded once; - with `output_gate`, then gated in fp32 and rounded again. - - `query_positions` and `compress_ratio` are ignored: on a prefix-packed - selection the prefill kernel's bound reads exactly the valid entries. - Walked 32 rows at a time to bound the gathered K/V. - """ - rows, num_q_heads, head_dim = q.shape - num_pages, num_kv_heads = k_cache.shape[0], k_cache.shape[1] - groups = num_q_heads // num_kv_heads - num_requests, table_width = block_table.shape - requests = request_indices.long() - out = torch.empty(rows, num_q_heads, head_dim, dtype=torch.float32, device=q.device) - for start in range(0, rows, 32): - stop = min(start + 32, rows) - chunk = selected_tokens[start:stop].long() - logical_page = chunk.clamp_min(0) // tokens_per_block - in_page = chunk.clamp_min(0) % tokens_per_block - request = requests[start:stop] - owned = (request >= 0) & (request < num_requests) - physical = block_table[ - request.clamp(0, num_requests - 1)[:, None].expand_as(logical_page), - logical_page.clamp(max=table_width - 1), - ].long() - resident = ( - (chunk >= 0) - & owned[:, None] - & (logical_page < table_width) - & (physical >= 0) - & (physical < num_pages) - ) - physical = physical.clamp(0, num_pages - 1) - keys = k_cache[physical, :, in_page, :].float() - values = v_cache[physical, :, in_page, :].float() - grouped = q[start:stop].reshape(-1, num_kv_heads, groups, head_dim).float() - scores = torch.einsum("bhgd,bkhd->bhgk", grouped, keys) * softmax_scale - scores = scores.masked_fill(~resident[:, None, None, :], float("-inf")) - # A row with nothing resident softmaxes over all -inf: NaN, i.e. zeros. - weights = torch.nan_to_num(torch.softmax(scores, dim=-1), nan=0.0) - pooled = torch.einsum("bhgk,bkhd->bhgd", weights, values) - out[start:stop] = pooled.reshape(stop - start, num_q_heads, head_dim) - if output_gate is None: - return out.to(q.dtype) - gate = output_gate.reshape(rows, num_q_heads, head_dim).float() - return (out.to(q.dtype).float() * torch.sigmoid(gate)).to(q.dtype) - - -qsa_paged_sparse_gqa = QsaPagedSparseGqa() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rollback_planes.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rollback_planes.py deleted file mode 100644 index 144df3bd24fc..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rollback_planes.py +++ /dev/null @@ -1,144 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The speculative rollback of every attention layer's QSA index-K plane in one launch each way: -gather the verify rows' slots of all planes before the trunk writes them, and after the -draft loop write the gathered values back into the slots of the rows the sampler rejected. - -Tier 5 (Triton written for this target). The planes are separate buffers of the cache -manager (one per attention layer, each `[pages, tokens_per_block, heads, dim]` with a free -page stride and a contiguous tail); the kernels reach them through a table of their base -addresses and page strides, one program per (plane, row). The composition this replaces ran -one gather per layer during the trunk and a gather, a `torch.where` and an indexed write per -layer after the draft loop. -""" - -from __future__ import annotations - -import torch -import triton -import triton.language as tl - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -#: This checkpoint's index plane row: one index KV head of 128 bf16. -ROW = 128 - - -@triton.jit -def _gather_kernel(ptrs, page_strides, pages, within, out, rows, ROW: tl.constexpr): - plane = tl.program_id(0) - r = tl.program_id(1) - base = tl.load(ptrs + plane).to(tl.pointer_type(tl.bfloat16)) - stride = tl.load(page_strides + plane) - page = tl.load(pages + r) - tok = tl.load(within + r) - offs = tl.arange(0, ROW) - vals = tl.load(base + page * stride + tok * ROW + offs) - tl.store(out + (plane.to(tl.int64) * rows + r) * ROW + offs, vals) - - -@triton.jit -def _restore_kernel(ptrs, page_strides, pages, within, rejected, old, rows, ROW: tl.constexpr): - plane = tl.program_id(0) - r = tl.program_id(1) - if tl.load(rejected + r) != 0: - base = tl.load(ptrs + plane).to(tl.pointer_type(tl.bfloat16)) - stride = tl.load(page_strides + plane) - page = tl.load(pages + r) - tok = tl.load(within + r) - offs = tl.arange(0, ROW) - vals = tl.load(old + (plane.to(tl.int64) * rows + r) * ROW + offs) - tl.store(base + page * stride + tok * ROW + offs, vals) - - -def plane_table(planes: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: - """`(base addresses, page strides)` of the planes, int64 on their device (built once per - cache-manager bind: the planes do not move).""" - dev = planes[0].device - for t in planes: - assert t.dtype == torch.bfloat16 and t.dim() == 4, "index planes [pages, tpb, H, D] bf16" - assert t.shape[2] * t.shape[3] == ROW, "one 128-wide row a token" - assert t.stride(1) == ROW and t.stride(2) == t.shape[3] and t.stride(3) == 1, ( - "contiguous [tpb, H, D] tail" - ) - ptrs = torch.tensor([t.data_ptr() for t in planes], dtype=torch.int64, device=dev) - strides = torch.tensor([t.stride(0) for t in planes], dtype=torch.int64, device=dev) - return ptrs, strides - - -class QsaRollbackGather(OpWrapper): - """`torch.stack([plane[pages, within] for plane in planes])` as `[L, R, 128]` bf16, with - the planes given as `plane_table(planes)`: every row's slot of every plane, in plane - order. Nothing is written but the output. - - Fusion boundary: the gather of all planes; the slots (`pages`, `within`) are the caller's. - """ - - ARCHS = frozenset({Arch.SM_103}) - CELLS: tuple[Cell, ...] = ( - Cell(why="12 planes, 4 verify rows (MTP3, one request)", spec=dict(planes=12, rows=4)), - Cell(why="12 planes, 512 verify rows (MTP3, 128 requests)", spec=dict(planes=12, rows=512)), - ) - note = """ - Metadata consumed: none. Call site: `SpeculativeRollback.begin`, before the verify step's - trunk writes any plane. One program per (plane, row), one 256-byte row each. - Preconditions: `pages` / `within` int64 `[R]`; the planes as `plane_table` accepts them. - """ - - def __call__(self, ptrs, page_strides, pages, within, num_planes: int) -> torch.Tensor: - rows = pages.shape[0] - out = torch.empty([num_planes, rows, ROW], dtype=torch.bfloat16, device=pages.device) - if rows: - _gather_kernel[(num_planes, rows)]( - ptrs, page_strides, pages, within, out, rows, ROW=ROW, num_warps=1 - ) - return out - - def reference(self, ptrs, page_strides, pages, within, num_planes: int) -> torch.Tensor: - raise NotImplementedError("the cell harness builds planes and compares against torch") - - def is_valid(self, ptrs, page_strides, pages, within, num_planes: int) -> None: - assert ptrs.dtype == torch.int64 and ptrs.numel() == num_planes - assert page_strides.dtype == torch.int64 and page_strides.numel() == num_planes - assert pages.dtype == torch.int64 and within.dtype == torch.int64 - assert pages.shape == within.shape and pages.dim() == 1 - - -class QsaRollbackRestore(OpWrapper): - """For every row `r` with `rejected[r]` and every plane `l`: `plane_l[pages[r], within[r]] - = old[l, r]` -- the inverse of `qsa_rollback_gather` over the rejected rows. Rows not - rejected are not touched. - - Fusion boundary: the restore of all planes; `rejected` is the caller's. - """ - - ARCHS = frozenset({Arch.SM_103}) - CELLS: tuple[Cell, ...] = ( - Cell(why="12 planes, 4 verify rows, two rejected", spec=dict(planes=12, rows=4)), - Cell(why="12 planes, 512 verify rows", spec=dict(planes=12, rows=512)), - ) - note = """ - Metadata consumed: none. Call site: `SpeculativeRollback.commit_speculative_states`, - after the draft loop. `rejected` int32/int8/bool `[R]` (nonzero = restore). - """ - - def __call__(self, ptrs, page_strides, pages, within, rejected, old) -> None: - num_planes, rows, _ = old.shape - if rows: - _restore_kernel[(num_planes, rows)]( - ptrs, page_strides, pages, within, rejected, old, rows, ROW=ROW, num_warps=1 - ) - - def reference(self, ptrs, page_strides, pages, within, rejected, old) -> None: - raise NotImplementedError("the cell harness builds planes and compares against torch") - - def is_valid(self, ptrs, page_strides, pages, within, rejected, old) -> None: - assert old.dtype == torch.bfloat16 and old.dim() == 3 and old.shape[2] == ROW - assert old.is_contiguous() - assert ptrs.numel() == old.shape[0] and page_strides.numel() == old.shape[0] - assert pages.shape == (old.shape[1],) and within.shape == (old.shape[1],) - assert rejected.shape == (old.shape[1],) - - -qsa_rollback_gather = QsaRollbackGather() -qsa_rollback_restore = QsaRollbackRestore() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py deleted file mode 100644 index ae062832dcb6..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_expand_decode.py +++ /dev/null @@ -1,651 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""QSA's block top-k and its expansion to token indices, one launch per batch of rows. - -Tier 5 (CuTe DSL written for this target). The decode path ran `qsa_topk_select` -(the CUDA indexer top-k, one CTA a row) and then `qsa_expand_block_indices` (one -more launch) per attention layer. Here one CTA per row holds four consecutive scores a -thread and narrows to the k-th largest by radix passes over the scores' bit patterns -(shared-memory histograms of 11 bits), starting below the bits every visible score -shares. It stops as soon as the bin holding the threshold has at most `CAP` keys and -ranks those exactly by pairwise comparison (ties to the higher block, the canonical -selection's rule). Every selected block is then placed in block order by one block-wide -scan, and expanded to token indices with the causal tail, as the expansion entry does. - -Scores are a sum of ReLUs times a positive scale, so non-negative: their fp32 bit -patterns, read as int32, order like their values (a negative zero is taken as zero, a -NaN ranks above everything). -""" - -from __future__ import annotations - -import cutlass -import cutlass.cute as cute -import torch -from cutlass import const_expr -from cutlass.cute.runtime import from_dlpack - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl -from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr - -from .qsa_expand_block_indices import qsa_expand_block_indices as _expand - -#: Consecutive scores a thread holds (one 16-byte load). -KPT = 4 -#: Score columns a row may have: `KPT` a thread, at most 1024 threads a CTA. -MAX_WIDTH = KPT * 1024 -#: CTAs an SM the register allocation is held to. At the 2592-wide score table a row is 672 -#: threads, which the per-subpartition register file holds once an SM at the compiler's own -#: 48 registers a thread; at three (32 a thread) rows overlap each other's barriers and -#: shared-memory round trips (GB300, a mixed step's 16256 context rows: 426 -> 236 us on -#: half-zero rows, 486 -> 284 on relu-sparse ones; 128 decode rows 8.2 us both ways). -MIN_BLOCKS = 3 -#: Radix passes at most (11 bits each; 31 magnitude bits). -_PASSES = 3 -_HIST = 2048 -_I32_MAX = 0x7FFFFFFF -#: At most this many keys in the threshold's bin are ranked pairwise instead of -#: narrowed by another pass. -CAP = 128 -#: Threads that own the 2048 bins of a scan, four each. -_SCAN_THREADS = _HIST // 4 -#: Consecutive threads that rank one candidate of the threshold's bin, each over every -#: `_RANK_LANES`-th candidate (`_SCAN_THREADS` threads cover `CAP` candidates). -_RANK_LANES = _SCAN_THREADS // CAP - - -def _threads(width: int) -> int: - return max(_SCAN_THREADS, -(-width // (KPT * 32)) * 32) - - -class _TopkExpandKernel: - def __init__(self, *, width, key_stride, block_topk, compress, token_topk, pdl): - self.width, self.key_stride, self.k, self.cr, self.token_topk, self.pdl = ( - width, - key_stride, - block_topk, - compress, - token_topk, - pdl, - ) - self.final = token_topk + compress - 1 - self.threads = _threads(width) - self.warps = self.threads // 32 - self.flag_words = -(-width // 32) - - @cute.jit - def __call__( - self, - keys_t: cute.Tensor, - end_ptr: cute.Pointer, - pos_ptr: cute.Pointer, - len_ptr: cute.Pointer, - topk_ptr: cute.Pointer, - sel_ptr: cute.Pointer, - rows: cutlass.Int32, - stream, - ): - # The keys arrive as a tensor, as the TVM-FFI entry takes its stream from a GPU - # tensor; their layout is rebuilt with the static width and row stride the 16-byte - # loads need. The row count is the call's. - mKeys = cute.make_tensor( - keys_t.iterator, cute.make_layout((rows, self.width), stride=(self.key_stride, 1)) - ) - per_row = cute.make_layout((rows,), stride=(1,)) - mEnd = cute.make_tensor(end_ptr, per_row) - mPos = cute.make_tensor(pos_ptr, per_row) - mLen = cute.make_tensor(len_ptr, per_row) - mTopk = cute.make_tensor(topk_ptr, cute.make_layout((rows, self.k), stride=(self.k, 1))) - mSel = cute.make_tensor( - sel_ptr, cute.make_layout((rows, self.final), stride=(self.final, 1)) - ) - ld_keys = cute.make_copy_atom( - cute.nvgpu.CopyG2ROp(), cutlass.Int32, num_bits_per_copy=32 * KPT - ) - self.kernel(mKeys, mEnd, mPos, mLen, mTopk, mSel, ld_keys).launch( - grid=[rows, 1, 1], - block=[self.threads, 1, 1], - smem=(_HIST + self.flag_words + 2 * CAP + 64 + 16) * 4 + 128, - stream=stream, - use_pdl=self.pdl, - min_blocks_per_mp=MIN_BLOCKS, - ) - - @cute.kernel - def kernel(self, mKeys, mEnd, mPos, mLen, mTopk, mSel, ld_keys: cute.CopyAtom): - tidx, _, _ = cute.arch.thread_idx() - row, _, _ = cute.arch.block_idx() - warp = cute.arch.warp_idx() - lane = cute.arch.lane_idx() - nt: cutlass.Constexpr = self.threads - nw: cutlass.Constexpr = self.warps - k: cutlass.Constexpr = self.k - cr: cutlass.Constexpr = self.cr - smem = cutlass.utils.SmemAllocator() - hist = smem.allocate_tensor( - cutlass.Int32, cute.make_layout((_HIST,), stride=(1,)), byte_alignment=16 - ) - flags = smem.allocate_tensor( - cutlass.Int32, cute.make_layout((self.flag_words,), stride=(1,)), byte_alignment=16 - ) - cand_key = smem.allocate_tensor( - cutlass.Int32, cute.make_layout((CAP,), stride=(1,)), byte_alignment=16 - ) - cand_idx = smem.allocate_tensor( - cutlass.Int32, cute.make_layout((CAP,), stride=(1,)), byte_alignment=16 - ) - wa = smem.allocate_tensor( - cutlass.Int32, cute.make_layout((32,), stride=(1,)), byte_alignment=16 - ) - wb = smem.allocate_tensor( - cutlass.Int32, cute.make_layout((32,), stride=(1,)), byte_alignment=16 - ) - state = smem.allocate_tensor( - cutlass.Int32, cute.make_layout((16,), stride=(1,)), byte_alignment=16 - ) - - # 0. Shared state, before the wait (nothing global is read here). - for b in cutlass.range_constexpr(-(-_HIST // nt)): - if tidx + b * nt < _HIST: - hist[tidx + b * nt] = cutlass.Int32(0) - for b in cutlass.range_constexpr(-(-self.flag_words // nt)): - if tidx + b * nt < self.flag_words: - flags[tidx + b * nt] = cutlass.Int32(0) - if tidx == 0: - state[5] = cutlass.Int32(0) # candidates compacted - # The row's metadata, host-filled: read before the wait too. - n = mEnd[row] - seqlen = mLen[row] - visible = cutlass.Int32(mPos[row]) + 1 - if const_expr(self.pdl): - cute.arch.griddepcontrol_wait() - cute.arch.griddepcontrol_launch_dependents() - - # 1. Four consecutive keys a thread (a column past the window is no candidate: - # -1); the largest and smallest key in the window. - e0 = tidx * KPT - keys = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32) - keys.fill(-1) - if e0 < self.width: - cute.copy(ld_keys, cute.logical_divide(mKeys, (None, KPT))[row, (None, tidx)], keys) - kmax = cutlass.Int32(-1) - kmin = cutlass.Int32(_I32_MAX) - for c in cutlass.range_constexpr(KPT): - b = keys[c] - if b < 0: # -0.0 (or a negative): rank as zero - b = cutlass.Int32(0) - if e0 + c >= n: - b = cutlass.Int32(-1) - keys[c] = b - kmax = cutlass.max(kmax, b) - if b >= 0: - kmin = cutlass.min(kmin, b) - kmax = cute.arch.warp_reduction_max(kmax) - kmin = cute.arch.warp_reduction(kmin, lambda x, y: cutlass.min(x, y)) - if lane == 0: - wa[warp] = kmax - wb[warp] = kmin - cute.arch.sync_threads() - gmax = cutlass.Int32(-1) - gmin = cutlass.Int32(_I32_MAX) - if lane < nw: - gmax = wa[lane] - gmin = wb[lane] - gmax = cute.arch.warp_reduction_max(gmax) - gmin = cute.arch.warp_reduction(gmin, lambda x, y: cutlass.min(x, y)) - - # 2. Every key in the window shares the bits above the highest bit where the - # largest and the smallest differ (`hb`); radix passes walk down from there, 11 bits - # at a time, until the threshold's bin holds at most CAP keys or no bits are left. - diff = gmax ^ gmin - hb = cutlass.Int32(-1) - if diff != 0: - hb = 31 - cutlass.Int32(cute.arch.clz(diff)) - low = (cutlass.Int64(1) << cutlass.Int64(hb + 1)) - 1 - mask0 = cutlass.Int32(cutlass.Int64(_I32_MAX) & (cutlass.Int64(-1) ^ low)) - if tidx == 0: - state[0] = gmax & mask0 # threshold prefix - state[1] = mask0 # its mask - state[2] = cutlass.Int32(k) # still needed at the threshold - state[3] = n # keys in the threshold's bin - state[4] = cutlass.Int32(1) # narrowing done - state[6] = cutlass.max(hb - 10, 0) # the next pass's shift - if n > k: - if hb >= 0: - state[4] = cutlass.Int32(0) - cute.arch.sync_threads() - for _p in cutlass.range_constexpr(_PASSES): - if state[4] == 0: - prefix = state[0] - mask = state[1] - need = state[2] - shift = state[6] - for c in cutlass.range_constexpr(KPT): - if keys[c] >= 0: - if (keys[c] & mask) == prefix: - digit = (keys[c] >> shift) & (_HIST - 1) - cute.arch.atomic_add( - hist.iterator + digit, cutlass.Int32(1), scope="cta" - ) - cute.arch.sync_threads() - # Thread t < 512 owns bins 2047-4t .. 2044-4t, counted from the top; each is - # cleared for the next pass once read. - h = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), cutlass.Int32) - h.fill(0) - own = cutlass.Int32(0) - if tidx < _SCAN_THREADS: - for j in cutlass.range_constexpr(4): - h[j] = hist[_HIST - 1 - 4 * tidx - j] - hist[_HIST - 1 - 4 * tidx - j] = cutlass.Int32(0) - own = own + h[j] - incl = own - for off in (1, 2, 4, 8, 16): - up = cute.arch.shuffle_sync_up(incl, offset=off) - if lane >= off: - incl = incl + up - if lane == 31: - wa[warp] = incl - cute.arch.sync_threads() - above = cutlass.Int32(0) - if lane < warp: - above = wa[lane] - before = cute.arch.warp_reduction_sum(above) + incl - own # keys in the bins above - if tidx < _SCAN_THREADS: - if before < need: - if need <= before + own: - run = before - for j in cutlass.range_constexpr(4): - if run < need: - if need <= run + h[j]: - state[0] = prefix | ((_HIST - 1 - 4 * tidx - j) << shift) - state[1] = mask | ((_HIST - 1) << shift) - state[2] = need - run - state[3] = h[j] - state[4] = cutlass.Int32( - cutlass.select_((h[j] <= CAP) | (shift == 0), 1, 0) - ) - state[6] = cutlass.max(shift - 11, 0) - run = run + h[j] - cute.arch.sync_threads() - - prefix = state[0] - mask = state[1] - need = state[2] - at_thr = state[3] - # 3. At most CAP keys in the threshold's bin: rank them exactly, ties to the - # higher block, and flag the `need` best. - if n > k: - if at_thr <= CAP: - for c in cutlass.range_constexpr(KPT): - if keys[c] >= 0: - if (keys[c] & mask) == prefix: - slot = cute.arch.atomic_add( - state.iterator + 5, cutlass.Int32(1), scope="cta" - ) - cand_key[slot] = keys[c] - cand_idx[slot] = e0 + c - cute.arch.sync_threads() - # Candidate `ci` is ranked by `_RANK_LANES` consecutive lanes, lane `sub` - # counting the better ones among candidates sub, sub + _RANK_LANES, ...; - # butterfly shuffles add the counts up. - ci = tidx // _RANK_LANES - sub = tidx % _RANK_LANES - rank = cutlass.Int32(0) - it = cutlass.Int32(0) - if ci < at_thr: - kt = cand_key[ci] - it = cand_idx[ci] - for j in cutlass.range(sub, at_thr, _RANK_LANES, unroll=4): - kj = cand_key[j] - ij = cand_idx[j] - better = (kj > kt) | ((kj == kt) & (ij > it)) - rank = rank + cutlass.Int32(cutlass.select_(better, 1, 0)) - for off in cutlass.range_constexpr(_RANK_LANES.bit_length() - 1): - rank = rank + cute.arch.shuffle_sync_bfly(rank, offset=1 << off) - if ci < at_thr: - if sub == 0: - if rank < need: - cute.arch.atomic_or( - flags.iterator + (it >> 5), - cutlass.Int32(1) << (it & 31), - scope="cta", - ) - cute.arch.sync_threads() - - # 4. Block order: one scan of (selected, tied at the threshold) counts. A tied key - # is left only when the passes ran out of bits with more than CAP keys tied; the - # last `need` of those in block order are taken. - take = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32) - tied = cute.make_rmem_tensor(cute.make_layout((KPT,), stride=(1,)), cutlass.Int32) - own = cutlass.Int32(0) - for c in cutlass.range_constexpr(KPT): - t_c = cutlass.Int32(0) - q_c = cutlass.Int32(0) - if keys[c] >= 0: - if n <= k: - t_c = cutlass.Int32(1) - else: - mk = keys[c] & mask - if mk > prefix: - t_c = cutlass.Int32(1) - if mk == prefix: - if at_thr > CAP: - q_c = cutlass.Int32(1) - else: - e = e0 + c - t_c = (flags[e >> 5] >> (e & 31)) & 1 - take[c] = t_c - tied[c] = q_c - own = own + t_c + (q_c << 16) - incl = own - for off in (1, 2, 4, 8, 16): - up = cute.arch.shuffle_sync_up(incl, offset=off) - if lane >= off: - incl = incl + up - if lane == 31: - wb[warp] = incl - cute.arch.sync_threads() - above = cutlass.Int32(0) - if lane < warp: - above = wb[lane] - pre = cute.arch.warp_reduction_sum(above) + incl - own - taken_before = pre & 0xFFFF - tied_before = pre >> 16 - # The tied keys left out, the lowest-index ones (none unless the threshold's bin - # holds more keys than are still needed). - skip = cutlass.max(at_thr - need, 0) - for c in cutlass.range_constexpr(KPT): - e = e0 + c - sel = cutlass.Int32(0) - pos = cutlass.Int32(0) - if tied[c] == 1: - if tied_before >= skip: - sel = cutlass.Int32(1) - pos = taken_before + tied_before - skip - tied_before = tied_before + 1 - if take[c] == 1: - sel = cutlass.Int32(1) - pos = taken_before + cutlass.max(tied_before - skip, 0) - taken_before = taken_before + 1 - if sel == 1: - mTopk[row, pos] = e - for j in cutlass.range_constexpr(cr): - tok = e * cr + j - val = cutlass.Int32(-1) - if tok < seqlen: - val = tok - mSel[row, pos * cr + j] = val - - # 5. Padding of the block list; the causal tail after the selected tokens, then -1. - total = n - if n > k: - total = cutlass.Int32(k) - for c in cutlass.range_constexpr(-(-k // nt)): - s = tidx + c * nt - if s < k: - if s >= total: - mTopk[row, s] = cutlass.Int32(-1) - tail_start = (visible // cr) * cr - tail_count = visible - tail_start - body = total * cr - if body > self.token_topk: - body = cutlass.Int32(self.token_topk) - for c in cutlass.range_constexpr(-(-self.final // nt)): - col = tidx + c * nt - if col < self.final: - if col >= body: - off = col - body - val = cutlass.Int32(-1) - if off < cr - 1: - if off < tail_count: - if tail_start + off < seqlen: - val = tail_start + off - mSel[row, col] = val - - -#: One compiled launch per score-table geometry, k, compress ratio, token budget, index -#: dtypes and PDL; the row count is a runtime argument. -_COMPILED: dict = {} -_INDEX_TYPES = {torch.int32: cutlass.Int32, torch.int64: cutlass.Int64} - - -def launch( - scores, - row_ends, - query_positions, - sequence_lengths, - topk_indices, - compress_ratio, - token_topk, - pdl=True, - out=None, -): - """The launch behind `QsaTopkExpandDecode.__call__`: writes `topk_indices` and returns the - token selection `[rows, token_topk + compress_ratio - 1]` int32 (`out` when given).""" - rows, width = scores.shape - block_topk = topk_indices.shape[1] - sel = out - if sel is None: - sel = torch.empty( - (rows, token_topk + compress_ratio - 1), dtype=torch.int32, device=scores.device - ) - if rows == 0: - return sel - pdl = bool(pdl and get_env_enable_pdl()) - keys = scores.view(torch.int32) - key = ( - width, - scores.stride(0), - block_topk, - compress_ratio, - token_topk, - query_positions.dtype, - sequence_lengths.dtype, - pdl, - ) - ptrs = ( - row_ends.data_ptr(), - query_positions.data_ptr(), - sequence_lengths.data_ptr(), - topk_indices.data_ptr(), - sel.data_ptr(), - ) - fn = _COMPILED.get(key) - if fn is None: - gm = cute.AddressSpace.gmem - fn = _COMPILED[key] = cute.compile( - _TopkExpandKernel( - width=width, - key_stride=scores.stride(0), - block_topk=block_topk, - compress=compress_ratio, - token_topk=token_topk, - pdl=pdl, - ), - from_dlpack(keys, assumed_align=16).mark_layout_dynamic(leading_dim=1), - make_ptr(cutlass.Int32, ptrs[0], gm, assumed_align=4), - make_ptr( - _INDEX_TYPES[query_positions.dtype], - ptrs[1], - gm, - assumed_align=query_positions.element_size(), - ), - make_ptr( - _INDEX_TYPES[sequence_lengths.dtype], - ptrs[2], - gm, - assumed_align=sequence_lengths.element_size(), - ), - make_ptr(cutlass.Int32, ptrs[3], gm, assumed_align=4), - make_ptr(cutlass.Int32, ptrs[4], gm, assumed_align=4), - cutlass.Int32(rows), - cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), - options="--enable-tvm-ffi", - ) - fn(keys, *ptrs, rows) - return sel - - -def supports(scores) -> bool: - """Whether a score table this wide is one this entry serves (see `is_valid`).""" - return scores.shape[1] % KPT == 0 and scores.shape[1] <= MAX_WIDTH - - -class QsaTopkExpandDecode(OpWrapper): - """`qsa_topk_select` (zero row starts, `row_ends` the visible block counts), - `qsa_canonical_selection` and `qsa_expand_block_indices`, in one launch: one CTA a row, for - decode rows and a context chunk's rows alike. - - Per row: the `block_topk` largest of `scores[row, :row_ends[row]]` (non-negative fp32), - ties at the threshold to the higher block, written to `topk_indices[row]` in block order - and padded with -1 -- the composition's selection after `qsa_canonical_selection`; the - returned token selection is the expansion entry's: each selected group's members below - `sequence_lengths[row]`, packed at the front, then the incomplete group at the causal - frontier (`query_positions[row] + 1` visible tokens), then -1. `topk_indices` is written - in place, and the selection into `out` when given (a contiguous int32 - `[rows, token_topk + compress_ratio - 1]`, e.g. a row range of the step's whole - selection); nothing else is written. - - Fusion boundary: the selection and the expansion. Scoring is the caller's. - """ - - ARCHS = frozenset({Arch.SM_103}) - - CELLS: tuple[Cell, ...] = ( - Cell( - why="decode, 4 rows across the budget, one below it", - spec=dict(rows=4, pos=(8700, 4095, 1500, 9900)), - ), - # The target writes a context chunk's rows through `out` into its rows of the - # step's selection; relu-sparse scores put the zero class (negative zeros among it) - # across the budget, the ties the canonical rule decides. - Cell( - why="context chunk: 2048 rows of one request (positions 6144..8191), relu-sparse " - "scores, written through `out` into rows of a step-sized selection", - spec=dict(rows=2048, chunk=(6144, 2048), relu_sparse=True, out_margin=64), - ), - Cell( - why="decode, 1 row of scores within 5% of each other", - spec=dict(rows=1, pos=(8700,), narrow=True), - ), - # Serving's decode rows at 8K: positive, log-normally spread over ~2 octaves, a - # zero or two; the first pass leaves tens to ~CAP candidates at the threshold. - Cell( - why="decode, 16 rows of log-normal scores with a few zeros: threshold bins of " - "up to CAP candidates ranked pairwise, ties included", - spec=dict(rows=16, pos=tuple(8192 + 37 * r for r in range(16)), lognormal=True), - ), - ) - - note = """ - Metadata consumed: `qsa_visible_blocks` (`row_ends`), the rows' logical positions - and sequence lengths. Preconditions: zero row starts; `scores` fp32 `[rows, W]`, unit - column stride, 16-byte aligned rows, `W % 4 == 0` and `W <= 4096` (four keys a thread, - at most 1024 threads); `topk_indices` int32 `[rows, K]` contiguous; `row_ends` int32, - the positions and lengths int32 or int64, each a contiguous `[rows]`. Ties: the higher - block at the threshold, as `qsa_canonical_selection` rewrites the CUDA top-k's pick. - - Compiled once per score-table width and row stride (the block table's, fixed for an - engine), k, compress ratio, token budget, index dtypes and PDL: the row count is a - runtime argument, so an eager step's context rows, any count, reuse the kernel (read - from the launch, not pinned by a test). - - Certified: decode batches of 1 and 4 rows over a 2560-wide score table, a 2048-row - context chunk written through `out`, and a replay under a CUDA graph that reads the - scores after its wait. - """ - - def __call__( - self, - scores, - row_ends, - query_positions, - sequence_lengths, - topk_indices, - compress_ratio, - token_topk, - out=None, - ): - return launch( - scores, - row_ends, - query_positions, - sequence_lengths, - topk_indices, - compress_ratio, - token_topk, - out=out, - ) - - def reference( - self, - scores, - row_ends, - query_positions, - sequence_lengths, - topk_indices, - compress_ratio, - token_topk, - out=None, - ): - """`(topk_indices_after, selection)`: a stable descending sort per row (ties to the - higher block), the k first in block order, then the expansion entry's reference.""" - rows, _ = scores.shape - k = topk_indices.shape[1] - out = torch.full_like(topk_indices, -1) - for r in range(rows): - n = int(row_ends[r]) - s = scores[r, :n].float().clamp(min=0.0) - # A stable sort of the reversed row puts the higher block first among equals. - order = (n - 1) - torch.sort(torch.flip(s, [0]), descending=True, stable=True).indices[ - :k - ] - chosen = torch.sort(order).values - out[r, : chosen.numel()] = chosen.to(torch.int32) - sel = _expand.reference( - out, - query_positions, - sequence_lengths, - compress_ratio=compress_ratio, - token_topk=token_topk, - ) - return out, sel - - def is_valid( - self, - scores, - row_ends, - query_positions, - sequence_lengths, - topk_indices, - compress_ratio, - token_topk, - out=None, - ): - assert scores.dtype == torch.float32 and scores.stride(1) == 1, ( - "fp32 scores, unit column stride" - ) - assert scores.shape[1] % KPT == 0 and scores.shape[1] <= MAX_WIDTH, "W % 4 == 0, W <= 4096" - assert scores.stride(0) % KPT == 0 and scores.data_ptr() % 16 == 0, "16-byte aligned rows" - assert topk_indices.dtype == torch.int32 and topk_indices.is_contiguous(), ( - "int32 topk_indices" - ) - rows = scores.shape[0] - assert row_ends.dtype == torch.int32 and row_ends.shape == (rows,), "int32 row_ends [rows]" - for t in (query_positions, sequence_lengths): - assert t.dtype in _INDEX_TYPES and t.shape == (rows,), "int32/int64 [rows]" - assert all(t.is_contiguous() for t in (row_ends, query_positions, sequence_lengths)), ( - "contiguous row_ends, positions and lengths" - ) - assert token_topk == topk_indices.shape[1] * compress_ratio, ( - "token_topk = k * compress_ratio" - ) - assert out is None or ( - out.dtype == torch.int32 - and out.is_contiguous() - and out.shape == (scores.shape[0], token_topk + compress_ratio - 1) - ), "out: contiguous int32 [rows, token_topk + compress_ratio - 1]" - - -qsa_topk_expand_decode = QsaTopkExpandDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py deleted file mode 100644 index 6e5d284f768b..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py +++ /dev/null @@ -1,262 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Vision-tower 2-D axial rotary embedding (patch row / patch column), no KV -cache. - -Not the text tower's mRoPE: that rotates 64 of a 256-wide head against three -interleaved position rows (`mrope_section`, `rope_theta` 1e7); this rotates the -whole 72-wide vision head against two spatial rows from the image grid, 18 row -angles then 18 column angles, `theta` 1e4. Neither reads the other's config. - -Tier-3 composed from native torch (angle table, index gather, cos/sin, two -multiply-adds). Every registered rope op takes 1-D positions; -`flashinfer_apply_rope_with_cos_sin_cache_inplace` with an explicit per-token -table would sidestep that, but at `head_size=72` it raises "misaligned -address" and poisons the CUDA context. -""" - -from __future__ import annotations - -from typing import List, Sequence, Tuple - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -#: This checkpoint: 1152 / 16 heads = a 72-wide head, spatial_merge_size 2, -#: theta 1e4 (`model.visual.blocks.N.attn.qkv.weight` is [3456, 1152]). -NUM_HEADS = 16 -HEAD_DIM = 72 -MERGE = 2 -THETA = 10000.0 - -#: (frames, height, width) in patch units, from the runtime-value dump of the -#: mmmu gate (58.33): 170 segment sizes from 256 to 12544 per frame, calls of up -#: to 22144 rows. 62x90 is a grid this processor emits. -SERVING_FRAME = ((1, 62, 90),) - - -def _rotate_half(x: torch.Tensor, half: int) -> torch.Tensor: - """NeoX-style half rotation: `[x0, x1] -> [-x1, x0]` over the last dim.""" - return torch.cat((-x[..., half:], x[..., :half]), dim=-1) - - -def _spatial_merge_positions(grid: Sequence[Sequence[int]], merge: int) -> List[Tuple[int, int]]: - """Per-token `(row, col)` by four nested loops, in spatial-merge order.""" - positions: List[Tuple[int, int]] = [] - for frames, height, width in grid: - frame = [ - (block_row * merge + r, block_col * merge + c) - for block_row in range(height // merge) - for block_col in range(width // merge) - for r in range(merge) - for c in range(merge) - ] - positions.extend(frame * frames) - return positions - - -class VisionRope2d(OpWrapper): - """Rotate every head of `q` and `k` by its patch's 2-D axial angles; return - fresh `(q_out, k_out)` in the inputs' shapes and dtypes. - - rope_dim = head_dim // 2; axis_dim = rope_dim // 2 - inv_freq[j] = theta ** (-2j / rope_dim) j < axis_dim - angle[t] = [row(t) * inv_freq, col(t) * inv_freq] height half first - out[.., i] = x[i] cos a_i - x[i + rope_dim] sin a_i - out[.., i + rope_dim] = x[i + rope_dim] cos a_i + x[i] sin a_i - - NeoX half-split pairing `(i, i + rope_dim)`, not GPT-J `(2j, 2j+1)`. Tokens - are in spatial-merge order (whole `merge x merge` blocks row-major, - row-major inside a block) -- what the patch merger consumes; a raster walk - is silently wrong. A video's frames repeat the spatial grid: no temporal - axis. The angle table is built in fp64 and rounded once to fp32; the - rotation runs in fp32 and rounds once on the store. - - Fusion boundary: the angle table and the rotation. The qkv projection, the - unrotated `v`, the per-frame attention segmentation and the attention are - the caller's. - """ - - ARCHS = frozenset({Arch.SM_103}) - - # spec: grid, dtype; `heads` / `heads_k` / `head_dim` (default 16 / heads / - # 72); `theta` (default 1e4); `grid_dtype` / `grid_device` for `grid_thw` - # (default int32 on the host); `layout` "separate" (default) or - # "packed_qkv" -- q and k as the `split` views of one [T, 3*heads, D] - # projection, which is how the tower calls it. - CELLS: tuple[Cell, ...] = ( - Cell( - why="call site: q/k split views of packed qkv, 62x90", - spec=dict(grid=SERVING_FRAME, dtype=torch.bfloat16, layout="packed_qkv"), - ), - Cell( - why="call site: packed qkv views, video + image (two grid rows)", - spec=dict( - grid=((3, 8, 8), (1, 16, 24)), - dtype=torch.bfloat16, - layout="packed_qkv", - grid_dtype=torch.int64, - ), - ), - Cell( - why="call site: packed qkv views, video at the 32768-row ceiling", - spec=dict( - grid=((64, 16, 32),), - dtype=torch.bfloat16, - layout="packed_qkv", - grid_dtype=torch.int64, - ), - ), - ) - - note = """ - Metadata consumed: none. `grid_thw` is read with `.tolist()`: a device - tensor forces a sync (int32/int64, host/device all bitwise equal). `q` and - `k` need last-dim stride 1 only; they are not mutated. - - `head_dim` must be divisible by 4; otherwise the angle table's width - mismatches the head and the broadcast raises (loud, so not guarded; not - pinned by a test). - - Accuracy vs fp64 angles over the cells: bf16 <= 1.56e-2 (one ulp at - unit-normal operands). Measured, not certified: fp32 <= 4.77e-7, flat in the - grid side up to 6272 -- the fp64 table is why (an fp32 `inv_freq` times fp32 - position moved an fp32 output 1.475e-5 at a 64 side, outside the default - gate); wrong operators at 1024 bf16 tokens land 8.5-9.5 away: GPT-J pairing, - axes swapped, raster order, `inv_freq` over `head_dim`. - - Host-bound: ~353 us/call at 1024 and 4096 tokens (`.tolist()` + small - launches), x27 blocks per image. The table depends only on the grid, and the - result does not depend on the grid's placement (host or device) or integer - dtype (int32 or int64). - - Not certified: float16 or float32; mixed q/k dtypes or head counts; head - geometries other than 16 x 72; float8; `spatial_merge_size` != 2; `theta` - other than 1e4; grid sides past 90 or calls past 32768 rows; CUDA-graph - capture (it would bake in the grid). - """ - - def __call__( - self, - q: torch.Tensor, - k: torch.Tensor, - grid_thw: torch.Tensor, - spatial_merge_size: int, - theta: float = 10000.0, - ) -> Tuple[torch.Tensor, torch.Tensor]: - num_tokens, _, head_dim = q.shape - # Angles per token: half the head (the other half is each pair's - # partner); each axis owns half of those, height first. - rope_dim = head_dim // 2 - - grid: List[List[int]] = grid_thw.tolist() - - device = q.device - # fp64 then one rounding to fp32: the largest angle is (side - 1) - # radians, and fp32 throughout rounds twice there. - inv_freq = 1.0 / ( - theta ** (torch.arange(0, rope_dim, 2, dtype=torch.float64, device=device) / rope_dim) - ) - max_side = max(max(height, width) for _, height, width in grid) - freq_table = torch.outer( - torch.arange(max_side, dtype=torch.float64, device=device), inv_freq - ).to(torch.float32) - - merge = spatial_merge_size - intra = torch.arange(merge, device=device) - pos_ids = torch.empty((num_tokens, 2), dtype=torch.long, device=device) - offset = 0 - for frames, height, width in grid: - blocks_h, blocks_w = height // merge, width // merge - rows = ( - ( - torch.arange(blocks_h, device=device)[:, None, None, None] * merge - + intra[None, None, :, None] - ) - .expand(blocks_h, blocks_w, merge, merge) - .reshape(-1) - ) - cols = ( - ( - torch.arange(blocks_w, device=device)[None, :, None, None] * merge - + intra[None, None, None, :] - ) - .expand(blocks_h, blocks_w, merge, merge) - .reshape(-1) - ) - coords = torch.stack((rows, cols), dim=-1) - if frames > 1: - coords = coords.repeat(frames, 1) - pos_ids[offset : offset + coords.shape[0]] = coords - offset += coords.shape[0] - - # [tokens, 2, axis_dim] -> [tokens, rope_dim]: height angles, then width. - angles = freq_table[pos_ids].flatten(1) - duplicated = torch.cat((angles, angles), dim=-1) - cos = duplicated.cos().unsqueeze(-2) - sin = duplicated.sin().unsqueeze(-2) - - q_fp32, k_fp32 = q.float(), k.float() - q_out = (q_fp32 * cos + _rotate_half(q_fp32, rope_dim) * sin).to(q.dtype) - k_out = (k_fp32 * cos + _rotate_half(k_fp32, rope_dim) * sin).to(k.dtype) - return q_out, k_out - - def reference( - self, - q: torch.Tensor, - k: torch.Tensor, - grid_thw: torch.Tensor, - spatial_merge_size: int, - theta: float = 10000.0, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """Positions from nested loops, angles from Python-float `inv_freq` times - integer positions in fp64, rotation as a complex multiply in fp32.""" - rope_dim = q.shape[-1] // 2 - axis_dim = rope_dim // 2 - positions = _spatial_merge_positions(grid_thw.tolist(), spatial_merge_size) - inverse = torch.tensor( - [theta ** (-(2.0 * j) / rope_dim) for j in range(axis_dim)], dtype=torch.float64 - ) - pos = torch.tensor(positions, dtype=torch.float64).reshape(-1, 2) - angles = (pos[:, :, None] * inverse).flatten(1).to(device=q.device, dtype=torch.float32) - rotor = torch.polar(torch.ones_like(angles), angles).unsqueeze(-2) - - def rotate(x: torch.Tensor) -> torch.Tensor: - fp32 = x.float() - rotated = torch.complex(fp32[..., :rope_dim], fp32[..., rope_dim:]) * rotor - return torch.cat((rotated.real, rotated.imag), dim=-1).to(x.dtype) - - return rotate(q), rotate(k) - - def is_valid( - self, - q: torch.Tensor, - k: torch.Tensor, - grid_thw: torch.Tensor, - spatial_merge_size: int, - theta: float = 10000.0, - ) -> None: - # Each of these returns a full-size result without an error: a side - # that does not divide fills only the whole blocks' rows of the - # position table and leaves the rest uninitialized (four tokens of a - # 4x5 grid read back as unrotated); a grid describing fewer tokens than - # q carries leaves the tail uninitialized the same way; a one-token k - # broadcasts against every token's angles. - assert q.dim() == 3 and k.dim() == 3, ( - f"q and k must be [tokens, heads, head_dim]; got q {tuple(q.shape)}, k {tuple(k.shape)}" - ) - assert q.shape[0] == k.shape[0] and q.shape[-1] == k.shape[-1], ( - f"q and k must agree on tokens and head_dim; got q {tuple(q.shape)}, k {tuple(k.shape)}" - ) - total = 0 - for frames, height, width in grid_thw.tolist(): - assert height % spatial_merge_size == 0 and width % spatial_merge_size == 0, ( - f"grid sides must be divisible by spatial_merge_size={spatial_merge_size}; " - f"got height={height}, width={width}" - ) - total += frames * height * width - assert total == q.shape[0], f"grid_thw describes {total} tokens but q/k carry {q.shape[0]}" - - -vision_rope_2d = VisionRope2d() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/gated_ngram_injection.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/gated_ngram_injection.py deleted file mode 100644 index af6161baa6ee..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/gated_ngram_injection.py +++ /dev/null @@ -1,187 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""PLE gating: one shared n-gram value injected into every hyper-connection -stream, each scaled by its own key-query gate. - -Tier-3 composed from native torch (two upcasts, one reduction, the signed -square root, one broadcast multiply); it launches no TensorRT-LLM kernel. -Nothing registered under ``torch.ops.trtllm`` computes a dot-product gate or a -signed square root (``fused_sigmoid_mul_inplace`` takes a per-element gate and -could at most stand in for the last multiply). ``Engram.forward`` in -``tensorrt_llm/_torch/modules/engram/engram.py`` computes this gate as four -statements inside a forward that also runs the projections, both norms (plain -``x * w``, not this checkpoint's grouped gemma ``(1 + w)``) and a stateless -short conv -- not an atomic callable, and a different composite. The built-in -Qwen4-Exp implementation has its own, not wrapped here; the semantics follow -the HuggingFace ``Qwen4ExpTextPLELayer``. -""" - -from __future__ import annotations - -import math - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -#: This checkpoint's PLE geometry and its gate floor (`_PLE_CLAMP_MIN`). -CKPT_HC_COUNT = 4 -CKPT_HIDDEN = 2560 -CKPT_CLAMP_MIN = 1e-6 - - -def gate_ratio(out: torch.Tensor, ref: torch.Tensor) -> float: - """Worst element's distance from `ref` as a fraction of the test's 16-bit - band: `rtol = eps`, `atol = tiny * eps` (one subnormal spacing). Below 1 is - inside.""" - info = torch.finfo(ref.dtype) - ref64 = ref.double() - return ( - ((out.double() - ref64).abs() / (info.tiny * info.eps + info.eps * ref64.abs())) - .max() - .item() - ) - - -class GatedNgramInjection(OpWrapper): - """For token `t`, stream `s`, `D = value.shape[-1]`, with `key_normed` and - `query_normed` packed stream-major as `hc_count` blocks of `D`:: - - dot[t, s] = sum_i key_normed[t, s*D + i] * query_normed[t, s*D + i] - g = dot / sqrt(D) - signed = sign(g) * sqrt(max(|g|, clamp_min)) - out[t, s*D + i] = sigmoid(signed[t, s]) * value[t, i] - - Returns a fresh contiguous `[..., hc_count * D]` in **`value`'s dtype** - (a bf16 key/query with an fp32 value returns fp32). The value is shared: - every stream carries a differently scaled copy of the same vector. - `sign(0) = 0`, so an exactly zero dot gives `sigmoid(0) = 0.5` in every - dtype; `clamp_min` only replaces a small nonzero magnitude. The whole - chain runs in fp32 and rounds once. - - Fusion boundary: the per-stream dot, scale, floor, signed root, sigmoid - and broadcast multiply. The `key_proj` / `value_proj` GEMMs, both grouped - gemma norms producing the operands, the short conv over the result and - the add into the stream are the caller's. - """ - - ARCHS = frozenset({Arch.SM_103}) - - # spec: kind "ckpt" (default) -- operands with this checkpoint's norm-gamma - # and value statistics: leading, value_dtype (bf16), clamp_min. kind - # "single_channel" -- one nonzero channel per stream so each dot is an exact - # product under the test's control: dots, value_dtype, clamp_min. - CELLS: tuple[Cell, ...] = ( - Cell( - why="checkpoint 4x2560 bf16, 33 tokens / MTP verify 8 x 4 + 1", - spec=dict(leading=(33,)), - ), - Cell( - why="checkpoint 4x2560 bf16, prefill 8192 tokens, largest measured", - spec=dict(leading=(8192,)), - ), - Cell( - why="dots 0, +-5e-8 (inside the floor), 0.05, fp32 value", - spec=dict( - kind="single_channel", - dots=(0.0, 5e-8, -5e-8, 0.05), - value_dtype=torch.float32, - clamp_min=CKPT_CLAMP_MIN, - ), - ), - ) - - note = """ - Metadata consumed: none; every token row and stream independent. Call - site: once per forward, on the PLE layer, positional `clamp_min` = - `_PLE_CLAMP_MIN` (1e-6), `key_normed` / `query_normed` `[T, 10240]` and - `value` `[T, 2560]`, all bf16. - - Preconditions: `key_normed.shape[-1] == hc_count * value.shape[-1]` (the - unflatten raises otherwise, which also catches a wrong `hc_count`). Leading - dims must match exactly (`is_valid`); measured, a single query or value - row broadcasts silently over every key row. - - fp32 rather than the HF reference's bf16 chain: against the unrounded - float64 the fp32 chain sat at 0.498 of a one-epsilon bf16 band, the - literal bf16 chain at 4.693 (4.81 against `reference`), and 31.5% of the - output's bf16 elements differ between the two -- a bitwise comparison - against a bf16 implementation will not match. - - The floor is invisible in bf16: a sub-floor dot moves the sigmoid by ~5e-4 - of its value, measurable in fp32 and below half a bf16 ulp - (`clamp_min` 1e-6 and 0.0 bitwise identical with a bf16 `value`). - - Every near-miss tried (no root, unsigned root, no / misplaced / `1/D` - scale, sigmoid before the root, tanh, stream-minor packing) stays bounded - by `|value|`; they land 105x..1115x outside the band. - - CUDA-graph capture certified, bitwise equal to eager; deterministic. - Cost, not gated: 0.114 ms at 8 tokens, 0.128 at 512, 0.495 at 4096, - 0.956 at 8192 (launch-bound below ~512); transient allocation 3.56x the - operand bytes (the fp32 upcasts). - - Hard limits: not probed (largest operand 8192 x 10240 = 2^26.3 elements). - """ - - def __call__( - self, - key_normed: torch.Tensor, - query_normed: torch.Tensor, - value: torch.Tensor, - hc_count: int, - clamp_min: float = 1e-6, - ) -> torch.Tensor: - activation_dtype = value.dtype - hidden_size = value.shape[-1] - - keys = key_normed.unflatten(-1, (hc_count, hidden_size)).float() - queries = query_normed.unflatten(-1, (hc_count, hidden_size)).float() - - gate = (keys * queries).sum(dim=-1, keepdim=True) / math.sqrt(hidden_size) - gate = gate.abs().clamp_min(clamp_min).sqrt() * gate.sign() - - gated = torch.sigmoid(gate) * value.float().unsqueeze(-2) - return gated.flatten(-2).to(activation_dtype) - - def reference( - self, - key_normed: torch.Tensor, - query_normed: torch.Tensor, - value: torch.Tensor, - hc_count: int, - clamp_min: float = 1e-6, - ) -> torch.Tensor: - """float64 -- the entry already computes in fp32, so an fp32 reference - would be the same arithmetic -- rounded once to `value`'s dtype.""" - hidden = value.shape[-1] - keys = key_normed.double().unflatten(-1, (hc_count, hidden)) - queries = query_normed.double().unflatten(-1, (hc_count, hidden)) - gate = (keys * queries).sum(dim=-1, keepdim=True) / math.sqrt(hidden) - gate = gate.abs().clamp_min(clamp_min).sqrt() * gate.sign() - out = torch.sigmoid(gate) * value.double().unsqueeze(-2) - return out.flatten(-2).to(value.dtype) - - def is_valid( - self, - key_normed: torch.Tensor, - query_normed: torch.Tensor, - value: torch.Tensor, - hc_count: int, - clamp_min: float = 1e-6, - ) -> None: - # Both measured: a [1, width] query beside [64, width] keys, or a - # [1, D] value, returns 64 rows computed from the wrong row. - assert key_normed.shape == query_normed.shape, ( - f"key_normed {tuple(key_normed.shape)} and query_normed " - f"{tuple(query_normed.shape)} must have the same shape; a mismatch " - f"broadcasts silently instead of raising" - ) - assert value.shape[:-1] == key_normed.shape[:-1], ( - f"value {tuple(value.shape)} must carry the same leading dims as " - f"key_normed {tuple(key_normed.shape)}; a mismatch broadcasts silently " - f"instead of raising" - ) - - -gated_ngram_injection = GatedNgramInjection() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/dequantize_e4m3_per_tensor.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/dequantize_e4m3_per_tensor.py deleted file mode 100644 index 7f3922dea06a..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/dequantize_e4m3_per_tensor.py +++ /dev/null @@ -1,94 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Per-tensor FP8 dequantization: float8_e4m3fn -> bf16/fp16/fp32, multiplied by one scale.""" - -from __future__ import annotations - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — loads the extension registering torch.ops.tensorrt_llm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -#: The PLE n-gram table's per-tensor dequantization scale in this checkpoint. -PLE_SCALE = 0.00019931793212890625 - -#: Scale values the whole-domain cells sweep: the checkpoint's, unity, a power -#: of two, an arbitrary one, and the smallest e4m3 subnormal. -_DOMAIN_SCALES = (PLE_SCALE, 1.0, 0.5, 0.037, 2.0**-9) - - -def _ple(tokens: int) -> dict: - """`_ple_block`'s call: 16 gathered 160-wide rows per token, one bf16 scale.""" - return dict(shape=(tokens, 16, 160), dtype=torch.bfloat16, scales=(PLE_SCALE,), source="gather") - - -class DequantizeE4m3PerTensor(OpWrapper): - """`out = to(scales.dtype)(float(input) * float(scales[0]))`, one rounding. - - Wraps `torch.ops.tensorrt_llm.dequantize_e4m3_per_tensor` -- the - `tensorrt_llm` namespace; there is no `trtllm::` alias. - - It MULTIPLIES: pass a checkpoint's `weight_scale` straight, never - reciprocated (the inverse, `static_quantize_e4m3_per_tensor`, divides). The - output dtype follows `scales`, not `input`: a bf16 scale yields bf16 with no - cast. Shape is preserved exactly. `scales` has `input`'s rank with every dim - 1 (`[1, 1, 1]` for the PLE's rank-3 gather). - - Fusion boundary: scale-and-widen only. The table gather before it and the - flatten to `[T, 2560]` after it (a view) are the caller's. - """ - - ARCHS = frozenset({Arch.SM_103}) - - CELLS: tuple[Cell, ...] = ( - Cell(why="_ple_block decode: [1, 16, 160] gathered rows, bf16 scale", spec=_ple(1)), - Cell( - why="_ple_block prefill at max_num_tokens: [16384, 16, 160] through the grid-stride " - "loop", - spec=_ple(16384), - ), - Cell( - why="every e4m3 byte into bf16, five scale values, NaN codes included", - spec=dict( - shape=(1, 256), dtype=torch.bfloat16, scales=_DOMAIN_SCALES, source="all_codes" - ), - ), - ) - - note = """ - Metadata consumed: none; stateless. One kernel per call, - `scaleMatrix` at a fixed 1024-block - grid with a grid-stride loop; there is no vectorized variant, so alignment - is not an axis. - - Preconditions, enforced by the op (`symmetric_dequantize_per_tensor` in - thop/fp8Op.cpp): `input` float8_e4m3fn, CUDA, contiguous, rank >= 2, - non-empty (zero rows raise -- an empty batch must branch first; the test - holds this rejection); `scales` CUDA, bf16 / fp16 / fp32, of `input`'s rank - with every dim 1. - - NaN codes stay NaN; `448 * PLE_SCALE` reads back as bf16 0.08935546875 (the - all-codes cell). - - CUDA-graph capture is certified: replay picks up a new `input` and a new - `scales` value written in place (the scale is read from memory at replay). - - Hard limits - numel certified to 41,943,040 (16384 tokens); probed to 167,772,160; the - int32 boundary 2^31 not probed - - Not certified: fp16 and fp32 outputs (the target's scale is bf16), ranks - other than 2 and 3, degenerate scale values (0, negative, inf, NaN) and - overflow, the op's other loud rejections (dtype, layout, rank, device), - `TRTLLM_ENABLE_PDL=1`, `scales` on a different device than `input`. - """ - - def __call__(self, input: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: - return torch.ops.tensorrt_llm.dequantize_e4m3_per_tensor(input, scales) - - def reference(self, input: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: - """One fp32 multiply, one conversion to the scales' dtype.""" - return (input.float() * scales.float().flatten()[0]).to(scales.dtype) - - -dequantize_e4m3_per_tensor = DequantizeE4m3PerTensor() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_fwd.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_fwd.py deleted file mode 100644 index af83f235c2f7..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_fwd.py +++ /dev/null @@ -1,250 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Causal depthwise conv1d over a whole (ragged) batch of sequences, with each -sequence's convolution state read from and published to a paged state pool in -the same call.""" - -from __future__ import annotations - -from typing import Optional - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -#: This checkpoint's GDN mixer: `[q|k|v]` is 10240 wide, the conv is width 4, -#: so each pool slot holds a 3-column window. -CKPT_DIM = 10240 -#: The largest conv-pool slot count the cache manager was seen to allocate for -#: this checkpoint (10, 35, 366, 368, 396, 399, 432 and 461 across twelve -#: runtime-value dumps, while `max_num_sequences` never exceeds 32). The target -#: hands the pool over unsliced, so 461 tops the certified interval. -ENGINE_SLOTS = 461 - - -class CausalConv1dFwd(OpWrapper): - """Depthwise causal conv1d (+ optional SiLU) of every sequence in a batch, - against its carried window, publishing the new window -- one launch. - - For sequence `s` over tokens `[b, e)` and `slot = cache_indices[s]` (or `s`): - `xpad = cat([conv_states[slot] if has_initial_state[s] else 0, x[:, b:e]])`, - `out[c, b + t] = act(sum_k weight[c, k] * xpad[c, t + k] + bias[c])` with the - taps accumulated in fp32 and rounded once, and - `conv_states[slot] = xpad[:, L : L + W - 1]` -- the raw input tail, not the - activation, exact in every dtype. - - Fusion boundary: the state read, the convolution, the bias, the SiLU and the - state publish are inside. The projection that produces `x`, and slicing the - result into q/k/v, are the caller's. - - Writes the activation into `out`, or into `x` in place when `out is None`, - and returns that tensor. `conv_states` is advanced in place for every - non-padded sequence and is not returned. - - Two kernels sit behind the op, selected from `x`'s layout: unit stride on - the token axis is channel-major (in place allowed; `out`, if given, must - have `x`'s exact strides); unit stride on the channel axis with a token - stride > 1, `dim > 1` and at least one token is channel-last (`out` - required, channel-last, disjoint from `x`). The two agree bitwise. - """ - - ARCHS = frozenset({Arch.SM_103}) - - # spec keys: dim, lengths (varlen), pool_slots, slots (cache_indices), - # initial (has_initial_state); bf16, W 4, no bias, SiLU, pad_slot_id -1, - # channel-major in place throughout. - CELLS: tuple[Cell, ...] = ( - Cell( - why="checkpoint prefill at the engine's 461-slot pool: slots 460, 0, 230", - spec=dict( - dim=CKPT_DIM, - lengths=(37, 1, 128), - pool_slots=ENGINE_SLOTS, - slots=(ENGINE_SLOTS - 1, 0, ENGINE_SLOTS // 2), - initial=(True, False, True), - ), - ), - Cell( - why="pad_slot_id -1 (the engine's sentinel): sequence and slot skipped", - spec=dict( - dim=1024, - lengths=(16, 8, 16), - pool_slots=4, - slots=(0, -1, 2), - initial=(True, True, True), - ), - ), - Cell( - why="one 2048-token sequence: the vectorised load, 128-thread block", - spec=dict(dim=CKPT_DIM, lengths=(2048,), pool_slots=1, slots=(0,), initial=(True,)), - ), - ) - - note = """ - Metadata consumed: none. Batch layout (`query_start_loc`, int32), pool - (`conv_states`, `[slots, dim, W - 1]`), slot map (`cache_indices`, int32) - and carry flags (`has_initial_state`, bool) are explicit arguments. - - Call site: the GDN prefill half -- `x` the `[10240, num_prefill_tokens]` - channel-major output of `extract_transpose_prefill_slice`, W 4, no bias, - SiLU on, the layer's whole conv pool unsliced, `pad_slot_id` -1, `out` None. - The cells are that call (bf16, W 4, no bias, SiLU, channel-major in - place): a three-sequence prefill at the engine's 461-slot pool, which - tops every width the cache manager was seen to allocate (the 64-thread - block with scalar loads); one 2048-token sequence (the 128-thread block - with vectorised loads); and the pad sentinel on a small pool. The - 128-thread block's scalar-load form (seen at a single 2047-token - sequence) is not certified. - - Preconditions the op checks itself (all raise): dtypes in {bf16, fp16, - fp32}, `conv_states`/`bias` dtype, int32 `query_start_loc`/`cache_indices`, - bool `has_initial_state` of length batch, `out` dtype/shape/layout for the - selected kernel, channel-last in place, `x` rank. - - Not checked, not guarded (content, not metadata, or an illegal access - rather than a wrong number): a `cache_indices` value outside `[0, slots)`; - `has_initial_state[s]` True with `conv_states=None` (reads through null); - `conv_states`' shape (only its dtype is validated). Sequences shorter - than `W - 1` are legal; one (a single token without carried state) is - certified, one with carried state is not. A zero-length sequence, which - the target never schedules, is not certified. Unnamed slots are bitwise - untouched. - - Only the channel-major kernel (the target's layout) is certified; the - channel-last kernel and `TRTLLM_CAUSAL_CONV1D_SILU_MODE`, which selects - faster SiLU forms on it, are not. - - Hard limits - x token/channel strides uint32 stride fields boundary 2^32 elements - not probed. - """ - - def __call__( - self, - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor], - conv_states: Optional[torch.Tensor], - query_start_loc: Optional[torch.Tensor], - cache_indices: Optional[torch.Tensor], - has_initial_state: Optional[torch.Tensor], - silu_activation: bool, - pad_slot_id: int, - out: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - torch.ops.trtllm.causal_conv1d_fwd( - x, - weight, - bias, - conv_states, - query_start_loc, - cache_indices, - has_initial_state, - silu_activation, - pad_slot_id, - out, - ) - return x if out is None else out - - def reference( - self, - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor], - conv_states: Optional[torch.Tensor], - query_start_loc: Optional[torch.Tensor], - cache_indices: Optional[torch.Tensor], - has_initial_state: Optional[torch.Tensor], - silu_activation: bool, - pad_slot_id: int, - out: Optional[torch.Tensor] = None, - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - """`(activation, conv_states)` as the call leaves them; mutates nothing. - - The activation starts from `out` (or `x`) so a padded sequence's - columns keep whatever was there, as the op leaves them. - """ - width = weight.shape[-1] - act = (x if out is None else out).clone() - states = None if conv_states is None else conv_states.clone() - if query_start_loc is not None: - starts = query_start_loc.tolist() - spans = [(slice(None), slice(starts[s], starts[s + 1])) for s in range(len(starts) - 1)] - else: - spans = [(s,) for s in range(x.shape[0])] - slots = None if cache_indices is None else cache_indices.tolist() - carry = None if has_initial_state is None else has_initial_state.tolist() - for s, where in enumerate(spans): - slot = s if slots is None else slots[s] - if slots is not None and slot == pad_slot_id: - continue - tokens = x[where].float() # [dim, L] - length = tokens.shape[1] - if conv_states is not None and carry is not None and carry[s]: - prefix = conv_states[slot].float() - else: - prefix = tokens.new_zeros(tokens.shape[0], width - 1) - padded = torch.cat([prefix, tokens], dim=1) - acc = torch.zeros_like(tokens) - for k in range(width): - acc = acc + weight[:, k].float().unsqueeze(1) * padded[:, k : k + length] - if bias is not None: - acc = acc + bias.float().unsqueeze(1) - if silu_activation: - acc = acc * torch.sigmoid(acc) - if length: - act[where] = acc.to(act.dtype) - if states is not None: - states[slot] = padded[:, length : length + width - 1].to(states.dtype) - return act, states - - def is_valid( - self, - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor], - conv_states: Optional[torch.Tensor], - query_start_loc: Optional[torch.Tensor], - cache_indices: Optional[torch.Tensor], - has_initial_state: Optional[torch.Tensor], - silu_activation: bool, - pad_slot_id: int, - out: Optional[torch.Tensor] = None, - ) -> None: - # Widths outside 2..4 fall off the dispatcher with no launch at all: - # measured at W 1 and W 5, destination and pool bitwise unchanged. - assert weight.shape[-1] in (2, 3, 4), ( - f"width {weight.shape[-1]} has no kernel; only 2, 3 and 4 are dispatched, " - "and any other width is a silent no-op" - ) - # The dispatch reads weight (and bias) as x's type without converting: - # an fp32 weight with bf16 x moved the result by 11.5. - assert weight.dtype == x.dtype, ( - "weight must have x's dtype; the kernel reads it as x's type without converting" - ) - varlen = query_start_loc is not None - c_axis, l_axis = (0, 1) if varlen else (1, 2) - if x.dim() != (2 if varlen else 3): - return # the op rejects a wrong rank itself - channel_last = ( - x.stride(c_axis) == 1 - and x.stride(l_axis) > 1 - and x.shape[c_axis] > 1 - and x.shape[l_axis] > 0 - ) - if not channel_last: - # The channel-major kernel walks the token axis and the state's last - # axis contiguously, ignoring both strides: token stride 2 landed - # 16.3 off; a state last-dim stride 2 landed 7.9 off (5.4 in the - # published state). The channel-last kernel honours every stride. - assert x.shape[l_axis] == 0 or x.stride(l_axis) == 1, ( - "the channel-major kernel loads the token axis contiguously; " - "x must have stride 1 along it" - ) - assert conv_states is None or conv_states.stride(-1) == 1, ( - "the channel-major kernel ignores conv_states' last-dim stride; it must be 1" - ) - - -causal_conv1d_fwd = CausalConv1dFwd() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/chunk_gated_delta_rule.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/chunk_gated_delta_rule.py deleted file mode 100644 index 701fa42c82ed..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/chunk_gated_delta_rule.py +++ /dev/null @@ -1,268 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Chunked gated delta rule: the Gated DeltaNet prefill core. - -Wraps the in-tree module callable -``tensorrt_llm._torch.modules.fla.chunk.chunk_gated_delta_rule`` -- the -**vendored Triton** implementation (tier-2 sourcing: the gated delta rule has -no ``torch.ops.trtllm`` surface). Its interior is a fixed sequence of Triton -steps (gate cumsum, scaled K K^T, triangular solve, WY recompute, chunked state -recurrence, output epilogue) no caller invokes on its own: one fused op. - -The GDN mixer resolves between this and a FlashInfer-backed sibling of the same -name (``tensorrt_llm._torch.modules.fla.flashinfer_chunk``) through -``gdn_mixer._resolve_chunk_gated_delta_rule``, which picks FlashInfer on sm_103 -unless ``TLLM_USE_FLASHINFER_GDN_PREFILL=0``. This entry names the Triton one -directly, so neither the resolver nor the variable changes what it runs. -""" - -from __future__ import annotations - -from typing import Optional, Tuple - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.modules.fla.chunk import chunk_gated_delta_rule as _chunk_gated_delta_rule - -#: This checkpoint's Gated DeltaNet geometry at tp1. -NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM = 16, 48, 128, 128 -#: A recurrent-pool slot count of a full-size engine run. The kernels address a slot by -#: its index alone, so a pool of another size (a bf16 pool holds more slots in the same -#: memory) takes the same path. The target passes the layer's whole pool unsliced. -ENGINE_SLOTS = 461 - - -def _engine_slots(num_seqs: int) -> tuple[int, ...]: - """`num_seqs` distinct slots of the engine pool, slot 0 and the last live.""" - return ((0, ENGINE_SLOTS - 1) + tuple(i * 13 + 4 for i in range(num_seqs)))[:num_seqs] - - -class ChunkGatedDeltaRule(OpWrapper): - """Run the gated delta rule over whole sequences, returning `(o, final_state)`. - - Per token `t` and value head, `h` the `[V, K]` state, `q_h`/`k_h` the key - heads broadcast by `repeat_interleave` (value head `i` reads key head - `i // (H // Hg)`), fp32 accumulation: - - h *= exp(g[t]); u = (v[t] - h k_h[t]) * beta[t]; h += u k_h[t]^T - o[t] = h (scale * q_h[t]) - - computed 64 tokens at a time over a WY representation rounded to - `q.dtype` -- the same recurrence, a different sum. `scale=None` is - `K ** -0.5`; `g` is in log space. `use_qk_l2norm_in_kernel` normalizes q/k - as `x / sqrt(sum(x^2) + 1e-6)` first. - - With `cu_seqlens` the batch is one packed row (`q.shape[0] == 1`) split at - those offsets; without it each row is a sequence. Initial state: zeros - (`initial_state=None`); `initial_state[n]`, read-only; or, with - `initial_state_indices` and `inplace_indexed_state_update=True`, - `initial_state[indices[n]]` with the final state **written back into that - slot**. `output_final_state` only decides whether a fresh fp32 - `[N, H, V, K]` final state is also returned (on the indexed path, a copy - gathered after the write-back). - - Fusion boundary: gate cumsum, optional q/k norm, scale, recurrence, state - read and write, output epilogue. The conv, q/k/v/g/beta preparation, slot - allocation and the gated norm on `o` are the caller's. - - Returns `(o, final_state)` unchanged: `o` `[B, T, H, V]` in `q.dtype` - (`output` itself when given), `final_state` or `None`. The pool write is in - place and not returned. - """ - - ARCHS = frozenset({Arch.SM_103}) - - # spec: seqlens (packed with cu_seqlens); the initial state is the - # `ENGINE_SLOTS`-row pool at `indices`, written in place; pool_dtype (fp32 | - # bf16); output: None | "buffer" (a fresh `output=` tensor). - CELLS: tuple[Cell, ...] = ( - Cell( - why="fp32 pool (a verify step's context rows), 461 slots, output= buffer", - spec=dict(seqlens=(200, 37, 91), indices=_engine_slots(3), output="buffer"), - ), - Cell( - why="bf16 pool, 461-slot pool, seqlens [1024]", - spec=dict(seqlens=(1024,), indices=_engine_slots(1), pool_dtype=torch.bfloat16), - ), - Cell( - why="bf16 pool, mixed step: a context and four 1-token rows [200, 1, 1, 1, 1]", - spec=dict( - seqlens=(200, 1, 1, 1, 1), indices=_engine_slots(5), pool_dtype=torch.bfloat16 - ), - ), - ) - - note = """ - Metadata consumed: none. - - Call site (`_gdn_block`): `(q, k, v, g, beta, None, ssm_states, si_all, - True, False, qsl_long, False, None)` after `fused_gdn_post_conv` (q/k - already normalized) -- the layer's whole pool `[slots, 48, 128, 128]`, fp32 - or bf16 as `mamba_ssm_cache_dtype` allocated it, int32 slots, int64 - `cu_seqlens`; the MTP-verify context half passes the `si_p` / `qsl_long_p` - prefixes, and a mixed batch's one-token decode rows ride in the same call. - Certified at 461 slots: three ragged context rows over an fp32 pool into - an `output=` buffer, one 1024-token sequence over a bf16 pool, and a - mixed step (a context and four one-token rows) over a bf16 pool. - - Over a bf16 pool the initial state is read as fp32 and the final state is stored - rounded once to bf16, once per call and sequence; the accumulation is unchanged. - - The first indexed in-place call for an autotune key in a process is as correct as a - warm one: the state kernel's autotuner snapshots and restores the indexed state slots - its benchmark launches write (`chunk_delta_h.py`, `_save_autotune_state`). The key - includes the pool dtype, so a bf16 pool has its own. Certified in a fresh interpreter - at this geometry over a bf16 pool (the fp32 pool's key takes the same snapshot and - restore). The positional path never aliased its state. - - Loud, from the callable: fp32 or mixed-dtype q/k/v, a 2-D `beta`, - `q.shape[0] != 1` with `cu_seqlens`, a sequence-count mismatch of indices - or initial states, indices without `inplace_indexed_state_update`. - - Not checked: slots must not alias (not probed); q/k must be unit-norm, - by the caller or the in-kernel flag (raw q/k with the flag off return - finite values ~1e25). `head_first` is pinned False. - - Sequence lengths are unconstrained: certified at 1, 37, 91, 200 and 1024, - off and on the 64 boundary, mixed ragged. Memory: a `[B, NT, H, K, V]` - per-chunk state scratch in `q.dtype`, 24 MiB per 1024 tokens here. - Deterministic. - - Not certified: the FlashInfer sibling, fp16 (accepted), bf16 `g`/`beta`, an - fp16 pool, `K = 256` (accepted, unmeasured), empty `T` or zero-length spans, - and what the target never passes: `initial_state=None` or a positional - initial state, `output_final_state`, `use_qk_l2norm_in_kernel`, an explicit - `scale`, no `cu_seqlens`. - - Hard limits - K state kernel asserts K <= 256 (read, not probed) - """ - - def __call__( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: Optional[float] = None, - initial_state: Optional[torch.Tensor] = None, - initial_state_indices: Optional[torch.Tensor] = None, - inplace_indexed_state_update: bool = False, - output_final_state: bool = False, - cu_seqlens: Optional[torch.Tensor] = None, - use_qk_l2norm_in_kernel: bool = False, - output: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - return _chunk_gated_delta_rule( - q=q, - k=k, - v=v, - g=g, - beta=beta, - scale=scale, - initial_state=initial_state, - initial_state_indices=initial_state_indices, - inplace_indexed_state_update=inplace_indexed_state_update, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - head_first=False, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - output=output, - ) - - def reference( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: Optional[float] = None, - initial_state: Optional[torch.Tensor] = None, - initial_state_indices: Optional[torch.Tensor] = None, - inplace_indexed_state_update: bool = False, - output_final_state: bool = False, - cu_seqlens: Optional[torch.Tensor] = None, - use_qk_l2norm_in_kernel: bool = False, - output: Optional[torch.Tensor] = None, - ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: - """`(o, final_state, initial_state as the call leaves it)`. - - fp64 token-at-a-time recurrence, rounded once: `o` to `q.dtype`, the - states to fp32 / the pool's dtype. Mutates nothing. `output` is ignored - (it only chooses where `o` lands). - """ - dt = torch.float64 - batch, num_tokens, num_k_heads, head_k = q.shape - num_v_heads, head_v = v.shape[2], v.shape[3] - repeats = num_v_heads // num_k_heads - scale = head_k**-0.5 if scale is None else scale - if cu_seqlens is None: - spans = [(b, 0, num_tokens) for b in range(batch)] - else: - bounds = cu_seqlens.tolist() - spans = [(0, bounds[n], bounds[n + 1]) for n in range(len(bounds) - 1)] - slots = None if initial_state_indices is None else initial_state_indices.tolist() - out = torch.zeros(batch, num_tokens, num_v_heads, head_v, dtype=dt, device=q.device) - finals = [] - for n, (row, bos, eos) in enumerate(spans): - if initial_state is None: - state = torch.zeros(num_v_heads, head_v, head_k, dtype=dt, device=q.device) - else: - state = initial_state[n if slots is None else slots[n]].to(dt) - for t in range(bos, eos): - q_t, k_t = q[row, t].to(dt), k[row, t].to(dt) - if use_qk_l2norm_in_kernel: - q_t = q_t / torch.sqrt(q_t.pow(2).sum(-1, keepdim=True) + 1e-6) - k_t = k_t / torch.sqrt(k_t.pow(2).sum(-1, keepdim=True) + 1e-6) - q_t = q_t.repeat_interleave(repeats, dim=0) * scale - k_t = k_t.repeat_interleave(repeats, dim=0) - state = state * g[row, t].to(dt).exp()[:, None, None] - u = (v[row, t].to(dt) - (state * k_t[:, None, :]).sum(-1)) * beta[row, t].to(dt)[ - :, None - ] - state = state + u[:, :, None] * k_t[:, None, :] - out[row, t] = (state * q_t[:, None, :]).sum(-1) - finals.append(state) - final_state = torch.stack(finals).float() if output_final_state else None - pool = None if initial_state is None else initial_state.clone() - if slots is not None and inplace_indexed_state_update: - for slot, state in zip(slots, finals): - pool[slot] = state.to(pool.dtype) - return out.to(q.dtype), final_state, pool - - def is_valid( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: Optional[float] = None, - initial_state: Optional[torch.Tensor] = None, - initial_state_indices: Optional[torch.Tensor] = None, - inplace_indexed_state_update: bool = False, - output_final_state: bool = False, - cu_seqlens: Optional[torch.Tensor] = None, - use_qk_l2norm_in_kernel: bool = False, - output: Optional[torch.Tensor] = None, - ) -> None: - if output is None: - return - # Measured: a head-slice view (strides (1179648, 12288, 128, 1)) was - # accepted and came back 9.9e-2 off where correct is ~4e-4; an fp32 - # buffer was filled correctly but the returned bf16 tensor no longer - # aliased it. A token-window slice that reports contiguous is fine. - assert output.is_contiguous(), ( - f"output must be contiguous; got strides {tuple(output.stride())} -- the " - "epilogue addresses it through strides derived from its shape" - ) - assert output.dtype == q.dtype, ( - f"output must have q's dtype {q.dtype}; got {output.dtype} -- the result " - "is cast to q.dtype and stops aliasing the buffer" - ) - - -chunk_gated_delta_rule = ChunkGatedDeltaRule() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/dilated_short_conv_paged.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/dilated_short_conv_paged.py deleted file mode 100644 index 8ff846801270..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/dilated_short_conv_paged.py +++ /dev/null @@ -1,292 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Dilated depthwise short convolution over a packed, ragged serving batch, with -each sequence's convolution state read from and published to a paged pool. - -Composed from native torch ops (sourcing tier 3): no TensorRT-LLM kernel takes -a dilation. The ladder, as measured on this machine in this paged form: - -* tier 1 -- no registered op schema mentions `dilation`; the three convolutions - (`causal_conv1d_fwd`, `causal_conv1d_update`, `mamba_conv1d`) are the - non-dilated family. `causal_conv1d_fwd` fits the batch layout exactly, and a - zero-interleaved width-10 weight (taps at 0, 3, 6, 9) over the very - `[slots, channels, 9]` pool **returns normally and computes nothing** -- - activation and pool bitwise unchanged -- because its dispatcher has no branch - past width 4 and the width check sits on the update entry point only. - `causal_conv1d_update` raises on the same weight. -* tier 2 -- the vendored Triton `causal_conv1d_fn` does not compile at width - 10; engram's `ShortConv` is dilated but rectangular and stateless. -* tier 3 -- `F.conv1d(dilation=, groups=)` over a packed history plus index - arithmetic. This entry. - -Not taken: `d` phase-separated non-dilated convolutions through -`causal_conv1d_fwd` (three engine calls plus gathers -- not one entry), and a -rectangular `[num_seqs, channels, state_len + max_tokens]` buffer (5.4 GB at 32 -x 8192 tokens). A fused dilation-aware paged kernel is the faster path. -""" - -from __future__ import annotations - -from typing import Optional - -import torch -import torch.nn.functional as F - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -#: `hc_count * hidden_size = 4 * 2560` channels, kernel 4, dilation `ngram_size` -#: 3, hence the 9-column state the mamba cache manager allocates per slot. -CKPT_CHANNELS = 10240 -CKPT_DILATION = 3 -#: A PLE-conv-pool slot count of a full-size engine run. The kernel addresses a slot by its -#: index alone, so a pool of another size (the pool shares the recurrent pool's slot space, -#: and a bf16 recurrent pool holds more slots in the same memory) takes the same path. The -#: target hands the whole pool over unsliced. -ENGINE_SLOTS = 461 - - -def _scattered(num_seqs: int) -> tuple[int, ...]: - """`num_seqs` distinct slots over the whole engine pool, both ends in.""" - slots = ((0, ENGINE_SLOTS - 1) + tuple(i * 13 + 4 for i in range(num_seqs)))[:num_seqs] - assert len(set(slots)) == num_seqs and max(slots) < ENGINE_SLOTS - return slots - - -def _ckpt(why: str, lengths: tuple[int, ...], **spec) -> Cell: - return Cell( - why=why, - spec=dict(channels=CKPT_CHANNELS, lengths=lengths, pool_slots=ENGINE_SLOTS, **spec), - ) - - -class DilatedShortConvPaged(OpWrapper): - """For sequence `s` over columns `[b, e)` of the packed channel-major `x`, - `slot = cache_indices[s]` (or `s`), `state_len = (kernel_size - 1) * - dilation`: `window = [conv_states[slot] if has_initial_state[s] else 0, - x[:, b:e]]`, `out[c, b + t] = silu(round(sum_k weight[c, 0, k] * - window[c, t + k * dilation] + bias[c]))` -- taps in fp32, rounded by the - convolution, then the SiLU reads that back and rounds again -- and - `conv_states[slot] = window[:, L : L + state_len]`, the raw input tail. - - The highest tap index multiplies the token being emitted. A padded - sequence (`cache_indices[s] == pad_slot_id`) publishes nothing and still - gets output columns -- its tokens against a zero carry, defined and - meaningless. Two sequences naming one slot: the last live one's state - survives (defined by construction, not a race; uncertified, as the - scheduler never names one slot twice). - - Fusion boundary: state read, zero fill for a fresh sequence, the dilated - depthwise conv, bias, SiLU and the state publish are inside. The - pre-conv norm producing `x`, slot allocation and reset, and the residual - add are the caller's. - - Returns a fresh contiguous `[channels, num_tokens]` activation; `x`, - `weight` and `bias` are read-only. `conv_states` is advanced in place and - not returned. - """ - - ARCHS = frozenset({Arch.SM_103}) - - # spec keys: lengths, slots, initial (default all True). Every cell is the - # call site's own form: a bare .t() of the token-major [T, 10240] - # activation, kernel 4, dilation 3, bf16, no bias, SiLU, - # query_start_loc_long (int64), int32 slots, the whole 461-slot pool. - CELLS: tuple[Cell, ...] = ( - _ckpt("decode: 32 one-token rows", (1,) * 32, slots=_scattered(32)), - _ckpt( - "mixed: a 2047-token context + 3 one-token rows", (2047, 1, 1, 1), slots=_scattered(4) - ), - _ckpt( - "MTP verify: a fresh 37-token context + 8 rows of 3 drafts", - (37,) + (3,) * 8, - slots=_scattered(9), - initial=(False,) + (True,) * 8, - ), - ) - - note = """ - Metadata consumed: none -- the same explicit argument set as - `causal_conv1d_fwd` (packed batch, pool, int32/int64 slot map, bool carry). - - Call site: the PLE short conv on one layer -- `x` a bare `.t()` of the - `[T, 10240]` normed activation, weight `[10240, 1, 4]` bf16, no bias, - dilation 3 (state 9), SiLU on, the whole `[slots, 10240, 9]` pool, int32 - slots, `query_start_loc_long` (int64, rebuilt from `seq_lens` under MTP), - carry flags covering the whole batch, `pad_slot_id` -1. Certified in that - form at the 461-slot pool: 32 one-token decode rows, a 2047-token context - mixed with three one-token rows, and an MTP verify step (a fresh 37-token - context and 8 requests of 3 tokens). A padded row is not certified. - - Preconditions: `query_start_loc[-1] == x.shape[1]`; a short one indexes - past the internal history and fires a device-side assert that kills the - CUDA context (loud, not guarded). An out-of-range slot reaches the same - bounds check (reasoned, not probed). Pool, weight and bias share `x`'s - dtype (each mismatch raises); `weight` is 3-D depthwise (2-D raises). - Only the checkpoint's kernel size 4 and dilation 3 are certified; there is - no width dispatch. - - CUDA-graph capture is certified (at 32 requests of 8 tokens): no host - sync, all batch reads on device, a captured graph stays correct when the - batch tensors' values change. - Launch-bound on decode (8 rows: 0.461 ms eager vs 0.029 ms of conv; 0.094 - ms replayed from a graph). - - Not measured: pools above 461 slots (a bf16 recurrent state gives the engine - larger ones; a slot is addressed by its index alone), token counts above 8192, - `torch.compile`, a pool dtype unlike the model's. - """ - - def __call__( - self, - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor], - conv_states: torch.Tensor, - query_start_loc: torch.Tensor, - cache_indices: Optional[torch.Tensor], - has_initial_state: Optional[torch.Tensor], - dilation: int, - silu_activation: bool, - pad_slot_id: int, - ) -> torch.Tensor: - channels, num_tokens = x.shape - kernel_size = weight.shape[-1] - state_len = (kernel_size - 1) * dilation - num_slots = conv_states.shape[0] - num_seqs = query_start_loc.numel() - 1 - device = x.device - - # Device-resident arithmetic only: no `.item()`, no boolean indexing, - # no data-dependent shape -- which is what makes the call capturable. - starts = query_start_loc[:-1].to(torch.int64) - ends = query_start_loc[1:].to(torch.int64) - rows = torch.arange(num_seqs, device=device) - columns = torch.arange(num_tokens, device=device) - taps = torch.arange(state_len, device=device) - - # The first sequence whose end is strictly past the column; skips - # zero-token sequences by construction. - seq_of_column = torch.bucketize(columns, ends.contiguous(), right=True) - - if cache_indices is None: - slots = rows - live = torch.ones(num_seqs, dtype=torch.bool, device=device) - else: - live = cache_indices != pad_slot_id - slots = torch.where(live, cache_indices.to(torch.int64), torch.zeros_like(rows)) - carries = ( - torch.zeros_like(live) if has_initial_state is None else (has_initial_state & live) - ) - - # Each sequence's carry placed right ahead of its own tokens, so one - # convolution over the buffer never reads across a sequence boundary. - history = x.new_zeros((channels, num_tokens + num_seqs * state_len)) - history.index_copy_(1, columns + state_len * (seq_of_column + 1), x) - - previous = conv_states.index_select(0, slots) - history.index_copy_( - 1, - ((starts + rows * state_len).unsqueeze(1) + taps).reshape(-1), - torch.where(carries.reshape(num_seqs, 1, 1), previous, torch.zeros_like(previous)) - .permute(1, 0, 2) - .reshape(channels, -1), - ) - - convolved = F.conv1d( - history, - weight, - bias=bias, - stride=1, - padding=0, - dilation=dilation, - groups=channels, - ) - if silu_activation: - convolved = F.silu(convolved) - out = convolved.index_select(1, columns + state_len * seq_of_column) - - published = ( - history.index_select(1, ((ends + rows * state_len).unsqueeze(1) + taps).reshape(-1)) - .view(channels, num_seqs, state_len) - .permute(1, 0, 2) - ) - - # Skip padded rows without a host sync: every row writes, but writes - # the bytes of the highest-numbered live row owning that slot (or the - # slot's own contents when no live row does), so duplicate - # destinations carry identical bytes and landing order cannot matter. - writer = torch.full((num_slots,), -1, dtype=torch.int64, device=device) - writer.scatter_reduce_( - 0, - slots, - torch.where(live, rows, torch.full_like(rows, -1)), - reduce="amax", - include_self=True, - ) - owner = writer.index_select(0, slots) - conv_states.index_copy_( - 0, - slots, - torch.where( - (owner >= 0).reshape(num_seqs, 1, 1), - published.index_select(0, owner.clamp_min(0)), - previous, - ), - ) - return out - - def reference( - self, - x: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor], - conv_states: torch.Tensor, - query_start_loc: torch.Tensor, - cache_indices: Optional[torch.Tensor], - has_initial_state: Optional[torch.Tensor], - dilation: int, - silu_activation: bool, - pad_slot_id: int, - ) -> tuple[torch.Tensor, torch.Tensor]: - """`(out, conv_states)` as the call leaves them, one sequence at a time - as an explicit fp32 tap-sum (no convolution op); mutates nothing. - - Rounds where the entry's two calls round: once after the taps, again - after the SiLU. Publishes in sequence order, so a duplicated slot ends - with the last live sequence's state. - """ - channels = x.shape[0] - kernel_size = weight.shape[-1] - state_len = (kernel_size - 1) * dilation - starts = query_start_loc.tolist() - slots = None if cache_indices is None else cache_indices.tolist() - carry = None if has_initial_state is None else has_initial_state.tolist() - w = weight.float() - out = torch.zeros(channels, x.shape[1], dtype=x.dtype, device=x.device) - pool = conv_states.clone() - for s in range(len(starts) - 1): - begin, end = starts[s], starts[s + 1] - length = end - begin - slot = s if slots is None else slots[s] - live = slots is None or slot != pad_slot_id - if live and carry is not None and carry[s]: - prefix = conv_states[slot] - else: - prefix = x.new_zeros(channels, state_len) - window = torch.cat([prefix, x[:, begin:end]], dim=1) - xf = window.float() - acc = torch.zeros(channels, length, dtype=torch.float32, device=x.device) - for k in range(kernel_size): - acc = acc + w[:, 0, k].reshape(-1, 1) * xf[:, k * dilation : k * dilation + length] - if bias is not None: - acc = acc + bias.float().reshape(-1, 1) - y = acc.to(x.dtype) - if silu_activation: - y = F.silu(y.float()).to(x.dtype) - out[:, begin:end] = y - if live: - pool[slot] = window[:, length : length + state_len] - return out, pool - - -dilated_short_conv_paged = DilatedShortConvPaged() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py deleted file mode 100644 index 434e48ce4b55..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py +++ /dev/null @@ -1,129 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Channel-major extraction of a prefill row/column slice, the operand a causal -conv1d over a token-major projection needs. - -Wraps the in-tree Triton callable -``tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops.extract_transpose_prefill_slice`` -(one kernel launch, ``_extract_transpose_prefill_kernel``). A module-plane -callable rather than a ``torch.ops.trtllm.*`` op: there is no op boundary -around this kernel. -""" - -from __future__ import annotations - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops import ( - extract_transpose_prefill_slice as _extract_transpose_prefill_slice, -) - -#: The GDN mixer's packed `[q|k|v]` width for this checkpoint; the target hands -#: this op the `[q|k|v]` column view of its one input-projection output. -GDN_MIXED_QKV_WIDTH = 10240 -#: That output's width: `[q|k|v | z | a | b]` = 10240 + 6144 + 48 + 48. -GDN_IN_PROJ_WIDTH = 16480 - - -class ExtractTransposePrefillSlice(OpWrapper): - """`out[c, t] = src[t, start_col + c]` for `t < num_prefill_tokens`, - `c < width`: `src[:n, a:a+w].transpose(0, 1).contiguous()` in one pass. - - Fusion boundary: the row crop, the column crop and the transpose. No - convolution, no conv state, no channel split -- the caller hands the result - to `causal_conv1d_fwd`. Returns a fresh contiguous `[width, - num_prefill_tokens]` tensor in `src`'s dtype; `src` is read-only. - """ - - ARCHS = frozenset({Arch.SM_103}) - - # `num_rows`/`num_cols` build the [T, 16480] projection; `row_stop` and - # `col_slice` turn it into the view the call site passes (its q|k|v - # columns, and in a mixed step the prefill rows of those). - CELLS: tuple[Cell, ...] = ( - Cell( - why="mixed batch: prefill rows of the q|k|v view of a [512, 16480] in_proj", - spec=dict( - num_rows=512, - num_cols=GDN_IN_PROJ_WIDTH, - row_stop=400, - col_slice=(0, GDN_MIXED_QKV_WIDTH), - n=400, - start_col=0, - width=GDN_MIXED_QKV_WIDTH, - dtype=torch.bfloat16, - ), - ), - *( - Cell( - why=f"GDN prefill, the q|k|v column view of the in_proj, {n} tokens", - spec=dict( - num_rows=n, - num_cols=GDN_IN_PROJ_WIDTH, - col_slice=(0, GDN_MIXED_QKV_WIDTH), - n=n, - start_col=0, - width=GDN_MIXED_QKV_WIDTH, - dtype=torch.bfloat16, - ), - ) - for n in (33, 8192) - ), - ) - - note = """ - Metadata consumed: none. Stateless. - - Certified bitwise in bf16 at the call site's operand, the q|k|v column - view of the `[T, 16480]` input projection: 33 and 8192 tokens, and a mixed - step's 400 prefill rows of a 512-row projection. - - Preconditions: `src` 2-D on CUDA with `src.stride(1) == 1`; the row stride - is free. A non-unit column stride (`src.t()`, `src[:, ::2]`) is refused - loudly by the callable's own bare `assert`, before any launch. - - The kernel masks its loads only against the requested `num_prefill_tokens` - and `width`, never against `src`'s extent: 128 rows of a 64-row `src`, or - columns `[200, 328)` of a 256-column one, were accepted and returned a - full tensor read from out-of-bounds memory. `is_valid` bounds both. - - Hard limits - width * num_prefill_tokens int32 destination offset boundary 2^31 - probed at width 16384: n 131000 (2.146e9) bitwise exact; n 140000 - (2.294e9) raises `CUDA error: an illegal memory access` (loud, so not - a guard). The source offset is widened to int64 in-kernel: a source - offset of 2.62e9 was bitwise exact. At width 10240 the destination - boundary is n > 209715, within reach of a large `max_num_tokens`. - - Not certified: dtypes other than bf16, `start_col` other than 0, a row crop - shorter than `src`, empty extents. - """ - - def __call__( - self, src: torch.Tensor, num_prefill_tokens: int, start_col: int, width: int - ) -> torch.Tensor: - return _extract_transpose_prefill_slice(src, num_prefill_tokens, start_col, width) - - def reference( - self, src: torch.Tensor, num_prefill_tokens: int, start_col: int, width: int - ) -> torch.Tensor: - return src[:num_prefill_tokens, start_col : start_col + width].transpose(0, 1).contiguous() - - def is_valid( - self, src: torch.Tensor, num_prefill_tokens: int, start_col: int, width: int - ) -> None: - # A 3-D src with a unit dim-1 stride passes the callable's own assert - # and is read as 2-D through stride(0). - assert src.dim() == 2, f"src must be 2D, got {src.dim()}D" - assert 0 <= num_prefill_tokens <= src.shape[0], ( - f"num_prefill_tokens {num_prefill_tokens} outside src rows [0, {src.shape[0]}]; " - "the kernel reads past src rather than failing" - ) - assert 0 <= start_col and 0 <= width and start_col + width <= src.shape[1], ( - f"column range [{start_col}, {start_col + width}) outside src columns " - f"[0, {src.shape[1]}); the kernel reads past src rather than failing" - ) - - -extract_transpose_prefill_slice = ExtractTransposePrefillSlice() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_gating.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_gating.py deleted file mode 100644 index bc996ad6dd0d..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_gating.py +++ /dev/null @@ -1,153 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Gated-DeltaNet decay gate: g = -exp(A_log) * softplus(a + dt_bias), per token -and per value head, in one Triton kernel. - -Wraps the in-tree module callable -``tensorrt_llm._torch.modules.mamba.gdn_mixer.fused_gdn_gating`` (tier-2 -sourcing: no op boundary exists). Its body allocates the fp32 output and -launches ``fused_gdn_gating_kernel`` once. -""" - -from __future__ import annotations - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.modules.mamba.gdn_mixer import fused_gdn_gating as _fused_gdn_gating - -#: This checkpoint's GDN value-head count at tp1, and the `[b | a]` width the -#: stock mixer slices `a` out of. -NUM_V_HEADS = 48 -BA_WIDTH = 2 * NUM_V_HEADS - -#: Read off this checkpoint's 48 layers of `linear_attn.A_log` / -#: `linear_attn.dt_bias` (bf16, 48 each): exp(A_log) spans [3.839e-03, 1.394e+02], -#: and dt_bias reaching 19.25 against the default threshold of 20.0 makes the -#: softplus's linear branch a live path, not a corner case. -A_LOG_RANGE = (-5.5625, 4.9375) -DT_BIAS_RANGE = (-8.9375, 19.25) - -#: Relative error of the SASS approximate transcendentals (`ex2.approx.f32`, -#: `lg2.approx.f32`) Triton lowers `tl.exp` / `tl.log` to. -APPROX_ULP = 2.0**-22 - - -class FusedGdnGating(OpWrapper): - """`g[t, h] = -exp(A_log[h]) * softplus(a[t, h] + dt_bias[h])`, computed in - fp32 whatever the operand dtypes, into a fresh contiguous fp32 - `[batch, num_heads]` tensor. `batch, num_heads = a.shape`. - - `softplus(x) = log(1 + exp(beta * x)) / beta` while `beta * x <= threshold`, - else `x` -- the unstabilised form: below `x = -24 ln 2 = -16.6355` the fp32 - sum drops `exp(x)` and `g` is exactly 0 (absolute error <= 6e-8 times - exp(A_log), relative error 1.0). Harmless for a log-domain decay. - - Fusion boundary: bias add, softplus with its threshold branch, `exp(A_log)` - and the product are inside. The projection producing `a`, the - recurrence's other gate `beta = sigmoid(b)`, and the recurrence are not. - No input is modified. - """ - - ARCHS = frozenset({Arch.SM_103}) - - # spec: rows; layout "contiguous" | "target_rows" (the decode rows of a - # contiguous [T, 48] `a` projection, past T - rows prefill rows); fill - # (constant a). `a` bf16, softplus beta 1.0, threshold 20 throughout. - CELLS: tuple[Cell, ...] = ( - Cell(why="target verify call: 32 seqs x 2 tokens, contiguous a", spec=dict(rows=64)), - Cell( - why="target mixed verify: decode rows of a [T, 48] projection", - spec=dict(rows=16, layout="target_rows"), - ), - Cell( - why="a = 200: linear branch; exp(x) is evaluated, inf never leaks", - spec=dict(rows=4, fill=200.0), - ), - ) - - note = """ - Metadata consumed: none. Stateless and per (token, head). - - Call site (`_gdn_verify`): `(A_log, a_d, dt_bias, 1.0, 20.0)` with `a_d` - the `[n * p, 48]` bf16 decode rows of the `a` projection (all of it on a - pure verify step) and fp32 `A_log` / `dt_bias`; `g` is then reshaped to - `[n, p, 48]` for `fused_recurrent_gated_delta_rule_update`. Certified at - that call: 32 sequences of 2 tokens, the 16 decode rows of a mixed step's - `[T, 48]` projection, and `a = 200` (the softplus's linear branch, where - `exp` overflows). - - Raise loudly, so not guarded: a non-2-D `a` (`ValueError` unpacking - `a.shape`), CPU tensors (`ValueError: Pointer argument cannot be accessed - from Triton`). - - `a`'s row stride is honoured (a row is read at `row * a.stride(0)`; the - target passes contiguous rows); everything else is addressed without - strides. `is_valid` guards the three violations measured to return a finite wrong `g`: a stride-2 A_log - view moved g by 6.031e+01 and a stride-2 dt_bias by 3.832e-01; 40-element - A_log / dt_bias against 48 heads moved it by 1.009e+02 / 1.010e+02 (the - per-head load is masked against `a.shape[1]` only); an `a` with head - stride 2 by 1.153e+00. - - Deterministic: two launches agree bitwise. - - Hard limits - none probed. - """ - - def __call__( - self, - A_log: torch.Tensor, - a: torch.Tensor, - dt_bias: torch.Tensor, - beta: float = 1.0, - threshold: float = 20.0, - ) -> torch.Tensor: - return _fused_gdn_gating(A_log, a, dt_bias, beta, threshold) - - def reference( - self, - A_log: torch.Tensor, - a: torch.Tensor, - dt_bias: torch.Tensor, - beta: float = 1.0, - threshold: float = 20.0, - ) -> torch.Tensor: - """fp64 with the stabilised `log1p(exp(x))`, rounded once to fp32. - - The stabilised form makes the reference the accurate side; the - kernel's unstabilised tail differs inside the test's absolute floor. - """ - num_heads = a.shape[-1] - x = a.double() + dt_bias[:num_heads].double() - softplus = torch.where(beta * x <= threshold, torch.log1p(torch.exp(beta * x)) / beta, x) - return (-torch.exp(A_log[:num_heads].double()) * softplus).float() - - def is_valid( - self, - A_log: torch.Tensor, - a: torch.Tensor, - dt_bias: torch.Tensor, - beta: float = 1.0, - threshold: float = 20.0, - ) -> None: - # `shape[-1]` / `stride(-1)` so a non-2-D `a` still reaches the op's own - # ValueError rather than being preempted here. - num_heads = a.shape[-1] - assert a.stride(-1) == 1, ( - f"`a`'s head axis must be unit-stride; got strides {a.stride()} -- the kernel " - "reads a row as `row * a.stride(0) + head`, so another column stride returns " - "a finite wrong g" - ) - assert A_log.stride(0) == 1 and dt_bias.stride(0) == 1, ( - f"A_log and dt_bias must be contiguous 1-D; got strides {A_log.stride()} and " - f"{dt_bias.stride()} -- both are read as `ptr + head`" - ) - assert A_log.numel() >= num_heads and dt_bias.numel() >= num_heads, ( - f"A_log and dt_bias must hold at least a.shape[-1]={num_heads} elements; got " - f"{A_log.numel()} and {dt_bias.numel()} -- the per-head load is masked against " - "a.shape[-1] only, so a short tensor is read out of bounds" - ) - - -fused_gdn_gating = FusedGdnGating() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py deleted file mode 100644 index e34c3fdb77f1..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py +++ /dev/null @@ -1,232 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Gated-DeltaNet post-convolution preparation: the q/k/v split of the causal -conv output with Q/K L2 normalization, plus the decay gate and beta. - -Wraps the in-tree module callable -``tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops.fused_gdn_post_conv`` -(tier-2 sourcing: no op boundary exists). One Triton launch, -``_fused_gdn_post_conv_kernel``, after allocating its five outputs. -""" - -from __future__ import annotations - -from typing import Optional - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops import ( - fused_gdn_post_conv as _fused_gdn_post_conv, -) - -#: This checkpoint's linear-attention geometry at tp1. -NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM = 16, 128, 48, 128 -QKV_DIM = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240 - -#: Read off this checkpoint's `linear_attn.A_log` / `linear_attn.dt_bias` -#: (bf16, 48 per layer): A_log spans [-5.5625, 3.5625] and dt_bias -#: [-8.9375, 19.25]. dt_bias reaching 19.25 makes the softplus threshold branch -#: a live path, not a corner case. -A_LOG_RANGE = (-5.5625, 3.5625) -DT_BIAS_RANGE = (-8.9375, 19.25) -SOFTPLUS_THRESHOLD = 20.0 - - -class FusedGdnPostConv(OpWrapper): - """Split a causal-conv output into q/k/v and build the GDN gates. - - The token axis is the `num_prefill_tokens` columns of the dim-major - `prefill` `[>= qkv_dim, n_p]`, then the rows of the token-major `decode` - `[n_d, >= qkv_dim]` (`None` for a prefill-only batch), with - `qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim` packed - `[q | k | v]`. Per token and head, in fp32: - - q, k = x / sqrt(sum(x^2) + l2_norm_eps) (per head) - v = copied bit-exact - g = -exp(A_log) * softplus(a + dt_bias) (fp32 always) - beta = sigmoid(b) (fp32, rounded once to beta_dtype) - - `softplus(x)` is the unstabilised `log(1 + exp(x))` below - `softplus_threshold` and `x` above it. - - Fusion boundary: split, Q/K norm, `g` and `beta` are inside. The conv is - not (`prefill`/`decode` are already its output), and the recurrence that - consumes these must not normalize q/k again (`use_qk_l2norm_in_kernel=False`). - - Returns `(q, k, v, g, beta)`, five fresh contiguous tensors: - `[1, T, num_k_heads, head_k_dim]` x2 and `[1, T, num_v_heads, head_v_dim]` - in `prefill.dtype`, then `[1, T, num_v_heads]` fp32 and `beta_dtype`. - No input is modified. - """ - - ARCHS = frozenset({Arch.SM_103}) - - # spec: num_prefill, num_decode; dt_bias_fill. Every cell is the target's - # one [T, 16480] input projection: `decode` the q|k|v column view of its - # decode rows, `a`/`b` its column views, `prefill` the transposing - # extraction of its prefill rows. - CELLS: tuple[Cell, ...] = ( - *( - Cell( - why=f"target in_proj views, {p} prefill + {d} decode tokens", - spec=dict(num_prefill=p, num_decode=d), - ) - for p, d in ((200, 32), (2048, 0)) - ), - Cell( - why="every head above softplus_threshold: the linear branch", - spec=dict(num_prefill=64, num_decode=0, dt_bias_fill=SOFTPLUS_THRESHOLD + 8.0), - ), - ) - - note = """ - Metadata consumed: none. Stateless and per token; it never needs sequence - boundaries. - - Call sites (`_gdn_block`): prefill-only `(mq_t, None, a, b, ...)`, mixed - `(mq_p_t, mq_d, a, b, ...)` and the MTP-verify context half - `(mq_p_t, None, a_p, b_p, ...)`, all with `l2_norm_eps=1e-6`, - `softplus_threshold=20.0` and the fp32 `beta_dtype` default; `prefill` is - `extract_transpose_prefill_slice`'s contiguous output; `decode` and `a`/`b` - column views (row stride 16480) of the layer's one `[T, 16480]` - input-projection output. The stock mixer's verify call passes - `beta_dtype=b.dtype` (not certified here). Certified in that form at a - mixed step of 200 prefill and 32 decode tokens, a 2048-token prefill-only - step, and a 64-token step with every head past the softplus threshold. - - Preconditions: `prefill.shape[0]` and `decode.shape[1]` >= qkv_dim (extra - trailing features ignored; a short one is read out of bounds); `decode` - has `prefill`'s dtype. `a`/`b` may have extra columns (the in_proj - case) and are read through their own strides. `head_k_dim`/`head_v_dim` - are rounded up to a power of two and masked. - - The op raises on none of the metadata violations probed; `is_valid` - guards the four the kernel reads out of bounds or through the wrong - addresses, measured: stride-2 A_log moved g by 2.8e+02, stride-2 dt_bias by - 3.2e+02; A_log, dt_bias, a, b short by rows/elements moved g by 4.9e+12, - 2.9e+02, 3.2e+37 and made beta NaN. - - The softplus is unstabilised: for very negative `x` the fp32 `1 + exp(x)` - loses `exp(x)` and `g` rounds to exactly 0. Absolute error stays ~1 fp32 - ulp of the log argument times `exp(A_log)` (4.8e-07 measured over - x in [-9, 5]); relative error in that tail reaches 1.0. Inconsequential - for a log-domain decay. `exp(x)` is evaluated on the linear side too - (`tl.where` is a select); harmless below fp32 overflow at x ~ 88, which no - checkpoint `a + dt_bias` reaches. - - Not certified: fp16/fp32 activations, other head geometries, a bf16 - `beta_dtype`, bf16 `A_log`/`dt_bias`, a non-contiguous `prefill`, contiguous - `decode`/`a`/`b` (not views of the projection), single-token halves, an - empty batch. - - Hard limits - none probed. - """ - - def __call__( - self, - prefill: torch.Tensor, - decode: Optional[torch.Tensor], - a: torch.Tensor, - b: torch.Tensor, - A_log: torch.Tensor, - dt_bias: torch.Tensor, - num_k_heads: int, - head_k_dim: int, - num_v_heads: int, - head_v_dim: int, - l2_norm_eps: float = 1e-6, - softplus_threshold: float = 20.0, - beta_dtype: torch.dtype = torch.float32, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - return _fused_gdn_post_conv( - prefill, - decode, - a, - b, - A_log, - dt_bias, - num_k_heads, - head_k_dim, - num_v_heads, - head_v_dim, - l2_norm_eps, - softplus_threshold, - beta_dtype, - ) - - def reference( - self, - prefill: torch.Tensor, - decode: Optional[torch.Tensor], - a: torch.Tensor, - b: torch.Tensor, - A_log: torch.Tensor, - dt_bias: torch.Tensor, - num_k_heads: int, - head_k_dim: int, - num_v_heads: int, - head_v_dim: int, - l2_norm_eps: float = 1e-6, - softplus_threshold: float = 20.0, - beta_dtype: torch.dtype = torch.float32, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """q/k normalized in fp32 and rounded once; `g` in fp64 with the - stabilised `log1p(exp(x))` (the kernel's unstabilised form differs - inside the test's floor, not around it), rounded to fp32; `beta` is - the fp32 sigmoid rounded once to `beta_dtype`.""" - blocks = [prefill.t()[:, : 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim]] - if decode is not None: - blocks.append(decode[:, : blocks[0].shape[1]]) - mixed = torch.cat(blocks, dim=0) - n = mixed.shape[0] - kd = num_k_heads * head_k_dim - - def l2norm(x): - x = x.reshape(n, num_k_heads, head_k_dim).float() - return (x / torch.sqrt(x.pow(2).sum(-1, keepdim=True) + l2_norm_eps)).to(prefill.dtype) - - q, k = l2norm(mixed[:, :kd]), l2norm(mixed[:, kd : 2 * kd]) - v = mixed[:, 2 * kd :].reshape(n, num_v_heads, head_v_dim).contiguous() - x = a[:n, :num_v_heads].double() + dt_bias[:num_v_heads].double() - softplus = torch.where(x <= softplus_threshold, torch.log1p(torch.exp(x)), x) - g = (-torch.exp(A_log[:num_v_heads].double()) * softplus).float() - beta = torch.sigmoid(b[:n, :num_v_heads].float()).to(beta_dtype) - return tuple(t.unsqueeze(0) for t in (q, k, v, g, beta)) - - def is_valid( - self, - prefill: torch.Tensor, - decode: Optional[torch.Tensor], - a: torch.Tensor, - b: torch.Tensor, - A_log: torch.Tensor, - dt_bias: torch.Tensor, - num_k_heads: int, - head_k_dim: int, - num_v_heads: int, - head_v_dim: int, - l2_norm_eps: float = 1e-6, - softplus_threshold: float = 20.0, - beta_dtype: torch.dtype = torch.float32, - ) -> None: - # The kernel indexes A_log/dt_bias as `ptr + head` with no stride and - # no mask, and does not mask a/b's token axis against their extent. - num_tokens = prefill.shape[1] + (0 if decode is None else decode.shape[0]) - assert A_log.stride(0) == 1 and dt_bias.stride(0) == 1, ( - f"A_log and dt_bias must be contiguous 1-D; got strides {A_log.stride()} " - f"and {dt_bias.stride()} -- read as `ptr + head`, a strided view gives a " - "finite wrong g" - ) - assert A_log.numel() >= num_v_heads and dt_bias.numel() >= num_v_heads, ( - f"A_log and dt_bias must hold at least num_v_heads={num_v_heads} elements; " - f"got {A_log.numel()} and {dt_bias.numel()} -- the load is unmasked" - ) - assert a.shape[0] >= num_tokens and b.shape[0] >= num_tokens, ( - f"a and b must have at least num_tokens={num_tokens} rows (prefill + decode); " - f"got {a.shape[0]} and {b.shape[0]} -- the token axis is unmasked" - ) - - -fused_gdn_post_conv = FusedGdnPostConv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py deleted file mode 100644 index 072f874f6d4e..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py +++ /dev/null @@ -1,298 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Gated DeltaNet recurrent core over a slot pool, several tokens per sequence: -the speculative verify step's recurrence. - -Wraps the in-tree module callable -``tensorrt_llm._torch.modules.fla.fused_recurrent.fused_recurrent_gated_delta_rule_update`` -(tier-2 sourcing: the gated delta rule has no op boundary). One launch of the -vendored Triton kernel ``fused_recurrent_gated_delta_rule_update_fwd_kernel`` -- -``triton.heuristics`` + ``triton.jit``, no autotuner. -""" - -from __future__ import annotations - -from typing import Optional - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.modules.fla.fused_recurrent import ( - fused_recurrent_gated_delta_rule_update as _fused_recurrent_gated_delta_rule_update, -) - -#: This checkpoint's GDN geometry at tp1. -NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM = 16, 48, 128, 128 - -#: Rows of the engine's per-step buffer (`max_num_sequences`), as the -#: runtime-value dumps record `intermediate_ssm_states`: [32, p, 48, 128, 128]. -ENGINE_VERIFY_ROWS = 32 -#: Default pool and buffer rows of the test's builds. -SLOTS = 8 - - -def _verify(n, p, why, **spec): - return Cell(why=why, spec=dict(mode="verify", n=n, p=p, **spec)) - - -class FusedRecurrentGatedDeltaRuleUpdate(OpWrapper): - """Advance each sequence's `[K, V]` recurrent state by *all* its tokens, in - order. Per token and v-head `hv` (q/k head `hv // (HV // H)`): - - q, k = q / sqrt(q.q + 1e-6), k / sqrt(k.k + 1e-6) # if use_qk_l2norm_in_kernel - h *= exp(g); u = (v - h^T k) * beta; h += k u^T; o = h^T (scale * q) - - in fp32 throughout. `h` starts at `initial_state_source[idx]` (`[S, HV, V, - K]`, K innermost), or zero when `idx < 0` or there is no pool. `scale=None` - is `K ** -0.5`. `beta` with `v`'s rank is read per channel. - - Two independent sinks for `h`: - `disable_state_update=False` stores the final state back into slot `idx`; - `intermediate_states_buffer` (`[S', cache_steps, HV, V, K]`) receives the - state after every token at `[idx, step]`. The verify call is the second - alone: the pool must not move until the sampler has accepted a prefix. - - Fusion boundary: L2 norm, scale, the indexed gather and both scatters are - inside. The gates are *not*: `g = -exp(A_log) * softplus(a + dt_bias)` and - `beta = sigmoid(b)` arrive evaluated. The conv and q/k/v split are the - caller's too. - - Returns the `[N, P, HV, V]` output (`[1, T, HV, V]` with `cu_seqlens`) in - `output`'s dtype, else `q`'s; it *is* `output` when one is given. The pool - and the per-step buffer are written in place and not returned. - """ - - ARCHS = frozenset({Arch.SM_103}) - - # spec: mode ("verify": disable_state_update + buffer with cache_steps = p - # | "write_back"), n sequences x p tokens, pool_slots, cache_rows, slots - # (None: arange(n)), pool_dtype / cache_dtype, beta ("fp32" | "bf16": - # sigmoid of a bf16 projection, as the target passes it), qkv ("views" - # over a packed [T, 10240] row block | "contiguous", as the target passes - # them). - CELLS: tuple[Cell, ...] = tuple( - _verify( - n, - p, - f"{label}target verify: n={n} p={p}, n-row gather, 32-row buffer", - pool_slots=n, - cache_rows=ENGINE_VERIFY_ROWS, - beta="bf16", - qkv="contiguous", - pool_dtype=dtype, - cache_dtype=dtype, - ) - for dtype, label, n, p in ( - (torch.float32, "", 32, 2), - (torch.float32, "", 8, 4), - (torch.bfloat16, "bf16 state: ", 32, 2), - ) - ) - - note = """ - Metadata consumed: none; every piece of state is an argument. - - Call site (`_gdn_verify`, the legacy verify path of `TRTLLM_USE_GDN_REPLAY=0`, `p` up - to 4), all positional: `(q, k, v, g, beta, src, int_idx, None, True, None, True, False, - int_ssm, p, None)` -- q/k/v the contiguous `[n, p, 16|48, 128]` reshapes of - `pack_gdn_decode_qkv`'s outputs (so the in-kernel L2 norm is required), `g` fp32 from - `fused_gdn_gating`, `beta = sigmoid(b)` bf16, `src` an `[n, 48, 128, 128]` gather of - the live slots and `int_ssm` the layer's `[32, p, 48, 128, 128]` per-step buffer, both - at the pool's dtype (fp32 or bf16, as `mamba_ssm_cache_dtype` allocated it; the buffer's - states are rounded once to it), `int_idx = intermediate_state_indices[:n]` (int32 - arange). The buffer's rows are the batch rows `update_mamba_states` reads, not slot - ids, and its step axis is `tokens_per_gen_step`: `p` (as `cache_steps`) must equal it. - - Preconditions not checkable from metadata: `HV % H == 0` (q/k heads are - repeat_interleaved, not tiled); indices in `[0, S)` and `[0, S')` or - negative (the device check is `idx >= 0` only; out of range is an - out-of-bounds access, not probed); distinct indices when writing back -- - measured at two sequences on one slot: outputs exact, the slot a - per-(v-tile, head) mixture 0.83 .. 0.89 from a sequential application; in the - varlen form `cache_steps >= max(seq_len)`; q/k unit-norm-conditioned (pass - `use_qk_l2norm_in_kernel=True` or normalize upstream). `g` is log-space. - - Raise loudly, so not guarded: `cu_seqlens` with `q.shape[0] > 1` and an - index count other than `len(cu_seqlens) - 1` (ValueError); a `None` pool - with the write-back on (triton CompilationError). - - `is_valid` guards what is accepted and silently wrong, each measured: - a strided `output` (the one tensor not made contiguous) is written with - packed arithmetic -- 7.04 off and the neighbouring half overwritten; - `output` with `disable_output_calculation=True` is never written and the - return is a [1, 1, 1, 1] scratch; a strided pool or per-step buffer is - silently replaced by a contiguous copy, so the writes land in the copy and - the caller's tensor comes back bitwise unchanged; `cache_steps` below the - token count (or None, i.e. 0) overwrites the next row's steps -- 0.850 off - at cache_steps 2, P 4; `cache_steps` other than the buffer's step axis - misplaces every row past 0 against the buffer's own indexing. - - The first call in a fresh process is bitwise the second (measured on the - write-back sink; the kernel has no autotuner). - - Certified: the call site's form alone -- the per-step buffer sink with the pool - untouched, at 32 sequences of P = 2 and 8 of P = 4 over an fp32 pool and buffer, and - 32 of P = 2 over a bf16 one. Not certified: the write-back sink - (`disable_state_update=False`), alone or with the buffer; the varlen form - (`cu_seqlens`); negative, scattered or int64 indices; no pool; fp32 or per-channel - `beta`; q/k/v as views of the packed rows; an `output` buffer; - `disable_output_calculation=True`; `K` or `V` other than 128; fp16/fp32 q/k/v; P = 3 - and P outside 2..4; N above 32. - - Hard limits - N * HV gridDim.z boundary 65535 probed at N=1300 (62400): runs; - N=1400 (67200): RuntimeError "invalid argument", context survives - """ - - def __call__( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - initial_state_source: Optional[torch.Tensor], - initial_state_indices: torch.Tensor, - scale: Optional[float] = None, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: Optional[torch.Tensor] = None, - disable_state_update: bool = False, - disable_output_calculation: bool = False, - intermediate_states_buffer: Optional[torch.Tensor] = None, - cache_steps: Optional[int] = None, - output: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - return _fused_recurrent_gated_delta_rule_update( - q=q, - k=k, - v=v, - g=g, - beta=beta, - scale=scale, - initial_state_source=initial_state_source, - initial_state_indices=initial_state_indices, - cu_seqlens=cu_seqlens, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - disable_state_update=disable_state_update, - disable_output_calculation=disable_output_calculation, - intermediate_states_buffer=intermediate_states_buffer, - cache_steps=cache_steps, - output=output, - ) - - def reference( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - initial_state_source: Optional[torch.Tensor], - initial_state_indices: torch.Tensor, - scale: Optional[float] = None, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: Optional[torch.Tensor] = None, - disable_state_update: bool = False, - disable_output_calculation: bool = False, - intermediate_states_buffer: Optional[torch.Tensor] = None, - cache_steps: Optional[int] = None, - output: Optional[torch.Tensor] = None, - ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: - """`(out, initial_state_source, intermediate_states_buffer)` as the call - leaves them; mutates nothing. Token-sequential, fp32 throughout, each - store rounded once. `disable_output_calculation` and `cache_steps` are - ignored (the output is always computed; the buffer is indexed by its - own step axis).""" - H, K = k.shape[2], k.shape[3] - HV, V = v.shape[2], v.shape[3] - rep = HV // H - scale = K**-0.5 if scale is None else scale - if cu_seqlens is None: - spans = [(b, 0, q.shape[1]) for b in range(q.shape[0])] - else: - cu = cu_seqlens.tolist() - spans = [(0, cu[i], cu[i + 1]) for i in range(len(cu) - 1)] - out = torch.zeros(v.shape, dtype=torch.float32, device=v.device) - pool = None if initial_state_source is None else initial_state_source.clone() - cache = None if intermediate_states_buffer is None else intermediate_states_buffer.clone() - per_channel = beta.dim() == v.dim() - for i, (b, start, end) in enumerate(spans): - slot = int(initial_state_indices[i]) - if slot >= 0 and initial_state_source is not None: - h = initial_state_source[slot].float().transpose(1, 2) # [HV, K, V] - else: - h = torch.zeros(HV, K, V, dtype=torch.float32, device=v.device) - for step, t in enumerate(range(start, end)): - q_t = q[b, t].float().repeat_interleave(rep, dim=0) # [HV, K] - k_t = k[b, t].float().repeat_interleave(rep, dim=0) - if use_qk_l2norm_in_kernel: - q_t = q_t / torch.sqrt(q_t.square().sum(-1, keepdim=True) + 1e-6) - k_t = k_t / torch.sqrt(k_t.square().sum(-1, keepdim=True) + 1e-6) - q_t = q_t * scale - h = h * torch.exp(g[b, t].float())[:, None, None] - u = v[b, t].float() - (h * k_t[:, :, None]).sum(1) # [HV, V] - u = u * (beta[b, t].float() if per_channel else beta[b, t].float()[:, None]) - h = h + k_t[:, :, None] * u[:, None, :] - out[b, t] = (h * q_t[:, :, None]).sum(1) - if cache is not None and slot >= 0: - cache[slot, step] = h.transpose(1, 2).to(cache.dtype) - if pool is not None and not disable_state_update and slot >= 0: - pool[slot] = h.transpose(1, 2).to(pool.dtype) - return out.to(q.dtype if output is None else output.dtype), pool, cache - - def is_valid( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - initial_state_source: Optional[torch.Tensor], - initial_state_indices: torch.Tensor, - scale: Optional[float] = None, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: Optional[torch.Tensor] = None, - disable_state_update: bool = False, - disable_output_calculation: bool = False, - intermediate_states_buffer: Optional[torch.Tensor] = None, - cache_steps: Optional[int] = None, - output: Optional[torch.Tensor] = None, - ) -> None: - assert not (output is not None and disable_output_calculation), ( - "output and disable_output_calculation=True are mutually exclusive: with the " - "output calculation off the buffer is never written and the return value is a " - "[1, 1, 1, 1] scratch tensor" - ) - if output is not None: - assert output.is_contiguous(), ( - f"output must be contiguous; got strides {output.stride()} for shape " - f"{tuple(output.shape)} -- the kernel writes it with packed [B, T, HV, V] " - "arithmetic, misplacing values and clobbering whatever lies between" - ) - if initial_state_source is not None and not disable_state_update: - assert initial_state_source.is_contiguous(), ( - f"initial_state_source must be contiguous when disable_state_update is False; " - f"got strides {initial_state_source.stride()} -- a strided pool is copied " - "before the call and the final state is written into the copy" - ) - if intermediate_states_buffer is not None: - assert intermediate_states_buffer.is_contiguous(), ( - f"intermediate_states_buffer must be contiguous; got strides " - f"{intermediate_states_buffer.stride()} -- a strided buffer is copied before " - "the call and every per-step state is written into the copy" - ) - if cu_seqlens is None: - assert cache_steps is not None and cache_steps >= q.shape[1], ( - f"cache_steps must be at least the per-sequence token count when an " - f"intermediate_states_buffer is given; got cache_steps={cache_steps} for " - f"{q.shape[1]} tokens per sequence -- a smaller stride makes each " - "sequence's tail steps overwrite the next row's" - ) - assert cache_steps == intermediate_states_buffer.shape[1], ( - f"cache_steps={cache_steps} must equal intermediate_states_buffer's step axis " - f"({intermediate_states_buffer.shape[1]}) -- the kernel strides rows by " - "cache_steps, so any other value misplaces every row past 0" - ) - - -fused_recurrent_gated_delta_rule_update = FusedRecurrentGatedDeltaRuleUpdate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py deleted file mode 100644 index 19e5cfdc4d9a..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py +++ /dev/null @@ -1,297 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Gated DeltaNet recurrent core with the gating fused in, over a slot pool: -the GDN decode step. - -Wraps the in-tree module callable -``tensorrt_llm._torch.modules.fla.fused_sigmoid_gating_recurrent.fused_sigmoid_gating_delta_rule_update`` -(tier-2 sourcing: the gated delta rule has no ``torch.ops.trtllm`` surface). -""" - -from __future__ import annotations - -from typing import Optional - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.modules.fla.fused_sigmoid_gating_recurrent import ( - fused_sigmoid_gating_delta_rule_update as _fused_sigmoid_gating_delta_rule_update, -) - -#: A recurrent-pool slot count of a full-size engine run. The kernels address a slot by -#: its index alone, so a pool of another size (a bf16 pool holds more slots in the same -#: memory) takes the same path. The target passes the layer's whole pool unsliced. -ENGINE_SLOTS = 461 - - -def _scattered(n_seqs: int) -> tuple[int, ...]: - """`n_seqs` distinct engine-pool slots, both ends first, then a stride-13 walk over the - interior (13 is coprime to its `ENGINE_SLOTS - 2` slots).""" - rest = tuple((i * 13 + 4) % (ENGINE_SLOTS - 2) + 1 for i in range(max(n_seqs - 2, 0))) - return ((0, ENGINE_SLOTS - 1) + rest)[:n_seqs] - - -class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper): - """Advance the gated delta rule over every sequence, gates computed in-kernel. - - Sequence `n` (tokens `cu_seqlens[n]:cu_seqlens[n+1]`, or row `n` without - `cu_seqlens`), v-head `hv` reading k/q head `hv // (HV // H)`, state `h` - `[K, V]` loaded from `initial_state_source[initial_state_indices[n]]` - (stored `[V, K]`, K innermost), per token, all in fp32: - - g = -exp(A_log) * softplus_b(a + dt_bias) beta = sigmoid(b) - q, k = q / sqrt(q.q + 1e-6), k / sqrt(k.k + 1e-6) (only with use_qk_l2norm_in_kernel) - h *= exp(g); u = (v - h^T k) * beta; h += k u^T; o = h^T (scale * q) - - `softplus_b(x) = log(1 + exp(beta_s x)) / beta_s` below `softplus_threshold` - (in `beta_s x`), `x` above; `scale=None` is `K ** -0.5`, bitwise. The final - `h` is written back to the same slot. A negative index means no slot: zero - initial state, nothing written. `initial_state_source=None`: zero state, - nothing written. - - Fusion boundary: both gates, the optional q/k norm, scale, recurrence, - indexed gather and scatter are inside. The conv, the q/k/v split (plain - views) and the gated RMS norm on `o` are the caller's. - - Returns `o` `[B, T, HV, V]` in `output`'s dtype (it *is* `output` when - given), else `q`'s. The pool is advanced in place and not returned. - """ - - ARCHS = frozenset({Arch.SM_103}) - - # spec: seq_lens, slots (pool rows), indices (None -> arange), pool_dtype - # (fp32), ab ("separate": two contiguous [T, 48] | "in_proj": q/k/v/a/b all - # column views of one [T, 16480] projection, as the target's one input GEMM - # makes them), layout ("varlen": [1, T, ...] with int64 cu_seqlens | - # "rows": cu_seqlens None, one token per row, q/k/v [N, 1, heads, 128] - # views of the packed [N, 10240] block whose token stride is the row - # stride, and a/b [N, 1, 48] -- the target's bf16 call). In-kernel q/k - # norm on, no output buffer. - CELLS: tuple[Cell, ...] = ( - Cell( - why="target bf16 decode call, rows form N=128, 461-slot pool", - spec=dict( - seq_lens=(1,) * 128, - slots=ENGINE_SLOTS, - indices=_scattered(128), - ab="separate", - layout="rows", - pool_dtype=torch.bfloat16, - ), - ), - Cell( - why="bf16 rows form, pad rows (-1) between slot 0 and 460", - spec=dict( - seq_lens=(1, 1, 1, 1), - slots=ENGINE_SLOTS, - indices=(0, -1, ENGINE_SLOTS - 1, -1), - ab="separate", - layout="rows", - pool_dtype=torch.bfloat16, - ), - ), - Cell( - why="fp32 pool, in_proj column views, N=128: the block test's reference", - spec=dict(seq_lens=(1,) * 128, slots=128, ab="in_proj"), - ), - ) - - note = """ - Metadata consumed: none. - - Call site (`_gdn_block` decode over a bf16 pool), one token per sequence, - int32 slots, the layer's whole pool `[slots, 48, 128, 128]`: the rows form, - `cu_seqlens` None -- q/k/v `[T, 1, heads, 128]` split from the whole - `[T, 10240]` block reshaped to `[T, 1, 80, 128]`, so the token stride is the - row stride (`is_valid` refuses other strides), `a`/`b` `[T, 1, 48]`. The - cells drive it at 128 rows of a 461-slot pool, and with CUDA-graph pad rows - (-1) at 4 rows. One more cell certifies the packed form the GDN block test - uses as its fp32-pool reference: `(a_log, a, dt_bias, 1.0, 20.0, q, k, v, - b, ssm_states, si_all, None, True, qsl_long, None)`, q/k/v `[1, T, heads, - 128]` and `a`/`b` column views of one `[T, 16480]` input-projection output - (row stride 16480), int64 `cu_seqlens`, at 128 rows. - - Branch: the callable takes a FlashInfer CuTe-DSL kernel - (`gdn_decode_bf16_state.gated_delta_rule`) when the pool is bf16, - `cu_seqlens` is given and `T == N` (head dims 128, supported arch); - everything else runs the vendored Triton kernel, and only the Triton - branch is certified. The target's rows form (no `cu_seqlens`) and the reference's - packed form over an fp32 pool both stay on it. The bf16 call takes - the rows form because FlashInfer is not batch-invariant: at 48 value - heads it runs one kernel up to N = 10 and another from N = 11, and one - fixed row's single-step output and state differ between the two, so a - row's tokens would depend on its batch. The Triton rows are one program - per (row, head); the test shows a fixed row's output and state bitwise - equal across batch sizes 1..32. FlashInfer also differs from `reference` - in reading and writing a negative slot as slot 0 (Triton gives a zero - state and no write) and in its L2-norm epsilon. - - Preconditions not checked by anything: - - `HV % H == 0`. - - Pool trailing axes `(V, K)`, K innermost. The callable asserts the - strides `(V*K, K, 1)`, but at K == V a transposed `[S, HV, K, V]` pool - has the same strides: accepted, measured 2.15e-2 off (vs 3.8e-6 - correct). The slot stride is free. - - Indices distinct (a shared slot races on write-back: outputs fine, the - slot holds one writer's state) and in `[0, S)` or negative - (out-of-range is an unchecked device access -- `tl.device_assert` is - compiled out -- not probed). - - q/k unit-norm conditioned: in-kernel norm on, or normalized by the - caller. Raw multi-token q/k with the norm off drifted the fp32 state - 1.7e-2 relative from an fp32 sequential reference: domain, not defect. - - `output` contiguous, `[B, T, HV, V]`: guarded by `is_valid`. Its dtype - wins (an fp32 `output` with bf16 q returns fp32). - - Deterministic: repeated calls are bitwise identical. Grid - `(1, ceil(V/32), min(N*HV, 65535))`, grid-striding past the cap (read, - not measured: N > 1365 at this geometry). - - Not certified: the FlashInfer branch, K or V other than 128, fp16/fp32 - operands, an fp16 pool, `softplus_beta`/`softplus_threshold` other than - 1.0/20.0, `B > 1` with `cu_seqlens`; the rows form over an fp32 pool or above - 128 rows; CUDA-graph replay against the eager call; the packed form over a - bf16 pool, with more than one token per sequence, with an `output` buffer, - with `l2norm` off, without a pool or on a layer view of one. - - Hard limits - none probed. - """ - - def __call__( - self, - A_log: torch.Tensor, - a: torch.Tensor, - dt_bias: torch.Tensor, - softplus_beta: float, - softplus_threshold: float, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - b: torch.Tensor, - initial_state_source: Optional[torch.Tensor], - initial_state_indices: torch.Tensor, - scale: Optional[float] = None, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: Optional[torch.Tensor] = None, - output: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - return _fused_sigmoid_gating_delta_rule_update( - A_log=A_log, - a=a, - dt_bias=dt_bias, - softplus_beta=softplus_beta, - softplus_threshold=softplus_threshold, - q=q, - k=k, - v=v, - b=b, - initial_state_source=initial_state_source, - initial_state_indices=initial_state_indices, - scale=scale, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - cu_seqlens=cu_seqlens, - output=output, - ) - - def reference( - self, - A_log: torch.Tensor, - a: torch.Tensor, - dt_bias: torch.Tensor, - softplus_beta: float, - softplus_threshold: float, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - b: torch.Tensor, - initial_state_source: Optional[torch.Tensor], - initial_state_indices: torch.Tensor, - scale: Optional[float] = None, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: Optional[torch.Tensor] = None, - output: Optional[torch.Tensor] = None, - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - """`(o, pool_after)`: sequential over tokens, fp32 throughout like the - kernel, `o` rounded once to the result dtype and the state once to the - pool's. `output` contributes only its dtype; nothing is mutated.""" - batch, total, num_k_heads, k_dim = q.shape - hv, v_dim = v.shape[2], v.shape[3] - repeat = hv // num_k_heads - scale = k_dim**-0.5 if scale is None else scale - if cu_seqlens is None: - spans = [(n, 0, total) for n in range(batch)] - else: - bounds = cu_seqlens.tolist() - spans = [(0, bounds[n], bounds[n + 1]) for n in range(len(bounds) - 1)] - a3, b3 = (a, b) if a.dim() == 3 else (a.unsqueeze(0), b.unsqueeze(0)) - out = torch.zeros(batch, total, hv, v_dim, device=q.device, dtype=torch.float32) - pool = None if initial_state_source is None else initial_state_source.clone() - neg_exp_a = -torch.exp(A_log.float()) - slots = initial_state_indices.tolist() - for n, (row, bos, eos) in enumerate(spans): - slot = slots[n] - if pool is not None and slot >= 0: - state = initial_state_source[slot].float().transpose(1, 2) - else: - state = torch.zeros(hv, k_dim, v_dim, device=q.device, dtype=torch.float32) - for t in range(bos, eos): - q_t = q[row, t].float().repeat_interleave(repeat, dim=0) - k_t = k[row, t].float().repeat_interleave(repeat, dim=0) - x = a3[row, t].float() + dt_bias.float() - bx = softplus_beta * x - softplus = torch.where( - bx <= softplus_threshold, torch.log(1.0 + torch.exp(bx)) / softplus_beta, x - ) - beta = torch.sigmoid(b3[row, t].float()) - if use_qk_l2norm_in_kernel: - q_t = q_t / torch.sqrt(q_t.pow(2).sum(-1, keepdim=True) + 1e-6) - k_t = k_t / torch.sqrt(k_t.pow(2).sum(-1, keepdim=True) + 1e-6) - state = state * torch.exp(neg_exp_a * softplus)[:, None, None] - u = (v[row, t].float() - (state * k_t[:, :, None]).sum(1)) * beta[:, None] - state = state + k_t[:, :, None] * u[:, None, :] - out[row, t] = (state * (q_t * scale)[:, :, None]).sum(1) - if pool is not None and slot >= 0: - pool[slot] = state.transpose(1, 2).to(pool.dtype) - return out.to(q.dtype if output is None else output.dtype), pool - - def is_valid( - self, - A_log: torch.Tensor, - a: torch.Tensor, - dt_bias: torch.Tensor, - softplus_beta: float, - softplus_threshold: float, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - b: torch.Tensor, - initial_state_source: Optional[torch.Tensor], - initial_state_indices: torch.Tensor, - scale: Optional[float] = None, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: Optional[torch.Tensor] = None, - output: Optional[torch.Tensor] = None, - ) -> None: - # Measured on the Triton branch: a last-dim slice of a doubly wide - # buffer and a token-strided view were both accepted, came back 5.3e-2 - # and 4.8e-2 off at an output absmax of 4.1e-2, and overwrote elements - # outside the view. - assert output is None or output.is_contiguous(), ( - f"output must be contiguous; got strides {output.stride()} for shape " - f"{tuple(output.shape)} -- the kernel writes it with packed " - "[NK, B, T, HV, V] arithmetic, misplacing values and clobbering the gaps" - ) - if cu_seqlens is None: - # q/k/v reshaped to [N, 1, heads, 128] from column slices of the packed - # [N, 10240] block carry stride(1) = heads * 128, not the row stride: they are - # accepted and read other rows' columns. - for name, t in (("q", q), ("k", k), ("v", v), ("a", a), ("b", b)): - assert t.dim() < 3 or t.stride(0) == t.shape[1] * t.stride(1), ( - f"{name} {tuple(t.shape)} has strides {t.stride()}: without " - "cu_seqlens the kernel finds row n's tokens at n * T * stride(1), " - "so stride(0) must equal T * stride(1)" - ) - - -fused_sigmoid_gating_delta_rule_update = FusedSigmoidGatingDeltaRuleUpdate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py deleted file mode 100644 index 4110469820d5..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_in_proj_conv1d_pdl.py +++ /dev/null @@ -1,432 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The Gated-DeltaNet input projection at decode rows with the causal conv1d step -as its epilogue, one PDL launch. - -Tier 5 (CuTe DSL written for this target). The decode GDN block projects into one -`[T, q|k|v | z | a | b]` buffer and then runs a separate depthwise conv kernel -over the q|k|v columns. That conv touches 3 cached values and 4 taps per column: -far too little work for a launch of its own. Here every CTA computes `OUTS` -output columns over the whole K, and for a column below `conv_dim` it finishes -the conv step on the projected value before storing it. - -Order of work inside a CTA: - -1. its weight rows into registers, the conv taps of its columns, the row's slot - and the cached window of each of its conv columns -- static, host-filled, or - last written by this kernel one step ago -- so before `griddepcontrol.wait` - (PDL), while the kernel ahead of this one is still running; -2. `griddepcontrol.wait`; -3. the activation, the products, a warp then a cross-warp reduction; -4. per column: the projection rounded to bf16 (the value the unfused path - stores), and for a conv column of a non-pad row the conv in fp32 over - `[window, projected]`, SiLU, rounded once; the window advances to - `[window[1:], projected]`. - -Element `k` of a row is held by thread `(k // VEC) % THREADS` as its vector -`(k // VEC) // THREADS`, so each k-tile of loads is contiguous across a warp. -""" - -from __future__ import annotations - -import cuda.bindings.driver as cuda_drv -import cutlass -import cutlass.cute as cute -import torch -from cutlass import const_expr -from cutlass._mlir.dialects import llvm -from cutlass.cute import experimental as cute_ext -from cutlass.cute.runtime import from_dlpack -from cutlass.cutlass_dsl import dsl_user_op - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( - Arch, - Cell, - OpWrapper, - true_fp32_matmul, -) -from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl - -VEC = 8 -#: Before its PDL wait a CTA also has the weight rows of CTAs `bidx + j * L2_DIST` -#: (j = 1 .. L2_AHEAD) brought into L2: the kernel ahead of this one is latency-bound and -#: only a few of these CTAs fit beside it, so the later waves' weights would otherwise -#: start streaming only after it. -L2_AHEAD = 2 -L2_DIST = 1216 - - -@dsl_user_op -def _bulk_prefetch_l2(src, nbytes, *, loc=None, ip=None): - llvm.inline_asm( - None, - [ - cutlass.Int64(src.toint(loc=loc, ip=ip)).ir_value(loc=loc, ip=ip), - cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), - ], - "cp.async.bulk.prefetch.L2.global [$0], $1;", - "l,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -#: Largest decode row count this entry serves (every row re-reads the weight from registers). -MAX_ROWS = 2 -#: (threads, outputs per CTA) per K; swept on GB300 at the checkpoint's 16480 x 2560. -TILING = {2560: (64, 2)} -WIDTH = 4 # conv taps -_COMPILE_OPTIONS = "--ptxas-options -maxrregcount=64" - - -class _InProjConvKernel: - def __init__(self, *, rows, k, threads, outs, conv_dim, pad_slot, pdl): - assert k % (threads * VEC) == 0 and threads % 32 == 0 - self.rows, self.threads, self.outs = rows, threads, outs - self.k_tiles = k // (threads * VEC) - self.warps = threads // 32 - self.conv_dim, self.pad_slot, self.pdl = conv_dim, pad_slot, pdl - - @cute.jit - def __call__(self, mA, mW, mOut, mState, mConvW, mSlot, stream: cuda_drv.CUstream): - n = cute.size(mW, mode=[0]) - load_a = cute.make_copy_atom( - cute.nvgpu.CopyG2ROp(), - cutlass.BFloat16, - num_bits_per_copy=16 * VEC, - load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, - ) - load_w = cute.make_copy_atom( - cute.nvgpu.CopyG2ROp(), - cutlass.BFloat16, - num_bits_per_copy=16 * VEC, - load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, - ) - self.kernel(mA, mW, mOut, mState, mConvW, mSlot, load_a, load_w).launch( - grid=[cute.ceil_div(n, self.outs), 1, 1], - block=[self.threads, 1, 1], - smem=self.rows * self.outs * self.warps * 4, - stream=stream, - use_pdl=self.pdl, - min_blocks_per_mp=1, - ) - - @cute.kernel - def kernel( - self, - mA, - mW, - mOut, - mState, - mConvW, - mSlot, - load_a: cute.CopyAtom, - load_w: cute.CopyAtom, - ): - tidx, _, _ = cute.arch.thread_idx() - bidx, _, _ = cute.arch.block_idx() - warp = cute.arch.warp_idx() - rows: cutlass.Constexpr = self.rows - outs: cutlass.Constexpr = self.outs - threads: cutlass.Constexpr = self.threads - k_tiles: cutlass.Constexpr = self.k_tiles - warps: cutlass.Constexpr = self.warps - n0 = bidx * outs - my_col = n0 + tidx - - # (row, (VEC, (threads, k_tiles))): thread t's vectors of a row, k-tile by k-tile. - w_vec = cute.logical_divide(cute.logical_divide(mW, (None, VEC)), (None, (None, threads))) - a_vec = cute.logical_divide(cute.logical_divide(mA, (None, VEC)), (None, (None, threads))) - - # 1. Static operands, before the wait. - w_reg = cute.make_rmem_tensor( - cute.make_layout((outs, k_tiles, VEC), stride=(k_tiles * VEC, VEC, 1)), cutlass.BFloat16 - ) - for j in cutlass.range_constexpr(outs): - w_row = w_vec[n0 + j, (None, (tidx, None))] - for kt in cutlass.range_constexpr(k_tiles): - cute.copy(load_w, w_row[None, kt], w_reg[j, kt, None]) - tap = cute.make_rmem_tensor(cute.make_layout((WIDTH,), stride=(1,)), cutlass.BFloat16) - win = cute.make_rmem_tensor( - cute.make_layout((rows, WIDTH - 1), stride=(WIDTH - 1, 1)), cutlass.BFloat16 - ) - if tidx < outs: - if my_col < self.conv_dim: - for j4 in cutlass.range_constexpr(WIDTH): - tap[j4] = mConvW[my_col, j4] - for r in cutlass.range_constexpr(rows): - slot = mSlot[r] - if slot != self.pad_slot: - for j3 in cutlass.range_constexpr(WIDTH - 1): - win[r, j3] = mState[slot, my_col, j3] - - if const_expr(self.pdl): - n_rows = cute.size(mW, mode=[0]) - k_len = cute.size(mW, mode=[1]) - if tidx < L2_AHEAD: - tgt = bidx + (tidx + 1) * L2_DIST - if (tgt + 1) * outs <= n_rows: - _bulk_prefetch_l2(mW.iterator + tgt * outs * k_len, outs * k_len * 2) - - # 2. The producer's output only after the wait. - if const_expr(self.pdl): - cute.arch.griddepcontrol_launch_dependents() - cute.arch.griddepcontrol_wait() - - # 3. Products and reductions. - acc = cute.make_rmem_tensor( - cute.make_layout((rows, outs), stride=(outs, 1)), cutlass.Float32 - ) - acc.fill(0.0) - a_reg = cute.make_rmem_tensor( - cute.make_layout((k_tiles, VEC), stride=(VEC, 1)), cutlass.BFloat16 - ) - for r in cutlass.range_constexpr(rows): - a_row = a_vec[r, (None, (tidx, None))] - for kt in cutlass.range_constexpr(k_tiles): - cute.copy(load_a, a_row[None, kt], a_reg[kt, None]) - for kt in cutlass.range_constexpr(k_tiles): - for v in cutlass.range_constexpr(VEC): - av = a_reg[kt, v].to(cutlass.Float32) - for j in cutlass.range_constexpr(outs): - acc[r, j] = acc[r, j] + av * w_reg[j, kt, v].to(cutlass.Float32) - for r in cutlass.range_constexpr(rows): - for j in cutlass.range_constexpr(outs): - acc[r, j] = cute.arch.warp_reduction_sum(acc[r, j]) - smem = cutlass.utils.SmemAllocator() - part = smem.allocate_tensor( - cutlass.Float32, - cute.make_layout((rows, outs, warps), stride=(outs * warps, warps, 1)), - byte_alignment=16, - ) - with cute.arch.elect_one(): - for r in cutlass.range_constexpr(rows): - for j in cutlass.range_constexpr(outs): - part[r, j, warp] = acc[r, j] - cute.arch.sync_threads() - - # 4. Thread j < outs finishes column n0 + j. - if tidx < outs: - for r in cutlass.range_constexpr(rows): - total = cutlass.Float32(0.0) - for w in cutlass.range_constexpr(warps): - total = total + part[r, tidx, w] - proj = total.to(cutlass.BFloat16) - res = proj - if my_col < self.conv_dim: - slot = mSlot[r] - if slot != self.pad_slot: - c = ( - tap[0].to(cutlass.Float32) * win[r, 0].to(cutlass.Float32) - + tap[1].to(cutlass.Float32) * win[r, 1].to(cutlass.Float32) - + tap[2].to(cutlass.Float32) * win[r, 2].to(cutlass.Float32) - + tap[3].to(cutlass.Float32) * proj.to(cutlass.Float32) - ) - res = (c / (1.0 + cute.math.exp(-c))).to(cutlass.BFloat16) - mState[slot, my_col, 0] = win[r, 1] - mState[slot, my_col, 1] = win[r, 2] - mState[slot, my_col, 2] = proj - mOut[r, my_col] = res - - -def _dl(t: torch.Tensor, align: int): - return from_dlpack(t, assumed_align=align) - - -_COMPILED: dict = {} - - -def _operands(x, mat_b, out, conv_state, conv_weight, conv_state_indices, stream): - return ( - _dl(x, 32), - _dl(mat_b.t(), 32), - _dl(out, 32), - _dl(conv_state, 2), - _dl(conv_weight, 8), - _dl(conv_state_indices, 4), - stream, - ) - - -def launch(x, mat_b, conv_state, conv_weight, conv_state_indices, conv_dim, pad_slot_id, pdl): - """The launch behind `GdnInProjConv1dPdl.__call__`. Returns the fresh `[M, N]` output.""" - rows, k = x.shape - n = mat_b.shape[1] - dev = x.device - pdl = bool(pdl and get_env_enable_pdl()) - out = torch.empty(rows, n, dtype=x.dtype, device=dev) - stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream) - operands = (x, mat_b, out, conv_state, conv_weight, conv_state_indices, stream) - key = ( - rows, - n, - k, - conv_dim, - int(pad_slot_id), - pdl, - tuple(conv_state.shape), - conv_state.stride(), - ) - fn = _COMPILED.get(key) - if fn is None: - # Compiled from the call's own operands (a stand-in for the pool would span every - # layer of the coalesced conv cache). - threads, outs = TILING[k] - kern = _InProjConvKernel( - rows=rows, - k=k, - threads=threads, - outs=outs, - conv_dim=conv_dim, - pad_slot=int(pad_slot_id), - pdl=pdl, - ) - fn = _COMPILED[key] = cute_ext.compile( - kern, *_operands(*operands), options=_COMPILE_OPTIONS - ) - fn(*_operands(*operands)) - return out - - -def _cell(why: str, rows: int, slots) -> Cell: - return Cell(why=why, spec=dict(rows=rows, slots=slots, pdl=True, pool=461)) - - -class GdnInProjConv1dPdl(OpWrapper): - """`out = x @ mat_b`, then for `c < conv_dim` and a row whose slot is not - `pad_slot_id` the causal conv step on the projected value, one launch. - - `x` is the dense row-major `[M, K]` activation (`M <= MAX_ROWS`), `mat_b` the - dense column-major `[K, N]` view of the stored `[N, K]` projection (its - `.t()`). Row `r`, slot `s = conv_state_indices[r]`, column `c < conv_dim`: - `p = round_bf16(sum_k x[r, k] * W[c, k])` (fp32 products and sums), - `out[r, c] = round_bf16(silu(sum_j conv_weight[c, j] * [conv_state[s, c, :], p][j]))` - in fp32, and `conv_state[s, c, :] = [conv_state[s, c, 1:], p]`. Every other - column, and every column of a pad row, is `p`; a pad row's slot is not read or - written. Returns the fresh `[M, N]` output; `conv_state` is advanced in place. - - Equivalent to `low_m_bf16_mm` followed by `causal_conv1d_update(x=out[:, :conv_dim], - conv_state, conv_weight, bias=None, silu_activation=True, cache_seqlens=None, - conv_state_indices, pad_slot_id)` up to the GEMV's summation order. - - Fusion boundary: the projection, the conv step, SiLU and the window's advance. - The split into q|k|v, z, a, b and everything after are the caller's. - """ - - ARCHS = frozenset({Arch.SM_103}) - - # spec: rows, slots (tuple per row; -1 is the pad slot), pdl, pool (slots in the pool). - CELLS: tuple[Cell, ...] = ( - _cell("decode 1 row, slot 37 of the 461-slot pool", 1, (37,)), - _cell("decode 2 rows, two slots", 2, (5, 211)), - _cell("decode 2 rows, one a pad row (-1)", 2, (17, -1)), - ) - - note = """ - Metadata consumed: the decode rows' conv slots (`conv_state_indices`, int32 - `[M]`). Preconditions: bf16 `x` dense `[M, K]`, `1 <= M <= MAX_ROWS`; `mat_b` - the dense column-major `[K, N]` view of a `[N, K]` weight with K in `TILING` - and `N` a multiple of the tile's outputs; `conv_state` `[slots, conv_dim, 3]` - bf16 (any strides; they are compiled in), `conv_weight` `[conv_dim, 4]` bf16 - contiguous, `conv_dim <= N`, no bias, SiLU on, `cache_seqlens` None (the - window shifts; not a ring). - - PDL: the weight rows, the conv taps, the slots and the pool's windows are read - before `griddepcontrol.wait` (the pool row is this layer's own, written by this - kernel one step ago; nothing between writes it); the activation after it. The - launch triggers its dependents right before its own wait. - - Compilation: CuTe DSL, JIT on the first call per (M, N, K, conv_dim, - pad_slot_id, pdl, pool shape and strides); that call must come before - CUDA-graph capture -- the engine's warmup forward does it. - - Numerics: the projection is fp32 products and sums rounded once to bf16 (the - value stored for non-conv columns and pushed into the window); the conv is - fp32 over the bf16 window, taps and projection, rounded once. Only the - GEMV's summation order differs from `low_m_bf16_mm` + `causal_conv1d_update`. - - Certified with PDL on, over the 461-slot pool: 1 and 2 rows, and 2 rows with a pad - row; two consecutive steps under CUDA-graph replay (1 row). Not certified: `pdl=False`. - - Hard limits - M rows per launch boundary MAX_ROWS=2 driven at 1, 2 (cells) - K a TILING key 2560 the checkpoint's hidden size - pool any slot count and strides 461 driven (64 under graph replay) - """ - - def __call__( - self, - x: torch.Tensor, - mat_b: torch.Tensor, - conv_state: torch.Tensor, - conv_weight: torch.Tensor, - conv_state_indices: torch.Tensor, - conv_dim: int, - pad_slot_id: int, - pdl: bool = True, - ) -> torch.Tensor: - return launch( - x, mat_b, conv_state, conv_weight, conv_state_indices, conv_dim, pad_slot_id, pdl - ) - - def reference( - self, - x: torch.Tensor, - mat_b: torch.Tensor, - conv_state: torch.Tensor, - conv_weight: torch.Tensor, - conv_state_indices: torch.Tensor, - conv_dim: int, - pad_slot_id: int, - pdl: bool = True, - ) -> tuple[torch.Tensor, torch.Tensor]: - """`(out, conv_state)` as the call leaves them, in fp32 then rounded; mutates nothing.""" - with true_fp32_matmul(): - proj = (x.float() @ mat_b.float()).to(x.dtype) - out = proj.clone() - state = conv_state.clone() - for r, s in enumerate(conv_state_indices.tolist()): - if s == pad_slot_id: - continue - p = proj[r, :conv_dim].float() - window = torch.cat([conv_state[s].float(), p[:, None]], dim=1) - c = (window * conv_weight.float()).sum(1) - out[r, :conv_dim] = (c * torch.sigmoid(c)).to(x.dtype) - state[s] = window[:, 1:].to(conv_state.dtype) - return out, state - - def is_valid( - self, - x: torch.Tensor, - mat_b: torch.Tensor, - conv_state: torch.Tensor, - conv_weight: torch.Tensor, - conv_state_indices: torch.Tensor, - conv_dim: int, - pad_slot_id: int, - pdl: bool = True, - ) -> None: - assert x.dtype == mat_b.dtype == conv_state.dtype == conv_weight.dtype == torch.bfloat16 - assert x.dim() == 2 and x.is_contiguous() and 1 <= x.shape[0] <= MAX_ROWS, tuple(x.shape) - assert mat_b.shape[0] == x.shape[1] and x.shape[1] in TILING, ( - tuple(mat_b.shape), - tuple(x.shape), - ) - assert mat_b.stride(0) == 1 and mat_b.stride(1) == mat_b.shape[0], ( - "mat_b must be dense column-major [K, N] (e.g. weight.t())" - ) - assert mat_b.shape[1] % TILING[x.shape[1]][1] == 0, ( - "N must be a multiple of the tile's outputs" - ) - assert conv_state.dim() == 3 and tuple(conv_state.shape[1:]) == (conv_dim, WIDTH - 1) - assert tuple(conv_weight.shape) == (conv_dim, WIDTH) and conv_weight.is_contiguous() - assert conv_dim <= mat_b.shape[1] - assert conv_state_indices.dtype == torch.int32 and tuple(conv_state_indices.shape) == ( - x.shape[0], - ) - - -gdn_in_proj_conv1d_pdl = GdnInProjConv1dPdl() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/pack_gdn_decode_qkv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/pack_gdn_decode_qkv.py deleted file mode 100644 index 8fe4119dbaaa..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/pack_gdn_decode_qkv.py +++ /dev/null @@ -1,135 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Gated-DeltaNet decode-phase q/k/v packing: the plain split of the causal -conv-update output into per-head q, k and v, with no normalization and no gates. - -Wraps the in-tree module callable -``tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops.pack_gdn_decode_qkv`` -(tier-2 sourcing: no op boundary exists). One Triton launch, -``_pack_gdn_decode_qkv_kernel``, after allocating the three outputs. -""" - -from __future__ import annotations - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops import ( - pack_gdn_decode_qkv as _pack_gdn_decode_qkv, -) - -#: This checkpoint's linear-attention geometry at tp1. -NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM = 16, 128, 48, 128 -CKPT_HEADS = (NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM) -QKV_DIM = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240 - - -class PackGdnDecodeQkv(OpWrapper): - """Cut a token-major `[num_tokens, >= qkv_dim]` conv output, packed - `[q | k | v]` along features with - `qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim`, into - - q, k [1, num_tokens, num_k_heads, head_k_dim] - v [1, num_tokens, num_v_heads, head_v_dim] - - Every element is a bit-exact copy: no Q/K L2 norm, no gate, no dtype - change. Columns at or past `qkv_dim` are dropped. - - Fusion boundary: only the split and per-head reshape are inside. The conv - (and its activation) is the caller's, and so is the Q/K normalization -- - the recurrence consuming these must run `use_qk_l2norm_in_kernel=True`. - That is the one way to misuse this op and still get plausible output. - - Row-wise, with no notion of a sequence: a verify step of `B` sequences x - `P` tokens is passed as `B * P` rows, the token axis keeps the input row - order, and `q.reshape(B, P, ...)` is a view, not a re-gather. - - Returns three fresh contiguous allocations in `mixed_qkv.dtype`; - `mixed_qkv` is read through both of its strides and is not modified. - """ - - ARCHS = frozenset({Arch.SM_103}) - - # spec: num_tokens, or (num_decodes, draft) for a verify batch, whose rows - # are the token-major `[n * p, qkv_dim]` view of the conv's `[n, C, p]` output. - CELLS: tuple[Cell, ...] = tuple( - Cell( - why=f"target verify call: {b} seqs x {p} tokens, token-major rows", - spec=dict(num_decodes=b, draft=p), - ) - for b, p in ((32, 2), (8, 4)) - ) + (Cell(why="127 tokens, contiguous: a partial final block", spec=dict(num_tokens=127)),) - - note = """ - Metadata consumed: none. Stateless; sequence boundaries and draft grouping - are carried by the caller's row order alone. - - Call site (`_gdn_verify`): `mixed_qkv` the `[n * p, 10240]` token-major - conv output (contiguous: the conv result carries its input's strides, so - transposing back is a view), heads `16, 128, 48, 128`, bf16. Certified bitwise - at that call for 32 sequences of 2 tokens and 8 of 4, and at 127 contiguous - rows (a partial final block). - - Raise loudly, so not guarded: `head_k_dim` and `head_v_dim` rounding up to - different powers of two (`128 / 64`: triton `CompilationError` - "Mismatched type for dim_mask", before any launch); a 1-D input - (`IndexError` reading `stride(1)`). - - The feature axis is never masked against the input's own extent; the - token axis is. `is_valid` guards that: `[32, 9216]` at qkv_dim 10240 was - accepted and returned finite values (83.3% of `v` bit-equal to the block - it should have read, the rest up to 6.25 off); over-declaring - `num_v_heads` likewise; a 3-D `[4, 2, 10240]` input faulted with an - illegal memory access that poisons the process. - - Hard limits - none probed. - """ - - def __call__( - self, - mixed_qkv: torch.Tensor, - num_k_heads: int, - head_k_dim: int, - num_v_heads: int, - head_v_dim: int, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - return _pack_gdn_decode_qkv(mixed_qkv, num_k_heads, head_k_dim, num_v_heads, head_v_dim) - - def reference( - self, - mixed_qkv: torch.Tensor, - num_k_heads: int, - head_k_dim: int, - num_v_heads: int, - head_v_dim: int, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """`[q | k | v]` cut with `narrow` and reshaped per head; views where torch allows.""" - n = mixed_qkv.shape[0] - qk_width, v_width = num_k_heads * head_k_dim, num_v_heads * head_v_dim - q = mixed_qkv.narrow(1, 0, qk_width).reshape(n, num_k_heads, head_k_dim) - k = mixed_qkv.narrow(1, qk_width, qk_width).reshape(n, num_k_heads, head_k_dim) - v = mixed_qkv.narrow(1, 2 * qk_width, v_width).reshape(n, num_v_heads, head_v_dim) - return q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0) - - def is_valid( - self, - mixed_qkv: torch.Tensor, - num_k_heads: int, - head_k_dim: int, - num_v_heads: int, - head_v_dim: int, - ) -> None: - # `shape[1]` rather than a rank check: a 1-D input still reaches the - # op's own IndexError, and a 3-D one (whose shape[1] is unrelated) is - # refused here instead of faulting. - qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim - assert mixed_qkv.shape[1] >= qkv_dim, ( - f"mixed_qkv must be at least qkv_dim={qkv_dim} wide " - f"(2 * {num_k_heads} * {head_k_dim} + {num_v_heads} * {head_v_dim}); got " - f"shape {tuple(mixed_qkv.shape)} -- the feature axis is not masked against " - "the input's extent, so a short row is read out of bounds" - ) - - -pack_gdn_decode_qkv = PackGdnDecodeQkv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py deleted file mode 100644 index b61b97d3bb67..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py +++ /dev/null @@ -1,163 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Zero the recurrent (and optional convolution) pool rows of fresh requests. - -Wraps the in-tree module callable -``tensorrt_llm._torch.modules.mamba.recurrent_state_cache.reset_recurrent_state_rows`` -(tier-2 sourcing: no ``torch.ops.trtllm`` op exists). One Triton launch, -``_reset_recurrent_state_rows_kernel``; the rest of the callable is shape and -stride validation. -""" - -from __future__ import annotations - -from typing import Optional - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.modules.mamba import recurrent_state_cache as _recurrent_state_cache - -#: This checkpoint's per-layer pool rows: a [48, 128, 128] recurrent state and a -#: [conv_dim, d_conv - 1] = [10240, 3] convolution window. -REC_SHAPE = (48, 128, 128) -CONV_SHAPE = (10240, 3) -#: A slot count of a full-size engine run, shared by the recurrent and conv pools. The -#: kernel addresses a slot by its index alone, so a pool of another size (a bf16 recurrent -#: pool holds more slots in the same memory) takes the same path. The target hands both -#: pools over unsliced. -ENGINE_SLOTS = 461 - - -def _engine_case(num_seqs: int, flip: bool) -> dict: - """`num_seqs` distinct slots over the engine pool, slot 0 and 460 live; - `flip` moves which of the two is the one reset.""" - slots = (0, ENGINE_SLOTS - 1) + tuple(i * 13 + 4 for i in range(num_seqs - 2)) - flags = tuple((i % 2 == 1) != flip for i in range(num_seqs)) - return dict(slots=ENGINE_SLOTS, indices=slots, flags=flags) - - -class ResetRecurrentStateRows(OpWrapper): - """For each request `i` with `slot = state_indices[i]`: when - `has_initial_states[i] == 0` and `0 <= slot < recurrent_states.shape[0]`, - fill row `slot` of `recurrent_states` -- and of `conv_states`, when given -- - with +0.0. Nothing is read from either pool. - - Untouched, bitwise: every unnamed slot, every named slot whose flag is - nonzero (the row a resumed or chunked request keeps), every negative or - out-of-range slot (skipped, not clamped), padding between rows, the whole - conv pool when it is `None`, both index tensors. A slot named twice is - zeroed if any of its entries asks. - - Fusion boundary: the call decides nothing. Computing `has_initial_states`, - assigning slots, and slicing both vectors to the context prefix are the - caller's. - - Returns `(recurrent_states, conv_states)` -- the pools it was handed, - written in place (`conv_states` is `None` when not passed). - """ - - ARCHS = frozenset({Arch.SM_103}) - - # spec: slots (pool rows), indices (int32), flags (bool), rec_dtype (fp32), - # layout: "plain" | "layer_view" (pool[1] of a [3, slots, ...] allocation). - # Both pools are passed throughout. - CELLS: tuple[Cell, ...] = ( - Cell( - why="engine pool 461, 32 requests, slot 460 reset", - spec=_engine_case(32, flip=True), - ), - Cell( - why="bf16 recurrent pool 461, 32 requests, slot 0 reset", - spec=dict(_engine_case(32, flip=False), rec_dtype=torch.bfloat16), - ), - Cell( - why="pool[layer] of a [layers, slots, ...] allocation", - spec=dict(slots=4, indices=(2, 0), flags=(False, True), layout="layer_view"), - ), - ) - - note = """ - Metadata consumed: none. At the call site `state_indices` / - `has_initial_states` are the context prefix of `Mamba2Metadata`'s int32 / - bool buffers, and the pools are the layer's views of the cache manager's - `[layers, slots, ...]` allocations, passed unsliced. Certified at 461 slots - over an fp32 and a bf16 recurrent pool, and at 4 slots over one layer's view - of a 3-layer allocation. Not certified: negative or out-of-range slots, - rows with padding between them, `conv_states=None`. - - Loud, from the callable itself (ValueError): an empty recurrent pool, conv - and recurrent slot counts differing (the in-range test uses the recurrent - pool's), flag/index shape mismatch, non-contiguous or overlapping rows in - either pool, any CPU tensor. - - Silent, so `is_valid` guards it: both index tensors are read as flat - contiguous buffers of `state_indices.shape[0]` elements. - - Not guarded (data, not metadata): an index dtype too narrow for its slot - wraps -- an int8 130 arrives as -126 and is skipped as invalid. int32 is - what the runtime uses. - - Launch: grid `(num_requests, cdiv(max(rec_row, conv_row), 256))`, the two - row widths masked independently; ordered on the current stream. - - Hard limits - slot * stride(0) int64 row base probed at 460 * 786432 = 3.6e8 exact - within-row offset int32 boundary 2^31 elements per row not probed - (this checkpoint's largest row is 786432) - """ - - def __call__( - self, - recurrent_states: torch.Tensor, - state_indices: torch.Tensor, - has_initial_states: torch.Tensor, - conv_states: Optional[torch.Tensor] = None, - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - _recurrent_state_cache.reset_recurrent_state_rows( - recurrent_states, - state_indices, - has_initial_states, - conv_states, - ) - return recurrent_states, conv_states - - def reference( - self, - recurrent_states: torch.Tensor, - state_indices: torch.Tensor, - has_initial_states: torch.Tensor, - conv_states: Optional[torch.Tensor] = None, - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - """Both pools as the call leaves them: clone, zero the rows the rule names.""" - slots = recurrent_states.shape[0] - rec = recurrent_states.clone() - conv = None if conv_states is None else conv_states.clone() - for slot, flag in zip(state_indices.tolist(), has_initial_states.tolist()): - if flag or not 0 <= slot < slots: - continue - rec[slot] = 0 - if conv is not None: - conv[slot] = 0 - return rec, conv - - def is_valid( - self, - recurrent_states: torch.Tensor, - state_indices: torch.Tensor, - has_initial_states: torch.Tensor, - conv_states: Optional[torch.Tensor] = None, - ) -> None: - # Measured: a stride-2 view of a 6-element index buffer reset rows - # [0, 1] instead of [0, 1, 2] (read through the base pointer); a - # stride-2 flag view reset row 1 instead of nothing; a [2, 3] - # index/flag pair considered only its first 2 of 6 requests. - assert state_indices.dim() == 1 and state_indices.is_contiguous(), ( - "state_indices must be a contiguous 1-D tensor" - ) - assert has_initial_states.dim() == 1 and has_initial_states.is_contiguous(), ( - "has_initial_states must be a contiguous 1-D tensor" - ) - - -reset_recurrent_state_rows = ResetRecurrentStateRows() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py deleted file mode 100644 index 16b691b4008a..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py +++ /dev/null @@ -1,141 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Vision patch embedding: one 3-D convolution per pixel patch. - -Tier-3 composed: one `torch.nn.functional.conv3d` call between two views that -launch nothing. No registered op is a 3-D convolution, and the CuTe DSL -implicit-GEMM conv3d kernels in `visual_gen/kernels/blackwell/conv/` refuse this -checkpoint's three input channels (bf16 needs `in_channels % 8 == 0`; -`can_implement` swallows the refusal and returns False). -""" - -from __future__ import annotations - -from typing import Optional - -import torch -import torch.nn.functional as F - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( - Arch, - Cell, - OpWrapper, - true_fp32_matmul, -) - -#: This checkpoint: in_channels 3, kernel (temporal_patch_size 2, patch_size -#: 16, 16), out_channels 1152 -- weight [1152, 3, 2, 16, 16], a 1536-pixel row. -IN_CHANNELS = 3 -KERNEL = (2, 16, 16) -EMBED_DIM = 1152 - - -class PatchEmbedConv3d(OpWrapper): - """`out[n, k] = bias[k] + sum_{c,t,i,j} pixel[n, c, t, i, j] * weight[k, c, t, i, j]`. - - One whole `[in_channels, kt, kh, kw]` patch per row of `pixel_values`, flat - `[N, C*kt*kh*kw]` (what the image processor emits, channel-slowest) or - explicit `[N, C, kt, kh, kw]`; the kernel extent is read from `weight`. - Cross-correlation (no kernel flip), like `nn.Conv3d`. Returns a fresh - contiguous `[N, out_channels]` in the input dtype. - - Fusion boundary: the patch projection and its bias. Patchification, the - position embedding added after it, and the row order (rows pass through in - the order they arrive -- the merge-block order the merger needs) are the - caller's. - """ - - ARCHS = frozenset({Arch.SM_103}) - - # spec: num_patches, dtype. The checkpoint's geometry, flat contiguous - # rows and a bias: what the vision tower passes. - CELLS: tuple[Cell, ...] = tuple( - Cell( - why=f"checkpoint geometry, {n} patches bf16", - spec=dict(num_patches=n, dtype=torch.bfloat16), - ) - for n in (64, 4096) - ) + ( - Cell( - why="tower call ceiling: 8192 tokens x merge 2x2 = 32768 patches bf16", - spec=dict(num_patches=32768, dtype=torch.bfloat16), - ), - ) - - note = """ - Metadata consumed: none. `pixel_values` is not mutated. All three operands - share one dtype and device -- torch raises "... should be the same" - otherwise. Several images, and a video's frame groups, are just more rows; - there is no batch axis. - - fp32 runs in TF32 unless the caller sets `torch.backends.cudnn.allow_tf32 = - False` (global state the caller owns): measured 5.84e-4 from an exact fp32 - matmul with it on (99.99% of outputs differ), 3.58e-6 with it off. bf16 and - fp16 are unaffected. - - At this shape the op is a GEMM: one output position per row, and the - `[1152, 1536]` reshape of the filter is a view of the same bytes. Not timed. - CUDA-graph capturable; the output is a fresh allocation. - - Not batch-invariant: cuDNN picks its kernel by row count, so a row's last bits - depend on how many rows share the call. Check rows against `reference`, never - against a call of a different size. - - Not certified: dtypes other than bf16 (the fp32 TF32 behavior above is - measured, not certified); other geometries, `bias=None` and non-contiguous - pixel rows; patch counts below 64, between the cells, and past 32768 - (cuDNN picks its kernel by row count). - """ - - def __call__( - self, - pixel_values: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - out_channels, in_channels, kt, kh, kw = weight.shape - patches = pixel_values.view(-1, in_channels, kt, kh, kw) - return F.conv3d(patches, weight, bias).view(-1, out_channels) - - def reference( - self, - pixel_values: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """True-fp32 matmul of the flattened rows against the flattened filter.""" - rows = pixel_values.reshape(pixel_values.shape[0], -1).float() - filters = weight.reshape(weight.shape[0], -1).float() - with true_fp32_matmul(): - out = rows @ filters.t() - if bias is not None: - out = out + bias.float() - return out.to(pixel_values.dtype) - - def is_valid( - self, - pixel_values: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor] = None, - ) -> None: - # The view reinterprets anything whose element count divides: a whole - # [1, 3, 2, 64, 64] volume becomes 16 "patches" -- a [16, 1152] result - # of the right shape from the wrong pixels, 2.52 from the strided conv - # the caller meant -- and [N, 3072] rows become 2N, without an error. - _, in_channels, kt, kh, kw = weight.shape - patch_numel = in_channels * kt * kh * kw - if pixel_values.dim() == 2: - one_patch_per_row = pixel_values.shape[1] == patch_numel - elif pixel_values.dim() == 5: - one_patch_per_row = tuple(pixel_values.shape[1:]) == (in_channels, kt, kh, kw) - else: - one_patch_per_row = False - assert one_patch_per_row, ( - f"each row must be exactly one patch: pixel_values must be " - f"[num_patches, {patch_numel}] or " - f"[num_patches, {in_channels}, {kt}, {kh}, {kw}] for this weight; got " - f"{tuple(pixel_values.shape)}" - ) - - -patch_embed_conv3d = PatchEmbedConv3d() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/pos_embed_interpolate.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/pos_embed_interpolate.py deleted file mode 100644 index b616a48a86de..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/pos_embed_interpolate.py +++ /dev/null @@ -1,243 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Vision position-embedding bilinear interpolation onto each image's patch -grid, emitted in spatial-merge order. - -Tier-3 composed: index and corner-weight arithmetic on `[h]`, `[w]` and -`[h, w]` tensors, then one `torch.nn.functional.embedding_bag` for the whole -batch. No registered op interpolates anything, and `modules/` has no vision -position-embedding callable. The exact-math sibling, -`_triton_pos_embed_interpolate` in `models/modeling_qwen3vl.py`, is a private -symbol of another model's modeling file; the other in-tree resamplers are -bicubic and/or `align_corners=False`, a different function. -""" - -from __future__ import annotations - -import math -from typing import Sequence - -import torch -import torch.nn.functional as F - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -#: This checkpoint: num_position_embeddings 2304 (a 48 x 48 grid), hidden 1152, -#: spatial_merge_size 2. -NUM_GRID = 48 -HIDDEN = 1152 -MERGE = 2 - - -def _merge_order(x: torch.Tensor, height: int, width: int, merge: int) -> torch.Tensor: - """`[h, w, ...]` raster -> `[h*w, ...]` in (h_block, w_block, row, col) order.""" - tail = x.shape[2:] - return ( - x.reshape(height // merge, merge, width // merge, merge, *tail) - .permute(0, 2, 1, 3, *range(4, 4 + len(tail))) - .reshape(height * width, *tail) - ) - - -class PosEmbedInterpolate(OpWrapper): - """Resample a square row-major `[N*N, hidden]` table onto each `(t, h, w)` - patch grid; return `[sum(t*h*w), hidden]` in the table's dtype. - - Corner-aligned bilinear: `pos = linspace(0, N - 1, n)` per axis, `lo = - floor(pos)`, `hi = min(lo + 1, N - 1)`; `h` indexes table rows (stride `N`), - `w` columns. Rows come out in spatial-merge order -- whole `merge x merge` - blocks row-major, row-major inside a block -- which is what the patch - merger's 4-to-1 reshape consumes; raster order is silently wrong. `t` - repeats the frame's rows (no temporal axis); images concatenate in order. - The four-term sum accumulates in fp32 and rounds once. - - Fusion boundary: the position embeddings alone. Adding them to the patch - embeddings, the 2-D rope and the attention segmentation are the caller's. - """ - - ARCHS = frozenset({Arch.SM_103}) - - # spec: grid (tuple of (t, h, w)), merge, dtype, rows (table rows), hidden. - CELLS: tuple[Cell, ...] = ( - Cell( - why="a 4-image batch incl. the 62x90 serving grid", - spec=dict( - grid=((1, 4, 8), (2, 6, 4), (1, 48, 48), (1, 62, 90)), - merge=MERGE, - dtype=torch.bfloat16, - rows=NUM_GRID**2, - hidden=HIDDEN, - ), - ), - Cell( - why="video: 3 frames repeat one 6x10 frame", - spec=dict( - grid=((3, 6, 10),), - merge=MERGE, - dtype=torch.bfloat16, - rows=NUM_GRID**2, - hidden=HIDDEN, - ), - ), - # The largest grid one call carries (32768 patches). - Cell( - why="checkpoint table bf16, grid [[1, 128, 256]]", - spec=dict( - grid=((1, 128, 256),), - merge=MERGE, - dtype=torch.bfloat16, - rows=NUM_GRID**2, - hidden=HIDDEN, - ), - ), - ) - - note = """ - Metadata consumed: none. `grid_thw` is host ints (a device tensor works but - `int()` syncs per element). The table is not mutated; the result is fresh. - - Loud failures, left to torch: an `h` or `w` not divisible by - `spatial_merge_size` (or a merge larger than the grid) raises "shape ... is - invalid for input of size ..."; a non-2-D table raises "weight has to be a - 2D Tensor". An empty grid side returns `[0, hidden]` rather than raising. - - At `h = w = N` the samples are exactly `arange(N)` and the result is - bitwise the table in merge order. The bf16 result is within one bf16 - half-ulp of the fp64 value (0.0147 at 62x90); a bf16 multiply-add chain - (a literal HF transcription) lands up to one bf16 ulp (0.031) from this. - - CUDA graphs: capturable, but the grid is a host argument, so it is frozen - at capture. Cost is ~225 us of launch-bound index construction per image, - flat in token count; the gather is the only part that scales. - - Certified at the checkpoint's bf16 table on a 4-image batch (4x8, a 2x6x4 - video, the 48x48 exact fit and the 62x90 serving grid), a 3-frame 6x10 - video and the 128x256 grid one call carries at most. Not certified: a CPU - table (it runs and returns on CPU); fp16 and fp32 tables; a - `spatial_merge_size` other than the checkpoint's 2. - """ - - def __call__( - self, - pos_embed_weight: torch.Tensor, - grid_thw: Sequence[Sequence[int]], - spatial_merge_size: int, - ) -> torch.Tensor: - num_grid_per_side = math.isqrt(pos_embed_weight.shape[0]) - device = pos_embed_weight.device - merge = spatial_merge_size - idx_parts = [] - weight_parts = [] - - for frames, height, width in grid_thw: - frames, height, width = int(frames), int(height), int(width) - - # Corner-aligned; non-negative positions, so the int64 cast floors. - h_pos = torch.linspace( - 0, num_grid_per_side - 1, height, device=device, dtype=torch.float32 - ) - w_pos = torch.linspace( - 0, num_grid_per_side - 1, width, device=device, dtype=torch.float32 - ) - h_lo = h_pos.to(torch.int64) - w_lo = w_pos.to(torch.int64) - h_hi = (h_lo + 1).clamp(max=num_grid_per_side - 1) - w_hi = (w_lo + 1).clamp(max=num_grid_per_side - 1) - dh = (h_pos - h_lo).unsqueeze(1) - dw = (w_pos - w_lo).unsqueeze(0) - - row_lo = (h_lo * num_grid_per_side).unsqueeze(1) - row_hi = (h_hi * num_grid_per_side).unsqueeze(1) - col_lo = w_lo.unsqueeze(0) - col_hi = w_hi.unsqueeze(0) - - # [h, w, 4]: corners in (lo,lo) (lo,hi) (hi,lo) (hi,hi) order. - idx = torch.stack( - (row_lo + col_lo, row_lo + col_hi, row_hi + col_lo, row_hi + col_hi), dim=-1 - ) - weights = torch.stack( - ((1 - dh) * (1 - dw), (1 - dh) * dw, dh * (1 - dw), dh * dw), dim=-1 - ) - - # Spatial-merge reorder; `reshape` is loud when a side does not divide. - merged_h, merged_w = height // merge, width // merge - idx = ( - idx.reshape(merged_h, merge, merged_w, merge, 4) - .permute(0, 2, 1, 3, 4) - .reshape(height * width, 4) - ) - weights = ( - weights.reshape(merged_h, merge, merged_w, merge, 4) - .permute(0, 2, 1, 3, 4) - .reshape(height * width, 4) - ) - - if frames > 1: - idx = idx.repeat(frames, 1) - weights = weights.repeat(frames, 1) - idx_parts.append(idx) - weight_parts.append(weights) - - idx = idx_parts[0] if len(idx_parts) == 1 else torch.cat(idx_parts) - weights = weight_parts[0] if len(weight_parts) == 1 else torch.cat(weight_parts) - - # One launch: each row of `idx` is a 4-element bag scaled by `weights` - # and summed, fp32 table and fp32 weights. - table = pos_embed_weight.float() - out = F.embedding_bag(idx, table, mode="sum", per_sample_weights=weights) - return out.to(pos_embed_weight.dtype) - - def reference( - self, - pos_embed_weight: torch.Tensor, - grid_thw: Sequence[Sequence[int]], - spatial_merge_size: int, - ) -> torch.Tensor: - """The explicit four-corner weighted sum in fp64, rounded once. - - Shares the entry's sample positions (fp32 `linspace`) but not its - gather or accumulation; the test cross-checks the structure separately - against `F.interpolate`. - """ - num_grid = math.isqrt(pos_embed_weight.shape[0]) - hidden = pos_embed_weight.shape[1] - grid = pos_embed_weight.double().view(num_grid, num_grid, hidden) - device = pos_embed_weight.device - outs = [] - for frames, height, width in grid_thw: - frames, height, width = int(frames), int(height), int(width) - h_pos = torch.linspace(0, num_grid - 1, height, device=device) - w_pos = torch.linspace(0, num_grid - 1, width, device=device) - h_lo, w_lo = h_pos.to(torch.int64), w_pos.to(torch.int64) - h_hi = (h_lo + 1).clamp(max=num_grid - 1) - w_hi = (w_lo + 1).clamp(max=num_grid - 1) - d_h = (h_pos - h_lo).double().unsqueeze(1) - d_w = (w_pos - w_lo).double().unsqueeze(0) - raster = ( - grid[h_lo][:, w_lo] * ((1 - d_h) * (1 - d_w))[..., None] - + grid[h_lo][:, w_hi] * ((1 - d_h) * d_w)[..., None] - + grid[h_hi][:, w_lo] * (d_h * (1 - d_w))[..., None] - + grid[h_hi][:, w_hi] * (d_h * d_w)[..., None] - ) - rows = _merge_order(raster, height, width, spatial_merge_size) - outs.append(rows.repeat(frames, 1) if frames > 1 else rows) - out = torch.cat(outs) if len(outs) > 1 else outs[0] - return out.to(pos_embed_weight.dtype) - - def is_valid( - self, - pos_embed_weight: torch.Tensor, - grid_thw: Sequence[Sequence[int]], - spatial_merge_size: int, - ) -> None: - # With 2300 rows isqrt truncates to 47, every gather index stays in - # range, and the result is a correctly shaped resample of the wrong grid. - num_rows = pos_embed_weight.shape[0] - side = math.isqrt(num_rows) - assert side * side == num_rows, ( - f"pos_embed_weight must hold a square grid; got {num_rows} rows, " - f"which is not a perfect square" - ) - - -pos_embed_interpolate = PosEmbedInterpolate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py deleted file mode 100644 index fbe80d324baf..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py +++ /dev/null @@ -1,1422 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Weight table and loader: Qwen3.8-Flash-Next-NVFP4 / sm_103 / tp1. - -One `W` per weight role (shape, checkpoint source, dtype, load-time transform); -`Qwen4ExpTp1Weights` unrolls the table into declaration, manifest and load. tp1 -means no per-rank slicing. Storage is HF `[out, in]`; GEMM views, grouped-norm -slices, fp32 gating params and NVFP4 scalars are derived in `modeling.py` after load. - -Transforms: - -* trunk attention input projection: `q_proj`, `k_proj`, `v_proj` and the indexer's - `index_qk_proj` read the same activation, so they are one - `[q | k | v | gate | index | pad, hidden]` operand and one GEMM (`_attn_in_w`). `q_proj` - interleaves q and the output gate per head; its q halves go first so the GEMM output's - leading columns are the packed `[q | k | v]` the fused norm+rope rewrites in place, then - the gate halves, the index projection and `index_head_dim` zero rows (a row is a whole - number of `head_dim` heads). The draft (MTP) layer keeps its separate projections; -* fused shared-expert gate|up: `[gate ; up]` (the half order `flashinfer_silu_and_mul` reads); -* NVFP4 routed-expert stacks, laid out for the trunk's routed backend - (`core.trunk_moe_backend`): `[up ; gate]`, then 128x4 scale swizzle. trtllm-gen also - interleaves FC1 row by row and 32-row block shuffles weight and scale bytes (FC2: - shuffle only); CuteDSL interleaves FC1 in 64-row up/gate groups and leaves FC2 in - checkpoint row order. Under trtllm-gen a second copy in the CuteDSL layout (`fc*_cd`) - serves the prefill-sized launches (`modeling._CUTE_DSL_MOE_MIN_ROWS`) when the core keeps - it (`core.moe_prefill_copy`); -* n-gram table: 128 row chunks concatenated into one `[rows, 160]` buffer (the hash - emits global row indices). - -The n-gram table is not in the table: its height is a checkpoint fact, so `load` -registers it with `torch.empty` and adds its rows to the manifest; -`post_load_weights` reconciles the height. Residency (device or host) is set by -`TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD` (`_host_resident_empty`); the manifest is the -same either way. Offloading frees ~47.7 GiB, which the KV cache manager turns into -more slots; several entries certify at most 461 slots, so raising -`kv_cache_fraction` with the offload on needs those cells re-checked. - -Not loaded: `mtp.*` when there is no `speculative_config`. With one, the whole MTP module -is loaded, including the draft layer's own QSA indexer (`_index_branch`). MTP experts are -plain row-major FP8 128x128: `[up ; gate]` concat, bf16 -`weight_scale_inv` cast to fp32. `lm_head` is the shell's parameter. - -`load(model, weights)` fills every declared parameter once and asserts coverage -both ways: every key is consumed or in the non-load set. -""" - -import os -from dataclasses import dataclass -from types import SimpleNamespace -from typing import Any, Callable - -import torch -from cuda.bindings import runtime as cudart -from torch import nn - -from tensorrt_llm._utils import CUASSERT -from tensorrt_llm.llmapi.utils import get_numa_aware_cpu_affinity - -#: `layers=ALL_LAYERS` is the common case and the default. A sentinel rather -#: than a lambda so the table reads as data, and rather than `True` so the -#: field has one type of answer: which layers. -ALL_LAYERS = object() - - -@dataclass(frozen=True) -class W: - """One weight role: what it is, where its bytes come from, what happens on the way. - - `shape` takes the dimension bundle and returns the declared shape; every - shape in these models is derived from configuration rather than constant. - - `src` says where the bytes come from, in one of three forms: - - * a checkpoint key template, with `{p}` standing for the layer prefix; - * a sequence of `(template, slice_fn)` pairs, when one parameter is - assembled from several checkpoint tensors side by side -- `slice_fn` - takes the dimension bundle and returns an index into the destination; - * a callable `(d, i) -> [(key, index), ...]`, taking the dimension bundle - and the layer index. The first two forms are special cases of it; this - one exists for sources that are not a fixed template -- a per-expert - row, or a prefix that differs between layer families. A `key` may itself - be a tuple, meaning the transform is handed those tensors together. - - `layers` says which layers carry this weight: `ALL_LAYERS` for every one, - `None` for a weight that is not per-layer at all, or a callable - `d -> iterable[int]` when only some layers have it. - - `transform` runs on the source tensor before the copy, on the - destination's device. `None` means the checkpoint tensor is already the - declared shape. - """ - - name: str - shape: Callable[[Any], tuple[int, ...]] - src: str | tuple[tuple[str, Callable], ...] | Callable - dtype: torch.dtype | None = None - transform: Callable | None = None - layers: object | None = ALL_LAYERS - - def layer_indices(self, d: SimpleNamespace) -> list[int] | None: - """The layers this weight exists on, or None if it is not per-layer.""" - if self.layers is None: - return None - if self.layers is ALL_LAYERS: - return list(range(d.num_layers)) - return list(self.layers(d)) - - def sources(self, d: SimpleNamespace, i: int | None, prefix: str) -> list[tuple[Any, Any]]: - """`[(checkpoint key or key tuple, index into the param | None), ...]`.""" - if callable(self.src): - return list(self.src(d, i)) - if isinstance(self.src, str): - return [(self.src.format(p=prefix), None)] - return [(template.format(p=prefix), slice_fn(d)) for template, slice_fn in self.src] - - -# ----------------------------------------------------------- host residency -- -# -# Allocates the one PLE n-gram table. The table stays a `cuda` tensor of the same -# shape/dtype. Registered, device-mapped pages (not a host gather) because the gather's -# row indices are produced on device inside the captured decode graph. - -#: `cudaMemoryType` values; the core asserts on them after load. -CUDA_MEMORY_TYPE_HOST = 1 -CUDA_MEMORY_TYPE_DEVICE = 2 - - -def memory_kind(ptr: int) -> int: - """The driver's own answer to "what kind of memory is this pointer". - - Returns a `cudaMemoryType`. This is the residency question asked of the - process rather than of `nvidia-smi`, and it is the only one a - working-but-resident implementation cannot answer the way an offloaded one - does. - """ - (attributes,) = CUASSERT(cudart.cudaPointerGetAttributes(ptr)) - return int(attributes.type) - - -#: `cudaHostRegister` flags (`driver_types.h`): mapped, so a kernel may dereference the -#: pages, and portable, for every context in the process. -_CUDA_HOST_REGISTER_PORTABLE = 0x01 -_CUDA_HOST_REGISTER_MAPPED = 0x02 - - -def _local_cpus(device) -> set | None: - """The CPUs of `device`'s NUMA affinity, as NVML reports it (`get_numa_aware_cpu_affinity`), - that this process may run on, or None when NVML cannot tell.""" - import pynvml - - index = torch.device(device).index - try: - local = get_numa_aware_cpu_affinity(torch.cuda.current_device() if index is None else index) - except pynvml.NVMLError: - return None - return (set(local) & os.sched_getaffinity(0)) or None - - -class _RegisteredHostBuffer: - """A pageable CPU allocation registered (pinned and mapped) for the device, and - unregistered before its storage is released: at this size glibc serves the - allocation with `mmap` and returns it with `munmap`, so freeing registered pages - would leave the driver mapping memory the OS may hand out again. - - Not torch's pinned allocator: it rounds the 47.7 GiB table up to 64 GiB, and on - GB300 the device's random row gathers from its pages run an order of magnitude - slower than from registered pageable memory.""" - - def __init__(self, nbytes: int, device=None) -> None: - # The registration faults the pages in on this thread, so under the default local - # allocation policy they land on the NUMA node of the CPU it runs on. Pinned to the - # CPUs local to `device` meanwhile, the table sits next to the GPU that gathers from - # it: on GB300, random row gathers from the other socket's memory run ~1.7x slower, - # and with a process per GPU half of them would otherwise land there. - local = _local_cpus(device) if device is not None else None - previous = os.sched_getaffinity(0) if local else None - if local: - os.sched_setaffinity(0, local) - try: - self._storage = torch.empty(nbytes, dtype=torch.uint8) - self.pointer = self._storage.data_ptr() - CUASSERT( - cudart.cudaHostRegister( - self.pointer, nbytes, _CUDA_HOST_REGISTER_PORTABLE | _CUDA_HOST_REGISTER_MAPPED - ) - ) - finally: - if previous: - os.sched_setaffinity(0, previous) - # Bound now: at interpreter shutdown the module may already be torn down. - self._unregister = cudart.cudaHostUnregister - - def __del__(self) -> None: - pointer, self.pointer = self.pointer, 0 - if pointer: - self._unregister(pointer) - - -def _host_resident_empty( - rows: int, cols: int, dtype, device -) -> tuple[torch.Tensor, _RegisteredHostBuffer]: - """`[rows, cols]` of `dtype` on `device`, whose bytes live in host memory. - - Returns the tensor and the registered CPU allocation that owns the pages. The - caller **must** keep the second alive for as long as the first: the tensor - is a view through a device address and owns nothing. - """ - itemsize = torch.empty(0, dtype=dtype).element_size() - assert itemsize == 1, ( - f"the host-resident allocation reinterprets a `|u1` buffer, so its " - f"element must be one byte; {dtype} is {itemsize}" - ) - assert device.type == "cuda", f"host offload needs a cuda device to map into; got {device}" - - host = _RegisteredHostBuffer(rows * cols, device) - (device_pointer,) = CUASSERT(cudart.cudaHostGetDevicePointer(host.pointer, 0)) - - # No public torch API builds a tensor over a pointer the caching allocator - # did not hand out; `__cuda_array_interface__` is the numba/cupy path. - # `|u1` because the protocol has no float8 typestr; reinterpreted below. - interface = SimpleNamespace( - __cuda_array_interface__={ - "shape": (rows, cols), - "typestr": "|u1", - "data": (int(device_pointer), False), - "version": 3, - "strides": None, - } - ) - table = torch.as_tensor(interface, device=device).view(dtype) - # `torch.as_tensor` copies when it decides the buffer is not already on the - # requested device; a copy here would be a silent 48 GiB device allocation. - assert table.data_ptr() == int(device_pointer), ( - "torch.as_tensor copied the host-mapped buffer instead of viewing it, " - "so the table would be device-resident after all" - ) - assert table.shape == (rows, cols) and table.dtype == dtype, (table.shape, table.dtype) - return table, host - - -# ------------------------------------------------------------------ the table -- - -_LM = "model.language_model" -_VISUAL = "model.visual" -_MTP = "mtp" -_MTP_LAYER = "mtp.layers.0" - -# NVFP4 block size: one e4m3 scale per 16 contiguous elements along K. -SF_VEC = 16 - -# The FP8 block-scale family's block edge, on both axes of the MTP experts' -# weight scales and along the activation row. -FP8_BLOCK = 128 - -# Parameter-key suffixes after which to release transform scratch. The expert -# relayouts allocate per-expert scratch 512 times per parameter and the table -# fill stages one shard at a time; without this, peak load memory is the whole -# model rather than one parameter. -_RELEASE_AFTER = ( - "_fc2_s", - "_fc1_s", - "_fc2_s_cd", - "_fc1_s_cd", - "ple_table", - "mtp_fc1_w", - "mtp_fc2_w", -) - - -def _block32_perm(rows: int, device) -> torch.Tensor: - """Gather index of the 32-row block shuffle: inside each aligned block of - 32 rows, source row `4u + v` (0<=u<8, 0<=v<4) lands at destination row - `8v + u`.""" - assert rows % 32 == 0, rows - src = torch.arange(32) - dst = (src % 4) * 8 + src // 4 - idx = torch.empty(32, dtype=torch.long) - idx[dst] = src - blocks = rows // 32 - return (idx.repeat(blocks) + torch.arange(blocks).repeat_interleave(32) * 32).to(device) - - -def _interleave_perm(rows: int, device) -> torch.Tensor: - """Gather index that interleaves the `[up | gate]` halves of a `2*I` row - stack into up0, gate0, up1, gate1, ...""" - p = torch.empty(rows, dtype=torch.long) - p[0::2] = torch.arange(0, rows // 2) - p[1::2] = torch.arange(rows // 2, rows) - return p.to(device) - - -#: The group the CuteDSL FC1 epilogue pairs up and gate rows in: rows `[128k, 128k + 64)` -#: of the fused operand are up rows and the next 64 their gate rows. -_CUTEDSL_GATE_UP_GROUP = 64 - - -def _group_interleave_perm(rows: int, group: int, device) -> torch.Tensor: - """Gather index that interleaves the `[up | gate]` halves of a `2*I` row stack in - `group`-row blocks: up[0:g], gate[0:g], up[g:2g], gate[g:2g], ...""" - assert rows % (2 * group) == 0, (rows, group) - idx = torch.arange(rows, device=device).view(2, rows // (2 * group), group) - return idx.transpose(0, 1).reshape(-1) - - -def _fc1_perm(rows: int, device, core) -> torch.Tensor: - """The FC1 row gather of the trunk's routed backend. trtllm-gen: interleave then - 32-row block shuffle, composed into one gather. CuteDSL: the 64-row group - interleave alone.""" - if core.trunk_moe_backend == "CUTEDSL": - return _group_interleave_perm(rows, _CUTEDSL_GATE_UP_GROUP, device) - return _interleave_perm(rows, device)[_block32_perm(rows, device)] - - -def _swizzle(x: torch.Tensor) -> torch.Tensor: - """128x4 block-scale swizzle of an `[M, C]` uint8 scale matrix, returned - flat (`M % 128 == 0` and `C % 4 == 0` hold at this geometry, so the result - is exactly `M * C` bytes).""" - return torch.ops.trtllm.block_scale_interleave(x.contiguous()) - - -#: Rows after the shared expert's gate|up rows in its fused operand: the per-token gate's -#: one row, then zeros, so the GEMM's output row stays a whole number of 16-byte vectors. -SH_GATE_ROWS = 8 - - -def _cat_rows(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: - """Fused shared-expert weight bytes: **gate** rows first, the half order - `flashinfer_silu_and_mul` reads, then up, then the per-token gate's row and - `SH_GATE_ROWS - 1` zero rows -- one GEMM forms gate|up and the gate logit.""" - gate, up, sh_gate = t - pad = sh_gate.new_zeros(SH_GATE_ROWS - 1, sh_gate.shape[-1]) - return torch.cat((gate, up, sh_gate, pad), dim=0).contiguous() - - -def _attn_in_w(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: - """(q_proj, k_proj, v_proj, index_qk_proj) -> the fused `[q | k | v | gate | index | - pad]` input projection of one trunk attention layer (see the module docstring). - - `q_proj`'s rows are `[heads, 2, head_dim]`: head `h`'s q rows, then its gate rows. - """ - q, k, v, idx = t - heads, head_dim, hidden = core.heads_q, core.head_dim, q.shape[-1] - qg = q.reshape(heads, 2, head_dim, hidden) - pad = q.new_zeros(core.index_head_dim, hidden) - rows = (qg[:, 0].reshape(heads * head_dim, hidden), k, v) - rows += (qg[:, 1].reshape(heads * head_dim, hidden), idx, pad) - return torch.cat(rows, dim=0).contiguous() - - -def _expert_fc1_w(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: - """(up, gate) packed nibbles -> one expert's FC1 operand.""" - up, gate = t - w = torch.cat([up, gate], dim=0) - return torch.index_select(w, 0, _fc1_perm(w.shape[0], w.device, core)).contiguous() - - -def _expert_fc1_s(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: - """(up, gate) e4m3 scales -> one expert's FC1 scale operand: the weight's own row - permutation, then the swizzle.""" - up, gate = t - s = torch.cat([up.view(torch.uint8), gate.view(torch.uint8)], dim=0) - rows, cols = s.shape - s = torch.index_select(s, 0, _fc1_perm(rows, s.device, core)) - return _swizzle(s).reshape(rows, cols) - - -def _expert_fc2_w(t: torch.Tensor, core) -> torch.Tensor: - """`down_proj` packed nibbles -> one expert's FC2 operand (no interleave): the - 32-row block shuffle on trtllm-gen, checkpoint row order on CuteDSL.""" - if core.trunk_moe_backend == "CUTEDSL": - return t.contiguous() - return torch.index_select(t, 0, _block32_perm(t.shape[0], t.device)).contiguous() - - -def _expert_fc2_s(t: torch.Tensor, core) -> torch.Tensor: - """`down_proj` e4m3 scales -> one expert's FC2 scale operand: the weight's own row - order, then the swizzle.""" - s = t.view(torch.uint8) - rows, cols = s.shape - if core.trunk_moe_backend != "CUTEDSL": - s = torch.index_select(s, 0, _block32_perm(rows, s.device)) - return _swizzle(s).reshape(rows, cols) - - -#: The layout argument of the second, CuteDSL-layout expert copy (`fc*_cd`). -_CUTEDSL_LAYOUT = SimpleNamespace(trunk_moe_backend="CUTEDSL") - - -def _expert_fc1_w_cd(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: - """(up, gate) packed nibbles -> one expert's FC1 operand in the CuteDSL layout.""" - return _expert_fc1_w(t, _CUTEDSL_LAYOUT) - - -def _expert_fc1_s_cd(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: - """(up, gate) e4m3 scales -> one expert's FC1 scale operand in the CuteDSL layout.""" - return _expert_fc1_s(t, _CUTEDSL_LAYOUT) - - -def _expert_fc2_s_cd(t: torch.Tensor, core) -> torch.Tensor: - """`down_proj` e4m3 scales -> one expert's FC2 scale operand in the CuteDSL layout.""" - return _expert_fc2_s(t, _CUTEDSL_LAYOUT) - - -def _mtp_fc1_w(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: - """(up, gate) e4m3 codes -> one draft expert's FC1 operand, `[up ; gate]` on - the row axis and nothing else. The FP8 block-scale runner reads the - checkpoint's plain row-major layout.""" - up, gate = t - return torch.cat([up, gate], dim=0).contiguous() - - -def _mtp_fc1_s(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: - """(up, gate) 128x128 block scales -> one draft expert's FC1 scale operand: - the same `[up ; gate]` concatenation, cast bf16 -> fp32.""" - up, gate = t - return torch.cat([up, gate], dim=0).to(torch.float32).contiguous() - - -def _mtp_fc2_s(t: torch.Tensor, core) -> torch.Tensor: - """`down_proj.weight_scale_inv` -> the FC2 scale operand: a bf16 -> fp32 - cast, no relayout.""" - return t.to(torch.float32).contiguous() - - -#: The layout argument of the draft experts' NVFP4 copy: the draft MoE runs it on trtllm-gen -#: whatever the trunk's routed backend is. -_TRTLLM_GEN_LAYOUT = SimpleNamespace(trunk_moe_backend="TRTLLM") - -#: `448 * 6`: an NVFP4 global scale is this over the tensor's amax (modelopt's recipe). -NVFP4_AMAX_TO_GLOBAL = 448.0 * 6.0 - - -def _fp8_block_dequant(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: - """`[E, N, K]` e4m3 codes times their `[E, N/128, K/128]` fp32 block scales, in fp32.""" - s = scales.repeat_interleave(FP8_BLOCK, 1).repeat_interleave(FP8_BLOCK, 2) - return codes.float() * s - - -def mtp_experts_to_nvfp4(fc1_w, fc1_s, fc2_w, fc2_s, chunk: int = 32) -> tuple: - """Requantize the draft layer's FP8 128x128 block-scaled experts to NVFP4 in the - trunk's trtllm-gen layout (`_expert_fc1_w/_s`, `_expert_fc2_w/_s`). - - Per expert and projection (FC1's up|gate share one, as the fused FC1's alpha - requires): dequantize with the block scales, take the global scale - `448 * 6 / amax`, then `fp4_quantize` (16-element blocks, e4m3 block scales, - linear order), which then gets the trunk's row permutation and 128x4 swizzle. - Returns `(fc1_w [E, 2I, H/2], fc1_s [E, 2I, H/16], fc2_w [E, H, I/2], - fc2_s [E, H, I/16])` uint8 and the fp32 `[E]` weight global scales `(g1, g2)`; - a weight's real value is its NVFP4 reconstruction over its global scale. - `fc1_w` / `fc1_s` are the `[up ; gate]` stacks `_mtp_fc1_w` / `_mtp_fc1_s` build.""" - experts, two_i, hidden = fc1_w.shape - inter = two_i // 2 - dev = fc1_w.device - u8 = torch.uint8 - out = ( - torch.empty(experts, two_i, hidden // 2, dtype=u8, device=dev), - torch.empty(experts, two_i, hidden // SF_VEC, dtype=u8, device=dev), - torch.empty(experts, hidden, inter // 2, dtype=u8, device=dev), - torch.empty(experts, hidden, inter // SF_VEC, dtype=u8, device=dev), - ) - g1 = torch.empty(experts, dtype=torch.float32, device=dev) - g2 = torch.empty(experts, dtype=torch.float32, device=dev) - layout = _TRTLLM_GEN_LAYOUT - for e0 in range(0, experts, chunk): - e1 = min(e0 + chunk, experts) - w1 = _fp8_block_dequant(fc1_w[e0:e1], fc1_s[e0:e1]) - w2 = _fp8_block_dequant(fc2_w[e0:e1], fc2_s[e0:e1]) - g1[e0:e1] = NVFP4_AMAX_TO_GLOBAL / w1.abs().amax(dim=(1, 2)).clamp_min(1e-30) - g2[e0:e1] = NVFP4_AMAX_TO_GLOBAL / w2.abs().amax(dim=(1, 2)).clamp_min(1e-30) - w1, w2 = w1.to(torch.bfloat16), w2.to(torch.bfloat16) - for j, e in enumerate(range(e0, e1)): - q, s = torch.ops.trtllm.fp4_quantize(w1[j], g1[e : e + 1], SF_VEC, False, False) - s = s.view(two_i, hidden // SF_VEC) - out[0][e] = _expert_fc1_w((q[:inter], q[inter:]), layout) - out[1][e] = _expert_fc1_s((s[:inter], s[inter:]), layout) - q, s = torch.ops.trtllm.fp4_quantize(w2[j], g2[e : e + 1], SF_VEC, False, False) - out[2][e] = _expert_fc2_w(q, layout) - out[3][e] = _expert_fc2_s(s.view(hidden, inter // SF_VEC), layout) - del w1, w2 - return (*out, g1, g2) - - -def _materialize(entry) -> torch.Tensor: - """Realize one checkpoint entry: `[:]` on a lazy safetensors slice, but a - 0-dim tensor (every NVFP4 per-tensor scalar on this checkpoint) rejects that - index.""" - return entry if getattr(entry, "ndim", 1) == 0 else entry[:] - - -# --- which layers carry a weight --- - - -def _attn_layers(d): - """The full-attention layers (`i % 4 == 3` on this checkpoint).""" - return [i for i, a in enumerate(d.is_attn) if a] - - -def _gdn_layers(d): - """The Gated-DeltaNet (linear-attention) layers.""" - return [i for i, a in enumerate(d.is_attn) if not a] - - -#: The Gated-DeltaNet input projections, in the row order of the concatenated -#: `in_proj_w` (and so the column order of its GEMM output). -_IN_PROJ_PARTS = ("in_proj_qkv", "in_proj_z", "in_proj_a", "in_proj_b") - - -def _in_proj_rows(d, part: int) -> tuple[slice]: - """The row range of `in_proj_w` that input projection `part` fills.""" - start = sum(d.in_proj_split[:part]) - return (slice(start, start + d.in_proj_split[part]),) - - -def _cute_dsl_copy_layers(d): - """Every layer when the core keeps the second, CuteDSL-layout copy that prefill-sized - launches read under trtllm-gen (`core.moe_prefill_copy`); none otherwise, CuteDSL among - them, whose one copy is it.""" - return range(d.num_layers) if d.moe_prefill_copy else () - - -def _vision_blocks(d): - """The vision tower's blocks. Keyed `l{b}_vis_*`: block `b`, not decoder layer `b`.""" - return range(d.v_depth) - - -def _mtp_only(d): - """The MTP module, at its own layer index `num_hidden_layers`, only when it is live.""" - return [d.num_layers] if d.mtp_enabled else [] - - -# --- where the bytes come from --- - - -def _ple(suffix): - """A PLE weight. The PLE layer's index is config, so the key is computed.""" - return lambda d, i: [(f"{_LM}.layers.{d.ple_layer}.ple.{suffix}", None)] - - -def _ngram_prefix(ple_layer: int) -> str: - """The quantized n-gram table submodule: the 128 shards and the one - per-tensor scale that dequantizes them both hang off this, one level below - `ple_embedding` itself.""" - return f"{_LM}.layers.{ple_layer}.ple.ple_embedding.ngram_embedding" - - -def _visual(suffix): - """A vision-tower weight that is not per-block.""" - return lambda d, i: [(f"{_VISUAL}.{suffix}", None)] - - -def _vision_block(suffix): - """One vision block's weight.""" - return lambda d, i: [(f"{_VISUAL}.blocks.{i}.{suffix}", None)] - - -def _mtp_key(key): - """One MTP-module tensor, at its own checkpoint key (the module's prefix is not the - trunk's layer prefix).""" - return lambda d, i: [(key, None)] - - -def _experts(prefix, *suffixes): - """One source per routed expert, written at index `e`; a pair of suffixes hands the - transform both tensors together.""" - - def build(d, i): - base = prefix(i) - out = [] - for e in range(d.num_experts): - keys = tuple(f"{base}.experts.{e}.{s}" for s in suffixes) - out.append((keys if len(keys) > 1 else keys[0], (e,))) - return out - - return build - - -def _trunk_mlp(i): - return f"{_LM}.layers.{i}.mlp" - - -def _mtp_mlp(i): - return f"{_MTP_LAYER}.mlp" - - -#: The six per-tensor NVFP4 scalars every routed expert carries, as -#: (parameter suffix, checkpoint suffix). The checkpoint stores modelopt -#: reciprocals; the core folds them into the runner's scalars after load. -_EXPERT_SCALARS = ( - ("e_isc1", "gate_proj.input_scale"), - ("e_isc1_up", "up_proj.input_scale"), - ("e_ws2_1", "gate_proj.weight_scale_2"), - ("e_ws2_1_up", "up_proj.weight_scale_2"), - ("e_isc2", "down_proj.input_scale"), - ("e_ws2_2", "down_proj.weight_scale_2"), -) - - -#: (parameter tag, checkpoint module) of the two per-layer hyper-connection gates. -_HC_MODULES = (("ahc", "attn_hyper_connection"), ("mhc", "mlp_hyper_connection")) - - -def _hyper_connection(tag, src, layers): - """The four weights of one gated-residual hyper-connection; `src(name)` names the source.""" - return ( - W(f"{tag}_norm", shape=lambda d: (d.stream,), src=src("hc_norm"), layers=layers), - W( - f"{tag}_down", - shape=lambda d: (d.hc_lowrank, d.stream), - src=src("input_mix_weight_down"), - layers=layers, - ), - W( - f"{tag}_up", - shape=lambda d: (d.stream, d.hc_lowrank), - src=src("input_mix_weight_up"), - layers=layers, - ), - W( - f"{tag}_inject", - shape=lambda d: (d.hc_count, d.stream), - src=src("block_inject_weight"), - layers=layers, - ), - ) - - -def _mixer(tag, prefix, layers): - """A stack-final hyper-connection mixer: a layer's gate minus the injection weight.""" - src = _mtp_key if layers is not None else (lambda key: key) - return ( - W( - f"{tag}_norm", - shape=lambda d: (d.stream,), - src=src(f"{prefix}.hc_norm.weight"), - layers=layers, - ), - W( - f"{tag}_down", - shape=lambda d: (d.hc_lowrank, d.stream), - src=src(f"{prefix}.input_mix_weight_down.weight"), - layers=layers, - ), - W( - f"{tag}_up", - shape=lambda d: (d.stream, d.hc_lowrank), - src=src(f"{prefix}.input_mix_weight_up.weight"), - layers=layers, - ), - ) - - -def _attn_in_fused(tag, src, layers): - """The trunk attention layers' one input projection: q_proj, k_proj, v_proj and the - QSA indexer's index_qk_proj, re-laid by `_attn_in_w`.""" - keys = ("q_proj.weight", "k_proj.weight", "v_proj.weight", "indexer.index_qk_proj.weight") - return ( - W( - f"{tag}attn_in_w", - shape=lambda d: (d.attn_in_rows, d.hidden), - src=lambda d, i: [(tuple(src(k).format(p=f"{_LM}.layers.{i}") for k in keys), None)], - transform=_attn_in_w, - layers=layers, - ), - ) - - -def _gated_attention(tag, src, layers, fused_in=False): - """The gated GQA block's six weights. `q_proj` is twice as wide as q: q and the output - gate interleave per head. With `fused_in`, q/k/v are not declared here: they are part - of `_attn_in_fused`'s one operand.""" - return ( - *_attn_in_separate(tag, src, layers, fused_in), - *_attn_out_and_norms(tag, src, layers), - ) - - -def _attn_in_separate(tag, src, layers, fused_in): - if fused_in: - return () - return ( - W( - f"{tag}q_w", - shape=lambda d: (2 * d.q_width, d.hidden), - src=src("q_proj.weight"), - layers=layers, - ), - W( - f"{tag}k_w", - shape=lambda d: (d.kv_width, d.hidden), - src=src("k_proj.weight"), - layers=layers, - ), - W( - f"{tag}v_w", - shape=lambda d: (d.kv_width, d.hidden), - src=src("v_proj.weight"), - layers=layers, - ), - ) - - -def _attn_out_and_norms(tag, src, layers): - return ( - W( - f"{tag}o_w", - shape=lambda d: (d.hidden, d.q_width), - src=src("o_proj.weight"), - layers=layers, - ), - W(f"{tag}q_norm", shape=lambda d: (d.head_dim,), src=src("q_norm.weight"), layers=layers), - W(f"{tag}k_norm", shape=lambda d: (d.head_dim,), src=src("k_norm.weight"), layers=layers), - ) - - -def _index_branch(tag, src, layers, fused_in=False): - """The QSA index branch's three weights: one replicated, unquantized projection - producing `[index_heads q | 1 k]`, and two per-head gemma norms. With `fused_in`, the - projection is part of `_attn_in_fused`'s one operand and only the norms are declared.""" - projection = W( - f"{tag}idx_w", - shape=lambda d: (d.index_width, d.hidden), - src=src("index_qk_proj.weight"), - layers=layers, - ) - return ( - *(() if fused_in else (projection,)), - W( - f"{tag}idx_qn", - shape=lambda d: (d.index_head_dim,), - src=src("q_layernorm.weight"), - layers=layers, - ), - W( - f"{tag}idx_kn", - shape=lambda d: (d.index_head_dim,), - src=src("k_layernorm.weight"), - layers=layers, - ), - ) - - -def _shared_expert(tag, prefix, layers): - """The bf16 shared expert beside the routed ones, its per-token gate in the same - operand as its gate|up (`_cat_rows`).""" - return ( - W( - f"{tag}sh_gu", - shape=lambda d: (2 * d.shared_inter + SH_GATE_ROWS, d.hidden), - src=lambda d, i: [ - ( - ( - f"{prefix(i)}.shared_expert.gate_proj.weight", - f"{prefix(i)}.shared_expert.up_proj.weight", - f"{prefix(i)}.shared_expert_gate.weight", - ), - None, - ) - ], - transform=_cat_rows, - layers=layers, - ), - W( - f"{tag}sh_dn", - shape=lambda d: (d.hidden, d.shared_inter), - src=lambda d, i: [(f"{prefix(i)}.shared_expert.down_proj.weight", None)], - layers=layers, - ), - ) - - -u8, f8, f32, i64 = torch.uint8, torch.float8_e4m3fn, torch.float32, torch.int64 - - -class Qwen4ExpTp1Weights: - """Qwen3.8-Flash-Next-NVFP4 / sm_103 / tp1. - - The table, and the passes over it: `declare` allocates every weight it names, - `manifest` maps checkpoint keys onto them, `load` fills them. `declare` and - `manifest` share one "which parameters exist" decision (`_targets`), so they - cannot disagree about a weight. - """ - - WEIGHTS: tuple[W, ...] = ( - # --- the two per-layer hyper-connection gates: every trunk layer --- - *( - w - for tag, module in _HC_MODULES - for w in _hyper_connection( - tag, lambda name, m=module: f"{{p}}.{m}.{name}.weight", ALL_LAYERS - ) - ), - # --- the gated GQA layers and their QSA index branch --- - # q/k/v and the indexer's replicated, unquantized [index_heads q | 1 k] - # projection as one operand, then o_proj, the q/k norms and the indexer's two - # per-head gemma norms. - *_attn_in_fused("", lambda s: f"{{p}}.self_attn.{s}", _attn_layers), - *_gated_attention("", lambda s: f"{{p}}.self_attn.{s}", _attn_layers, fused_in=True), - *_index_branch("", lambda s: f"{{p}}.self_attn.indexer.{s}", _attn_layers, fused_in=True), - # --- the Gated-DeltaNet layers --- - # The four input projections read the same activation, so they are one - # weight and one GEMM: rows `[q|k|v (conv_dim) | z (value_dim) | a | b]` - # (`d.in_proj_split`), each checkpoint tensor copied into its own row - # range. The forward splits the output into column views. - W( - "in_proj_w", - shape=lambda d: (sum(d.in_proj_split), d.hidden), - src=tuple( - (f"{{p}}.linear_attn.{name}.weight", lambda d, j=j: _in_proj_rows(d, j)) - for j, name in enumerate(_IN_PROJ_PARTS) - ), - layers=_gdn_layers, - ), - W( - "A_log", - shape=lambda d: (d.num_v_heads,), - src="{p}.linear_attn.A_log", - layers=_gdn_layers, - ), - W( - "dt_bias", - shape=lambda d: (d.num_v_heads,), - src="{p}.linear_attn.dt_bias", - layers=_gdn_layers, - ), - # Stored [conv_dim, 1, width]; the conv entries want [conv_dim, width], which is - # a zero-copy view the targets take once, after load. - W( - "conv_w", - shape=lambda d: (d.conv_dim, 1, d.conv_kernel), - src="{p}.linear_attn.conv1d.weight", - layers=_gdn_layers, - ), - W( - "gnorm", - shape=lambda d: (d.head_v_dim,), - src="{p}.linear_attn.norm.weight", - layers=_gdn_layers, - ), - W( - "out_w", - shape=lambda d: (d.hidden, d.value_dim), - src="{p}.linear_attn.out_proj.weight", - layers=_gdn_layers, - ), - # --- MoE on every layer: NVFP4 routed experts, a bf16 router and shared expert --- - W("router", shape=lambda d: (d.num_experts, d.hidden), src="{p}.mlp.gate.weight"), - W( - "fc1_w", - shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // 2), - dtype=u8, - src=_experts(_trunk_mlp, "up_proj.weight", "gate_proj.weight"), - transform=_expert_fc1_w, - ), - W( - "fc1_s", - shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // SF_VEC), - dtype=u8, - src=_experts(_trunk_mlp, "up_proj.weight_scale", "gate_proj.weight_scale"), - transform=_expert_fc1_s, - ), - W( - "fc2_w", - shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // 2), - dtype=u8, - src=_experts(_trunk_mlp, "down_proj.weight"), - transform=_expert_fc2_w, - ), - W( - "fc2_s", - shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // SF_VEC), - dtype=u8, - src=_experts(_trunk_mlp, "down_proj.weight_scale"), - transform=_expert_fc2_s, - ), - # Under trtllm-gen, the same experts again in the CuteDSL layout: prefill-sized - # launches run on them (`modeling._CUTE_DSL_MOE_MIN_ROWS`). - W( - "fc1_w_cd", - shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // 2), - dtype=u8, - src=_experts(_trunk_mlp, "up_proj.weight", "gate_proj.weight"), - transform=_expert_fc1_w_cd, - layers=_cute_dsl_copy_layers, - ), - W( - "fc1_s_cd", - shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // SF_VEC), - dtype=u8, - src=_experts(_trunk_mlp, "up_proj.weight_scale", "gate_proj.weight_scale"), - transform=_expert_fc1_s_cd, - layers=_cute_dsl_copy_layers, - ), - W( - "fc2_w_cd", - shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // 2), - dtype=u8, - src=_experts(_trunk_mlp, "down_proj.weight"), - layers=_cute_dsl_copy_layers, - ), - W( - "fc2_s_cd", - shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // SF_VEC), - dtype=u8, - src=_experts(_trunk_mlp, "down_proj.weight_scale"), - transform=_expert_fc2_s_cd, - layers=_cute_dsl_copy_layers, - ), - *( - W(name, shape=lambda d: (d.num_experts,), dtype=f32, src=_experts(_trunk_mlp, suffix)) - for name, suffix in _EXPERT_SCALARS - ), - *_shared_expert("", _trunk_mlp, ALL_LAYERS), - # --- not per-layer: the stack-final mixer, which is also the final norm --- - *_mixer("fhc", f"{_LM}.hyper_connection_mixer", None), - W( - "embed", - shape=lambda d: (d.vocab, d.hidden), - src=f"{_LM}.embed_tokens.weight", - layers=None, - ), - # --- PLE, on its one decoder layer (the n-gram table itself is declared at load) --- - W( - "ple_key_w", - shape=lambda d: (d.stream, d.hidden), - src=_ple("key_proj.weight"), - layers=None, - ), - W( - "ple_val_w", - shape=lambda d: (d.ple_embed_dim, d.hidden), - src=_ple("value_proj.weight"), - layers=None, - ), - W("ple_nk", shape=lambda d: (d.stream,), src=_ple("norm_key.weight"), layers=None), - W("ple_nq", shape=lambda d: (d.stream,), src=_ple("norm_query.weight"), layers=None), - W("ple_nc", shape=lambda d: (d.stream,), src=_ple("norm_conv.weight"), layers=None), - W( - "ple_conv_w", - shape=lambda d: (d.stream, 1, d.ple_conv_kernel), - src=_ple("conv1d.weight"), - layers=None, - ), - W( - "ple_mult", - shape=lambda d: (d.ngram_size,), - dtype=i64, - src=_ple("ple_embedding.layer_multipliers"), - layers=None, - ), - W( - "ple_vocab", - shape=lambda d: (d.ngram_heads,), - dtype=i64, - src=_ple("ple_embedding.ngram_heads_vocab_sizes"), - layers=None, - ), - W( - "ple_offs", - shape=lambda d: (d.ngram_heads,), - dtype=i64, - src=_ple("ple_embedding.ngram_heads_offsets"), - layers=None, - ), - # The table's per-tensor scale sits with the **table**, one level below the three - # metadata tensors above: the shards and this scale are the quantized - # `ngram_embedding` submodule, while the multipliers, offsets and vocabulary sizes - # describe the id space and stay on `ple_embedding` itself. The two groups also - # live in different safetensors files, which is the other way to notice the split. - W( - "ple_tscale", - shape=lambda d: (1,), - src=lambda d, i: [(f"{_ngram_prefix(d.ple_layer)}.weight_scale", None)], - layers=None, - ), - # --- the vision tower: all bf16, one-to-one, no transform --- - W( - "vis_patch_w", - shape=lambda d: (d.v_hidden, d.v_channels, d.v_tpatch, d.v_patch, d.v_patch), - src=_visual("patch_embed.proj.weight"), - layers=None, - ), - W( - "vis_patch_b", - shape=lambda d: (d.v_hidden,), - src=_visual("patch_embed.proj.bias"), - layers=None, - ), - W( - "vis_pos", - shape=lambda d: (d.v_pos_rows, d.v_hidden), - src=_visual("pos_embed.weight"), - layers=None, - ), - # LayerNorm **with a bias**, unlike every norm in the language tower. - W( - "vis_n1_w", - shape=lambda d: (d.v_hidden,), - src=_vision_block("norm1.weight"), - layers=_vision_blocks, - ), - W( - "vis_n1_b", - shape=lambda d: (d.v_hidden,), - src=_vision_block("norm1.bias"), - layers=_vision_blocks, - ), - # One fused `[q | k | v]` projection, MHA, with a bias. - W( - "vis_qkv_w", - shape=lambda d: (3 * d.v_hidden, d.v_hidden), - src=_vision_block("attn.qkv.weight"), - layers=_vision_blocks, - ), - W( - "vis_qkv_b", - shape=lambda d: (3 * d.v_hidden,), - src=_vision_block("attn.qkv.bias"), - layers=_vision_blocks, - ), - W( - "vis_proj_w", - shape=lambda d: (d.v_hidden, d.v_hidden), - src=_vision_block("attn.proj.weight"), - layers=_vision_blocks, - ), - W( - "vis_proj_b", - shape=lambda d: (d.v_hidden,), - src=_vision_block("attn.proj.bias"), - layers=_vision_blocks, - ), - W( - "vis_n2_w", - shape=lambda d: (d.v_hidden,), - src=_vision_block("norm2.weight"), - layers=_vision_blocks, - ), - W( - "vis_n2_b", - shape=lambda d: (d.v_hidden,), - src=_vision_block("norm2.bias"), - layers=_vision_blocks, - ), - # A **non-gated** MLP: nothing here is split in half, unlike every SwiGLU in the - # language tower. - W( - "vis_fc1_w", - shape=lambda d: (d.v_inter, d.v_hidden), - src=_vision_block("mlp.linear_fc1.weight"), - layers=_vision_blocks, - ), - W( - "vis_fc1_b", - shape=lambda d: (d.v_inter,), - src=_vision_block("mlp.linear_fc1.bias"), - layers=_vision_blocks, - ), - W( - "vis_fc2_w", - shape=lambda d: (d.v_hidden, d.v_inter), - src=_vision_block("mlp.linear_fc2.weight"), - layers=_vision_blocks, - ), - W( - "vis_fc2_b", - shape=lambda d: (d.v_hidden,), - src=_vision_block("mlp.linear_fc2.bias"), - layers=_vision_blocks, - ), - # The patch merger. Its norm is applied **before** the 4-to-1 shuffle, so it is - # `hidden` wide and not `merge_dim` wide. - W( - "vis_m_n_w", - shape=lambda d: (d.v_hidden,), - src=_visual("merger.norm.weight"), - layers=None, - ), - W("vis_m_n_b", shape=lambda d: (d.v_hidden,), src=_visual("merger.norm.bias"), layers=None), - W( - "vis_m_fc1_w", - shape=lambda d: (d.v_merge_dim, d.v_merge_dim), - src=_visual("merger.linear_fc1.weight"), - layers=None, - ), - W( - "vis_m_fc1_b", - shape=lambda d: (d.v_merge_dim,), - src=_visual("merger.linear_fc1.bias"), - layers=None, - ), - W( - "vis_m_fc2_w", - shape=lambda d: (d.hidden, d.v_merge_dim), - src=_visual("merger.linear_fc2.weight"), - layers=None, - ), - W( - "vis_m_fc2_b", - shape=lambda d: (d.hidden,), - src=_visual("merger.linear_fc2.bias"), - layers=None, - ), - # --- the MTP draft module, at layer index num_hidden_layers, only when live --- - # - # Structurally one more decoder layer of this architecture with a two-branch - # front end; every shape equals a trunk full-attention layer's except the front - # end and the expert quantization (FP8 128x128 block scales, not NVFP4). - W( - "mtp_pre_e_norm", - shape=lambda d: (d.hidden,), - src=_mtp_key(f"{_MTP}.pre_fc_norm_embedding.weight"), - layers=_mtp_only, - ), - # `[hc_count * hidden]`: the module's input is the trunk's four-stream residual. - W( - "mtp_pre_h_norm", - shape=lambda d: (d.stream,), - src=_mtp_key(f"{_MTP}.pre_fc_norm_hidden.weight"), - layers=_mtp_only, - ), - W( - "mtp_fc_e", - shape=lambda d: (d.hidden, d.hidden), - src=_mtp_key(f"{_MTP}.fc_embedding.weight"), - layers=_mtp_only, - ), - W( - "mtp_fc_h", - shape=lambda d: (d.hidden, d.hidden), - src=_mtp_key(f"{_MTP}.fc_hidden.weight"), - layers=_mtp_only, - ), - *( - w - for tag, module in _HC_MODULES - for w in _hyper_connection( - f"mtp_{tag}", - lambda name, m=module: _mtp_key(f"{_MTP_LAYER}.{m}.{name}.weight"), - _mtp_only, - ) - ), - *_gated_attention("mtp_", lambda s: _mtp_key(f"{_MTP_LAYER}.self_attn.{s}"), _mtp_only), - *_index_branch( - "mtp_", lambda s: _mtp_key(f"{_MTP_LAYER}.self_attn.indexer.{s}"), _mtp_only - ), - W( - "mtp_router", - shape=lambda d: (d.num_experts, d.hidden), - src=_mtp_key(f"{_MTP_LAYER}.mlp.gate.weight"), - layers=_mtp_only, - ), - # e4m3 codes stored row-major with one fp32 scale per 128x128 weight block, the - # FC1 halves concatenated `[up ; gate]` on the row axis. - W( - "mtp_fc1_w", - shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden), - dtype=f8, - src=_experts(_mtp_mlp, "up_proj.weight", "gate_proj.weight"), - transform=_mtp_fc1_w, - layers=_mtp_only, - ), - W( - "mtp_fc1_s", - shape=lambda d: ( - d.num_experts, - (2 * d.moe_inter) // FP8_BLOCK, - d.hidden // FP8_BLOCK, - ), - dtype=f32, - src=_experts(_mtp_mlp, "up_proj.weight_scale_inv", "gate_proj.weight_scale_inv"), - transform=_mtp_fc1_s, - layers=_mtp_only, - ), - # FC2 needs no relayout at all: the checkpoint's `[hidden, I]` e4m3 block is - # exactly the operand. - W( - "mtp_fc2_w", - shape=lambda d: (d.num_experts, d.hidden, d.moe_inter), - dtype=f8, - src=_experts(_mtp_mlp, "down_proj.weight"), - layers=_mtp_only, - ), - W( - "mtp_fc2_s", - shape=lambda d: (d.num_experts, d.hidden // FP8_BLOCK, d.moe_inter // FP8_BLOCK), - dtype=f32, - src=_experts(_mtp_mlp, "down_proj.weight_scale_inv"), - transform=_mtp_fc2_s, - layers=_mtp_only, - ), - *_shared_expert("mtp_", _mtp_mlp, _mtp_only), - # The module's own stack-final mixer, which is its `shared_head`. - *_mixer("mtp_fhc", f"{_MTP}.hyper_connection_mixer", _mtp_only), - ) - - def dims(self, core) -> SimpleNamespace: - """Shapes are read off the normalized text config and the composite's vision - config, which the core holds as `vision_config` before it declares anything.""" - cfg = core.model_config.pretrained_config - vcfg = core.vision_config - hidden = cfg.hidden_size - index_heads, index_kv_heads = cfg.indexer_n_heads, cfg.indexer_kv_heads - num_k_heads, num_v_heads = cfg.linear_num_key_heads, cfg.linear_num_value_heads - key_dim = num_k_heads * cfg.linear_key_head_dim - value_dim = num_v_heads * cfg.linear_value_head_dim - ngram_heads = cfg.heads_per_ngram * (cfg.ngram_size - 1) - v_hidden = vcfg.hidden_size - return SimpleNamespace( - num_layers=cfg.num_hidden_layers, - is_attn=[t == "full_attention" for t in cfg.layer_types], - hidden=hidden, - hc_count=cfg.hc_count, - hc_lowrank=cfg.hc_lowrank, - stream=cfg.hc_count * hidden, - head_dim=cfg.head_dim, - q_width=cfg.num_attention_heads * cfg.head_dim, - kv_width=cfg.num_key_value_heads * cfg.head_dim, - index_head_dim=cfg.indexer_head_dim, - index_width=(index_heads + index_kv_heads) * cfg.indexer_head_dim, - # [q | k | v | gate | index q, k | pad of one index head] - attn_in_rows=2 * cfg.num_attention_heads * cfg.head_dim - + 2 * cfg.num_key_value_heads * cfg.head_dim - + (index_heads + index_kv_heads + 1) * cfg.indexer_head_dim, - num_v_heads=num_v_heads, - head_v_dim=cfg.linear_value_head_dim, - conv_dim=2 * key_dim + value_dim, - value_dim=value_dim, - in_proj_split=(2 * key_dim + value_dim, value_dim, num_v_heads, num_v_heads), - conv_kernel=cfg.linear_conv_kernel_dim, - num_experts=cfg.num_experts, - moe_inter=cfg.moe_intermediate_size, - shared_inter=cfg.shared_expert_intermediate_size, - vocab=cfg.vocab_size, - ple_layer=list(cfg.ple_layer_ids)[0] - 1, # one-based in the config - ple_embed_dim=cfg.ple_embed_dim, - ple_conv_kernel=cfg.ple_conv_kernel_size, - ngram_size=cfg.ngram_size, - ngram_heads=ngram_heads, - v_depth=vcfg.depth, - v_hidden=v_hidden, - v_inter=vcfg.intermediate_size, - v_channels=vcfg.in_channels, - v_tpatch=vcfg.temporal_patch_size, - v_patch=vcfg.patch_size, - v_pos_rows=vcfg.num_position_embeddings, - v_merge_dim=v_hidden * vcfg.spatial_merge_size * vcfg.spatial_merge_size, - mtp_enabled=getattr(core.model_config, "spec_config", None) is not None, - trunk_moe_backend=core.trunk_moe_backend, - moe_prefill_copy=core.moe_prefill_copy, - dtype=core.model_config.torch_dtype, - ) - - def expected_unconsumed(self, core, weights) -> set: - """With MTP off, the whole `mtp.` subtree; with it on, nothing, so that any - unloaded `mtp.*` key fails the coverage assert instead of hiding.""" - if core.mtp_enabled: - return set() - return {k for k in weights if k.startswith(f"{_MTP}.")} - - def declare(self, d: SimpleNamespace) -> nn.ParameterDict: - """Allocate every weight the table declares. - - Meta-init intercepts `torch.empty` here -- real CUDA storage arrives - when the engine materializes the registry. - """ - w = nn.ParameterDict() - for entry in self.WEIGHTS: - dtype = entry.dtype or d.dtype - for key, _, _ in self._targets(entry, d): - w[key] = nn.Parameter( - torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False - ) - return w - - def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None, str]]: - """`[(param key, layer index | None, checkpoint layer prefix), ...]`. - - One place decides which parameters an entry declares, so `declare` and - `manifest` cannot disagree about whether a weight exists. - """ - layers = entry.layer_indices(d) - if layers is None: - return [(entry.name, None, "")] - return [(f"l{i}_{entry.name}", i, f"{_LM}.layers.{i}") for i in layers] - - def manifest(self, core) -> dict[str, list[tuple[Any, Any, Callable | None]]]: - """param key -> [(checkpoint key, index into the param | None, transform | None)].""" - d = self.dims(core) - rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {} - for entry in self.WEIGHTS: - for key, i, prefix in self._targets(entry, d): - rows[key] = [ - (ckpt, index, entry.transform) for ckpt, index in entry.sources(d, i, prefix) - ] - return rows - - def load(self, model, weights) -> None: - """Fill every declared parameter. The n-gram table is declared here first, - multi-key rows feed their transform several tensors at once, the per-tensor - scalars are 0-dim and reject `[:]`, and coverage is asserted in both - directions.""" - core = model.model - # The one load-time declaration, before the manifest is built: the table's - # height comes from the shards, and the manifest addresses it by row range. - shard_keys = _declare_ngram_table(core, weights) - manifest = self.manifest(core) - manifest["ple_table"] = [ - (key, (slice(s * core.ple_rows_per_shard, (s + 1) * core.ple_rows_per_shard),), None) - for s, key in enumerate(shard_keys) - ] - consumed: set = set() - - def fill(param: nn.Parameter, ckpt_key, index, transform) -> None: - keys = ckpt_key if isinstance(ckpt_key, tuple) else (ckpt_key,) - for key in keys: - assert key in weights, f"checkpoint key missing: {key}" - dst = param.data if index is None else param.data[index] - # Materialize where the destination lives: the expert relayouts are row - # gathers and the scale swizzle is a kernel, both on the device. - src_all = tuple( - _materialize(weights[key]).to(dst.device, non_blocking=True) for key in keys - ) - if transform is not None: - src = transform(src_all if isinstance(ckpt_key, tuple) else src_all[0], core) - else: - assert len(src_all) == 1, "a multi-key row needs a transform" - src = src_all[0] - assert dst.shape == src.shape, (ckpt_key, tuple(dst.shape), tuple(src.shape)) - assert src.dtype == dst.dtype, (ckpt_key, src.dtype, dst.dtype) - dst.copy_(src, non_blocking=True) - consumed.update(keys) - - assert set(manifest.keys()) == set(core.w.keys()), ( - "manifest/parameter drift", - set(manifest.keys()) ^ set(core.w.keys()), - ) - for param_key, sources in manifest.items(): - for ckpt_key, index, transform in sources: - fill(core.w[param_key], ckpt_key, index, transform) - if param_key.endswith(_RELEASE_AFTER): - torch.cuda.empty_cache() - - # Shell-registered exception: the base class owns lm_head (untied), and this - # checkpoint leaves it bf16, so it is a single dense weight. Asserted rather than - # adapted — a shape change is a different logits path. - assert tuple(model.lm_head.weight.shape) == (core.vocab, core.hidden), ( - f"the shell built lm_head {tuple(model.lm_head.weight.shape)} instead " - f"of {(core.vocab, core.hidden)}" - ) - fill(model.lm_head.weight, "lm_head.weight", None, None) - # Parameter-side coverage, the other direction: name every registered parameter - # outside the target's own ParameterDict instead of assuming lm_head is the only - # one. - shell = {n for n, _ in model.named_parameters() if not n.startswith("model.w.")} - assert shell == {"lm_head.weight"}, f"unfed shell parameters: {sorted(shell)}" - - torch.cuda.synchronize() - leftover = set(weights.keys()) - consumed - expected = self.expected_unconsumed(core, weights) - what = "empty" if core.mtp_enabled else "the MTP module" - assert leftover == expected, ( - f"checkpoint coverage: leftover keys are not exactly {what}; " - f"unexpected {sorted(leftover - expected)[:8]}, " - f"missed {sorted(expected - leftover)[:8]}" - ) - - -MODEL_WEIGHTS = Qwen4ExpTp1Weights() - - -def _declare_ngram_table(core, weights) -> list[str]: - """Register the n-gram table with the height its shards concatenate to, and return - the shard keys in concatenation order. - - The shards are **contiguous row chunks**, not an interleaving: global row `r` lives in - shard `r // rows_per_shard` at local row `r % rows_per_shard`. That is what makes the - flat buffer indexable by the hash's global rows. - - **This is the whole of PLE host offload.** `core.ple_host_offload`, resolved at - construction from `TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD`, picks between an ordinary - device allocation and host memory the device addresses. Nothing else changes: same - shape, same dtype, same `ParameterDict` key, same rows filling it shard by shard, same - gather in the forward. An offloaded tensor is **loaded, not skipped** — it stays - fully covered in both directions. What the manifest cannot record is *where* it - landed, so the core asserts that separately against the driver. - """ - prefix = _ngram_prefix(core.ple_layer) - keys = [f"{prefix}.shard_{s}.weight" for s in range(core.ple_shards)] - for key in keys: - assert key in weights, f"checkpoint key missing: {key}" - shapes = {tuple(weights[key].shape) for key in keys} - assert len(shapes) == 1, ( - f"the n-gram table's {len(keys)} shards do not share one shape: " - f"{sorted(shapes)}; they are concatenated as equal contiguous row " - f"chunks and an uneven one would shift every row above it" - ) - rows_per_shard, head_dim = shapes.pop() - assert head_dim == core.ple_head_dim, ( - f"the n-gram table is {head_dim} wide but {core.ngram_heads} heads must " - f"tile ple_embed_dim {core.ple_embed_dim}, i.e. {core.ple_head_dim}" - ) - dtype = weights[keys[0]].dtype - assert dtype == core.ple_table_dtype, ( - f"the n-gram table is stored {dtype}; this target declares " - f"{core.ple_table_dtype} and does not convert it" - ) - device = core.w["embed"].device - rows = core.ple_shards * rows_per_shard - if core.ple_host_offload: - table, core._ple_table_host = _host_resident_empty(rows, head_dim, dtype, device) - else: - table = torch.empty(rows, head_dim, dtype=dtype, device=device) - core._ple_table_host = None - core.w["ple_table"] = nn.Parameter(table, requires_grad=False) - # `nn.Parameter` is `_make_subclass`, which shares storage rather than copying — but - # "rather than copying" is the entire feature here, so it is asserted instead of - # assumed. - assert core.w["ple_table"].data_ptr() == table.data_ptr(), ( - "wrapping the n-gram table in nn.Parameter moved its storage" - ) - core.ple_rows_per_shard = rows_per_shard - return keys diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py index af8a50119783..cc8135690654 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py @@ -2,12 +2,16 @@ # SPDX-License-Identifier: Apache-2.0 """Routing for Qwen4ExpForConditionalGeneration. -`route` compares the facts that fix which code runs against the one deployment the -target was written for, and raises `NoTarget` stating them otherwise (`auto` falls +`route` compares the facts that fix which code runs against the deployments the +targets were written for, and raises `NoTarget` stating them otherwise (`auto` falls back to the built-in implementation; `require` raises). The checkpoint's published architecture name is kept by `load_pretrained_config`, so `route` registers under it. All criteria are read off `text_config` / `vision_config`; the composite parent's mirrored fields are unreliable. + +Qwen3.8-Flash-Next-NVFP4 on sm103 lands on `qwen3_8_flash_next_nvfp4__sm_103__default` +at one of its parallel layouts (`layout`): tp1 on one rank, and tp4, tep4, dep4 or dep2 in +one node. """ from __future__ import annotations @@ -25,31 +29,54 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm.llmapi.llm_args import TorchLlmArgs +# The GPU the target is certified on. +_SM = (10, 3) + # (layers, hidden, experts, attention heads, linear value heads, vocab). Shape-sniffing # also matches fine-tunes; the accuracy gate pins identity. _TEXT_SHAPE = (48, 2560, 512, 24, 48, 248320) # (depth, hidden, heads, merge, out_hidden, num_pos, deepstack_indexes, act). # `deepstack_visual_indexes == ()` is the key member: a deepstack sibling has a merger -# list this target lacks. +# list these targets lack. _VISION_SHAPE = (27, 1152, 16, 2, 2560, 2304, (), "gelu_pytorch_tanh") -# `moe_config.backend` requests the target implements (`_TRUNK_MOE_BACKENDS` in modeling.py). -_MOE_BACKENDS = ("AUTO", "TRTLLM", "CUTEDSL") - -# The longest MTP draft the target runs (`_MAX_VERIFY_WIDTH - 1` in modeling.py); the legacy -# per-step recurrent states (`_MAX_LEGACY_VERIFY_WIDTH - 1`) and the CUTEDSL experts -# (`_CUTEDSL_MAX_DRAFT_LEN`) are certified at shorter ones. +# The longest MTP draft the target runs (`_MAX_VERIFY_WIDTH - 1` in modeling.py). _MAX_DRAFT_LEN = 7 -_MAX_LEGACY_DRAFT_LEN = 3 + +# The engine choices the target implements at one rank (tp1), all in its modeling.py: the +# `moe_config.backend` requests (`_TRUNK_MOE_BACKENDS`), the longest draft of the CUTEDSL +# experts (`_CUTEDSL_MAX_DRAFT_LEN`) and the recurrent-state cache dtypes +# (`_ONE_RANK_SSM_STATE_DTYPES`). +_ONE_RANK_MOE_BACKENDS = ("AUTO", "TRTLLM", "CUTEDSL") _MAX_CUTEDSL_DRAFT_LEN = 3 +_ONE_RANK_SSM_STATE_DTYPES = (torch.float32, torch.bfloat16) + +# Their counterparts at every other layout: no CUTEDSL experts (`_TRUNK_MOE_BACKENDS` less +# CUTEDSL) and a bf16 recurrent state (`_SSM_STATE_DTYPES`). +_MOE_BACKENDS = ("AUTO", "TRTLLM") +_SSM_STATE_DTYPES = (torch.bfloat16,) + +# (world size, attention DP, MoE TP, MoE EP) -> the target's layout. Every layout spans its +# whole world in one node, attention TP over all of it unless attention DP replicates +# attention. One rank (tp1) holds everything. At four ranks the routed experts are split by +# intermediate (tp4) or by expert (tep4, dep4). At two ranks only attention DP with two +# 256-expert windows (dep2), where each rank's attention, Gated-DeltaNet and PLE geometry is +# dep4's. +_LAYOUTS = { + (1, False, 1, 1): "tp1", + (4, False, 4, 1): "tp4", + (4, False, 1, 4): "tep4", + (4, True, 1, 4): "dep4", + (2, True, 1, 2): "dep2", +} -# Recurrent-state cache dtypes the target serves (`_SSM_STATE_DTYPES` in modeling.py). -_SSM_STATE_DTYPES = (torch.float32, torch.bfloat16) +# The attention-DP layouts MTP is admitted at: those its draft acceptance was measured at. +_ATTENTION_DP_MTP_LAYOUTS = ("dep4",) class EngineChoices(NamedTuple): - """Engine choices that fix which code the target runs but are not in the checkpoint's + """Engine choices that fix which code a target runs but are not in the checkpoint's config, as the engine will resolve them.""" torch_dtype: torch.dtype # the dtype the engine runs the model in @@ -60,7 +87,7 @@ class EngineChoices(NamedTuple): def engine_choices(model_config: ModelConfig, llm_args: TorchLlmArgs) -> EngineChoices: - # Imported on the claimed branch only, as the target is. + # Imported on a claimed branch only, as the targets are. from tensorrt_llm._torch.pyexecutor.model_loader import effective_mamba_ssm_cache_dtype from tensorrt_llm._torch.utils import is_gdn_replay_enabled @@ -76,14 +103,57 @@ def engine_choices(model_config: ModelConfig, llm_args: TorchLlmArgs) -> EngineC ) +def layout(mapping) -> Optional[str]: + """The target's layout ``mapping`` describes (``tp1``, ``tp4``, ``tep4``, ``dep4``, + ``dep2``), or None. + + Pipeline and context parallelism, an attention TP narrower than the TP group, MoE + clusters, lm_head TP under attention DP and a group spanning nodes are other + forwards. + """ + if mapping.tp_size != mapping.world_size or (mapping.pp_size, mapping.cp_size) != (1, 1): + return None + if mapping.is_multi_node() or mapping.moe_cluster_size != 1: + return None + if mapping.enable_attention_dp: + if mapping.enable_lm_head_tp_in_adp: + return None + elif mapping.attn_tp_size != mapping.tp_size: + return None + return _LAYOUTS.get( + ( + mapping.world_size, + mapping.enable_attention_dp, + mapping.moe_tp_size, + mapping.moe_ep_size, + ) + ) + + +def _describe(mapping) -> str: + return ( + f"world_size={mapping.world_size} (tp={mapping.tp_size}, pp={mapping.pp_size}, " + f"cp={mapping.cp_size}, attention_dp={mapping.enable_attention_dp}, " + f"moe_tp={mapping.moe_tp_size}, moe_ep={mapping.moe_ep_size})" + ) + + def _check_engine( - spec_config: Optional[Any], kv_cache_quant_algo: Optional[Any], engine: EngineChoices + spec_config: Optional[Any], + kv_cache_quant_algo: Optional[Any], + engine: EngineChoices, + member: str, ) -> None: - """Decline what the target's constructor or first forward would refuse. + """Decline what the target's constructor or first forward would refuse at layout + `member`: at tp1 the CuteDSL experts and the fp32 recurrent state are served too. Without this ``auto`` claims the deployment and then fails mid-build instead of falling back to the built-in implementation. """ + if member == "tp1": + moe_backends, ssm_state_dtypes = _ONE_RANK_MOE_BACKENDS, _ONE_RANK_SSM_STATE_DTYPES + else: + moe_backends, ssm_state_dtypes = _MOE_BACKENDS, _SSM_STATE_DTYPES if spec_config is not None and ( getattr(spec_config, "draft_len_schedule", None) is not None or getattr(spec_config, "acceptance_rate_threshold", None) is not None @@ -100,15 +170,16 @@ def _check_engine( ) if engine.torch_dtype != torch.bfloat16: raise NoTarget(f"the engine runs the model in {engine.torch_dtype}; the target is bf16") - if str(engine.moe_backend).upper() not in _MOE_BACKENDS: + if str(engine.moe_backend).upper() not in moe_backends: raise NoTarget( - f"moe_config.backend={engine.moe_backend!r}; the target runs the trunk's routed " - f"experts on {', '.join(_MOE_BACKENDS)}" + f"moe_config.backend={engine.moe_backend!r}; at {member} the target runs the " + f"trunk's routed experts on {', '.join(moe_backends)}" ) - if engine.mamba_ssm_cache_dtype not in _SSM_STATE_DTYPES: + if engine.mamba_ssm_cache_dtype not in ssm_state_dtypes: + served = " or ".join(str(dtype).removeprefix("torch.") for dtype in ssm_state_dtypes) raise NoTarget( - f"a {engine.mamba_ssm_cache_dtype} recurrent-state cache; the target serves " - "float32 or bfloat16 (kv_cache_config.mamba_ssm_cache_dtype)" + f"a {engine.mamba_ssm_cache_dtype} recurrent-state cache; at {member} the target " + f"serves {served} (kv_cache_config.mamba_ssm_cache_dtype)" ) if engine.prefill_cuda_graph_backend == "piecewise": raise NoTarget( @@ -119,11 +190,11 @@ def _check_engine( if spec_config is None: return draft = spec_config.max_draft_len - if not engine.gdn_replay and draft > _MAX_LEGACY_DRAFT_LEN: + if not engine.gdn_replay: raise NoTarget( - f"max_draft_len={draft} with TRTLLM_USE_GDN_REPLAY=0: the legacy per-step recurrent " - f"states are certified at max_draft_len up to {_MAX_LEGACY_DRAFT_LEN}; leave " - "TRTLLM_USE_GDN_REPLAY unset for a longer draft" + "MTP with TRTLLM_USE_GDN_REPLAY=0: the target verifies its Gated-DeltaNet rows on " + "the cached replay only, not on the legacy per-step recurrent states; leave " + "TRTLLM_USE_GDN_REPLAY unset" ) if str(engine.moe_backend).upper() == "CUTEDSL" and draft > _MAX_CUTEDSL_DRAFT_LEN: raise NoTarget( @@ -132,6 +203,28 @@ def _check_engine( ) +def _check_attention_dp(member: str, spec_config: Optional[Any], engine: EngineChoices) -> None: + """Decline what the target does not run under attention DP (dep4, dep2), where + each rank serves its own requests and only the routed experts cross ranks. + + MTP is admitted at dep4, where each rank drafts for its own requests and the draft layer's + FP8 expert window runs the trunk's attention-DP round trip. Routing admits both MTP and + the breakable prefill graph at the attention-TP layouts (tp4, tep4). + """ + if spec_config is not None and member not in _ATTENTION_DP_MTP_LAYOUTS: + raise NoTarget( + f"speculative decoding with attention DP ({member}): no draft acceptance was " + "measured at this layout; MTP under attention DP is admitted at " + f"{', '.join(_ATTENTION_DP_MTP_LAYOUTS)}" + ) + if engine.prefill_cuda_graph_backend == "breakable": + raise NoTarget( + f"prefill_cuda_graph_backend='breakable' with attention DP ({member}): the engine " + "pads every rank to one prefill bucket, so a rank holding only generation requests " + "is handed more rows than its batch, which the target's decode step refuses" + ) + + # Qwen3.8-Flash-Next publishes ``Qwen4ExpForConditionalGeneration``, and the # config loader's composite branch writes the same name back. @routes("Qwen4ExpForConditionalGeneration") @@ -140,8 +233,8 @@ def route( ) -> type: """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits. - The target is imported where it is returned, so declining costs no import - of modeling code. It also declines the engine choices (`EngineChoices`) the + A target is imported where it is returned, so declining costs no import of + modeling code. Routing also declines the engine choices (`EngineChoices`) the target's constructor would refuse. """ pretrained_config = model_config.pretrained_config @@ -193,31 +286,34 @@ def route( and spec_config.max_total_draft_tokens == spec_config.max_draft_len ) - # Qwen3.8-Flash-Next-NVFP4 on one GB300 (sm103), MTP-Eagle or no speculation. - if ( - compute_capability == (10, 3) - and shape == _TEXT_SHAPE - and vshape == _VISION_SHAPE - and mapping.world_size == 1 - and mtp_eagle_or_none - ): - # One line on purpose: the claims suite reads the routed targets off - # `from ..modeling import ` lines. + checkpoint = compute_capability == _SM and shape == _TEXT_SHAPE and vshape == _VISION_SHAPE + # Qwen3.8-Flash-Next-NVFP4 on GB300s (sm103): on one at tp1, on four at tp4, tep4 or dep4, + # MTP-Eagle or no speculation; on two at dep2, without speculation. + member = layout(mapping) + if checkpoint and member is not None and mtp_eagle_or_none: + engine = engine_choices(model_config, llm_args) + if mapping.enable_attention_dp: + _check_attention_dp(member, spec_config, engine) + _check_engine(spec_config, kv_cache_quant_algo, engine, member) + # One line on purpose: the claims suite reads the routed target off the + # `from ..modeling import ` line. # fmt: off # isort: off - from .qwen3_8_flash_next_nvfp4__sm_103__tp1.modeling import ModelingV2Qwen38FlashNextNvfp4Sm103Tp1 + from .qwen3_8_flash_next_nvfp4__sm_103__default.modeling import ModelingV2Qwen38FlashNextNvfp4Sm103Default # isort: on # fmt: on - _check_engine(spec_config, kv_cache_quant_algo, engine_choices(model_config, llm_args)) - return ModelingV2Qwen38FlashNextNvfp4Sm103Tp1 + return ModelingV2Qwen38FlashNextNvfp4Sm103Default raise NoTarget( f"compute capability {compute_capability}, text shape {shape}, vision shape " - f"{vshape}, {mapping.world_size} ranks, speculative decoding " + f"{vshape}, {_describe(mapping)}, speculative decoding " f"{spec_config.spec_dec_mode.name if spec_config is not None else 'off'}" f"{f' max_draft_len={spec_config.max_draft_len}' if spec_config is not None else ''}, " - f"KV cache {kv_cache_quant_algo or 'bf16'}; " - "the one qwen4_exp target is Qwen3.8-Flash-Next-NVFP4 on sm103, tp1, with " - f"one-model MTP-Eagle (max_draft_len 1..{_MAX_DRAFT_LEN}) or no speculative decoding" + f"KV cache {kv_cache_quant_algo or 'bf16'}; the qwen4_exp target is " + "Qwen3.8-Flash-Next-NVFP4 on sm103 in one node: at 1 rank (tp1), and at 4 ranks tp4 " + "(attention TP4, experts split by intermediate), tep4 (attention TP4, experts EP4) or " + f"dep4 (attention DP, experts EP4), with one-model MTP-Eagle (max_draft_len " + f"1..{_MAX_DRAFT_LEN}) or no speculative decoding; at 2 ranks dep2 (attention DP, " + "experts EP2) without speculative decoding" ) diff --git a/tests/integration/defs/.test_durations b/tests/integration/defs/.test_durations index 8f5f9718a185..9f5741a6d2f9 100644 --- a/tests/integration/defs/.test_durations +++ b/tests/integration/defs/.test_durations @@ -20,12 +20,12 @@ "accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp4::test_gsm8k": 1247.1575, "accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp8::test_eagle3_gsm8k_and_acceptance": 998.465, "accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp8::test_gsm8k": 317.73875, - "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_breakable_prefill_matches_eager": 730.6153846153846, - "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_gsm8k": 643.2887001455605, - "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_gsm8k_mmmu_and_video": 969.7372977443608, - "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_cutedsl_fp8_kv_bf16_state": 658.6539230769231, - "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey": 492.2441946022727, - "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_long_context_passkey": 537.4055384615385, + "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_breakable_prefill_matches_eager": 730.6153846153846, + "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_gsm8k": 643.2887001455605, + "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp1_gsm8k_mmmu_and_video": 969.7372977443608, + "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp3_cutedsl_fp8_kv_bf16_state": 658.6539230769231, + "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp3_gsm8k_and_long_context": 492.2441946022727, + "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp7_long_context_passkey": 537.4055384615385, "accuracy/test_disaggregated_serving.py::TestDeepSeekR1::test_kv_cache_v2_nixl_python[cache_mgr_v2]": 606.7294444444445, "accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=0-overlap_scheduler=False]": 348.56913636363635, "accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=0-overlap_scheduler=True]": 207.74877272727272, @@ -1598,20 +1598,6 @@ "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_rank_agreement.py]": 23.462214285714285, "unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py": 23.585939655172414, "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py]": 23.585939655172414, - "unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_kv_quant.py": 29.04516883116883, - "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_kv_quant.py]": 29.04516883116883, - "unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_moe_scalars.py": 25.506426666666666, - "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_moe_scalars.py]": 25.506426666666666, - "unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_multimodal_embeds.py": 25.56098734177215, - "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_multimodal_embeds.py]": 25.56098734177215, - "unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_pool_contract.py": 25.537, - "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_pool_contract.py]": 25.537, - "unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_thop_attention_calls.py": 23.835025, - "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_thop_attention_calls.py]": 23.835025, - "unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_validate_request.py": 25.437897435897437, - "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_validate_request.py]": 25.437897435897437, - "unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_weight_table.py": 32.944, - "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_weight_table.py]": 32.944, "unittest/_torch/modeling_v2/moe/test_modeling_v2_mxe4m3_mxe2m1_block_scale_moe_runner.py": 24.025819444444444, "test_unittests.py::test_unittests_v2[unittest/_torch/modeling_v2/moe/test_modeling_v2_mxe4m3_mxe2m1_block_scale_moe_runner.py]": 24.025819444444444, "unittest/_torch/modeling_v2/moe/test_modeling_v2_noaux_tc_op.py::test_kimi_k3_config": 10.210739463601533, diff --git a/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py b/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py index 11761258ad08..a8edd1cdd25a 100644 --- a/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py +++ b/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py @@ -14,9 +14,16 @@ # limitations under the License. """Whole-model gates for the qwen4_exp modeling_v2 target. -Every test passes ``modeling_v2="require"``: under ``"auto"`` a configuration -that missed the target's routing would fall back to the built-in model, pass, -and report the built-in's numbers as the target's. +``TestModelingV2Qwen38FlashNextNvfp4Sm103Default``: Qwen3.8-Flash-Next-NVFP4 on sm_103 at every +layout the target serves: one rank (tp1, the ``test_tp1_*`` cases), four (tp4, tep4, dep4) and +two (dep2). + +Each test runs one engine (two for the breakable-graph comparison) and checks every feature +that engine carries, so a feature is covered without a test of its own. Every test passes +``modeling_v2="require"``: under ``"auto"`` a configuration that missed the target's routing +would fall back to the built-in model, pass, and report the built-in's numbers as the +target's. The checkpoint is also served disaggregated, the target on both sides of the hop: +one rank per engine with the AgentX engines, and dep2 on both. """ import random @@ -29,7 +36,7 @@ from tensorrt_llm import LLM from tensorrt_llm._torch._experimental.modeling_v2 import MODELING_V2_VALIDATE_ENV -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501 modeling as target_modeling, ) from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner import BreakableCUDAGraphRunner @@ -50,6 +57,7 @@ QSASparseAttentionConfig, SamplingParams, ) +from tensorrt_llm.quantization import QuantAlgo from ...conftest import llm_models_root from ..accuracy_core import ( @@ -59,6 +67,7 @@ assert_acceptance_length, compute_acceptance_length, ) +from ..test_disaggregated_serving import launch_disaggregated_llm, run_accuracy_test from ._env import scoped_env_overrides # The target is certified on sm_103 only. @@ -70,6 +79,32 @@ # engine prepared; the flag has to reach the worker rank. _VALIDATE = {MODELING_V2_VALIDATE_ENV: "1"} +# Every worker of a disaggregated test: the step-contract check, and what a worker spawning +# its second rank over MPI needs in a container running as root (docker/Dockerfile.multi). +_DISAGG_WORKER_ENV = { + **_VALIDATE, + "PRTE_ALLOW_RUN_AS_ROOT": "1", + "PRTE_ALLOW_RUN_AS_ROOT_CONFIRM": "1", +} + + +def _disagg_server_config(ctx_router=None, gen_router=None) -> dict: + """One context and one generation server on this host, with optional router types.""" + + def servers(port, router): + config = {"num_instances": 1, "urls": [f"localhost:{port}"]} + if router is not None: + config["router"] = {"type": router} + return config + + return { + "hostname": "localhost", + "port": 8000, + "backend": "pytorch", + "context_servers": servers(8001, ctx_router), + "generation_servers": servers(8002, gen_router), + } + def _target_core(llm: LLM): """The target's core module. @@ -306,6 +341,21 @@ def _synthetic_image(color: tuple, size: int = 256) -> Image.Image: return image +def _mmmu(llm, model_name: str) -> None: + """MMMU on ``llm`` against the reference row the engine resolves.""" + # MMMU adds <|endoftext|> to the stop token, as the built-in model's MMMU test does. + sampling_params = SamplingParams( + max_tokens=MMMU.MAX_OUTPUT_LEN, + truncate_prompt_tokens=MMMU.MAX_INPUT_LEN, + stop="<|endoftext|>", + ) + MMMU(model_name).evaluate( + llm, + sampling_params=sampling_params, + extra_evaluator_kwargs=dict(chat_template_kwargs=dict(enable_thinking=False)), + ) + + def _media_request(tokenizer, processor, text: str, items: list[tuple[str, object]]) -> dict: """One request carrying ``items``, ``(modality, data)`` pairs in prompt order. @@ -331,11 +381,37 @@ def _media_request(tokenizer, processor, text: str, items: list[tuple[str, objec return dict(prompt=prompt, multi_modal_data=media, mm_item_order=tracker.item_order()) -class TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1(LlmapiAccuracyTestHarness): - """qwen3.8-flash-next-nvfp4 / sm_103 / tp1: text, MTP and multimodal.""" +# The layouts routing admits, as LLM API parallel arguments. Each cuts the weights +# differently and crosses ranks differently, and a wrong cut still produces fluent text, so +# every layout runs one gsm8k against an anchor: tep4, tp4 and dep4 below, dep2 in the +# disaggregated case. +# Attention split by head over four ranks, the routed experts by expert. +_TEP4 = dict(tensor_parallel_size=4, moe_expert_parallel_size=4, enable_attention_dp=False) +# Attention split by head, every expert split by intermediate (160 of 640 rows a rank). +_TP4 = dict( + tensor_parallel_size=4, + moe_tensor_parallel_size=4, + moe_expert_parallel_size=1, + enable_attention_dp=False, +) +# Attention replicated, each rank serving its own requests; the routed experts by expert. +_DEP4 = dict(tensor_parallel_size=4, moe_expert_parallel_size=4, enable_attention_dp=True) + +# The quantized KV pools, one layout each: NVFP4 at tp4 with the breakable prefill graph +# (routed at attention TP only), FP8 at dep4 (MTP under attention DP, each rank drafting for +# its own requests). +_QUANTIZED_KV = [ + pytest.param(_TP4, "nvfp4", True, id="tp4_nvfp4_kv_breakable"), + pytest.param(_DEP4, "fp8", False, id="dep4_fp8_kv"), +] + + +class TestModelingV2Qwen38FlashNextNvfp4Sm103Default(LlmapiAccuracyTestHarness): + """qwen3.8-flash-next-nvfp4 / sm_103 / default: tp1, tp4, tep4, dep4 and dep2.""" # References are keyed by model name plus the quantization the engine - # resolves for this checkpoint (MIXED_PRECISION). + # resolves for this checkpoint (MIXED_PRECISION), and the MTP row by the + # speculative mode as well. MODEL_NAME = "Qwen/Qwen3.8-Flash-Next" MODEL_PATH = f"{llm_models_root()}/Qwen3.8-Flash-Next-NVFP4" @@ -351,17 +427,60 @@ class TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1(LlmapiAccuracyTestHarness): # The anchor's `max_tokens`; the class default (256) truncates this checkpoint's answers. GSM8K_MAX_OUTPUT_LEN = 512 - # `mamba_ssm_cache_dtype` is float32 because the checkpoint declares + # The checkpoint's bf16 K/V, a bf16 recurrent state (the one the target serves), no + # block reuse. + KV_CACHE = KvCacheConfig( + free_gpu_memory_fraction=0.6, + enable_block_reuse=False, + mamba_ssm_cache_dtype="bfloat16", + ) + + # The routed experts run on trtllm-gen, the MoE backend the target implements at every + # layout (CuteDSL is opt-in at tp1). It is also what the draft layer's FP8 block-scale + # experts need: the default CUTLASS MoE does not support them on sm_103. + MOE = MoeConfig(backend="TRTLLM") + + ENGINE = dict( + max_seq_len=9216, + max_batch_size=128, + max_num_tokens=8192, + enable_chunked_prefill=False, + kv_cache_config=KV_CACHE, + cuda_graph_config=CudaGraphConfig(enable_padding=True, max_batch_size=128), + moe_config=MOE, + ) + + # One-model MTP-Eagle with the checkpoint's single MTP module replayed three times, on + # the engine the MTP configuration is served with (a 16384-token sequence bound). + MTP3 = MTPDecodingConfig(max_draft_len=3) + ENGINE_MTP = dict(ENGINE, max_seq_len=16384) + + # The breakable prefill CUDA graph, up to max_num_tokens. + BREAKABLE_PREFILL = dict( + prefill_cuda_graph_backend="breakable", + prefill_capture_num_tokens=[512, 1024, 2048, 4096, 8192], + ) + + # A collapsed or miscomputed draft path accepts nothing and reads 1.0; every draft + # accepted reads 1 + max_draft_len. The midpoint of that range is the gsm8k MTP + # cases' tripwire: at or above it, more than half the drafts are accepted. + MTP3_COLLAPSE_FLOOR = 1.0 + MTP3.max_draft_len / 2 + + # The pass keys' acceptance-length floor in references/acceptance_length.yaml. + ACCEPTANCE_KEY_LONG_CONTEXT = "ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3_long_context" + + # ---- one rank (tp1) ---------------------------------------------------------------- + # At tp1 `mamba_ssm_cache_dtype` is float32 because the checkpoint declares # `mamba_ssm_dtype: float32` (the engine default is bfloat16); the NVFP4- and fp8-KV cases # set it to bfloat16, what `auto` resolves to for this checkpoint. `dtype` is left at the # checkpoint's bf16 KV, which is what the anchors were measured at; only the fp8 and NVFP4 # KV cases set it. - KV_CACHE = KvCacheConfig( + TP1_KV_CACHE = KvCacheConfig( free_gpu_memory_fraction=0.6, enable_block_reuse=True, mamba_ssm_cache_dtype="float32", ) - KV_CACHE_FP8_BF16_STATE = KvCacheConfig( + TP1_KV_CACHE_FP8_BF16_STATE = KvCacheConfig( dtype="fp8", free_gpu_memory_fraction=0.6, enable_block_reuse=True, @@ -369,113 +488,82 @@ class TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1(LlmapiAccuracyTestHarness): ) # Block reuse is on, which (with chunked prefill and the MTP drafts) is what lets the # attention op read an NVFP4 pool: its kernels exist for the paged-context path only. - KV_CACHE_NVFP4_BF16_STATE = KvCacheConfig( + TP1_KV_CACHE_NVFP4_BF16_STATE = KvCacheConfig( dtype="nvfp4", free_gpu_memory_fraction=0.6, enable_block_reuse=True, mamba_ssm_cache_dtype="bfloat16", ) # Block reuse is off for the vision path, which requires it. - KV_CACHE_MULTIMODAL = KvCacheConfig( + TP1_KV_CACHE_MULTIMODAL = KvCacheConfig( free_gpu_memory_fraction=0.6, enable_block_reuse=False, mamba_ssm_cache_dtype="float32", ) - # Required, not a tuning choice: the MTP draft layer's FP8 block-scale - # experts are not supported by the default CUTLASS MoE on sm_103. Set for - # every case so they share one engine configuration. - MOE = MoeConfig(backend="TRTLLM") - - ENGINE = dict( + # One rank: chunked prefill over an 8192-token sequence bound, 32 requests. + TP1_ENGINE = dict( max_seq_len=8192, max_batch_size=32, enable_chunked_prefill=True, - kv_cache_config=KV_CACHE, + kv_cache_config=TP1_KV_CACHE, cuda_graph_config=CudaGraphConfig(), moe_config=MOE, ) - # The breakable prefill CUDA graph is opt-in; these buckets cover the prompts of - # `test_breakable_prefill_matches_eager`. - BREAKABLE_PREFILL = dict( + # These buckets cover the prompts of `test_tp1_breakable_prefill_matches_eager`. + TP1_BREAKABLE_PREFILL = dict( prefill_cuda_graph_backend="breakable", prefill_capture_num_tokens=[128, 256, 512, 1024, 2048, 3072], ) - ENGINE_BREAKABLE = dict(ENGINE, **BREAKABLE_PREFILL) - # One-model MTP-Eagle with the checkpoint's single MTP module replayed - # `max_draft_len` times. + # One draft, and the longest the target runs. MTP1 = MTPDecodingConfig(max_draft_len=1) - MTP3 = MTPDecodingConfig(max_draft_len=3) MTP7 = MTPDecodingConfig(max_draft_len=7) - # Acceptance-length anchors in references/acceptance_length.yaml, keyed by - # target and variant. - ACCEPTANCE_KEY = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1" - ACCEPTANCE_KEY_LONG_CONTEXT = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_long_context" - ACCEPTANCE_KEY_OPT_IN = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_fp8_kv_bf16_state" - ACCEPTANCE_KEY_OPT_IN_LONG_CONTEXT = ( - "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_fp8_kv_bf16_state_long_context" + # The tp1 cases' acceptance-length anchors in references/acceptance_length.yaml. + TP1_ACCEPTANCE_KEY = "ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp1" + TP1_ACCEPTANCE_KEY_LONG_CONTEXT = ( + "ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp3_long_context" ) - ACCEPTANCE_KEY_MTP7_LONG_CONTEXT = ( - "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7_nvfp4_kv_bf16_state_long_context" + TP1_ACCEPTANCE_KEY_OPT_IN = ( + "ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp3_cutedsl_fp8_kv_bf16_state" + ) + TP1_ACCEPTANCE_KEY_OPT_IN_LONG_CONTEXT = ( + "ModelingV2Qwen38FlashNextNvfp4Sm103Default::" + "tp1_mtp3_cutedsl_fp8_kv_bf16_state_long_context" + ) + TP1_ACCEPTANCE_KEY_MTP7_LONG_CONTEXT = ( + "ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp7_nvfp4_kv_bf16_state_long_context" ) - def _gsm8k_and_acceptance(self, mocker, llm, key, label): - """gsm8k on ``llm``, then the same run's acceptance length against ``key``. + def _gsm8k(self, mocker, llm, label: str, key: str | None = None) -> None: + """gsm8k against the engine's reference row, then its acceptance length. - gsm8k is read against the reference row the engine resolves; the acceptance length comes - from the engine's iteration statistics, so the engine needs ``enable_iter_perf_stats``. + The acceptance length is held to the floor under `key` in + references/acceptance_length.yaml, or without one above the three-draft collapse + tripwire. Rejection sampling holds the emitted distribution to the target model's, so + a miscomputed draft layer leaves the score in place and only slows generation; the + acceptance length is what sees it. It comes from the engine's iteration statistics, + so the engine needs ``enable_iter_perf_stats``. """ mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN) GSM8K(self.MODEL_NAME).evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) acceptance_length = compute_acceptance_length(llm) print(f"[AL] {label} acceptance_length = {acceptance_length:.3f}") - assert_acceptance_length(key, acceptance_length) - - def _passkeys(self, llm, lengths, draft_len, key, label): - """Pass key retrieval on ``llm`` over prompts past the indexer's budget. + if key is not None: + assert_acceptance_length(key, acceptance_length) + return + assert acceptance_length >= self.MTP3_COLLAPSE_FLOOR, ( + f"{label}: acceptance length {acceptance_length:.3f} is below " + f"{self.MTP3_COLLAPSE_FLOOR}: the draft path is rejected more often than accepted" + ) - Every key must be stated back, and two identical requests far apart in one batch - (different slots) must give the same tokens, which fails if a request's selection leaks - into another slot. The acceptance length, from the requests' own metrics, is held to the - floor under ``key``. - """ + def _passkeys(self, llm, lengths, draft_len, key, label) -> None: + """Pass keys over prompts of `lengths` tokens, the acceptance length held to `key`.""" tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) prompts = _passkey_prompts(tokenizer, lengths) - texts = [prompt for prompt, _ in prompts] + [prompts[len(prompts) // 2][0]] - outputs = llm.generate( - texts, - SamplingParams(max_tokens=_PASSKEY_MAX_TOKENS, temperature=0, return_perf_metrics=True), - use_tqdm=False, - ) - missed = [ - k - for (_, k), output in zip(prompts, outputs) - if k not in output.outputs[0].text[:_PASSKEY_MAX_TOKENS] - ] - assert not missed, f"pass keys not stated back: {missed} of {[k for _, k in prompts]}" - twins = outputs[len(prompts) // 2], outputs[-1] - assert twins[0].outputs[0].token_ids == twins[1].outputs[0].token_ids - acceptance_length = _acceptance_length_from_requests(outputs, draft_len) - print(f"[AL] {label} acceptance_length = {acceptance_length:.3f}") - assert_acceptance_length(key, acceptance_length) - - def _mmmu(self, llm): - """MMMU on ``llm`` against the reference row the engine resolves.""" - # MMMU adds <|endoftext|> to the stop token, as the built-in model's - # MMMU test does. - sampling_params = SamplingParams( - max_tokens=MMMU.MAX_OUTPUT_LEN, - truncate_prompt_tokens=MMMU.MAX_INPUT_LEN, - stop="<|endoftext|>", - ) - MMMU(self.MODEL_NAME).evaluate( - llm, - sampling_params=sampling_params, - extra_evaluator_kwargs=dict(chat_template_kwargs=dict(enable_thinking=False)), - ) + assert_acceptance_length(key, self._passkey_acceptance(llm, prompts, label, draft_len)) def _video_checks(self, llm): """Synthetic videos through ``llm``. @@ -488,7 +576,7 @@ def _video_checks(self, llm): bound (a 3072x3072 image is 192x192 = 36864 patches) must fail as that request's own error, and the request after it must still complete; ``max_seq_len`` is raised so that image reaches the admission check rather than the sequence length check. The order of - the splice itself is pinned by ``test_modeling_v2_qwen4_exp_multimodal_embeds``. + the splice itself is pinned by the target's ``test_vision``. """ tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) processor = AutoProcessor.from_pretrained( @@ -522,8 +610,234 @@ def ask(text, items): answer = ask(_VIDEO_QUESTION, [("video", _synthetic_video("red_rightward"))]).lower() assert "red" in answer and "right" in answer, answer + def _long_context_passkey( + self, monkeypatch, speculative_config, engine, index_share, lengths, key, label, mocker=None + ): + """`_passkeys` with the executor in this process, checking the op calls too. + + With ``mocker``, gsm8k runs first on the same engine. The draft layer's IndexShare + reuse op must run exactly when IndexShare is on, and no prefill graph may replay. + """ + monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") + reuse_calls = [] + reuse = target_modeling.qsa_mtp_reuse_topk + + def counting_reuse(*args, **kwargs): + reuse_calls.append(1) + return reuse(*args, **kwargs) + + monkeypatch.setattr(target_modeling, "qsa_mtp_reuse_topk", counting_reuse) + replays = _count_prefill_graph_replays(monkeypatch) + engine = dict( + engine, + sparse_attention_config=QSASparseAttentionConfig( + index_share_for_mtp_iteration=index_share + ), + ) + with ( + scoped_env_overrides(_VALIDATE) as rank_env, + LLM( + self.MODEL_PATH, + speculative_config=speculative_config, + **engine, + modeling_v2="require", + env_overrides=rank_env, + ) as llm, + ): + assert _target_core(llm).mtp_index_share == index_share + if mocker is not None: + mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN) + GSM8K(self.MODEL_NAME).evaluate( + llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS + ) + self._passkeys(llm, lengths, speculative_config.max_draft_len, key, label) + assert not replays, replays + if index_share: + assert reuse_calls, "IndexShare is on and the reuse op never ran" + else: + assert not reuse_calls, "IndexShare is off and the reuse op ran" + + def _passkey_acceptance(self, llm, prompts, label, draft_len: int = 3) -> float: + """Pass-key retrieval on ``llm`` and the requests' mean acceptance length. + + Every key must be stated back, and two identical requests far apart in one batch + (different slots) must give the same tokens, which fails if a request's selection + leaks into another slot. + """ + texts = [prompt for prompt, _ in prompts] + [prompts[len(prompts) // 2][0]] + outputs = llm.generate( + texts, + SamplingParams(max_tokens=_PASSKEY_MAX_TOKENS, temperature=0, return_perf_metrics=True), + use_tqdm=False, + ) + missed = [ + key + for (_, key), output in zip(prompts, outputs) + if key not in output.outputs[0].text[:_PASSKEY_MAX_TOKENS] + ] + assert not missed, ( + f"{label}: pass keys not stated back: {missed} of {[key for _, key in prompts]}" + ) + twins = outputs[len(prompts) // 2], outputs[-1] + assert twins[0].outputs[0].token_ids == twins[1].outputs[0].token_ids, ( + f"{label}: the same prompt in two slots gave different tokens" + ) + acceptance_length = _acceptance_length_from_requests(outputs, draft_len) + print(f"[AL] {label} acceptance_length = {acceptance_length:.3f}") + return acceptance_length + + @skip_not_sm103 + @pytest.mark.skip_less_device(4) + def test_mtp3_gsm8k_and_long_context(self, mocker): + """tep4 with three drafts and IndexShare: gsm8k, then pass keys over 3k-6k tokens. + + gsm8k reads the MTP row of references/gsm8k.yaml, its acceptance length held above + the collapse tripwire. The pass keys are past the QSA indexer's 2048-token budget, so + the target's decode steps and its draft layer's take the sparse attention path, which + gsm8k never leaves, and IndexShare makes the draft loop reuse the selection of its + first step. A miscomputed draft path keeps the text correct and only lowers the + acceptance length, which is held to its floor in references/acceptance_length.yaml. + """ + tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) + prompts = _passkey_prompts(tokenizer, _PASSKEY_LENGTHS) + with ( + scoped_env_overrides(_VALIDATE) as rank_env, + LLM( + self.MODEL_PATH, + speculative_config=self.MTP3, + enable_iter_perf_stats=True, + sparse_attention_config=QSASparseAttentionConfig( + index_share_for_mtp_iteration=True + ), + **self.ENGINE_MTP, + **_TEP4, + modeling_v2="require", + env_overrides=rank_env, + ) as llm, + ): + self._gsm8k(mocker, llm, "tep4 mtp3 gsm8k") + acceptance_length = self._passkey_acceptance(llm, prompts, "tep4 mtp3 long context") + assert_acceptance_length(self.ACCEPTANCE_KEY_LONG_CONTEXT, acceptance_length) + @skip_not_sm103 - def test_gsm8k(self, mocker): + @pytest.mark.skip_less_device(4) + @pytest.mark.parametrize("parallel, kv_dtype, breakable", _QUANTIZED_KV) + def test_mtp3_gsm8k_quantized_kv(self, mocker, parallel, kv_dtype, breakable): + """gsm8k with three drafts over a quantized KV pool, block reuse and chunked prefill. + + gsm8k reads the MTP row of references/gsm8k.yaml for that KV dtype, its acceptance + length held above the collapse tripwire. + """ + kv_cache_config = KvCacheConfig( + dtype=kv_dtype, + free_gpu_memory_fraction=0.6, + enable_block_reuse=True, + mamba_ssm_cache_dtype="bfloat16", + ) + engine = dict(self.ENGINE_MTP, enable_chunked_prefill=True, kv_cache_config=kv_cache_config) + if breakable: + engine.update(self.BREAKABLE_PREFILL) + with ( + scoped_env_overrides(_VALIDATE) as rank_env, + LLM( + self.MODEL_PATH, + speculative_config=self.MTP3, + enable_iter_perf_stats=True, + **engine, + **parallel, + modeling_v2="require", + env_overrides=rank_env, + ) as llm, + ): + self._gsm8k(mocker, llm, f"mtp3 {kv_dtype} kv") + + @skip_not_sm103 + @pytest.mark.skip_less_device(4) + def test_mtp3_mmmu(self): + """MMMU at tep4 with three drafts, against the MTP row of references/mmmu.yaml. + + It runs the vision tower, replicated on every rank, and the mrope path. + """ + with ( + scoped_env_overrides(_VALIDATE) as rank_env, + LLM( + self.MODEL_PATH, + speculative_config=self.MTP3, + **self.ENGINE_MTP, + **_TEP4, + modeling_v2="require", + env_overrides=rank_env, + ) as llm, + ): + _mmmu(llm, self.MODEL_NAME) + + @skip_not_sm103 + @pytest.mark.timeout(5400) + @pytest.mark.skip_less_device(4) + @pytest.mark.skip_less_device_memory(145000) + @pytest.mark.skip_less_host_memory(204800) + def test_disaggregated_dep2(self, mocker): + """dep2 on both engines, one context and one generation server of two GPUs each. + + Attention DP over two ranks, the routed experts in two 256-expert windows; both + engines run one layout, so each rank's K/V, QSA index, Gated-DeltaNet and PLE pools + have one layout on both sides of the Python NIXL transceiver. + """ + # The core configuration on both engines: the checkpoint's bf16 K/V, a bf16 + # recurrent state, no block reuse, no chunked prefill, the trtllm-gen experts. + common_config = { + "modeling_v2": "require", + "tensor_parallel_size": 2, + "moe_expert_parallel_size": 2, + "enable_attention_dp": True, + "max_seq_len": 9216, + "max_batch_size": 128, + "max_num_tokens": 8192, + "enable_chunked_prefill": False, + "moe_config": {"backend": "TRTLLM"}, + "kv_cache_config": { + "enable_block_reuse": False, + "mamba_ssm_cache_dtype": "bfloat16", + "free_gpu_memory_fraction": 0.6, + }, + # The V2 hybrid manager's pools travel over the Python transceiver. + "cache_transceiver_config": {"backend": "NIXL", "transceiver_runtime": "PYTHON"}, + } + ctx_server_config = { + **common_config, + "disable_overlap_scheduler": True, + "cuda_graph_config": None, + } + gen_server_config = { + **common_config, + "disable_overlap_scheduler": False, + "cuda_graph_config": {"enable_padding": True, "max_batch_size": 128}, + } + + mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN) + with launch_disaggregated_llm( + _disagg_server_config(), + ctx_server_config, + gen_server_config, + self.MODEL_PATH, + max_workers=128, + extra_env=_DISAGG_WORKER_ENV, + # What every rank of both engines prints once it has bound the manager's pools. + assert_worker_log_contains="modeling_v2 qwen4_exp default target: rank", + ) as llm: + # The served quantization the reference is keyed by; the launcher reads NVFP4 + # off the checkpoint's name. + llm.args.quant_config.quant_algo = QuantAlgo.MIXED_PRECISION + run_accuracy_test( + llm, + self.MODEL_NAME, + ["GSM8K"], + extra_evaluator_kwargs={GSM8K: self.EXTRA_EVALUATOR_KWARGS}, + timeout=3000, + ) + + @skip_not_sm103 + def test_tp1_gsm8k(self, mocker): """The text model without speculation, against the gsm8k anchor.""" mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN) @@ -531,7 +845,7 @@ def test_gsm8k(self, mocker): scoped_env_overrides(_VALIDATE) as rank_env, LLM( self.MODEL_PATH, - **self.ENGINE, + **self.TP1_ENGINE, modeling_v2="require", env_overrides=rank_env, ) as llm, @@ -540,7 +854,7 @@ def test_gsm8k(self, mocker): task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) @skip_not_sm103 - def test_mtp1_gsm8k_mmmu_and_video(self, mocker): + def test_tp1_mtp1_gsm8k_mmmu_and_video(self, mocker): """gsm8k with its acceptance-length floor, MMMU and synthetic videos on one MTP engine. - gsm8k against the MTP row of references/gsm8k.yaml. A miscomputed draft layer only @@ -553,7 +867,9 @@ def test_mtp1_gsm8k_mmmu_and_video(self, mocker): Block reuse is off, as the vision path requires; it does not change the text numerics. ``max_seq_len`` is raised for the video checks' oversized image. """ - engine = dict(self.ENGINE, kv_cache_config=self.KV_CACHE_MULTIMODAL, max_seq_len=16384) + engine = dict( + self.TP1_ENGINE, kv_cache_config=self.TP1_KV_CACHE_MULTIMODAL, max_seq_len=16384 + ) with ( scoped_env_overrides(_VALIDATE) as rank_env, LLM( @@ -565,12 +881,12 @@ def test_mtp1_gsm8k_mmmu_and_video(self, mocker): env_overrides=rank_env, ) as llm, ): - self._gsm8k_and_acceptance(mocker, llm, self.ACCEPTANCE_KEY, "mtp1") - self._mmmu(llm) + self._gsm8k(mocker, llm, "tp1 mtp1", self.TP1_ACCEPTANCE_KEY) + _mmmu(llm, self.MODEL_NAME) self._video_checks(llm) @skip_not_sm103 - def test_mtp3_cutedsl_fp8_kv_bf16_state(self, mocker): + def test_tp1_mtp3_cutedsl_fp8_kv_bf16_state(self, mocker): """gsm8k and 3k-6k-token pass keys with three drafts over the opt-in choices. One engine runs gsm8k with an acceptance-length floor, then the pass keys. The trunk's @@ -585,9 +901,9 @@ def test_mtp3_cutedsl_fp8_kv_bf16_state(self, mocker): cannot use, which accuracy does not. """ engine = dict( - self.ENGINE, + self.TP1_ENGINE, moe_config=MoeConfig(backend="CUTEDSL"), - kv_cache_config=self.KV_CACHE_FP8_BF16_STATE, + kv_cache_config=self.TP1_KV_CACHE_FP8_BF16_STATE, sparse_attention_config=QSASparseAttentionConfig(index_share_for_mtp_iteration=False), ) with ( @@ -601,81 +917,41 @@ def test_mtp3_cutedsl_fp8_kv_bf16_state(self, mocker): env_overrides=rank_env, ) as llm, ): - self._gsm8k_and_acceptance(mocker, llm, self.ACCEPTANCE_KEY_OPT_IN, "mtp3 opt-in") + self._gsm8k(mocker, llm, "tp1 mtp3 opt-in", self.TP1_ACCEPTANCE_KEY_OPT_IN) self._passkeys( llm, _PASSKEY_LENGTHS, self.MTP3.max_draft_len, - self.ACCEPTANCE_KEY_OPT_IN_LONG_CONTEXT, - "mtp3 opt-in long context", + self.TP1_ACCEPTANCE_KEY_OPT_IN_LONG_CONTEXT, + "tp1 mtp3 opt-in long context", ) - def _long_context_passkey( - self, monkeypatch, speculative_config, engine, index_share, lengths, key, label - ): - """`_passkeys` with the executor in this process, checking the op calls too. - - The draft layer's IndexShare reuse op must run exactly when IndexShare is on, and no - prefill graph may replay. - """ - monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") - reuse_calls = [] - reuse = target_modeling.qsa_mtp_reuse_topk - - def counting_reuse(*args, **kwargs): - reuse_calls.append(1) - return reuse(*args, **kwargs) - - monkeypatch.setattr(target_modeling, "qsa_mtp_reuse_topk", counting_reuse) - replays = _count_prefill_graph_replays(monkeypatch) - engine = dict( - engine, - sparse_attention_config=QSASparseAttentionConfig( - index_share_for_mtp_iteration=index_share - ), - ) - with ( - scoped_env_overrides(_VALIDATE) as rank_env, - LLM( - self.MODEL_PATH, - speculative_config=speculative_config, - **engine, - modeling_v2="require", - env_overrides=rank_env, - ) as llm, - ): - assert _target_core(llm).mtp_index_share == index_share - self._passkeys(llm, lengths, speculative_config.max_draft_len, key, label) - assert not replays, replays - if index_share: - assert reuse_calls, "IndexShare is on and the reuse op never ran" - else: - assert not reuse_calls, "IndexShare is off and the reuse op ran" - @skip_not_sm103 - def test_mtp3_long_context_passkey(self, monkeypatch): - """Pass key retrieval over 3k-6k-token prompts with three drafts and IndexShare. - - The KV cache is bf16; the prompts are long enough to run the sparse attention path. The - gsm8k prompts never leave the dense path (they are under the indexer's 2048-token - budget), so the pass-key cases are the ones that run the target's QSA selection and the - draft layer's end to end; a pass key stated back correctly needs the right blocks - selected at every depth. IndexShare makes the draft loop reuse the selection of its - first step: the acceptance length must not fall below the floor taken from the same - prompts with the draft layer selecting on every step, and the reuse op must have run. + def test_tp1_mtp3_gsm8k_and_long_context(self, mocker, monkeypatch): + """gsm8k, then 3k-6k-token pass keys, with three drafts and IndexShare on one engine. + + The KV cache is bf16, the recurrent state fp32, with block reuse and chunked prefill. + gsm8k reads the MTP row of references/gsm8k.yaml; its prompts stay on the dense + attention path (they are under the indexer's 2048-token budget). The pass keys are + past it, so they run the target's QSA selection and the draft layer's end to end; a + pass key stated back correctly needs the right blocks selected at every depth. + IndexShare makes the draft loop reuse the selection of its first step: the pass keys' + acceptance length must not fall below the floor taken from the same prompts with the + draft layer selecting on every step, and the reuse op must have run. """ self._long_context_passkey( monkeypatch, self.MTP3, - self.ENGINE, + self.TP1_ENGINE, True, _PASSKEY_LENGTHS, - self.ACCEPTANCE_KEY_LONG_CONTEXT, - "mtp3 long context", + self.TP1_ACCEPTANCE_KEY_LONG_CONTEXT, + "tp1 mtp3 long context", + mocker=mocker, ) @skip_not_sm103 - def test_breakable_prefill_matches_eager(self, monkeypatch): + def test_tp1_breakable_prefill_matches_eager(self, monkeypatch): """Greedy continuations of the same prompts, with the breakable prefill graph and without. MTP with three drafts, a 3072-token chunk budget, and prompts of a natural document @@ -748,13 +1024,15 @@ def run(engine): long_tokens = long_future.result().outputs[0].token_ids return single, again, answers, long_tokens - single, again, answers, long_tokens = run(dict(self.ENGINE_BREAKABLE, max_num_tokens=3072)) + single, again, answers, long_tokens = run( + dict(dict(self.TP1_ENGINE, **self.TP1_BREAKABLE_PREFILL), max_num_tokens=3072) + ) assert {128, 256, 1024, 2048, 3072} <= set(replays), replays assert single == again, "the same engine repeated itself with a different answer" replays.clear() eager_single, _, eager_answers, eager_long_tokens = run( - dict(self.ENGINE, max_num_tokens=3072) + dict(self.TP1_ENGINE, max_num_tokens=3072) ) assert not replays, replays @@ -771,7 +1049,7 @@ def run(engine): assert long_tokens[:4] == eager_long_tokens[:4], (long_tokens, eager_long_tokens) @skip_not_sm103 - def test_mtp7_long_context_passkey(self, monkeypatch): + def test_tp1_mtp7_long_context_passkey(self, monkeypatch): """Pass key retrieval over 3k-12k-token prompts with seven drafts and IndexShare. The KV cache is NVFP4 and the recurrent state bf16. Seven drafts run the sparse path @@ -785,7 +1063,7 @@ def test_mtp7_long_context_passkey(self, monkeypatch): tokens bound that drift. 12k tokens past the 8k default needs ``max_seq_len`` 16384. """ engine = dict( - self.ENGINE, max_seq_len=16384, kv_cache_config=self.KV_CACHE_NVFP4_BF16_STATE + self.TP1_ENGINE, max_seq_len=16384, kv_cache_config=self.TP1_KV_CACHE_NVFP4_BF16_STATE ) self._long_context_passkey( monkeypatch, @@ -793,6 +1071,112 @@ def test_mtp7_long_context_passkey(self, monkeypatch): engine, True, _PASSKEY_LENGTHS_LONG, - self.ACCEPTANCE_KEY_MTP7_LONG_CONTEXT, - "mtp7 nvfp4 kv bf16 state long context", + self.TP1_ACCEPTANCE_KEY_MTP7_LONG_CONTEXT, + "tp1 mtp7 nvfp4 kv bf16 state long context", ) + + # The target's AgentX deployment environment at tp1 on both engines: PDL, eight K/V + # transfer threads and every opt-in projection of the target. + AGENTX_ENV = { + "TRTLLM_ENABLE_PDL": "1", + "TRTLLM_KV_TRANSFER_NUM_THREADS": "8", + "TRTLLM_QWEN4_EXP_PLE_HOST_OFFLOAD": "1", + "TRTLLM_MODELING_V2_DENSE_FP8": "1", + "TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4": "1", + "TRTLLM_MODELING_V2_MTP_MOE_NVFP4": "1", + "TRTLLM_MODELING_V2_MTP_HEAD_NVFP4": "1", + "TRTLLM_MODELING_V2_TARGET_HEAD_NVFP4": "1", + "TRTLLM_MODELING_V2_MTP_DENSE_FP8": "1", + } + + @skip_not_sm103 + @pytest.mark.timeout(5400) + @pytest.mark.skip_less_device(2) + @pytest.mark.skip_less_device_memory(145000) + @pytest.mark.skip_less_host_memory(204800) + def test_tp1_disaggregated_mtp2(self, mocker): + """One context and one generation GPU, the engines of the target's AgentX deployment. + + One rank each, so each runs the target at tp1. + + NVFP4 K/V, a bf16 recurrent state, MTP with two drafts, per-conversation block reuse + on the context engine and rejection sampling on the generation engine, over the + Python NIXL transceiver. Acceptance is real. + """ + common_config = { + "modeling_v2": "require", + "max_seq_len": 262144, + "max_num_tokens": 8192, + "enable_chunked_prefill": True, + "disable_mm_encoder": True, + "moe_config": {"backend": "TRTLLM"}, + "scheduler_config": {"capacity_scheduler_policy": "GUARANTEED_NO_EVICT"}, + "cache_transceiver_config": {"backend": "NIXL", "transceiver_runtime": "PYTHON"}, + } + kv_cache_config = { + "dtype": "nvfp4", + "free_gpu_memory_fraction": 0.85, + "mamba_ssm_cache_dtype": "bfloat16", + "use_kv_cache_manager_v2": True, + "tokens_per_block": 32, + } + # The context engine's breakable prefill graphs, up to max_num_tokens. + prefill_capture_num_tokens = ( + [2**i for i in range(10)] + list(range(768, 3328, 256)) + [4096, 6144, 8192] + ) + ctx_server_config = { + **common_config, + "max_batch_size": 32, + "kv_cache_config": { + **kv_cache_config, + "enable_block_reuse": True, + "block_reuse_config": {"policy": "per_conversation", "max_num_turns": 3}, + # One recurrent-state snapshot two tokens before each prompt's end. + "mamba_state_config": { + "periodic_snapshot_interval": 0, + "additional_snapshot_offsets_from_end": [2], + }, + }, + "cuda_graph_config": {"enable_padding": True, "max_batch_size": 32}, + "prefill_cuda_graph_backend": "breakable", + "prefill_capture_num_tokens": prefill_capture_num_tokens, + "speculative_config": {"decoding_type": "MTP", "max_draft_len": 2}, + } + gen_server_config = { + **common_config, + "max_batch_size": 128, + "kv_cache_config": {**kv_cache_config, "enable_block_reuse": False}, + "cuda_graph_config": { + "enable_padding": True, + "batch_sizes": list(range(1, 17)) + list(range(24, 129, 8)), + }, + "speculative_config": { + "decoding_type": "MTP", + "max_draft_len": 2, + "use_rejection_sampling": True, + }, + } + + mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN) + with launch_disaggregated_llm( + _disagg_server_config(ctx_router="conversation", gen_router="load_balancing"), + ctx_server_config, + gen_server_config, + self.MODEL_PATH, + max_workers=128, + extra_env={**_DISAGG_WORKER_ENV, **self.AGENTX_ENV}, + # What the target prints once it has packed the target head for decode rows. + assert_worker_log_contains="target head NVFP4", + ) as llm: + # The served quantization and draft the reference is keyed by; the launcher + # reads NVFP4 off the checkpoint's name. + llm.args.quant_config.quant_algo = QuantAlgo.MIXED_PRECISION + llm.args.quant_config.kv_cache_quant_algo = QuantAlgo.NVFP4 + llm.args.speculative_config = MTPDecodingConfig(max_draft_len=2) + run_accuracy_test( + llm, + self.MODEL_NAME, + ["GSM8K"], + extra_evaluator_kwargs={GSM8K: self.EXTRA_EVALUATOR_KWARGS}, + timeout=3000, + ) diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml index fc3eaad3ac9b..308420d9a5aa 100644 --- a/tests/integration/defs/accuracy/references/acceptance_length.yaml +++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml @@ -153,8 +153,9 @@ ModelingV2Glm52Nvfp4Sm103Tep4::mtp1: # self-referential; min_al is the automatic 95% of it. ref_al: 1.942036 min_al: 1.844935 -# Read by the acceptance-length cases of TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1 -# (test_qwen4_exp.py). Keyed by target and variant for the same reason as the +# Read by the tp1 acceptance-length cases of TestModelingV2Qwen38FlashNextNvfp4Sm103Default +# (test_qwen4_exp.py, the `test_tp1_*` cases: the target at one rank). Keyed by target, layout +# and variant for the same reason as the # ModelingV2DeepseekR10528Nvfp4Sm10xDefault::mtp3 entry above, and one key per variant: # TRTLLM_POPULATE_ACCEPTANCE_LENGTH=1 on one case must not overwrite another's floor. # @@ -165,8 +166,8 @@ ModelingV2Glm52Nvfp4Sm103Tep4::mtp1: # every draft accepted. Each min_al below is a hand-set collapse tripwire unless its # comment says it is 95% of ref_al. -# test_mtp1_gsm8k_mmmu_and_video, over gsm8k -ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1: +# test_tp1_mtp1_gsm8k_mmmu_and_video, over gsm8k +ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp1: # Smoke-scale, and labelled as such: measured on GB300 over the boot prompt # pair, STOCK 1.7403 against modeling_v2's 1.8063 on the same prompts. The # number below is the STOCK reading, so the gate is not self-referential. @@ -184,9 +185,9 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1: # measured, and tighten min_al then. min_al: 1.5 -# test_mtp3_cutedsl_fp8_kv_bf16_state, over gsm8k: three drafts with the trunk's experts on +# test_tp1_mtp3_cutedsl_fp8_kv_bf16_state, over gsm8k: three drafts with the trunk's experts on # CuteDSL over an fp8 KV cache and a bf16 recurrent state. -ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_fp8_kv_bf16_state: +ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp3_cutedsl_fp8_kv_bf16_state: # modeling_v2's own full-gsm8k reading on GB300 (1319 samples); there is no built-in reading # for this combination (the built-in model has no CuteDSL trunk). 3.705 over bf16 KV. ref_al: 3.509 @@ -194,11 +195,11 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_fp8_kv_bf16_state: # accepted fails it. min_al: 3.334 -# test_mtp3_long_context_passkey: three drafts with IndexShare on, bf16 KV, over -# 13 prompts of 3k-6k tokens (the sparse attention path). Most of the output is a copy of the -# document's section headings, which the draft predicts almost always, so the readings sit -# at the range's ceiling of 4.0. -ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_long_context: +# test_tp1_mtp3_gsm8k_and_long_context, over its pass keys: three drafts with IndexShare on, bf16 +# KV, over 13 prompts of 3k-6k tokens (the sparse attention path). Most of the output is a copy +# of the document's section headings, which the draft predicts almost always, so the readings +# sit at the range's ceiling of 4.0. +ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp3_long_context: # The same prompts with the draft layer selecting on every step (IndexShare off): 4.000, # equal to the IndexShare reading. ref_al: 4.000 @@ -207,20 +208,32 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_long_context: # 3.760, a 6% drop the automatic floor would pass. min_al: 3.9 -# test_mtp3_cutedsl_fp8_kv_bf16_state, over the pass keys of mtp3_long_context with IndexShare -# off. -ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp3_cutedsl_fp8_kv_bf16_state_long_context: +# test_tp1_mtp3_cutedsl_fp8_kv_bf16_state, over the pass keys of tp1_mtp3_long_context with +# IndexShare off. +ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp3_cutedsl_fp8_kv_bf16_state_long_context: # 3.962 here, against 4.000 over bf16 KV. ref_al: 3.962 - # 2.5% below ref_al, as for mtp3_long_context. + # 2.5% below ref_al, as for tp1_mtp3_long_context. min_al: 3.86 -# test_mtp7_long_context_passkey: seven drafts with IndexShare on, over an NVFP4 KV cache and a +# test_tp1_mtp7_long_context_passkey: seven drafts with IndexShare on, over an NVFP4 KV cache and a # bf16 recurrent state, over 12 prompts of 3k-12k tokens (the sparse attention path), whose # output mostly copies the document's section headings. The KV dtype sets the reading, and seven # chained drafts amplify it: 7.857 over bf16 KV, 7.500 over fp8, 6.867 over NVFP4 (6.961 with # IndexShare off, 7.003 over an fp32 state). -ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp7_nvfp4_kv_bf16_state_long_context: +ModelingV2Qwen38FlashNextNvfp4Sm103Default::tp1_mtp7_nvfp4_kv_bf16_state_long_context: ref_al: 6.867 - # 2.5% below ref_al, as for mtp3_long_context. + # 2.5% below ref_al, as for tp1_mtp3_long_context. min_al: 6.69 + +# Read by TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_and_long_context +# (test_qwen4_exp.py), over the pass keys of tp1_mtp3_long_context at tep4: three drafts with +# IndexShare on, bf16 KV. +ModelingV2Qwen38FlashNextNvfp4Sm103Default::mtp3_long_context: + # modeling_v2's own reading on GB300, at the range's ceiling as at tp1. The + # built-in implementation does not start this engine (its decode front's router GEMM + # asserts a power-of-two K), so there is no stock reading. + ref_al: 4.000 + # 2.5% below ref_al, as for tp1_mtp3_long_context, whose deliberately wrong + # reused selection scores 3.760. + min_al: 3.9 diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 6ee3a6d34211..a25cf5e0dc0b 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -778,6 +778,22 @@ test_e2e.py::test_relaxed_acceptance_quickstart_advanced_deepseek_r1_8gpus[DeepS test_e2e.py::test_trtllm_benchmark_serving[gpt_oss/gpt-oss-20b] test_e2e.py::test_trtllm_multimodal_benchmark_serving +# Qwen3.8-Flash-Next-NVFP4, modeling_v2 qwen4_exp target (sm_103 only): every accuracy test of +# the target, the CI ones included, and those that run only here (tp1 without speculation, +# tp1's breakable prefill graph against eager, mmmu at tep4). +accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_gsm8k +accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp3_gsm8k_and_long_context +accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp1_gsm8k_mmmu_and_video +accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp3_cutedsl_fp8_kv_bf16_state +accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp7_long_context_passkey +accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_breakable_prefill_matches_eager +accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_disaggregated_mtp2 TIMEOUT (90) +accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_and_long_context TIMEOUT (60) +accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[tp4_nvfp4_kv_breakable] TIMEOUT (60) +accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[dep4_fp8_kv] TIMEOUT (60) +accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_mmmu TIMEOUT (90) +accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_disaggregated_dep2 TIMEOUT (90) + # Dense FP4 MLA accuracy (SM107 only) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_mla accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_nvfp4_kv_cache_nixl diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 044a6b612c56..b29f7cc9fd62 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -203,14 +203,12 @@ l0_b300: - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_rms_norm_gated_token_major.py - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_modeling_v2_situ_and_mul.py - accuracy/modeling_v2/test_gpt_oss.py::TestModelingV2GptOss120bSm10xDefault::test_gsm8k[tp1] - # qwen4_exp target, pre-merge: gsm8k; one MTP engine for gsm8k with the acceptance-length - # floor (the only gate that sees a miscomputed draft layer), mmmu (the vision tower) and - # synthetic videos; and 3k-6k-token pass keys with three drafts and IndexShare over bf16 KV, - # the only prompts long enough to run the sparse attention path. The opt-in variants run - # post-merge (below). - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_gsm8k - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_gsm8k_mmmu_and_video - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_long_context_passkey + # qwen4_exp target at tp1, pre-merge: one engine with three drafts and IndexShare over bf16 KV, + # gsm8k and then 3k-6k-token pass keys, the only prompts long enough to run the sparse + # attention path, with their acceptance-length floor (what sees a miscomputed draft layer). + # The other engines run post-merge (below); no speculation and the breakable prefill graph's + # comparison with eager run in QA (llm_function_core). + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp3_gsm8k_and_long_context - condition: ranges: system_gpu_count: @@ -231,11 +229,11 @@ l0_b300: - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM" - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "CUTEDSL" - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "DEEPGEMM" - # qwen4_exp target, post-merge: the opt-in variants -- three drafts over CuteDSL experts, fp8 - # KV and a bf16 recurrent state in one engine (gsm8k and the pass keys); seven drafts over - # 3k-12k-token pass keys on NVFP4 KV and a bf16 state (the verify step's eight rows per - # request); and the breakable prefill CUDA graph replayed against eager greedy continuations - # at the bucket edges and with context and generation rows in one batch. - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp3_cutedsl_fp8_kv_bf16_state - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp7_long_context_passkey - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_breakable_prefill_matches_eager + # qwen4_exp target at tp1, post-merge: one draft with mmmu (the vision tower) and synthetic + # videos; three drafts over CuteDSL experts, fp8 KV and a bf16 recurrent state (gsm8k and + # the pass keys); seven drafts over 3k-12k-token pass keys on NVFP4 KV and a bf16 state (the + # verify step's eight rows per request). Its disaggregated engines run in + # l0_gb300_multi_gpus. + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp1_gsm8k_mmmu_and_video + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp3_cutedsl_fp8_kv_bf16_state + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_mtp7_long_context_passkey diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 964a02bd2295..f57116fa170b 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -57,6 +57,9 @@ l0_cpu: # GLM routing and draft workspace setup use CPU metadata and mocked allocations. - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/spec/test_modeling_v2_draft_exchange.py - unittest/_torch/modeling_v2/models/glm_moe_dsa/test_routing.py + # The qwen4_exp default target's cpu_only tests (its routing among them); this stage runs + # -m cpu_only, so the directory's GPU tests are left to l0_b300's modeling_v2 entry. + - unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default - unittest/_torch/pyexecutor - unittest/_torch/test_connector.py::test_a_partial_grouped_scheduler_is_rejected_at_bring_up - unittest/_torch/test_connector.py::test_a_grouped_scheduler_and_a_single_group_cache_are_accepted @@ -68,11 +71,8 @@ l0_cpu: - unittest/_torch/modeling # modeling_v2 routing and routing-table claims: config shapes and a monkeypatched # compute capability, no device. The default targets' shard tests: small CPU - # tensors and the CPU block_scale_interleave. The qwen4_exp entries are the target's - # tests that run on the source text, stubbed weights and meta tensors: its multimodal - # splice and admission check on a stub tower, the KV pool contract, its KV storage - # derivation and attention-call arguments, its routed-expert scalars and its weight table. - # The rest of unittest/_torch/modeling_v2 needs a GPU. + # tensors and the CPU block_scale_interleave. The rest of unittest/_torch/modeling_v2 + # needs a GPU. - unittest/_torch/modeling_v2/test_modeling_v2_claims.py - unittest/_torch/modeling_v2/test_modeling_v2_native_claims.py - unittest/_torch/modeling_v2/test_modeling_v2_native_loader.py @@ -95,13 +95,6 @@ l0_cpu: - unittest/_torch/modeling_v2/models/deepseek_v3/test_modeling_v2_deepseek_default_shards.py - unittest/_torch/modeling_v2/models/gpt_oss/test_modeling_v2_gpt_oss_default_shards.py - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_dense_backend_library_order - - unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_multimodal_embeds.py - - unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_moe_scalars.py - - unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_validate_request.py - - unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_pool_contract.py - - unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_kv_quant.py - - unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_thop_attention_calls.py - - unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_weight_table.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/integration/test_lists/test-db/l0_gb300.yml b/tests/integration/test_lists/test-db/l0_gb300.yml index 1de8159db276..a0e8fca03f24 100644 --- a/tests/integration/test_lists/test-db/l0_gb300.yml +++ b/tests/integration/test_lists/test-db/l0_gb300.yml @@ -121,6 +121,10 @@ l0_gb300: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV41::test_gsm8k_dspark[nvfp4-tp4] TIMEOUT (120) - accuracy/test_llm_api_pytorch.py::TestDeepSeekV41::test_gsm8k_dspark[nvfp4-dp4] TIMEOUT (120) - accuracy/test_glm53_flash.py::TestGLM53FlashFP8::test_mtp[tp_size=4-ep_size=4] TIMEOUT (60) + # qwen4_exp target, pre-merge: tep4 with three drafts and IndexShare, gsm8k and then + # 3k-6k-token pass keys with their acceptance-length floor. The other layouts and KV dtypes + # run post-merge (l0_gb300_multi_gpus). + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_and_long_context TIMEOUT (60) # K3 collectives and decode-path exchanges, four ranks in one NVLink domain. - condition: diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml index 97d9ebd7ac45..fcdf9456ff91 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml @@ -33,6 +33,17 @@ l0_gb300_multi_gpus: - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm10xDefault::test_gsm8k_chunked_prefill[dep4] - accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp4::test_gsm8k - accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp4::test_eagle3_gsm8k_and_acceptance + # qwen4_exp target, post-merge: three drafts over NVFP4 KV at tp4 with the + # breakable prefill graph and over FP8 KV at dep4 (MTP under attention DP), both with block + # reuse and chunked prefill; and dep2 context and generation engines on one node. tep4 runs + # pre-merge (l0_gb300); mmmu runs in QA (llm_function_core). + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[tp4_nvfp4_kv_breakable] TIMEOUT (60) + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_mtp3_gsm8k_quantized_kv[dep4_fp8_kv] TIMEOUT (60) + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_disaggregated_dep2 TIMEOUT (90) + # qwen4_exp target at tp1, post-merge: its AgentX engines (NVFP4 KV, MTP2 with rejection + # sampling, the breakable prefill graph, every opt-in projection) on one context and one + # generation GPU. + - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Default::test_tp1_disaggregated_mtp2 TIMEOUT (90) # Covers the GPU cases of tests/unittest/_torch/attention/ on GB300, kernels/ and # multi_gpu/ included; its cpu_only cases run on CPU (l0_cpu). The GPU cases of # kernels/parallel_hw_agnostic/ run here too, although l0_b200 and l0_b300 ignore that diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py index 4bc00766af36..0610b489bc1a 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py @@ -180,7 +180,7 @@ def _deviation(out: torch.Tensor, ref: torch.Tensor) -> str: ), ), ), - "qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1": ( + "qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default": ( Cell( why="qwen4_exp main rope, decode: one token, Gemma norm and interleaved mRoPE " "with three distinct position rows", diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py index ccbd73d80daa..b534b1a08b47 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py @@ -7550,7 +7550,7 @@ def _spread(batch: int, longest: int) -> tuple[int, ...]: ), ), ), - "qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1": ( + "qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default": ( Cell( why="qwen4_exp prefill: GQA 24q/2kv d256, bf16 pool, paged context; " "a fresh 70-row sequence beside 33 rows over a 48-token cached prefix", diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py new file mode 100644 index 000000000000..bb60f45d720c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the fused_sigmoid_mul_inplace catalog entry: every cell against +the entry's reference with nothing written past `attention_output`'s last row.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.fused_sigmoid_mul_inplace import ( # noqa: E501 + fused_sigmoid_mul_inplace as op, +) + +__extra_import_path__ = ["../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "fused_sigmoid_mul_inplace requires a CUDA device" + +#: Elements past `attention_output`'s last row, inside the buffer it views: a +#: masked lane that stored anyway would land here (one kernel block is 2048 wide). +_TAIL = 2048 + + +def _seed(text: str) -> int: + return zlib.crc32(text.encode()) + + +def _build(spec, seed): + """`(out, out_buffer, gate_buffer, gate)` for one spec. + + `out` is the compact `[T, hidden]` head of `out_buffer`, which runs `_TAIL` + elements past its last row. `gate` is a view of `gate_buffer`: the second + half of each `[T, heads, 2 * D]` q_proj head (`gate="interleaved"`), or the + 2-D gate window of the target's fused projection row (`gate="fused_row"`: + `row_stride` columns, the gate at `gate_col`; tp1's 14080 / 7168 where the + spec names neither). + """ + g = torch.Generator(device="cuda").manual_seed(seed) + t, heads, d, dtype = spec["tokens"], spec["heads"], spec["head_dim"], spec["dtype"] + hidden = heads * d + + def randn(*shape): + return torch.randn(*shape, generator=g, device="cuda").to(dtype) + + out_buffer = randn(t * hidden + _TAIL) + out = out_buffer[: t * hidden].view(t, hidden) + if spec["gate"] == "interleaved": + gate_buffer = randn(t, heads, 2 * d) + gate = gate_buffer[:, :, d:] + else: + assert spec["gate"] == "fused_row", spec["gate"] + # The target's fused input projection row [q | k | v | gate | index scratch]: + # the gate a 2-D [T, hidden] window at column qkv_width. + row_stride, gate_col = spec.get("row_stride", 14080), spec.get("gate_col", 7168) + gate_buffer = randn(t, row_stride) + gate = gate_buffer[:, gate_col : gate_col + hidden] + return out, out_buffer, gate_buffer, gate + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + out, out_buffer, gate_buffer, gate = _build(cell.spec, seed=_seed(cell.why)) + gate_before = gate_buffer.clone() + tail_before = out_buffer[out.numel() :].clone() + expected = op.reference(out, gate) + + result = op(out, gate) + + assert result.data_ptr() == out.data_ptr() and result.shape == out.shape + torch.testing.assert_close(result, expected) + assert torch.equal(gate_buffer, gate_before) + assert torch.equal(out_buffer[out.numel() :], tail_before), "a masked lane stored" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_erf.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_erf.py similarity index 96% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_erf.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_erf.py index 5255d4197ee8..ed7a1d95f5ac 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_erf.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_erf.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qwen4_exp gelu_erf entry. +"""GPU test for the gelu_erf catalog entry. The entry's identity is *which* GELU it is, and the wrong answer -- the tanh GELU this same tower's block MLPs use -- rounds to within about one bf16 ulp of @@ -14,11 +14,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.activation.gelu_erf import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.gelu_erf import ( # noqa: E501 gelu_erf as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 assert torch.cuda.is_available(), "gelu_erf requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_tanh.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_tanh.py similarity index 97% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_tanh.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_tanh.py index b09168720ce5..4f818110d2f6 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_tanh.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_gelu_tanh.py @@ -13,11 +13,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.activation.gelu_tanh import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.gelu_tanh import ( # noqa: E501 gelu_tanh as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 assert torch.cuda.is_available(), "gelu_tanh requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_silu_and_mul_rows.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_silu_and_mul_rows.py similarity index 65% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_silu_and_mul_rows.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_silu_and_mul_rows.py index e5291a35eb15..7a4874047e69 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_silu_and_mul_rows.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/activation/test_modeling_v2_silu_and_mul_rows.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the silu_and_mul_rows entry: the cells bitwise `flashinfer_silu_and_mul` on a contiguous copy of the gate|up columns (and so inside its reference band) with the fused -output's extra columns present, graph capture, the width guard.""" +output's extra columns present.""" import zlib @@ -12,11 +12,11 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.flashinfer_silu_and_mul import ( # noqa: E501 flashinfer_silu_and_mul, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.activation.silu_and_mul_rows import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.activation.silu_and_mul_rows import ( # noqa: E501 silu_and_mul_rows as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "silu_and_mul_rows requires a CUDA device" @@ -42,25 +42,3 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by torch.testing.assert_close(out, op.reference(x, INTER)) if x.shape[0]: assert torch.equal(out, flashinfer_silu_and_mul(x[:, : 2 * INTER].contiguous())) - - -def test_cuda_graph_capture_and_replay() -> None: - x = _build(dict(rows=128), 8) - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - op(x, INTER) - torch.cuda.current_stream().wait_stream(warmup) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - out = op(x, INTER) - x.copy_(_build(dict(rows=128), 9)) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(out, op(x, INTER)) - - -def test_guards() -> None: - x = _build(dict(rows=4), 5) - with validating(op), pytest.raises(AssertionError): - op(x[:, : 2 * INTER - 2], INTER) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_canonical_selection.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_canonical_selection.py new file mode 100644 index 000000000000..8701e8322c8c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_canonical_selection.py @@ -0,0 +1,282 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_canonical_selection entry. + +The rewrite fixes every slot of every row, so `_compare` is index equality against a +composite-key top-k, and against `_torch_rewrite`, the composed-torch rewrite the fused +kernel replaced, which also pins the rows outside the contract (NaN in a window, a +window cut by the table's width). The input rows come from an adversary that keeps a +random part of the tie class at the k-th value in a random order, or from the radix +engine the target calls. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_canonical_selection import ( # noqa: E501 + _SPLIT_ROWS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_canonical_selection import ( # noqa: E501 + qsa_canonical_selection as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_topk_select import ( # noqa: E501 + qsa_topk_select, +) + +__extra_import_path__ = ["../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "qsa_canonical_selection requires a CUDA device" + +DEV = "cuda" +COMPRESS_RATIO = 4 +BLOCK_TOPK = 512 +COLS = 2080 + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Index for index. The rule fixes every slot of every row, so nothing is tolerated: a + selection does no arithmetic.""" + assert out.dtype == torch.int32, out.dtype + assert out.shape == ref.shape, (out.shape, ref.shape) + wrong = (out != ref).any(1) + if bool(wrong.any()): + row = int(wrong.nonzero()[0, 0]) + raise AssertionError( + f"{int(wrong.sum())} of {out.shape[0]} rows differ from the canonical " + f"selection; first is row {row}" + ) + + +def _torch_rewrite(logits, row_starts, row_ends, indices) -> torch.Tensor: + """The rewrite as composed torch ops, in row tiles of 2^24 elements: the + implementation the fused kernel replaced, kept as a second oracle that is bitwise the + old rows wherever they disagree with `reference` (outside the contract).""" + num_rows, num_cols = logits.shape + index_topk = indices.shape[1] + tile = max(1, (1 << 24) // max(num_cols, 1)) + for lo in range(0, num_rows, tile): + hi = min(lo + tile, num_rows) + scores, out = logits[lo:hi], indices[lo:hi] + starts, ends = row_starts[lo:hi, None], row_ends[lo:hi, None] + columns = torch.arange(num_cols, dtype=torch.int32, device=logits.device)[None, :] + windowed = torch.where((columns >= starts) & (columns < ends), scores, float("nan")) + picked = torch.gather(scores, 1, (out + starts).clamp_(0, num_cols - 1).long()) + kth = picked.masked_fill_(out < 0, float("inf")).amin(1, keepdim=True) + at_least = windowed >= kth + tied = windowed == kth + surplus = at_least.sum(1, keepdim=True, dtype=torch.int32) - index_topk + dropped = tied & (torch.cumsum(tied, 1, dtype=torch.int32) <= surplus) + kept_seen = torch.cumsum(at_least ^ dropped, 1, dtype=torch.int32) + slots = torch.arange(1, index_topk + 1, dtype=torch.int32, device=logits.device) + slots = slots.repeat(hi - lo, 1) + torch.searchsorted(kept_seen, slots, out_int32=True, out=out) + out.sub_(starts) + out.masked_fill_(slots > kept_seen[:, -1:], -1) + return indices + + +def _logits(kind: str, rows: int, cols: int, gen: torch.Generator) -> torch.Tensor: + if kind == "relu": + return torch.relu(torch.randn(rows, cols, generator=gen, device=DEV)) * 3.0 + if kind == "relu_sparse": + x = torch.relu(torch.randn(rows, cols, generator=gen, device=DEV)) * 3.0 + return x * (torch.rand(rows, cols, generator=gen, device=DEV) < 0.2) + if kind == "relu_rare": + # About one positive score in 750: a long window still has fewer than k, so + # the k-th value is 0 and the tie class is most of the window, half of it -0.0. + x = torch.relu(torch.randn(rows, cols, generator=gen, device=DEV) - 3.0) + flip = torch.rand(rows, cols, generator=gen, device=DEV) < 0.5 + return torch.where(flip & (x == 0), -0.0, x) + if kind == "quantized": + # A handful of levels, so every k-th value is tied; rounding a small negative + # score gives -0.0. + return torch.round(torch.randn(rows, cols, generator=gen, device=DEV) * 2.0) / 2.0 + raise ValueError(kind) + + +def _windows(spec: dict, rows: int, cols: int, gen: torch.Generator): + kind = spec["windows"][0] + starts = torch.zeros(rows, dtype=torch.int32, device=DEV) + if kind == "chunk": + positions = spec["windows"][1] + torch.arange(rows, device=DEV) + else: + assert kind == "positions_random", kind + lo, hi = spec["windows"][1:] + positions = torch.randint(lo, hi + 1, (rows,), generator=gen, device=DEV) + ends = ((positions + 1) // COMPRESS_RATIO).clamp_max(cols).to(torch.int32) + return starts, ends + + +def _poison(logits, starts, ends, kind) -> None: + columns = torch.arange(logits.shape[1], device=DEV)[None, :] + outside = (columns < starts[:, None]) | (columns >= ends[:, None]) + if kind == "nan": + logits.masked_fill_(outside, float("nan")) + return + # Values that would win any comparison they were allowed into. + pattern = torch.tensor([float("nan"), float("inf"), 3e38, float("nan")], device=DEV) + logits.copy_(torch.where(outside, pattern[columns % 4].expand_as(logits), logits)) + + +def _order_key(logits: torch.Tensor) -> torch.Tensor: + """int64 that orders as the fp32 values do, -0.0 folded into +0.0.""" + bits = (logits.float() + 0.0).view(torch.int32).long() + return torch.where(bits < 0, bits ^ 0x7FFFFFFF, bits) + + +def _adversary(logits, starts, ends, index_topk, gen) -> torch.Tensor: + """An exact top-k of each window that keeps a random part of the tie class at the + k-th value, in a random order, `-1` last.""" + rows, cols = logits.shape + columns = torch.arange(cols, device=DEV)[None, :] + in_window = (columns >= starts.long()[:, None]) & (columns < ends.long()[:, None]) + # A random 32-bit low word: equal scores order at random (a collision only repeats + # an order), without sorting a `[rows, cols]` matrix. + tiebreak = torch.randint(0, 1 << 32, (rows, cols), generator=gen, device=DEV) + key = torch.where( + in_window, (_order_key(logits) << 32) | tiebreak, torch.iinfo(torch.int64).min + ) + width = min(index_topk, cols) + top = torch.topk(key, width, dim=1).indices + local = torch.where(torch.gather(in_window, 1, top), top - starts.long()[:, None], -1) + shuffle = torch.where(local >= 0, torch.rand(local.shape, generator=gen, device=DEV), 2.0) + local = torch.gather(local, 1, shuffle.argsort(1)) + out = torch.full((rows, index_topk), -1, dtype=torch.int32, device=DEV) + out[:, :width] = local.to(torch.int32) + return out + + +def _radix(logits, starts, ends) -> torch.Tensor: + out = torch.empty(logits.shape[0], BLOCK_TOPK, dtype=torch.int32, device=DEV) + return qsa_topk_select(logits, starts, ends, out, index_topk=BLOCK_TOPK) + + +def _build(spec: dict, rows: int, seed: int): + gen = torch.Generator(device=DEV).manual_seed(seed) + logits = _logits(spec["logits"], rows, spec["cols"], gen) + starts, ends = _windows(spec, rows, spec["cols"], gen) + _poison(logits, starts, ends, spec["tail"]) + if spec["selection"] == "random": + indices = _adversary(logits, starts, ends, spec["index_topk"], gen) + else: + indices = _radix(logits, starts, ends) + return logits, starts, ends, indices + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 + certifying(op) + spec = cell.spec + seed = sum(map(ord, cell.why)) + row_counts = spec["rows"] if isinstance(spec["rows"], tuple) else (spec["rows"],) + for rows in row_counts: + args = _build(spec, rows, seed + rows) + expected = op.reference(*args) + old = _torch_rewrite(args[0], args[1], args[2], args[3].clone()) + out = op(*args) + assert out is args[3], "the rewrite writes `indices` and returns it" + _compare(out, expected) + _compare(out, old) + + +def _random_windows(rows: int, cols: int, gen: torch.Generator): + """Windows anywhere in the table: a tenth empty, a fifth at most k wide, the rest of + any width the table leaves room for.""" + starts = torch.randint(0, cols + 1, (rows,), generator=gen, device=DEV) + u = torch.rand(rows, generator=gen, device=DEV) + short = torch.randint(1, BLOCK_TOPK + 1, (rows,), generator=gen, device=DEV) + wide = torch.randint(0, cols + 1, (rows,), generator=gen, device=DEV) + length = torch.where(u < 0.1, 0, torch.where(u < 0.3, short, wide)) + ends = (starts + length).clamp_max(cols) + return starts.to(torch.int32), ends.to(torch.int32) + + +def _chunk_windows(rows: int, cols: int, last_position: int): + """A prefill chunk's causal windows, the last row at `last_position`.""" + positions = last_position - rows + 1 + torch.arange(rows, device=DEV) + ends = ((positions + 1) // COMPRESS_RATIO).clamp(0, cols).to(torch.int32) + return torch.zeros(rows, dtype=torch.int32, device=DEV), ends + + +def _long_windows(kind: str, rows: int, cols: int, gen: torch.Generator): + if kind == "random": + return _random_windows(rows, cols, gen) + if kind == "chunk_start": + return _chunk_windows(rows, cols, rows - 1) + assert kind == "chunk_end", kind + return _chunk_windows(rows, cols, cols * COMPRESS_RATIO - 1) + + +#: (rows, cols, windows, logits) at long context: decode-sized to 16384-row prefill +#: chunks, 2080- to 65536-wide tables (65536 is the 256k context), on both sides of the +#: rows at which the op stops splitting a row into segments. +_LONG_CASES = ( + (1, 65536, "chunk_end", "relu_rare"), + (16, 32768, "random", "quantized"), + (256, 65536, "random", "relu_rare"), + (_SPLIT_ROWS - 1, 25000, "random", "quantized"), + (_SPLIT_ROWS, 25000, "random", "relu_rare"), + (4096, 8192, "random", "quantized"), + (4096, 32768, "chunk_end", "relu_sparse"), + (16384, 2080, "chunk_start", "relu_rare"), + (16384, 65536, "chunk_end", "quantized"), +) + + +@pytest.mark.parametrize( + "rows,cols,windows,kind", _LONG_CASES, ids=["-".join(map(str, c)) for c in _LONG_CASES] +) +def test_long_context_rows_match_reference_and_torch_rewrite(rows, cols, windows, kind) -> None: + """Tie classes of most of a window (relu zeros, half of them -0.0) or of a quantized + level, empty windows and windows of at most k, garbage past the windows.""" + gen = torch.Generator(device=DEV).manual_seed(rows * 131 + cols) + logits = _logits(kind, rows, cols, gen) + starts, ends = _long_windows(windows, rows, cols, gen) + _poison(logits, starts, ends, "poison") + raw = _adversary(logits, starts, ends, BLOCK_TOPK, gen) + expected = op.reference(logits, starts, ends, raw) + _compare(_torch_rewrite(logits, starts, ends, raw.clone()), expected) + _compare(op(logits, starts, ends, raw), expected) + + +@pytest.mark.parametrize("rows", (_SPLIT_ROWS - 1, 2048)) +def test_rows_outside_the_contract_match_the_torch_rewrite(rows) -> None: + """NaN and +-inf inside the windows, NaN picks, rows with no pick, picks past the + window, and windows cut short by a narrowed table: no canonical row to compare + against, so the rows are the torch rewrite's, bit for bit.""" + gen = torch.Generator(device=DEV).manual_seed(5) + cols, live = 9000, 6000 + clean = _logits("quantized", rows, cols, gen) + starts, ends = _random_windows(rows, cols, gen) + _poison(clean, starts, ends, "poison") + specials = torch.tensor([float("nan"), float("inf"), float("-inf")], device=DEV) + which = torch.randint(0, 3, (rows, cols), generator=gen, device=DEV) + sprinkle = torch.rand(rows, cols, generator=gen, device=DEV) < 0.002 + dirty = torch.where(sprinkle, specials[which], clean) + # Even rows were selected over the dirty scores (NaN sorts first: NaN picks), odd + # rows over the clean ones. + seen = torch.where((torch.arange(rows, device=DEV) % 2 == 0)[:, None], dirty, clean) + raw = _adversary(seen, starts, ends, BLOCK_TOPK, gen) + raw[::7] = -1 + raw[3::11, 0] = cols + 100 + raw[5::13, 1] = -5 + for scores in (dirty, dirty[:, :live]): + expected = _torch_rewrite(scores, starts, ends, raw.clone()) + _compare(op(scores, starts, ends, raw.clone()), expected) + + +def test_narrowing_the_table_to_a_bound_on_the_windows_changes_nothing() -> None: + """A caller that knows every window ends by column 1500 passes `logits[:, :1500]` of a + wide table: the rows are the full table's.""" + gen = torch.Generator(device=DEV).manual_seed(21) + rows, cols, bound = 4096, 8192, 1500 + logits = _logits("relu_sparse", rows, cols, gen) + starts = torch.zeros(rows, dtype=torch.int32, device=DEV) + ends = torch.randint(0, bound + 1, (rows,), generator=gen, device=DEV).to(torch.int32) + _poison(logits, starts, ends, "poison") + raw = _radix(logits, starts, ends) + wide = op(logits, starts, ends, raw.clone()) + narrow = op(logits[:, :bound], starts, ends, raw.clone()) + _compare(wide, op.reference(logits, starts, ends, raw)) + assert torch.equal(wide, narrow) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_expand_block_indices.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_expand_block_indices.py similarity index 80% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_expand_block_indices.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_expand_block_indices.py index 0c83dcf90793..f45b9470a169 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_expand_block_indices.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_expand_block_indices.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qsa_expand_block_indices catalog entry. +"""GPU test for the default target's qsa_expand_block_indices variant. Selections are built the way `qsa_topk_select` emits them: a random subset of the complete visible groups, prefix-packed, then `-1`. Every gate is bitwise. @@ -11,11 +11,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_expand_block_indices import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_expand_block_indices import ( # noqa: E501 qsa_expand_block_indices as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 assert torch.cuda.is_available(), "qsa_expand_block_indices requires a CUDA device" @@ -31,7 +31,10 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: DEV = "cuda" CR, TOPK = 4, 2048 +# The tp1 engine's max_seq_len: where `_serving_rows` places its continuation chunks. ENGINE_MAX_SEQ_LEN = 8192 +# The default target's longest sequence: every cell's bound. +ENVELOPE_MAX_SEQ_LEN = 262144 def _serving_rows(requests: int, context_lens: tuple) -> tuple[list[int], list[int]]: @@ -55,18 +58,34 @@ def _serving_rows(requests: int, context_lens: tuple) -> tuple[list[int], list[i return positions, lengths +def _decode_positions(rows: int, top: int) -> list[int]: + """One decode row per request, spread over [1, top): the first row sees one group, + the last is near `top`, and `i % 4` cycles the tail width through 0..3.""" + step = (top - 4) // (rows - 1) // CR * CR + positions = [1 + i * step + i % CR for i in range(rows)] + assert positions[-1] < top + return positions + + def _rows(batch) -> tuple[list[int], list[int]]: kind = batch[0] if kind == "sweep": positions = [4096 + 13 * i for i in range(batch[1])] + elif kind == "decode": + positions = _decode_positions(batch[1], batch[2]) + assert max(positions) + 1 <= ENVELOPE_MAX_SEQ_LEN elif kind == "serving": return _serving_rows(batch[1], batch[2]) + elif kind == "chunk": + start, rows = batch[1], batch[2] + assert start + rows <= ENVELOPE_MAX_SEQ_LEN + return list(range(start, start + rows)), [start + rows] * rows else: assert kind == "verify", kind starts, per_request = batch[1], batch[2] positions = [s + i for s in starts for i in range(per_request)] lengths = [s + per_request for s in starts for _ in range(per_request)] - assert max(lengths) <= ENGINE_MAX_SEQ_LEN + assert max(lengths) <= ENVELOPE_MAX_SEQ_LEN return positions, lengths return positions, [p + 1 for p in positions] @@ -108,17 +127,17 @@ def test_certified_cells(cell, certifying): # noqa: F811 assert out.dtype == torch.int32 and out.is_cuda and out.is_contiguous() assert tuple(out.shape) == (rows, final) _compare(out, op.reference(*args, **geometry)) - if cell.spec["batch"][0] == "serving": + if cell.spec["batch"][0] in ("serving", "chunk"): positions, lengths = args[1], args[2] assert bool((lengths.long() > positions.long() + 1).any()), ( - "a serving cell must drive sequence_lengths past position + 1" + "a context cell must drive sequence_lengths past position + 1" ) -def test_cuda_graph_capture_and_replay() -> None: - """Capturable, and a replay recomputes from whatever the buffers now hold: the 256 rows - a verify step of 32 requests carries at max_draft_len 7.""" - rows = 256 +@pytest.mark.parametrize("rows", [1, 128, 256, 1024]) +def test_cuda_graph_capture_and_replay(rows) -> None: + """Capturable at decode and verify row counts (both launch shapes), and a replay + recomputes from whatever the buffers now hold.""" positions = torch.tensor([4096 + 7 * i for i in range(rows)], device=DEV) lengths = (positions + 1).to(torch.int32) blocks = _packed_selection(positions, CR, TOPK, seed=31) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_mtp_reuse_topk.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_mtp_reuse_topk.py similarity index 83% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_mtp_reuse_topk.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_mtp_reuse_topk.py index 1f8e34122860..d689bfd91d89 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_mtp_reuse_topk.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_mtp_reuse_topk.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qsa_mtp_reuse_topk catalog entry. +"""GPU test for the default target's qsa_mtp_reuse_topk variant. Captured rows are built the way `qsa_topk_select` emits them: `min(K, visible)` distinct complete groups in an arbitrary order as a prefix, then `-1`. Inputs @@ -15,11 +15,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_mtp_reuse_topk import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_mtp_reuse_topk import ( # noqa: E501 qsa_mtp_reuse_topk as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 assert torch.cuda.is_available(), "qsa_mtp_reuse_topk requires a CUDA device" @@ -37,10 +37,10 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: CR = 4 #: The engine's `max_num_tokens`: height of the request-map and Top-K buffers. ENGINE_MAX_NUM_TOKENS = 8192 -ENGINE_MAX_SEQ_LEN = 8192 #: Captured positions for the first requests of a "mixed" batch: empty (0 #: groups), 1, 100, 255, 511 and 512 groups, rows that gain a group and rows -#: that do not, up to the last position the 8192-token envelope allows. +#: that do not, up to the last position an 8192-token sequence allows. A cell +#: with a longer `max_seq_len` also anchors its own last position. ANCHORS = (2, 3, 402, 2046, 2050, 8188, 2051, 2047, 1022, 4095, 5998, 7001, 37, 2044, 3001, 6143) @@ -53,21 +53,25 @@ def _captured_row(visible: int, block_topk: int, gen: torch.Generator) -> torch. return row -def _requests(contexts, drafted, gen) -> list[tuple[int, int]]: +def _requests(contexts, drafted, gen, max_seq_len: int) -> list[tuple[int, int]]: """(captured_visible, current_visible) per request.""" kind = contexts[0] if kind == "explicit": positions = list(contexts[1]) else: count = contexts[1] - low, high = {"mixed": (0, 8188), "full": (2047, 8188)}[kind] - fixed = {"mixed": ANCHORS, "full": (2047, 2050, 8188, 2051)}[kind] + top = max_seq_len - CR + low, high = {"mixed": (0, top), "full": (2047, top)}[kind] + fixed = { + "mixed": ANCHORS if top in ANCHORS else (*ANCHORS, top), + "full": (2047, 2050, top, 2051), + }[kind] positions = list(fixed[:count]) positions += torch.randint(low, high + 1, (count - len(positions),), generator=gen).tolist() out = [] for q, position in enumerate(positions): d = drafted[q % len(drafted)] - position = min(position, ENGINE_MAX_SEQ_LEN - 1 - d) + position = min(position, max_seq_len - 1 - d) out.append(((position + 1) // CR, (position + d + 1) // CR)) return out @@ -83,7 +87,7 @@ def _build(spec, seed: int, requests: int | None = None) -> dict: contexts = spec["contexts"] if requests is not None: contexts = (contexts[0], requests) - pairs = _requests(contexts, spec["drafted"], gen) + pairs = _requests(contexts, spec["drafted"], gen, spec["max_seq_len"]) num_requests = len(pairs) per_request = spec["rows_per_request"] rows = num_requests * per_request @@ -146,6 +150,21 @@ def _assert_legal_selection(out: torch.Tensor, kwargs: dict) -> None: assert bool((ids >= -1).all()) +def _assert_widest_displacement_driven(spec, kwargs: dict) -> None: + """The cell reaches the riskiest path its offsets allow: a full row (every slot a + valid id) gaining `(CR - 1 + max(drafted)) // CR` fresh groups, which displaces that + many of its oldest ids. Seeds are fixed, so this holds the cell's coverage, not luck.""" + block_topk = spec["block_topk"] + requests = kwargs["request_indices"].long() + filled = (kwargs["shared_indices"].long()[requests] >= 0).sum(dim=1) + captured = kwargs["shared_visible_blocks"].long()[requests] + fresh = (kwargs["visible_blocks"].long() - captured).clamp(0, block_topk) + widest = (CR - 1 + max(spec["drafted"])) // CR + assert bool(((filled == block_topk) & (fresh == widest)).any()), ( + f"no full row gains {widest} fresh groups: the cell misses its displacement path" + ) + + def _check_cell(spec, seed: int, requests: int | None = None) -> torch.Tensor: built = _build(spec, seed, requests) kwargs, buffers, rows = built["kwargs"], built["buffers"], built["rows"] @@ -156,6 +175,8 @@ def _check_cell(spec, seed: int, requests: int | None = None) -> torch.Tensor: assert out.dtype == torch.int32 and tuple(out.shape) == (rows, spec["block_topk"]) _compare(out, op.reference(**kwargs)) _assert_legal_selection(out, kwargs) + if requests is None: + _assert_widest_displacement_driven(spec, kwargs) for name, tensor in buffers.items(): if name == "output": assert torch.equal(tensor[rows:], before[name][rows:]), "rows past the view moved" @@ -203,12 +224,13 @@ def _check_graph(spec, seed: int) -> None: @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) def test_certified_cells(cell, certifying): # noqa: F811 + """Eager at the cell's own request count (bitwise, legal selection, nothing outside + the views moved), then each of its CUDA-graph request counts.""" certifying(op) seed = zlib.crc32(cell.why.encode()) + _check_cell(cell.spec, seed) if cell.spec["graph_batches"]: _check_graph(cell.spec, seed) - else: - _check_cell(cell.spec, seed) def _spec_of(prefix: str) -> dict: diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_nvfp4_pages_to_bf16.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_nvfp4_pages_to_bf16.py similarity index 96% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_nvfp4_pages_to_bf16.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_nvfp4_pages_to_bf16.py index 99c939511024..8adaf4bdce6e 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_nvfp4_pages_to_bf16.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_nvfp4_pages_to_bf16.py @@ -17,14 +17,14 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention import ( qsa_paged_sparse_gqa_nvfp4 as _gqa, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_nvfp4_pages_to_bf16 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_nvfp4_pages_to_bf16 import ( # noqa: E501 qsa_nvfp4_pages_to_bf16 as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 assert torch.cuda.is_available(), "qsa_nvfp4_pages_to_bf16 requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_scores.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_scores.py new file mode 100644 index 000000000000..84f02ccb594e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_scores.py @@ -0,0 +1,460 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp qsa_paged_index_scores entry. + +The paged state is real: a `[pages, tokens_per_block, 1, dim]` index cache +(a view of a wider page under the NVFP4 cells) sized from the pages the table maps, and an +int32 block table whose rows are request *slots* -- more of them than the batch has requests, assigned by a +seeded permutation -- with shuffled physical pages and `-1` past each +sequence on the serving layouts, which is what the engine's table carries. + +The four-rank layouts' cells add one rank's measured pool: every attention layer's +plane in one page (the last layer's plane taken), NaN wherever no request maps a +page, with the pool's first and last page among those the cell reads -- the last +one past 2^31 elements from the plane's base. + +The call has four kernels behind it: the entry's own visible-split kernel for +ungrouped decode-sized visible-only calls and request-group kernel for verify-shaped +ones, and the engine's row-wise and query-tile kernels. `_watch_launches` swaps the four Triton entry points +for recording proxies that still launch the real kernel, so every cell asserts +the branch it names. +""" + +import math + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention import ( + qsa_paged_index_scores as _entry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_index_scores import ( # noqa: E501 + qsa_paged_index_scores as op, +) +from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels + +__extra_import_path__ = ["../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "qsa_paged_index_scores requires a CUDA device" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Torch's default fp32 band on every specified column; `-inf` exactly. + + bf16 and fp16 products are exact in the kernel's fp32 accumulator, so + only summation order separates it from the reference, well inside the + default rtol 1.3e-6 / atol 1e-5 on every cell, the measured-pool ones + included. Wrong variants (no relu, mean over heads, first group member as + key, unscaled, 1/head_dim scale, block table ignored) land orders of + magnitude outside it. Columns the reference marks NaN are unspecified and + not read. + """ + assert out.dtype == ref.dtype == torch.float32, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + specified = ~torch.isnan(ref) + torch.testing.assert_close(out[specified], ref[specified]) + + +DEV = torch.device("cuda") +#: The largest batch the visible-split kernel serves. +DECODE_ROW_LIMIT = 256 +#: Index heads must divide this into more than one row for the query tile. +DOT_TILE_WIDTH = 128 +#: Index heads must divide this for the visible-split kernel. +SPLIT_DOT_WIDTH = 16 +#: A table width with columns enough for several of the request-group kernel's tiles. +_WIDTH_PAGES = 260 + + +class _LaunchRecorder: + """Transparent proxy over a Triton entry point: records the grid, then launches.""" + + def __init__(self, inner, name, log): + self._inner, self._name, self._log = inner, name, log + + def __getitem__(self, grid): + self._log.append((self._name, tuple(grid))) + return self._inner[grid] + + def __getattr__(self, item): + return getattr(self._inner, item) + + +class _watch_launches: + """Record which of the entry's four Triton kernels a call launches.""" + + _NAMES = ( + (_qsa_kernels, "_qsa_paged_index_scores_kernel", "row_wise"), + (_qsa_kernels, "_qsa_paged_index_scores_qtile_kernel", "query_tile"), + (_entry, "_qsa_index_scores_group_kernel", "request_group"), + (_entry, "_qsa_index_scores_visible_split_kernel", "visible_split"), + ) + + def __enter__(self): + self.log = [] + self._saved = [(module, name, getattr(module, name)) for module, name, _ in self._NAMES] + for (module, name, tag), (_, _, kernel) in zip(self._NAMES, self._saved): + setattr(module, name, _LaunchRecorder(kernel, tag, self.log)) + return self + + def __exit__(self, *exc): + for module, name, kernel in self._saved: + setattr(module, name, kernel) + return False + + @property + def branch(self): + assert len(self.log) == 1, f"expected exactly one kernel launch, saw {self.log}" + return self.log[0][0] + + @property + def grid(self): + assert len(self.log) == 1, f"expected exactly one kernel launch, saw {self.log}" + return self.log[0][1] + + +def _case(rows, requests, **overrides): + """A narrow-table spec for the non-cell tests, same geometry as the entry's own.""" + spec = dict( + layout="random", + rows=rows, + requests=requests, + slots=requests + 4, + pages=8, + heads=4, + head_dim=128, + tokens_per_block=32, + compress_ratio=4, + dtype=torch.bfloat16, + positions_dtype=torch.int32, + request_dtype=torch.int32, + q_layout="packed", + only_visible_blocks=False, + context_rows=False, + ) + spec.update(overrides) + return spec + + +def _rows_and_requests(spec, gen): + """Per-row positions, per-row request ordinal, and each request's mapped pages.""" + layout, tpb, pages = spec["layout"], spec["tokens_per_block"], spec["pages"] + + def pages_for(seq_len): + return math.ceil(seq_len / tpb) + + if layout == "random": + rows, requests = spec["rows"], spec["requests"] + req = torch.arange(rows, device=DEV) * requests // max(rows, 1) + positions = torch.randint(0, pages * tpb, (rows,), generator=gen, device=DEV) + return positions, req, [pages] * requests + + # (first position, row count) per request, in row order + if layout == "capture": + segments = [(0, n) for n in spec["lengths"]] + elif layout == "decode": + segments = [(p, spec["rows_per_request"]) for p in spec["gen_positions"]] + else: + raise AssertionError(f"unknown layout {layout!r}") + positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments]) + req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)]) + mapped = [pages_for(s + n) for s, n in segments] + assert max(mapped) <= pages + return positions, req, mapped + + +def _build(spec, seed): + """Real paged state for one spec, keyed by `__call__`'s parameter names.""" + gen = torch.Generator(device=DEV).manual_seed(seed) + heads, head_dim, dtype = spec["heads"], spec["head_dim"], spec["dtype"] + positions, req, mapped = _rows_and_requests(spec, gen) + rows, slots = positions.numel(), spec["slots"] + assert len(mapped) <= slots + + slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(mapped)] + total = sum(mapped) + pool_pages = spec.get("pool_pages", total + 3) + page_ids = torch.randperm(pool_pages, generator=gen, device=DEV) + if "pool_pages" in spec: + # A measured pool: its first and last page go to the first request's first two + # logical pages, so the cell reads both ends. + ends = torch.tensor([pool_pages - 1, 0], device=DEV) + page_ids = torch.cat([ends, page_ids[(page_ids != 0) & (page_ids != pool_pages - 1)]]) + page_ids = page_ids[:total].to(torch.int32) + block_table = torch.full((slots, spec["pages"]), -1, dtype=torch.int32, device=DEV) + cursor = 0 + for slot, count in zip(slot_of.tolist(), mapped): + block_table[slot, :count] = page_ids[cursor : cursor + count] + cursor += count + tpb = spec["tokens_per_block"] + if "page_pad" in spec: + # The trunk pool's page: every attention layer's index buffer (the fp8 and NVFP4 + # pools then the position page), so one layer's plane is a view at page stride + # buffers * tpb * dim + page_pad -- the last layer's here. Pages no request maps + # hold NaN: a read of one shows up in the scores. + buffers, buffer = spec["pool_buffers"], tpb * head_dim + page = torch.full( + (pool_pages, buffers * buffer + spec["page_pad"]), float("nan"), dtype=dtype, device=DEV + ) + index_cache = page[:, (buffers - 1) * buffer : buffers * buffer].view( + pool_pages, tpb, 1, head_dim + ) + index_cache[page_ids.long()] = torch.randn( + total, tpb, 1, head_dim, generator=gen, device=DEV + ).to(dtype) + assert index_cache.stride() == (page.stride(0), head_dim, head_dim, 1) + else: + index_cache = torch.randn(pool_pages, tpb, 1, head_dim, generator=gen, device=DEV).to(dtype) + + if spec["q_layout"] == "scratch_slice": + # [T, index heads | 1 k | 1 v, D], q is the first `heads` slots. + scratch = torch.randn(rows, heads + 2, head_dim, generator=gen, device=DEV).to(dtype) + q = scratch[:, :heads] + else: + q = torch.randn(rows, heads, head_dim, generator=gen, device=DEV).to(dtype) + if spec["q_layout"] != "packed" and rows > 1: # one row is trivially contiguous + assert not q.is_contiguous() + + return dict( + q=q, + index_cache=index_cache, + block_table=block_table, + query_positions=positions.to(spec["positions_dtype"]), + request_indices=slot_of.to(DEV)[req].to(spec["request_dtype"]), + tokens_per_block=spec["tokens_per_block"], + compress_ratio=spec["compress_ratio"], + only_visible_blocks=spec["only_visible_blocks"], + context_rows=spec["context_rows"], + use_pdl=spec.get("use_pdl", False), + prefetch=spec.get("prefetch", False), + rows_per_request=spec.get("call_rows_per_request", 1), + ) + + +def test_visible_split_metadata_offsets_reuse_the_compiled_kernel() -> None: + """Mixed-step metadata offsets preserve results without new alignment variants.""" + cell = next(c for c in op.CELLS if c.spec["branch"] == "visible_split") + args = _build(cell.spec, seed=81) + rows = args["q"].shape[0] + position_storage = torch.zeros(rows + 4, dtype=args["query_positions"].dtype, device=DEV) + request_storage = torch.zeros(rows + 4, dtype=args["request_indices"].dtype, device=DEV) + compiled = [] + for position_offset, request_offset in ((0, 0), (0, 1), (1, 1), (1, 3)): + positions = position_storage[position_offset : position_offset + rows] + requests = request_storage[request_offset : request_offset + rows] + positions.copy_(args["query_positions"]) + requests.copy_(args["request_indices"]) + shifted = {**args, "query_positions": positions, "request_indices": requests} + with _watch_launches() as watcher: + got = op(**shifted) + torch.cuda.synchronize() + assert watcher.branch == "visible_split" + _compare(got, op.reference(**shifted)) + cache = _entry._qsa_index_scores_visible_split_kernel.device_caches[0][0] + compiled.append(len(cache)) + assert len(set(compiled)) == 1, compiled + + +def _assert_reaches_the_pool_ends(args, spec) -> None: + """A measured-pool cell scores keys on the pool's first and last page, and the last page + lies past 2^31 elements from the plane's base, so a 32-bit offset would wrap. + """ + cache, pool_pages = args["index_cache"], spec["pool_pages"] + tpb, dim = spec["tokens_per_block"], spec["head_dim"] + page_stride = spec["pool_buffers"] * tpb * dim + spec["page_pad"] + assert cache.shape[0] == pool_pages and cache.stride() == (page_stride, dim, dim, 1) + assert (pool_pages - 1) * page_stride >= 2**31 + cr = args["compress_ratio"] + logical = (torch.arange(_columns(args), device=DEV) * cr + cr - 1) // tpb + pages = args["block_table"][args["request_indices"].long()][:, logical] + read = pages[_visible(args)] + assert bool((read == 0).any()) and bool((read == pool_pages - 1).any()), "a pool end is unread" + + +def _branch_rule(rows, spec): + heads = spec["heads"] + if ( + spec.get("call_rows_per_request", 1) == 1 + and spec["only_visible_blocks"] + and not spec.get("prefetch", False) + and rows <= DECODE_ROW_LIMIT + and SPLIT_DOT_WIDTH % heads == 0 + ): + return "visible_split" + packs = DOT_TILE_WIDTH % heads == 0 and DOT_TILE_WIDTH // heads > 1 + if spec["context_rows"] and rows > DECODE_ROW_LIMIT and packs: + return "query_tile" + grouped = ( + spec.get("call_rows_per_request", 1) > 1 + and spec["only_visible_blocks"] + and not spec.get("prefetch", False) + and heads in (2, 4, 8) + ) + return "request_group" if grouped else "row_wise" + + +def _columns(args): + return args["block_table"].shape[1] * args["tokens_per_block"] // args["compress_ratio"] + + +def _engine_scores(**args): + """The engine's own launcher on the entry's arguments; `rows_per_request` is the entry's.""" + args.pop("rows_per_request", None) + return _qsa_kernels.triton_qsa_paged_index_scores(**args) + + +def _visible(args): + blocks = torch.arange(_columns(args), device=DEV) + bound = (args["query_positions"].long() + 1) // args["compress_ratio"] + return blocks[None, :] < bound[:, None] + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=sum(map(ord, cell.why))) + rows = args["q"].shape[0] + assert cell.spec["branch"] == _branch_rule(rows, cell.spec), "the cell names the wrong kernel" + if "pool_pages" in cell.spec: + _assert_reaches_the_pool_ends(args, cell.spec) + with _watch_launches() as watcher: + out = op(**args) + torch.cuda.synchronize() + assert watcher.branch == cell.spec["branch"] + assert out.shape == (rows, _columns(args)) and out.dtype == torch.float32 + _compare(out, op.reference(**args)) + + +#: The decode and verify cells over a measured pool: the sites the step's CUDA graph captures. +_CAPTURED = tuple(c for c in op.CELLS if "pool_pages" in c.spec and not c.spec["context_rows"]) + + +#: The visible-split cells on the tune's table: their grid is fixed at capture, the windows not. +_SPLIT_CAPTURED = tuple( + c for c in _CAPTURED if c.spec["branch"] == "visible_split" and c.spec["pages"] > 516 +) + + +@pytest.mark.parametrize("cell", _SPLIT_CAPTURED, ids=[c.why[:44] for c in _SPLIT_CAPTURED]) +def test_captured_visible_split_follows_the_positions(cell) -> None: + """The visible-split grid is sized by the device, and each program reads its rows' + positions when it runs: capture at the cell's positions, then replay with every row's + window shrunk by a different amount (to empty for some) and grown back, and match the + reference on every specified column each time. + """ + args = _build(cell.spec, seed=sum(map(ord, cell.why)) + 2) + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + op(**args) + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(**args) + positions = args["query_positions"] + original = positions.clone() + gen = torch.Generator(device=DEV).manual_seed(5) + for trial in range(3): + if trial < 2: + keep = torch.rand(original.shape, generator=gen, device=DEV) ** (trial + 1) + positions.copy_((original.double() * keep).to(original.dtype)) + else: + positions.copy_(original) + graph.replay() + torch.cuda.synchronize() + _compare(captured, op.reference(**args)) + + +def test_prefetch_reads_the_fresh_anchor_key_and_q_after_the_wait() -> None: + """`prefetch` reads pages and keys before the PDL wait. Under graph replay a producer + rewrites q and, at rows whose position ends a group, the last visible block's anchor + key -- what the decode prologue writes -- right before the op; both must be seen. + The other keys are static, as they are in the decode step. Without prefetch the + entry takes its visible-split kernel, so the plain row-wise launch it is held to + bitwise is the engine callable's. + """ + args = _build(_case(4, 4, only_visible_blocks=True), seed=411) + cr, tpb = args["compress_ratio"], args["tokens_per_block"] + pos = args["query_positions"].clone() + assert bool((pos >= cr).all()) + ends = torch.tensor([True, False, True, True], device=DEV) + args["query_positions"] = torch.where(ends, pos - (pos + 1) % cr, pos).to(pos.dtype) + pos = args["query_positions"].long() + assert bool(((pos + 1) % cr == 0).eq(ends).all()) + pages = args["block_table"][args["request_indices"].long(), pos // tpb].long() + anchors = (pages[ends], (pos % tpb)[ends]) + visible = _visible(args) + plain = _engine_scores(**args) + assert torch.equal(op(**{**args, "use_pdl": True, "prefetch": True})[visible], plain[visible]) + src_q = args["q"].clone() + src_k = args["index_cache"][anchors].clone() + + def step(): + args["q"].copy_(src_q * 1.0) + args["index_cache"][anchors] = src_k * 1.0 + return op(**{**args, "use_pdl": True, "prefetch": True}) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + step() + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_out = step() + for trial in range(3): + fresh = _build(_case(4, 4, only_visible_blocks=True), seed=412 + trial) + src_q.copy_(fresh["q"]) + src_k.copy_(torch.randn(src_k.shape, device=DEV).to(src_k.dtype)) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(graph_out[visible], _engine_scores(**args)[visible]), f"replay {trial}" + _compare(graph_out, op.reference(**args)) + + +@pytest.mark.parametrize( + "group, rows, heads, mixed, use_pdl", + [ + (2, 16, 4, False, False), + (3, 13, 4, False, False), + (4, 64, 4, False, True), + (4, 62, 4, False, False), + (4, 64, 4, True, False), + (8, 40, 4, False, False), + (8, 36, 4, True, True), + (4, 16, 2, False, False), + (4, 16, 8, False, False), + (2, 16, 16, False, False), # 16 heads stay row-wise + (4, 256, 4, False, True), + ], +) +def test_request_group_is_bitwise_the_row_wise_kernel(group, rows, heads, mixed, use_pdl) -> None: + """`rows_per_request` picks the request-group kernel; its scores are the row-wise + kernel's bit patterns on every visible column. Runs of `group` rows of one request + (a partial last run when `group` does not divide the rows), or with `mixed` a + random request a row, so runs that name several requests gather row by row. + """ + spec = _case(rows, 8, heads=heads, pages=_WIDTH_PAGES, only_visible_blocks=True) + spec.update(q_layout="scratch_slice", use_pdl=use_pdl) + args = _build(spec, seed=97 * group + rows + heads) + slots = torch.unique(args["request_indices"]) + gen = torch.Generator(device=DEV).manual_seed(rows) + if mixed: + pick = torch.randint(0, slots.numel(), (rows,), generator=gen, device=DEV) + else: + pick = (torch.arange(rows, device=DEV) // group) % slots.numel() + args["request_indices"] = slots[pick].to(args["request_indices"].dtype) + # Every row sees at least one block (every request maps the whole table). + args["query_positions"] = args["query_positions"].clamp_min(args["compress_ratio"] - 1) + visible = _visible(args) + with _watch_launches() as watcher: + row_wise = _engine_scores(**args) + assert watcher.branch == "row_wise" + with _watch_launches() as watcher: + grouped = op(**{**args, "rows_per_request": group}) + assert watcher.branch == ("request_group" if heads in (2, 4, 8) else "row_wise") + assert torch.equal(grouped[visible].view(torch.int32), row_wise[visible].view(torch.int32)), ( + "the request-group kernel's bits differ from the row-wise kernel's" + ) + _compare(grouped, op.reference(**args)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_store.py similarity index 56% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_store.py index d6d7c5b03b47..3852fbcc9523 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_index_store.py @@ -8,18 +8,27 @@ its allocation, rows past the batch holding a previous step's plausible page ids. Both planes are pre-filled with a sentinel so a stray write shows. A `hole` cell appends one row that resolves to no slot. + +The default target's cells store into the stock manager's bf16 layout -- one +layer's plane of the 12 sparse layers' coalesced index pool, the other 11 +checked untouched -- at the measured tep4 pool, with the batch's last row on the +pool's last page. `test_the_stock_managers_planes_are_the_cells_layout` builds +the stock manager itself (through the target's pool-layout test) and stores +through its own views. """ +import importlib.util import zlib +from pathlib import Path import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_index_store import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_index_store import ( # noqa: E501 qsa_paged_index_store as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 assert torch.cuda.is_available(), "qsa_paged_index_store requires a CUDA device" @@ -80,8 +89,43 @@ def _shared_page_planes(spec, extras: dict): return index_cache, position_cache -def _build(spec, seed: int): - """Real planes, table and rows for one cell; returns the op's kwargs and extras.""" +def _coalesced_planes(spec, extras: dict): + """The bf16 index pool of the default target: the `layers` sparse layers' INDEX_KEY buffers + coalesced into one allocation, so the cell's plane (layer `layer`) is a strided view, and the + position plane its own contiguous pool. The whole allocation holds `KEY_FILL`; the pool is + recorded for the untouched check.""" + layers, layer, tpb, hkv, head_dim, axes, pages = ( + spec[k] for k in ("layers", "layer", "tokens_per_block", "hkv", "head_dim", "axes", "pages") + ) + pool = torch.full( + (pages, layers, tpb, hkv, head_dim), KEY_FILL, dtype=torch.bfloat16, device=DEV + ) + index_cache = pool[:, layer] + assert index_cache.stride() == (spec["page_stride"], hkv * head_dim, head_dim, 1) + position_cache = torch.full((pages, tpb, axes), POS_FILL, dtype=torch.int32, device=DEV) + extras["coalesced"] = (pool, layer) + return index_cache, position_cache + + +def _other_layers_untouched(pool: torch.Tensor, layer: int) -> None: + """Every other layer's plane in the coalesced pool still holds `KEY_FILL`, checked a band of + pages at a time so the check stays a fraction of the pool's size.""" + fill = int(_bits(torch.tensor([KEY_FILL], dtype=torch.bfloat16))[0]) + bits = pool.view(torch.int16) + for start in range(0, pool.shape[0], 16384): + band = bits[start : start + 16384] + for other in range(pool.shape[1]): + if other != layer: + assert bool((band[:, other] == fill).all()), ( + f"layer {other}'s index plane moved in pages {start}..{start + band.shape[0]}" + ) + + +def _build(spec, seed: int, planes=None): + """Real planes, table and rows for one cell; returns the op's kwargs and extras. + + `planes`, if given, is an `(index_cache, position_cache)` pair to store into instead of the + cell's own (a stock manager's views).""" gen = torch.Generator(device="cpu").manual_seed(seed) tpb, hkv, head_dim, axes = (spec[k] for k in ("tokens_per_block", "hkv", "head_dim", "axes")) slots, width, pages = spec["slots"], spec["width"], spec["pages"] @@ -99,7 +143,17 @@ def _build(spec, seed: int): table[r, :cols] = order[used : used + cols] used += cols allocated_cols.append(cols) - assert used <= pages and len(segments) < slots, "pool or table too small for the batch" + assert used <= pages and len(segments) <= slots, "pool or table too small for the batch" + if spec.get("straddle"): + # The batch's last row stores into the pool's last page; every other row keeps its page + # from the permutation, so the rows land on both sides of 2^31 elements of the plane. + last_start, last_rows = segments[-1] + column = (last_start + last_rows - 1) // tpb + live = table[: len(segments)] + holder = (live == pages - 1).nonzero() + if len(holder): + live[holder[0, 0], holder[0, 1]] = live[len(segments) - 1, column] + live[len(segments) - 1, column] = pages - 1 if spec["hole"] is not None: request, position = _hole(spec, table, allocated_cols) requests.append(request) @@ -113,9 +167,12 @@ def _build(spec, seed: int): token_k = projection[:, spec["index_heads"] :] if spec["key_layout"] == "fused_row": # The target's fused input projection row, [q | k | v | gate | index q, k | - # pad] = 110 heads of 128: the raw index key is head 108. - row = torch.randn(rows, 110, head_dim, generator=gen, dtype=torch.float32) - token_k = row.to(device=DEV, dtype=spec["key_dtype"])[:, 108 : 108 + hkv] + # pad of one index head], in heads of 128: the raw index key is the second to + # last (108 of 110 at 14080 wide; 32 of 34 at the default target's 4352). + heads = spec.get("row_width", 14080) // head_dim + row = torch.randn(rows, heads, head_dim, generator=gen, dtype=torch.float32) + token_k = row.to(device=DEV, dtype=spec["key_dtype"])[:, heads - 2 : heads - 2 + hkv] + assert token_k.stride(0) == heads * head_dim request_indices = torch.tensor(requests, dtype=spec["request_dtype"], device=DEV) logical_positions = torch.tensor(positions, dtype=spec["position_dtype"], device=DEV) @@ -131,8 +188,12 @@ def _build(spec, seed: int): extras = {} position_cache = None - if spec["cache_layout"] == "shared_page": + if planes is not None: + index_cache, position_cache = planes + elif spec["cache_layout"] == "shared_page": index_cache, position_cache = _shared_page_planes(spec, extras) + elif spec["cache_layout"] == "coalesced": + index_cache, position_cache = _coalesced_planes(spec, extras) else: index_cache = torch.full((pages, tpb, hkv, head_dim), KEY_FILL, device=DEV).to( torch.bfloat16 @@ -156,6 +217,9 @@ def _build(spec, seed: int): def test_certified_cells(cell, certifying): # noqa: F811 certifying(op) args, extras = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + if cell.spec.get("straddle"): + # The last row's page starts past 2^31 elements of the plane it is stored through. + assert (cell.spec["pages"] - 1) * args["index_cache"].stride(0) >= 2**31 ref = op.reference(**args) out = op(**args) assert out[0] is args["index_cache"] and out[1] is args["position_cache"] @@ -165,6 +229,8 @@ def test_certified_cells(cell, certifying): # noqa: F811 assert bool((extras["roles"].view(torch.int16) == fill).all()), ( "another index-K role of the shared page moved" ) + if "coalesced" in extras: + _other_layers_untouched(*extras["coalesced"]) #: The cell of 32 MTP7 verify blocks (256 rows): the segments, pool and table it runs over. @@ -178,6 +244,62 @@ def _small(seed: int, **overrides): return _build({**spec, **overrides}, seed=seed)[0] +def _pool_layout(): + """The target's pool-layout test, loaded by path: its `_manager` builds the stock manager the + engine builds on a rank of a member (core cache config, bf16 KV), and `ATTN_LAYERS` names the + sparse layers. Loaded rather than restated so the two cannot disagree on what the stock pool + is.""" + path = ( + Path(__file__).resolve().parents[1] + / "qwen3_8_flash_next_nvfp4__sm_103__default" + / "test_pool_layout.py" + ) + spec = importlib.util.spec_from_file_location("_qsa_index_store_pool_layout", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.mark.parametrize("member", ("tp4", "tep4", "dep4")) +def test_the_stock_managers_planes_are_the_cells_layout(member) -> None: + """The stock manager's own index and position planes at rank 0 of `member` have the layout + the member's cells model, and the op stores through them with the member's row width: four + rows, the last on the pool's last page, compared against `reference` over the whole plane, + every other sparse layer's plane untouched.""" + layout = _pool_layout() + cell = next(c.spec for c in op.CELLS if c.spec.get("member") == member) + tpb, hkv, head_dim, axes = (cell[k] for k in ("tokens_per_block", "hkv", "head_dim", "axes")) + with layout._manager(member, 0) as cm: + planes = [cm.get_index_k_buffer(i) for i in layout.ATTN_LAYERS] + index_cache, position_cache = planes[cell["layer"]], cm.get_qsa_position_buffer() + assert len(planes) == cell["layers"] + assert index_cache.stride() == (cell["page_stride"], hkv * head_dim, head_dim, 1) + assert position_cache.stride() == (tpb * axes, axes, 1) + assert position_cache.shape[0] == index_cache.shape[0] + plane_bytes = tpb * hkv * head_dim * index_cache.element_size() + offsets = [p.data_ptr() - planes[0].data_ptr() for p in planes] + assert offsets == [j * plane_bytes for j in range(len(planes))] + + for plane in planes: + plane.fill_(KEY_FILL) + position_cache.fill_(POS_FILL) + spec = dict( + cell, + segments=((0, 1), (31, 1), (33, 2)), + pages=index_cache.shape[0], + slots=8, + ) + args, _ = _build(spec, seed=71, planes=(index_cache, position_cache)) + assert int(args["block_table"][2, 1]) == index_cache.shape[0] - 1 + ref = op.reference(**args) + op(**args) + _compare((index_cache, position_cache), ref) + fill = int(_bits(torch.tensor([KEY_FILL], dtype=torch.bfloat16))[0]) + for j, plane in enumerate(planes): + if j != cell["layer"]: + assert bool((_bits(plane) == fill).all()), f"sparse layer {j}'s plane moved" + + def test_cuda_graph_capture_and_replay() -> None: """Captured once at the 32 MTP7 verify blocks (256 rows), replayed against new positions written into the same buffer.""" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store.py new file mode 100644 index 000000000000..dc4581afb8c9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store.py @@ -0,0 +1,346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_paged_kv_store entry. + +The paged state is real: an HND pool whose K and V role views are the caches, and the +engine's `[slots, width]` table -- a row per request slot, each live row's pages drawn from a +seeded permutation of the pool, `-1` past its allocation, and the rows past the batch holding a +previous step's plausible page ids. The pool is pre-filled with a sentinel so a stray write +shows. A cell's pool is one of two: + +- a contiguous `[pages, 2, Hkv, tpb, D]` pool, compared whole; +- with `pool="engine"`, the stock cache manager's layout at an engine's measured page count: + every attention layer's K and V planes in one page (24 role pages), the views at layer 47's + planes. At ~257k pages that spans ~100 GB, so it is a reserved virtual range with physical + memory mapped under the pages the rows touch only (`_EnginePool`); the pool's last page is + always one of them. The reference runs on a dense copy of those pages, and every other mapped + byte is checked unchanged. +""" + +import zlib + +import cuda.bindings.driver as cuda_drv +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_kv_store import ( # noqa: E501 + qsa_paged_kv_store as op, +) + +__extra_import_path__ = ["../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "qsa_paged_kv_store requires a CUDA device" + + +def _compare(out, ref) -> None: + """Bitwise over both role planes, compared as raw bits. + + A store is a move with one conversion, and `reference` performs the + same conversion, so there is no rounding to spend a tolerance on. In + the target both planes together are the whole pool allocation, so this + is also the stray-write check. + """ + for name, got, want in zip(("k_cache", "v_cache"), out, ref): + assert got.dtype == want.dtype, (name, got.dtype, want.dtype) + assert got.shape == want.shape, (name, got.shape, want.shape) + torch.testing.assert_close( + _bits(got), _bits(want), rtol=0, atol=0, msg=lambda m, n=name: f"{n}: {m}" + ) + + +DEV = "cuda" +FILL = -2.0 + + +def _bits(tensor: torch.Tensor) -> torch.Tensor: + as_int = {1: torch.uint8, 2: torch.int16, 4: torch.int32}[tensor.element_size()] + return tensor.contiguous().view(as_int) + + +def _check(result): + err, *value = result + assert err == cuda_drv.CUresult.CUDA_SUCCESS, err + return value[0] if len(value) == 1 else value + + +class _ArrayInterface: + def __init__(self, ptr: int, nbytes: int): + self.__cuda_array_interface__ = dict( + shape=(nbytes,), strides=None, typestr="|u1", data=(ptr, False), version=2 + ) + + +class _EnginePool: + """One attention layer's K and V views of the stock manager's QSA pool at `pages` pages. + + A page holds `role_pages` role planes of `[hkv, tpb, D]` (K then V of every attention + layer); the views are layer `layer_slot`'s. The range is reserved whole and physical memory + is mapped under the K and V planes of `touched` pages only, and under page 0's (a tensor's + base pointer). Reading an unmapped address faults, so only those pages are ever read. + """ + + def __init__(self, spec, touched): + dtype = spec["cache_dtype"] + elem = torch.empty((), dtype=dtype).element_size() + hkv, tpb, dim = spec["hkv"], spec["tokens_per_block"], spec["head_dim"] + self.pages, self.touched = spec["pages"], sorted(set(touched)) + self.role = hkv * tpb * dim + self.stride = spec["role_pages"] * self.role + self.own = spec["layer_slot"] * 2 * self.role + + def planes(page): + """This layer's K and V planes of `page`, in bytes of the range.""" + lo = (page * self.stride + self.own) * elem + return lo, lo + 2 * self.role * elem + + self.written = [planes(p) for p in self.touched] + + dev = torch.cuda.current_device() + prop = cuda_drv.CUmemAllocationProp() + prop.type = cuda_drv.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + prop.location.type = cuda_drv.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + prop.location.id = dev + gran = _check( + cuda_drv.cuMemGetAllocationGranularity( + prop, + cuda_drv.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED, + ) + ) + self.size = -(-self.pages * self.stride * elem // gran) * gran + self.va = int(_check(cuda_drv.cuMemAddressReserve(self.size, gran, 0, 0))) + merged = [] + for lo, hi in sorted({planes(0), *self.written}): + lo, hi = lo // gran * gran, min(-(-hi // gran) * gran, self.size) + if merged and lo <= merged[-1][1]: + merged[-1][1] = max(merged[-1][1], hi) + else: + merged.append([lo, hi]) + access = cuda_drv.CUmemAccessDesc() + access.location.type = cuda_drv.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + access.location.id = dev + access.flags = cuda_drv.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + self.maps, self.spans = [], merged + for lo, hi in merged: + handle = _check(cuda_drv.cuMemCreate(hi - lo, prop, 0)) + ptr = cuda_drv.CUdeviceptr(self.va + lo) + _check(cuda_drv.cuMemMap(ptr, hi - lo, 0, handle, 0)) + self.maps.append((ptr, hi - lo, handle)) + _check(cuda_drv.cuMemSetAccess(ptr, hi - lo, [access], 1)) + + self.raw = torch.as_tensor(_ArrayInterface(self.va, self.size), device=DEV) + pool = self.raw.view(dtype).as_strided( + (self.pages, 2, hkv, tpb, dim), + (self.stride, self.role, tpb * dim, dim, 1), + self.own, + ) + self.k, self.v = pool[:, 0], pool[:, 1] + self.index = torch.tensor(self.touched, dtype=torch.long, device=DEV) + self.fill() + + def fill(self) -> None: + """Every mapped byte to the sentinel's pattern, and remember it.""" + dtype = self.k.dtype + pattern = _bits(torch.full((2,), FILL, dtype=torch.bfloat16).to(dtype)) + for lo, hi in self.spans: + self.raw[lo:hi].view(pattern.dtype).fill_(int(pattern[0])) + torch.cuda.synchronize() + self.before = [self.raw[lo:hi].clone() for lo, hi in self.spans] + + def dense(self, table: torch.Tensor): + """`(k, v, table)`: copies of the touched pages, and `table` renumbered into them. + + An entry naming no touched page becomes `-1`. No row resolves to one, so the + reference reads the same pages it would over the whole pool. + """ + remap = torch.full((self.pages,), -1, dtype=torch.int32, device=DEV) + remap[self.index] = torch.arange(len(self.touched), dtype=torch.int32, device=DEV) + renamed = torch.where(table >= 0, remap[table.clamp_min(0).long()], -1).to(torch.int32) + return self.k[self.index].clone(), self.v[self.index].clone(), renamed.contiguous() + + def gathered(self): + return self.k[self.index], self.v[self.index] + + def assert_rest_unchanged(self) -> None: + """Every mapped byte outside the touched pages' K and V planes is what it was.""" + torch.cuda.synchronize() + for (lo, hi), before in zip(self.spans, self.before): + keep = torch.ones(hi - lo, dtype=torch.bool, device=DEV) + for p_lo, p_hi in self.written: + if p_hi > lo and p_lo < hi: + keep[max(p_lo, lo) - lo : min(p_hi, hi) - lo] = False + assert torch.equal(self.raw[lo:hi][keep], before[keep]), ( + f"bytes [{lo}, {hi}) of the pool changed outside the touched pages' planes" + ) + + def close(self) -> None: + torch.cuda.synchronize() + del self.k, self.v, self.raw + for ptr, size, handle in self.maps: + _check(cuda_drv.cuMemUnmap(ptr, size)) + _check(cuda_drv.cuMemRelease(handle)) + _check(cuda_drv.cuMemAddressFree(cuda_drv.CUdeviceptr(self.va), self.size)) + + +def _hole(spec, table, allocated_cols): + """One row that resolves to no slot: its (request, position).""" + assert spec["hole"] == "unallocated_page", spec["hole"] + assert int(table[0, allocated_cols[0]]) == -1 + return 0, allocated_cols[0] * spec["tokens_per_block"] + + +def _build(spec, seed: int): + """Real pool, table and rows for one cell: `(op kwargs, pool)`. + + `pool` is the contiguous pool tensor, or the `_EnginePool` under a `pool="engine"` cell. + """ + gen = torch.Generator(device="cpu").manual_seed(seed) + tpb, hkv, head_dim = spec["tokens_per_block"], spec["hkv"], spec["head_dim"] + slots, width, pages = spec["slots"], spec["width"], spec["pages"] + engine = spec.get("pool") == "engine" + segments = spec["segments"] + requests = [r for r, (_, n) in enumerate(segments) for _ in range(n)] + positions = [p for start, n in segments for p in range(start, start + n)] + + table = torch.randint(0, pages, (slots, width), generator=gen, dtype=torch.int32) + order = torch.randperm(pages, generator=gen).to(torch.int32) + if engine: + # The first row's page is the pool's last one, whose planes are its highest bytes. + first = segments[0][0] // tpb + last = int((order == pages - 1).nonzero()[0]) + order[[first, last]] = order[[last, first]] + used, allocated_cols = 0, [] + for r, (start, n) in enumerate(segments): + cols = (start + n - 1) // tpb + 1 + assert cols <= width, f"request {r} ends past the table's {width} columns" + table[r] = -1 + table[r, :cols] = order[used : used + cols] + used += cols + allocated_cols.append(cols) + assert used <= pages and len(segments) <= slots, "pool or table too small for the batch" + if spec["hole"] is not None: + request, position = _hole(spec, table, allocated_cols) + requests.append(request) + positions.append(position) + + rows = len(positions) + assert rows == spec.get("rows", rows), (rows, spec["rows"]) + qkv = torch.randn( + rows, + spec["heads_q"] + 2 * hkv + spec.get("row_pad_heads", 0), + head_dim, + generator=gen, + dtype=torch.float32, + ).to(device=DEV, dtype=spec["row_dtype"]) + assert qkv.stride(0) == spec.get("row_stride", qkv.stride(0)), qkv.stride(0) + k = qkv[:, spec["heads_q"] : spec["heads_q"] + hkv] + v = qkv[:, spec["heads_q"] + hkv : spec["heads_q"] + 2 * hkv] + touched = [int(table[r, p // tpb]) for r, p in zip(requests, positions)] + table = table.to(DEV) + + if engine: + pool = _EnginePool(spec, [t for t in touched if t >= 0]) + k_cache, v_cache = pool.k, pool.v + else: + pool = torch.full( + (pages, 2, hkv, tpb, head_dim), FILL, dtype=torch.bfloat16, device=DEV + ).to(spec["cache_dtype"]) + k_cache, v_cache = pool[:, 0], pool[:, 1] + args = dict( + k=k, + v=v, + k_cache=k_cache, + v_cache=v_cache, + request_indices=torch.tensor(requests, dtype=spec["request_dtype"], device=DEV), + logical_positions=torch.tensor(positions, dtype=spec["position_dtype"], device=DEV), + block_table=table, + tokens_per_block=tpb, + ) + return args, pool + + +def _engine_reference(args, pool: _EnginePool): + """The entry's reference over a dense copy of the touched pages.""" + dense_k, dense_v, renamed = pool.dense(args["block_table"]) + return op.reference(**{**args, "k_cache": dense_k, "v_cache": dense_v, "block_table": renamed}) + + +def _assert_reaches_past_int32(pool: _EnginePool) -> None: + """The highest touched page's planes lie past 2^32 elements of the view's base: a page + offset computed in 32 bits would land elsewhere.""" + top = max(pool.touched) + assert top == pool.pages - 1, (top, pool.pages) + assert top * pool.k.stride(0) >= 2**32, top * pool.k.stride(0) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args, pool = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + if isinstance(pool, _EnginePool): + try: + _assert_reaches_past_int32(pool) + before = pool.gathered() + before = (before[0].clone(), before[1].clone()) + ref = _engine_reference(args, pool) + out = op(**args) + assert out[0] is args["k_cache"] and out[1] is args["v_cache"] + got = pool.gathered() + _compare(got, ref) + pool.assert_rest_unchanged() + wrote = any(not torch.equal(_bits(g), _bits(b)) for g, b in zip(got, before)) + assert wrote, "a cell with rows must write" + print( + f"[{cell.why[:44]}] {len(pool.touched)} of {pool.pages} pages touched; " + f"top plane at byte {pool.written[-1][0]} of {pool.size}; " + f"{sum(hi - lo for lo, hi in pool.spans)} bytes mapped" + ) + finally: + pool.close() + return + before = pool.clone() + ref = op.reference(**args) + out = op(**args) + assert out[0] is args["k_cache"] and out[1] is args["v_cache"] + _compare(out, ref) + wrote = not torch.equal(_bits(pool), _bits(before)) + assert wrote == bool(cell.spec["segments"]), ( + "a cell with rows must write, an empty one must not" + ) + + +def _decode_args(seed: int, **overrides): + """The decode cell's layout over a bf16 pool, without its skipped row.""" + spec = {**op.CELLS[0].spec, "hole": None} + return _build({**spec, **overrides}, seed=seed) + + +#: The cell of 32 MTP7 verify blocks (256 rows): the segments, pool and table it runs over. +_MTP7_32 = next(c.spec for c in op.CELLS if c.why.startswith("fp8 decode MTP7: 32 requests")) +_MTP7_BATCH = {key: _MTP7_32[key] for key in ("segments", "pages", "slots")} + + +def test_cuda_graph_replay_follows_updated_positions() -> None: + """Captured once at 32 MTP7 verify blocks, replayed against new positions written + into the same buffer.""" + args, pool = _decode_args(seed=9, **_MTP7_BATCH) + positions = args["logical_positions"] + saved = positions.clone() + positions.zero_() + + def call(): + op(**args) + + for _ in range(3): # compile the Triton kernel before capture + call() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + call() + + pool.fill_(FILL) + positions.copy_(saved) + ref = op.reference(**args) + graph.replay() + torch.cuda.synchronize() + _compare((args["k_cache"], args["v_cache"]), ref) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py similarity index 88% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py index ce3ddec3389e..a03f5821bdfd 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_kv_store_nvfp4.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qsa_paged_kv_store_nvfp4 entry. +"""GPU test for the qsa_paged_kv_store_nvfp4 catalog entry. The paged state is real: one NVFP4 pool as the QSA cache manager serves it -- e2m1 data `[pages, 2, Hkv, tpb, D/2]` and e4m3 block scales `[pages, 2, Hkv, tpb, D/16]`, @@ -15,12 +15,12 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_kv_store_nvfp4 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_kv_store_nvfp4 import ( # noqa: E501 qsa_paged_kv_store_nvfp4 as op, ) from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 assert torch.cuda.is_available(), "qsa_paged_kv_store_nvfp4 requires a CUDA device" @@ -94,11 +94,6 @@ def _planes(args): return (args["k_cache"], args["v_cache"], args["k_scale_cache"], args["v_scale_cache"]) -def _reset(planes) -> None: - for plane, fill in zip(planes, (DATA_FILL, DATA_FILL, SCALE_FILL, SCALE_FILL)): - plane.fill_(fill) - - def _assert_other_roles_untouched(spec, pool) -> None: """Every role page of a slot but the call's K and V still holds its sentinel.""" data, scales = pool @@ -129,6 +124,16 @@ def _build(spec, seed: int, pool=None): used += cols allocated_cols.append(cols) assert used <= pages and len(segments) <= slots, "pool or table too small for the batch" + if spec.get("last_page") is not None: + # The named request's last page becomes the pool's last slot; the live entry that held + # it, if any, takes the swapped-out page, so the live pages stay distinct. + r = spec["last_page"] + column = allocated_cols[r] - 1 + live = table[: len(segments)] + holder = (live == pages - 1).nonzero() + if holder.numel(): + live[holder[0, 0], holder[0, 1]] = live[r, column] + live[r, column] = pages - 1 if spec["hole"] is not None: request, position = _hole(spec, table, allocated_cols) requests.append(request) @@ -169,60 +174,6 @@ def _snapshot(pool): return tuple(t.clone() for t in pool) -def _pools_equal(a, b) -> bool: - return all(torch.equal(x, y) for x, y in zip(a, b)) - - -def _graph_replay(spec, seed: int, graph_batches) -> None: - """Each G: capture one step of the first G requests, replay the next step. - - The next step lives in the captured buffers: new k/v rows, new pages in the - table, and every request's positions advanced by its rows, the way a decode - step moves. The replay must equal the reference and an eager call on the - same inputs, bit for bit. One pool serves the whole sweep, as the engine's does. - """ - pool = _pool(spec) - steps = {n for _, n in spec["segments"]} - assert len(steps) == 1, "a graph cell's requests all carry the same row count" - step = steps.pop() - for g in graph_batches: - sub = {**spec, "segments": spec["segments"][:g]} - args, _, _ = _build(sub, seed=seed + g, pool=pool) - following, _, _ = _build(sub, seed=seed + 1000 + g, pool=pool) - planes = _planes(args) - positions = args["logical_positions"] - next_positions = positions.clone() - positions.sub_(step) - assert int(positions.min()) >= 0 - - def call(): - op(**args) - - for _ in range(3): # compile the Triton kernel before capture - call() - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - call() - - for name in ("k", "v", "block_table"): - args[name].copy_(following[name]) - positions.copy_(next_positions) - _reset(planes) - op(**args) - eager = tuple(plane.clone() for plane in planes) - _reset(planes) - ref = op.reference(**args) - graph.replay() - torch.cuda.synchronize() - _compare(planes, ref) - assert all(torch.equal(p, e) for p, e in zip(planes, eager)), ( - f"G={g} x {step} rows: replay differs from eager" - ) - del graph - _assert_other_roles_untouched(spec, pool) - - def _thop_append(qkv, pool, table, requests, new, cached, num_contexts, p, scale_orig_quant): """The attention op's own append over the same rows (e4m3 output: the append does not depend on the output type).""" @@ -327,7 +278,31 @@ def _thop_equality(spec, seed: int) -> None: cached = [c + n for c, n in zip(cached, new)] -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def _wrote(planes) -> bool: + """Whether any byte of the call's four planes moved off its sentinel.""" + fills = (DATA_FILL, DATA_FILL, SCALE_FILL, SCALE_FILL) + return any(bool((plane != fill).any()) for plane, fill in zip(planes, fills)) + + +def _assert_straddles(spec, args) -> None: + """The request named by `last_page` stores into the pool's last slot, which lies past + 2^31 bytes of both the data and the scale plane: every offset has to be formed in 64 bits.""" + row = sum(n for _, n in spec["segments"][: spec["last_page"] + 1]) - 1 + column = int(args["logical_positions"][row]) // spec["tokens_per_block"] + request = int(args["request_indices"][row]) + assert int(args["block_table"][request, column]) == spec["pages"] - 1 + for name in ("k_cache", "k_scale_cache"): + assert (spec["pages"] - 1) * args[name].stride(0) >= 2**31, (name, args[name].stride()) + + +# A pool sized to a member's measured page count takes 47 GB of device memory. +_CELL_PARAMS = [ + pytest.param(c, marks=pytest.mark.high_cuda_memory) if c.spec.get("large") else c + for c in op.CELLS +] + + +@pytest.mark.parametrize("cell", _CELL_PARAMS, ids=[c.why[:44] for c in op.CELLS]) def test_certified_cells(cell, certifying): # noqa: F811 certifying(op) seed = zlib.crc32(cell.why.encode()) @@ -335,18 +310,77 @@ def test_certified_cells(cell, certifying): # noqa: F811 _thop_equality(cell.spec, seed=seed) return args, pool, _ = _build(cell.spec, seed=seed) - before = _snapshot(pool) + if cell.spec.get("last_page") is not None: + _assert_straddles(cell.spec, args) ref = op.reference(**args) out = op(**args) assert out[0] is args["k_cache"] and out[3] is args["v_scale_cache"] _compare(out, ref) - wrote = not _pools_equal(pool, before) - assert wrote == bool(cell.spec["segments"]), ( + assert _wrote(_planes(args)) == bool(cell.spec["segments"]), ( "a cell with rows must write, an empty one must not" ) _assert_other_roles_untouched(cell.spec, pool) +def _decode_args(seed: int, **overrides): + return _build({**op.CELLS[0].spec, **overrides}, seed=seed) + + +def _reset(planes) -> None: + for plane, fill in zip(planes, (DATA_FILL, DATA_FILL, SCALE_FILL, SCALE_FILL)): + plane.fill_(fill) + + +def _graph_replay(spec, seed: int, graph_batches) -> None: + """Each G: capture one step of the first G requests, replay the next step. + + The next step lives in the captured buffers: new k/v rows, new pages in the + table, and every request's positions advanced by its rows, the way a decode + step moves. The replay must equal the reference and an eager call on the + same inputs, bit for bit. One pool serves the whole sweep, as the engine's does. + """ + pool = _pool(spec) + steps = {n for _, n in spec["segments"]} + assert len(steps) == 1, "a graph cell's requests all carry the same row count" + step = steps.pop() + for g in graph_batches: + sub = {**spec, "segments": spec["segments"][:g]} + args, _, _ = _build(sub, seed=seed + g, pool=pool) + following, _, _ = _build(sub, seed=seed + 1000 + g, pool=pool) + planes = _planes(args) + positions = args["logical_positions"] + next_positions = positions.clone() + positions.sub_(step) + assert int(positions.min()) >= 0 + + def call(): + op(**args) + + for _ in range(3): # compile the Triton kernel before capture + call() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + call() + + for name in ("k", "v", "block_table"): + args[name].copy_(following[name]) + positions.copy_(next_positions) + _reset(planes) + op(**args) + eager = tuple(plane.clone() for plane in planes) + _reset(planes) + ref = op.reference(**args) + graph.replay() + torch.cuda.synchronize() + _compare(planes, ref) + assert all(torch.equal(p, e) for p, e in zip(planes, eager)), ( + f"G={g} x {step} rows: replay differs from eager" + ) + del graph + _assert_other_roles_untouched(spec, pool) + + def test_cuda_graph_replay_matches_eager() -> None: """Decode MTP under a graph: 8 requests x 8 rows in the prefill cell's trunk layout, captured once, replayed on the next step's rows, pages and positions.""" @@ -355,10 +389,6 @@ def test_cuda_graph_replay_matches_eager() -> None: _graph_replay(spec, seed=8, graph_batches=(8,)) -def _decode_args(seed: int, **overrides): - return _build({**op.CELLS[0].spec, **overrides}, seed=seed) - - def test_op_rejects_strided_request_indices() -> None: """The callable's own loud error -- why the entry carries no `is_valid` -- for the metadata a flat-pointer kernel would otherwise read at the wrong rows.""" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa.py new file mode 100644 index 000000000000..91ea1f705957 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa.py @@ -0,0 +1,353 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_paged_sparse_gqa catalog entry. + +The paged state is real: one `[pages, slot_pages, kv, tpb, D]` HND pool whose +layer planes are the call's K and V -- tp1's cells one layer's adjacent pair, +the default target's cells the engine's 24-role-page trunk, sized so the +requests' pages straddle 2^31 elements -- and the engine's slot table: a row +per request slot, assigned by a seeded permutation, `-1` past each sequence. +Selection rows are built the way QSA's top-k + expansion lay them out: +complete groups expanded to tokens, then the trailing partial group, then +`-1`, all packed at the front. + +The call has two kernels behind it (and a merge). The `launches` fixture +records which ones a call indexed with a grid, and whether each launched with +PDL, so every cell asserts its branch (and, where it names `splits`, the +split count). A cell marked `capture` is also captured in a CUDA graph, +refilled in place and replayed, as the engine's decode step is. +""" + +import math + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa import ( # noqa: E501 + qsa_paged_sparse_gqa as op, +) +from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels + +__extra_import_path__ = ["../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "qsa_paged_sparse_gqa requires a CUDA device" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """4 ulp of each (row, head)'s magnitude, 2 ulp relative RMS, in q's dtype. + + The kernel rounds its online-softmax probabilities to the cache dtype + before P.V, so the error tracks the magnitude of the values averaged, + not of the output: against an fp32 reference it measured 1.2e-3 on rows + averaging a thousand tokens (outputs 0.14 to 0.4) and 1.2e-2 on rows + averaging a handful (outputs to 5). Scaled by each head row's own + magnitude that is one band. Against this reference (fp32, rounded once) + every ungated cell measured at most 2.0 element ulp -- a single final-rounding + flip, since bf16's step is 2 of these units at the bottom of a binade -- + and at most 0.62 ulp RMS: 2x and 3.2x margin. Six wrong references (no + scale, interleaved heads, padding attended, identity pages, K/V + swapped, page/slot transposed) were each measured more than 0.1 off. + + The default target's cells (6/1 heads at 1 to 8192 rows, 24/2 over an + e4m3 pool) sit in the same band, and the table shifted by one page, the + softmax scale dropped and one selected token dropped a row each fall + outside it. + + A gated output is rounded twice, the attention and then the gated + product, each by at most one bf16 ulp of the value rounded, so its bound + is near the band's 4 element ulp. The gate not applied, or read one head + early, falls far outside it; one selected token dropped at the 1-row + gated cell is refused on the RMS gate. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + if out.numel() == 0: + return + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) + + +DEV = "cuda" + +_KERNELS = { + "_qsa_paged_sparse_gqa_kernel": "prefill", + "_qsa_paged_sparse_gqa_splitk_kernel": "splitk", + "_qsa_merge_splitk_kernel": "merge", +} + + +class _LaunchRecorder: + """Records a Triton kernel's grid indexing and whether it launched with PDL, then + delegates to the kernel.""" + + def __init__(self, inner, name, log): + self._inner, self._name, self._log = inner, name, log + + def __getitem__(self, grid): + entry = [self._name, tuple(grid), False] + self._log.append(entry) + launch = self._inner[grid] + + def run(*args, **kwargs): + entry[2] = bool(kwargs.get("launch_pdl", False)) + return launch(*args, **kwargs) + + return run + + +@pytest.fixture +def launches(monkeypatch): + """Which branch kernels the next call launches, in order: [name, grid, PDL].""" + log: list = [] + for name, tag in _KERNELS.items(): + inner = getattr(_qsa_kernels, name) + monkeypatch.setattr(_qsa_kernels, name, _LaunchRecorder(inner, tag, log)) + return log + + +# -------------------------------------------------------------------------- +# building real state from a cell spec +# -------------------------------------------------------------------------- + + +def _selected_tokens(positions, seed, compress_ratio, token_topk): + """QSA's expansion layout, built in one pass. + + A random `min(available, block_topk)`-subset of the complete visible + groups, each expanded to its `compress_ratio` tokens, then the incomplete + trailing group, then `-1` -- packed at the front, as the producer emits. + """ + block_topk = token_topk // compress_ratio + final_topk = token_topk + compress_ratio - 1 + gen = torch.Generator(device=DEV).manual_seed(seed) + rows = positions.numel() + visible = positions.long() + 1 + available = visible // compress_ratio + taken = available.clamp(max=block_topk) + width = max(block_topk, int(available.max().item()) if rows else 0) + + scores = torch.rand(rows, width, generator=gen, device=DEV) + group_ids = torch.arange(width, device=DEV)[None, :] + scores = torch.where(group_ids < available[:, None], scores, torch.inf) + chosen = scores.argsort(dim=1)[:, :block_topk] + + column = torch.arange(final_topk, device=DEV)[None, :] + group_of_column = (column // compress_ratio).clamp(max=block_topk - 1).expand(rows, -1) + body = torch.gather(chosen, 1, group_of_column) * compress_ratio + column % compress_ratio + in_body = column < (taken * compress_ratio)[:, None] + tail_column = column - (taken * compress_ratio)[:, None] + tail = (available * compress_ratio)[:, None] + tail_column + in_tail = ~in_body & (tail_column >= 0) & (tail < visible[:, None]) + selected = torch.where(in_body, body, torch.where(in_tail, tail, torch.full_like(body, -1))) + return selected.to(torch.int32) + + +#: NaN's bit pattern per pool dtype: filled through an integer view, since not every fill +#: kernel takes an e4m3 tensor. +_NAN_BITS = {torch.bfloat16: (torch.int16, 0x7FC0), torch.float8_e4m3fn: (torch.uint8, 0x7F)} + + +def _nan_pool(shape, dtype): + bits_dtype, nan = _NAN_BITS[dtype] + return torch.full(shape, nan, dtype=bits_dtype, device=DEV).view(dtype) + + +def _build(spec, seed): + """Real paged state for one spec, keyed by `__call__`'s parameter names. + + The pool is `[pool_pages, slot_pages, kv, tpb, D]` and the call's K and V are layer + `layer`'s two role planes: tp1's cells keep one layer's adjacent pair (2 role pages a + page), the default target's cells the engine's trunk (24). The requests sit on the + pool's last pages; every other page, and every other plane, is NaN, so an address that + strayed could not come back finite. + """ + gen = torch.Generator(device=DEV).manual_seed(seed) + tpb, kv, d, q_heads = ( + spec["tokens_per_block"], + spec["kv_heads"], + spec["head_dim"], + spec["q_heads"], + ) + segments = list(spec["segments"]) + positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments]) + req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)]) + rows, slots = positions.numel(), spec["slots"] + assert len(segments) <= slots + + slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(segments)] + mapped = [math.ceil((s + n) / tpb) for s, n in segments] + assert max(mapped) <= spec["pages"] + needed = sum(mapped) + 7 + pool_pages = spec.get("pool_pages", needed) + assert pool_pages >= needed + first = pool_pages - needed + page_ids = first + torch.randperm(needed, generator=gen, device=DEV).to(torch.int32) + block_table = torch.full((slots, spec["pages"]), -1, dtype=torch.int32, device=DEV) + cursor = 0 + for slot, count in zip(slot_of.tolist(), mapped): + block_table[slot, :count] = page_ids[cursor : cursor + count] + cursor += count + + slot_pages, layer = spec.get("slot_pages", 2), spec.get("layer", 0) + planes = slice(2 * layer, 2 * layer + 2) + pool = _nan_pool((pool_pages, slot_pages, kv, tpb, d), spec["cache_dtype"]) + pool[first:, planes] = torch.randn(needed, 2, kv, tpb, d, generator=gen, device=DEV).to( + spec["cache_dtype"] + ) + k_cache, v_cache = pool[:, 2 * layer], pool[:, 2 * layer + 1] + + q_dtype = spec["q_dtype"] + fused = spec["q_layout"] == "fused_row" + # The row q is read off: the packed `[q | k | v]` buffer, or the target's fused input + # projection row `[q | k | v | gate | index q, k | pad]` (`row_heads` heads of D). + row_heads = spec.get("row_heads", 55 if fused else q_heads + 2 * kv) + row = torch.randn(rows, row_heads * d, generator=gen, device=DEV).to(q_dtype) + q = row.view(rows, row_heads, d)[:, :q_heads] + + selected = _selected_tokens(positions, seed + 3, spec["compress_ratio"], spec["token_topk"]) + request_indices = slot_of.to(DEV)[req].to(spec["request_dtype"]) + + gate = None + if spec.get("gate") == "fused_row": + assert fused, "the gate is read off the fused row" + # The gate columns of the same row, after q, k and v. + start = (q_heads + 2 * kv) * d + gate = row[:, start : start + q_heads * d] + + args = dict( + q=q, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + selected_tokens=selected, + request_indices=request_indices, + tokens_per_block=tpb, + softmax_scale=d**-0.5, + ) + if spec["prefill_pair"]: + args.update(query_positions=positions, compress_ratio=spec["compress_ratio"]) + if gate is not None: + args.update(output_gate=gate) + return args + + +def _bound(positions, spec): + """The prefill kernel's per-row scan bound, from the positions alone.""" + cr = spec["compress_ratio"] + visible = positions.long() + 1 + return (visible // cr).clamp(max=spec["token_topk"] // cr) * cr + visible % cr + + +def _assert_straddles(args, edge: int) -> None: + """The pages the cell's rows map put the K and V planes' element offsets on both sides + of `edge`: what the cell claims about its pool.""" + table = args["block_table"] + pages = table[table >= 0].long() + for name in ("k_cache", "v_cache"): + plane = args[name] + ends = plane.storage_offset() + pages * plane.stride(0) + plane[0].numel() - 1 + assert bool((ends >= edge).any()), f"{name}: no mapped page lies past {edge}" + assert bool((ends < edge).any()), f"{name}: every mapped page lies past {edge}" + + +def _is_prefix_packed(table): + valid = table >= 0 + columns = torch.arange(table.shape[1], device=table.device)[None, :] + return (valid == (columns < valid.sum(1)[:, None])).all(1) + + +def _refill(args, spec, seed) -> None: + """The engine's next step, written into the call's own buffers: new q and gate values + (the fused row's columns), a new selection at the same positions, and new K/V in every + mapped page.""" + gen = torch.Generator(device=DEV).manual_seed(seed) + for name in ("q", "output_gate"): + if name in args: + view = args[name] + view.copy_(torch.randn(view.shape, generator=gen, device=DEV).to(view.dtype)) + positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in spec["segments"]]) + args["selected_tokens"].copy_( + _selected_tokens(positions, seed + 3, spec["compress_ratio"], spec["token_topk"]) + ) + table = args["block_table"] + pages = table[table >= 0].long() + for name in ("k_cache", "v_cache"): + plane = args[name] + bits = _NAN_BITS[plane.dtype][0] + fresh = torch.randn((pages.numel(), *plane.shape[1:]), generator=gen, device=DEV) + plane.view(bits)[pages] = fresh.to(plane.dtype).view(bits) + + +def _replay_in_a_graph(args, spec, seed, launches, first) -> None: + """The decode site's capture: the call captured into a CUDA graph (its split-K partials + allocated from the graph's pool), the inputs refilled in place, the graph replayed. + The replay must take the eager call's launches (`first` is that call's output), answer + the new inputs within the band, and match an eager call on them bitwise.""" + eager = [tuple(entry) for entry in launches] + del launches[:] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(**args) + assert [tuple(entry) for entry in launches] == eager, launches + _refill(args, spec, seed) + graph.replay() + torch.cuda.synchronize() + replayed = captured.clone() + assert not torch.equal(replayed, first), "the replay did not read the refilled inputs" + _compare(replayed, op.reference(**args)) + assert torch.equal(replayed, op(**args)), "the replay differs from an eager call" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811 + certifying(op) + spec = cell.spec + args = _build(spec, seed=sum(map(ord, cell.why))) + rows = args["q"].shape[0] + assert rows == spec.get("rows", rows) + if "row_heads" in spec: + # q, and the gate when fused, are views of one fused row `row_heads` heads wide. + q = args["q"] + assert q.stride(0) == spec["row_heads"] * spec["head_dim"] + if "output_gate" in args: + gate = args["output_gate"] + assert gate.stride(0) == q.stride(0) + assert gate.untyped_storage().data_ptr() == q.untyped_storage().data_ptr() + if "rows_per_request" in spec: + assert {n for _, n in spec["segments"]} == {spec["rows_per_request"]} + if "offset_edge" in spec: + _assert_straddles(args, spec["offset_edge"]) + if spec["prefill_pair"]: + # The precondition the op cannot check: front-packed, and exactly as + # many valid entries as the prefill kernel will scan. + selected = args["selected_tokens"] + assert bool(_is_prefix_packed(selected).all()) + valid = (selected >= 0).sum(1) + assert torch.equal(valid, _bound(args["query_positions"], spec)) + + out = op(**args) + assert [name for name, _, _ in launches] == list(spec["launches"]) + assert launches[0][1][:2] == (rows, spec["kv_heads"]) + if "splits" in spec: + assert launches[0][1][2] == spec["splits"], launches + assert all(pdl == spec.get("pdl", False) for _, _, pdl in launches), launches + if "output_gate" in args: + gated = (args["q"], args["k_cache"], args["selected_tokens"], args["output_gate"]) + assert op.fuses_output_gate(*gated) + assert out.shape == (rows, spec["q_heads"], spec["head_dim"]) + assert out.dtype == args["q"].dtype and out.is_contiguous() + assert not bool(out.isnan().any()) + _compare(out, op.reference(**args)) + if spec.get("capture"): + seed = sum(map(ord, cell.why)) + 1 + _replay_in_a_graph(args, spec, seed=seed, launches=launches, first=out) + + +#: The first cell is this checkpoint's own decode geometry; the tests below +#: start from it. +CKPT = op.CELLS[0].spec +assert (CKPT["q_heads"], CKPT["kv_heads"], CKPT["head_dim"]) == (24, 2, 256) + + +def _with(**overrides): + return {**CKPT, **overrides} diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py similarity index 87% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py index 7d24c80e8471..7ec9a3535fbc 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_decode.py @@ -2,22 +2,21 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the qsa_paged_sparse_gqa_decode entry: cells against the composition entry's fp32 reference in that entry's band (4 ulp of each head row's magnitude, 2 ulp -relative RMS), the selection's read-after-wait under graph replay with PDL, and the row -guard.""" +relative RMS), and the selection's read-after-wait under graph replay with PDL.""" import zlib import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_sparse_gqa_decode import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_decode import ( # noqa: E501 qsa_paged_sparse_gqa_decode as op, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_expand_decode import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_topk_expand_decode import ( # noqa: E501 qsa_topk_expand_decode, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import ( # noqa: E402,F401 — path above; fixture by name assert_within_ulp, certifying, @@ -109,10 +108,3 @@ def step(): return op(*args[:4], sel, *args[5:]) _replays_like_eager(step, lambda: scores.copy_(torch.rand(scores.shape, device=DEV))) - - -def test_guard_refuses_rows_past_max_rows() -> None: - args = _build(dict(pos=(8700,), gated=True), seed=5) - many = [a.expand(300, *a.shape[1:]) if i in (0, 4, 5) else a for i, a in enumerate(args[:6])] - with validating(op), pytest.raises(AssertionError): - op(*many, *args[6:8], None) # 300 rows: past MAX_ROWS diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py similarity index 84% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py index ab4313b61e21..f7db6f517d43 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qsa_paged_sparse_gqa_nvfp4 entry. +"""GPU test for the qsa_paged_sparse_gqa_nvfp4 catalog entry. The paged state is real: one NVFP4 pool as the QSA cache manager serves it -- e2m1 data `[pages, 2, kv, tpb, D/2]` and e4m3 block scales `[pages, 2, kv, tpb, @@ -15,20 +15,19 @@ """ import math -import re import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention import ( qsa_paged_sparse_gqa_nvfp4 as _entry, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_sparse_gqa_nvfp4 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4 import ( # noqa: E501 qsa_paged_sparse_gqa_nvfp4 as op, ) from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import triton_qsa_paged_sparse_gqa -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 assert torch.cuda.is_available(), "qsa_paged_sparse_gqa_nvfp4 requires a CUDA device" @@ -171,24 +170,37 @@ def _build(spec, seed): def _bf16_twin(args): - """QSA's bf16 sparse GQA over a bf16 pool holding the dequantized values.""" - pages = args["k_cache"].shape[0] - all_pages = torch.arange(pages, device=DEV)[:, None].expand(-1, args["tokens_per_block"]) - slots = torch.arange(args["tokens_per_block"], device=DEV)[None, :].expand(pages, -1) + """QSA's bf16 sparse GQA over a bf16 pool holding the dequantized values. + + The twin pool holds only the pages the table maps, renumbered in order, so a cell's + pool can reach past 2^31 bytes without the twin dequantizing all of it. + """ + table = args["block_table"] + mapped = torch.unique(table[table >= 0]).long() + renumber = torch.full((args["k_cache"].shape[0],), -1, dtype=torch.int32, device=DEV) + renumber[mapped] = torch.arange(mapped.numel(), dtype=torch.int32, device=DEV) + twin_table = torch.where(table >= 0, renumber[table.clamp_min(0).long()], table) + tpb = args["tokens_per_block"] + pages = mapped[:, None].expand(-1, tpb) + slots = torch.arange(tpb, device=DEV)[None, :].expand(mapped.numel(), -1) planes = [] for data, scales, index, grouped in ( (args["k_cache"], args["k_scale_cache"], 1, False), (args["v_cache"], args["v_scale_cache"], 2, True), ): values = _entry._dequantize_tokens( - data, scales, all_pages, slots, args["kv_scale_quant_orig"][index], grouped - ) # [pages, tpb, H, D] + data, scales, pages, slots, args["kv_scale_quant_orig"][index], grouped + ) # [mapped pages, tpb, H, D] planes.append(values.to(args["q"].dtype).transpose(1, 2).contiguous()) twin = { - name: args[name] - for name in ("q", "block_table", "selected_tokens", "request_indices", "tokens_per_block") + name: args[name] for name in ("q", "selected_tokens", "request_indices", "tokens_per_block") } - twin.update(softmax_scale=args["softmax_scale"], k_cache=planes[0], v_cache=planes[1]) + twin.update( + softmax_scale=args["softmax_scale"], + block_table=twin_table, + k_cache=planes[0], + v_cache=planes[1], + ) if "query_positions" in args: twin.update(query_positions=args["query_positions"], compress_ratio=args["compress_ratio"]) return triton_qsa_paged_sparse_gqa(**twin) @@ -200,6 +212,20 @@ def _bound(positions, spec): return (visible // cr).clamp(max=spec["token_topk"] // cr) * cr + visible % cr +def _assert_straddles(args, edge: int) -> None: + """The pages the cell's rows read put the scale plane's byte offsets on both sides of + `edge` and the data plane's past it: what the cell claims about its pool.""" + table = args["block_table"] + pages = table[table >= 0].long() + for name, side in (("k_scale_cache", "both"), ("v_scale_cache", "both"), ("k_cache", "past")): + plane = args[name] + base = plane.data_ptr() - plane.untyped_storage().data_ptr() + ends = base + pages * plane.stride(0) + plane[0].numel() - 1 + assert bool((ends >= edge).any()), f"{name}: no page the rows read lies past {edge}" + if side == "both": + assert bool((ends < edge).any()), f"{name}: every page lies past {edge}" + + def _is_prefix_packed(table): valid = table >= 0 columns = torch.arange(table.shape[1], device=table.device)[None, :] @@ -284,6 +310,8 @@ def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811 assert bool(_is_prefix_packed(selected).all()) valid = (selected >= 0).sum(1) assert torch.equal(valid, _bound(args["query_positions"], spec)) + if "offset_edge" in spec: + _assert_straddles(args, spec["offset_edge"]) out = op(**args) assert [name for name, _ in launches] == list(spec["launches"]) @@ -302,20 +330,3 @@ def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811 def _with(**overrides): return {**CKPT, **overrides} - - -def test_calls_are_bitwise_deterministic() -> None: - args = _build(_with(segments=tuple((4000 + 37 * i, 2) for i in range(4))), seed=8800) - assert torch.equal(op(**args), op(**args)) - - -def test_rejects_positions_without_a_compress_ratio() -> None: - """The op's own loud refusal -- the reason the entry carries no `is_valid`: half of - the prefill pair.""" - args = _build(_with(segments=((600, 1),) * 4), seed=9100) - args["selected_tokens"] = args["selected_tokens"][:, :64].contiguous() - args["query_positions"] = torch.zeros(4, dtype=torch.int64, device=DEV) - message = "pass both query_positions and compress_ratio, or neither" - with pytest.raises(ValueError, match=re.escape(message)): - with validating(op): - op(**args) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py similarity index 50% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py index d8ddeb55e29d..354a8469963d 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_decode.py @@ -1,9 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test: `qsa_paged_sparse_gqa_nvfp4_decode` against the float64 reference of -`qsa_paged_sparse_gqa_nvfp4` and QSA's bf16 kernel over the dequantized pool, inside the -NVFP4 op's band (4 ulp of each (row, head)'s max, 2 ulp relative RMS), over its cells with and -without the output gate; plus determinism and graph replay against eager.""" +"""GPU test for the `qsa_paged_sparse_gqa_nvfp4_decode` catalog entry: against the reference of +`qsa_paged_sparse_gqa_nvfp4` +and QSA's bf16 kernel over the dequantized pool, inside the NVFP4 op's band (4 ulp of each +(row, head)'s max, 2 ulp relative RMS), over its cells with and without the output gate; plus +an unaligned row slice of the selection. + +The tp4-heads cells run one tp4 / tep4 rank's 6 q heads over 1 KV head, q a view of the packed +`[q | k | v]` rows as the target hands it. The dep4 cells run at one rank's pool as the engine +allocates it under `kv_nvfp4`: the requests' pages scattered over all 211797 pages, both planes' +byte offsets on both sides of 2^31 and 2^32 (asserted), every scale outside the mapped pages +NaN, and each captured and replayed at its own compile variant. The bf16 twin dequantizes only +the mapped pages.""" import math import zlib @@ -11,14 +19,17 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention import ( qsa_paged_sparse_gqa_nvfp4 as _nvfp4_module, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_sparse_gqa_nvfp4_decode import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4_decode import ( # noqa: E501 qsa_paged_sparse_gqa_nvfp4_decode as op, ) from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import triton_qsa_paged_sparse_gqa +__extra_import_path__ = ["../../.."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + assert torch.cuda.is_available(), "qsa_paged_sparse_gqa_nvfp4_decode requires a CUDA device" DEV = "cuda" @@ -70,28 +81,43 @@ def build( global_v=1.0, request_dtype=torch.int32, pool_pages=None, + spread=False, + q_heads=24, + kv_heads=2, ): gen = torch.Generator(device=DEV).manual_seed(seed) - tpb, kv, d, q_heads = 32, 2, 256, 24 + tpb, kv, d = 32, kv_heads, 256 positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments]) req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)]) rows = positions.numel() slots = max(slots, len(segments)) slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(segments)] mapped = [math.ceil((s + n) / tpb) for s, n in segments] + assert max(mapped) <= pages, (mapped, pages) needed = sum(mapped) + 7 planes = slice(2 * layer, 2 * layer + 2) pool_pages = pool_pages or needed - first = pool_pages - needed - page_ids = first + torch.randperm(needed, generator=gen, device=DEV).to(torch.int32) data = torch.zeros(pool_pages, slot_pages, kv, tpb, d // 2, dtype=torch.uint8, device=DEV) + # Every scale outside the requests' pages is NaN, so a page address that strayed could + # not come back finite. scales = torch.full( (pool_pages, slot_pages, kv, tpb, d // 16), E4M3_NAN, dtype=torch.uint8, device=DEV ) - data[first:, planes] = torch.randint( + if spread: + # The requests' pages scattered over the whole pool; the pool's last page is the first + # request's first page, so it is always mapped. + picked = torch.randperm(pool_pages - 1, generator=gen, device=DEV)[: needed - 1] + last = torch.full((1,), pool_pages - 1, device=DEV) + page_ids = torch.cat([last, picked]).to(torch.int32) + resident = page_ids.long() + else: + first = pool_pages - needed + page_ids = first + torch.randperm(needed, generator=gen, device=DEV).to(torch.int32) + resident = slice(first, None) + data[resident, planes] = torch.randint( 0, 256, (needed, 2, kv, tpb, d // 2), generator=gen, device=DEV ).to(torch.uint8) - scales[first:, planes] = torch.randint( + scales[resident, planes] = torch.randint( 0x08, 0x31, (needed, 2, kv, tpb, d // 16), generator=gen, device=DEV ).to(torch.uint8) block_table = torch.full((slots, pages), -1, dtype=torch.int32, device=DEV) @@ -119,48 +145,114 @@ def build( def bf16_twin(args): - pages = args["k_cache"].shape[0] - all_pages = torch.arange(pages, device=DEV)[:, None].expand(-1, args["tokens_per_block"]) - slots = torch.arange(args["tokens_per_block"], device=DEV)[None, :].expand(pages, -1) + """QSA's bf16 sparse GQA over a bf16 pool holding the dequantized values of the pages the + table maps, renumbered in order, so a cell's pool can reach past 2^32 bytes without the + twin dequantizing all of it.""" + table = args["block_table"] + mapped = torch.unique(table[table >= 0]).long() + renumber = torch.full((args["k_cache"].shape[0],), -1, dtype=torch.int32, device=DEV) + renumber[mapped] = torch.arange(mapped.numel(), dtype=torch.int32, device=DEV) + twin_table = torch.where(table >= 0, renumber[table.clamp_min(0).long()], table) + tpb = args["tokens_per_block"] + pages = mapped[:, None].expand(-1, tpb) + slots = torch.arange(tpb, device=DEV)[None, :].expand(mapped.numel(), -1) planes = [] for data, scales, index, grouped in ( (args["k_cache"], args["k_scale_cache"], 1, False), (args["v_cache"], args["v_scale_cache"], 2, True), ): values = _nvfp4_module._dequantize_tokens( - data, scales, all_pages, slots, args["kv_scale_quant_orig"][index], grouped + data, scales, pages, slots, args["kv_scale_quant_orig"][index], grouped ) planes.append(values.to(args["q"].dtype).transpose(1, 2).contiguous()) - twin = { - n: args[n] - for n in ("q", "block_table", "selected_tokens", "request_indices", "tokens_per_block") - } - twin.update(softmax_scale=args["softmax_scale"], k_cache=planes[0], v_cache=planes[1]) + twin = {n: args[n] for n in ("q", "selected_tokens", "request_indices", "tokens_per_block")} + twin.update( + softmax_scale=args["softmax_scale"], + block_table=twin_table, + k_cache=planes[0], + v_cache=planes[1], + ) return triton_qsa_paged_sparse_gqa(**twin) +def _assert_straddles(args, edges) -> None: + """The pages the cell's rows read put both planes' byte offsets on both sides of every + edge: what the cell claims about its pool.""" + table = args["block_table"] + pages = table[table >= 0].long() + for name in ("k_cache", "v_cache", "k_scale_cache", "v_scale_cache"): + plane = args[name] + base = plane.data_ptr() - plane.untyped_storage().data_ptr() + ends = base + pages * plane.stride(0) + plane[0].numel() - 1 + for edge in edges: + assert bool((ends >= edge).any()), f"{name}: no page the rows read lies past {edge}" + assert bool((ends < edge).any()), f"{name}: every page the rows read lies past {edge}" + last = args["k_cache"].shape[0] - 1 + assert bool((pages == last).any()), "the pool's last page is not among the rows' pages" + + +#: Spec keys that describe the cell rather than the state `build` makes. +_CELL_ONLY = ("segments", "gated", "capture", "large", "offset_edges") + + def _args(spec, seed): - kw = {k: v for k, v in spec.items() if k not in ("segments", "gated")} + kw = {k: v for k, v in spec.items() if k not in _CELL_ONLY} args = build(spec["segments"], seed=seed, **kw) if spec.get("gated"): rows = args["q"].shape[0] g = torch.Generator(device=DEV).manual_seed(seed + 1) - args["output_gate"] = torch.randn(rows, 24 * 256, generator=g, device=DEV).to( - torch.bfloat16 - ) + width = args["q"].shape[1] * args["q"].shape[2] + args["output_gate"] = torch.randn(rows, width, generator=g, device=DEV).to(torch.bfloat16) return args -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell) -> None: - args = _args(cell.spec, zlib.crc32(cell.why.encode()) % (2**31)) +def _replay_in_place(args, seed) -> None: + """The call captured at the cell's own compile variant, fresh q (and gate) written into + the captured buffers, replayed: bitwise the eager call on the fresh values.""" + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + try: + with torch.cuda.graph(graph): + captured = op(**args) + g = torch.Generator(device=DEV).manual_seed(seed) + for name in ("q", "output_gate"): + if name in args: + args[name].copy_(torch.randn(args[name].shape, generator=g, device=DEV)) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(captured, op(**args)), "replay differs from eager" + finally: + del graph + torch.cuda.synchronize() + + +def _marks(cell): + return (pytest.mark.high_cuda_memory,) if cell.spec.get("large") else () + + +@pytest.mark.parametrize( + "cell", [pytest.param(c, marks=_marks(c), id=c.why[:44]) for c in op.CELLS] +) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture + certifying(op) + seed = zlib.crc32(cell.why.encode()) % (2**31) + args = _args(cell.spec, seed) + if "offset_edges" in cell.spec: + _assert_straddles(args, cell.spec["offset_edges"]) out = op(**args) assert out.dtype == torch.bfloat16 and not out.isnan().any() el, rms, ok = within_ulp(out, op.reference(**args)) assert ok, f"vs reference: {el:.2f} ulp element, {rms:.2f} ulp RMS" - if "output_gate" not in args: - el, rms, ok = within_ulp(out, bf16_twin(args)) - assert ok, f"vs QSA bf16 kernel: {el:.2f} ulp element, {rms:.2f} ulp RMS" + # QSA's bf16 kernel has no output gate, so a gated cell checks its ungated call over the + # same state against it. + gate = args.pop("output_gate", None) + plain = out if gate is None else op(**args) + el, rms, ok = within_ulp(plain, bf16_twin(args)) + assert ok, f"vs QSA bf16 kernel: {el:.2f} ulp element, {rms:.2f} ulp RMS" + if gate is not None: + args["output_gate"] = gate + if cell.spec.get("capture"): + _replay_in_place(args, seed + 7) def test_unaligned_row_slice_of_the_selection() -> None: @@ -174,36 +266,3 @@ def test_unaligned_row_slice_of_the_selection() -> None: aligned[name] = args[name].clone() assert args["selected_tokens"].data_ptr() % 16 assert torch.equal(op(**args), op(**aligned)) - - -def test_calls_are_bitwise_deterministic() -> None: - args = build(tuple((4000 + 37 * i, 2) for i in range(4)), seed=8800) - assert torch.equal(op(**args), op(**args)) - - -@pytest.mark.parametrize("rows", [1, 64]) -def test_replay_matches_eager(rows) -> None: - segments = ( - ((4000, 1),) if rows == 1 else tuple((100 + (499 * g) % 8000, 1) for g in range(rows)) - ) - args = build(segments, seed=77) - fresh = build(segments, seed=127) - op(**args) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured = op(**args) - for name in ( - "q", - "k_cache", - "v_cache", - "k_scale_cache", - "v_scale_cache", - "block_table", - "selected_tokens", - "request_indices", - ): - args[name].copy_(fresh[name]) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(captured, op(**fresh)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py similarity index 97% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py index 664da3193ba6..bbe42f210b46 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4_verify.py @@ -13,13 +13,13 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention import ( # noqa: E501 qsa_paged_sparse_gqa_nvfp4_verify as _verify_module, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_sparse_gqa_nvfp4_decode import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4_decode import ( # noqa: E501 nvfp4_sparse_gqa_decode as per_row, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_sparse_gqa_nvfp4_verify import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_sparse_gqa_nvfp4_verify import ( # noqa: E501 qsa_paged_sparse_gqa_nvfp4_verify as op, ) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_prefill_compress.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_prefill_compress.py similarity index 64% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_prefill_compress.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_prefill_compress.py index cac393b9865c..f2426eb8d1ec 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_prefill_compress.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_prefill_compress.py @@ -1,13 +1,19 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qsa_prefill_compress catalog entry. +"""GPU test for the default target's qsa_prefill_compress variant. The paged state is real: an index-key pool of random raw keys, a position plane holding `(t, h, w)` for every allocated position, and the engine's -`[slots, width]` table -- more slots than requests, each live row's pages drawn -from a seeded permutation of the pool, `-1` past its allocation, rows past the -batch holding a previous step's plausible page ids. The cos|sin table is built -the way the target builds it at load. +`[slots, width]` table -- at least as many slots as requests, each live row's +pages drawn from a seeded permutation of the pool, `-1` past its allocation, +rows past the batch holding a previous step's plausible page ids. The cos|sin +table is built the way the target builds it at load, `max_position_embeddings` +rows tall. + +The `coalesced` cells hold the default target's pool: one rank's measured page count, +the 12 sparse layers' index pages sharing each slot as KV cache manager V2 lays +them out at bf16, the other 11 layers' planes holding a sentinel that must +survive. Those pools are drawn on the device. A padding row at logical position -1 is deliberately absent: it is an illegal memory access that poisons the process for every test after it. @@ -18,11 +24,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_prefill_compress import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_prefill_compress import ( # noqa: E501 qsa_prefill_compress as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 assert torch.cuda.is_available(), "qsa_prefill_compress requires a CUDA device" @@ -40,24 +46,30 @@ def _compare(out: torch.Tensor, ref: torch.Tensor, flips: int = 0) -> None: kernel takes `tl.rsqrt`, the hardware approximation, where the reference takes torch's, and the two differ by an fp32 ulp, which moves the bf16 rounding of an element whose fp32 value sits on a rounding boundary. One element of its 8192 anchor elements is one bf16 - step off. + step off. The 8192-row tep4 prefill cell has five of 262016, each an exact bf16 tie in + the reference's fp32 value; the nearest wrong variant measured there (sequential bf16 + pooling) moves 94580. """ assert out.dtype == ref.dtype, (out.dtype, ref.dtype) assert out.shape == ref.shape, (out.shape, ref.shape) - if flips == 0: - torch.testing.assert_close(out.view(torch.int16), ref.view(torch.int16), rtol=0, atol=0) + # Only the differing elements are widened: a pool-sized int32 copy would not fit beside + # the target's pool. + a, b = out.view(torch.int16), ref.view(torch.int16) + differs = a != b + count = int(differs.sum()) + if count == 0: return - distance = (out.view(torch.int16).int() - ref.view(torch.int16).int()).abs() - assert int(distance.max()) <= 1 and int((distance != 0).sum()) <= flips, ( - f"{int((distance != 0).sum())} elements differ, by up to {int(distance.max())} bf16 " - f"steps; at most {flips} may differ by one" + distance = int((a[differs].int() - b[differs].int()).abs().max()) + assert distance <= 1 and count <= flips, ( + f"{count} elements differ, by up to {distance} bf16 steps; at most {flips} may differ by one" ) DEV = "cuda" ROTARY_DIM = 64 THETA = 1.0e7 -MAX_POSITIONS = 16384 +#: The target's table height: the checkpoint's `max_position_embeddings`. +MAX_POSITIONS = 262144 SENTINEL = 0x3C3C @@ -81,7 +93,17 @@ def _build(spec, seed: int): table = torch.randint(0, pages, (slots, width), generator=gen, dtype=torch.int32) order = torch.randperm(pages, generator=gen).to(torch.int32) - index_cache = torch.randn(pages, tpb, 1, head_dim, generator=gen).to(torch.bfloat16) + if spec.get("top_page"): + # Request 0's last page is the pool's last, so its anchors address the top of the pool. + last = (segments[0][0] + segments[0][1] - 1) // tpb + where = int((order == pages - 1).nonzero()) + order[[where, last]] = order[[last, where]] + coalesced = spec["layout"] == "coalesced" + index_cache = ( + None # drawn on the device below + if coalesced + else torch.randn(pages, tpb, 1, head_dim, generator=gen).to(torch.bfloat16) + ) position_cache = torch.zeros(pages, tpb, 3, dtype=torch.int32) used = 0 for r, (start, n) in enumerate(segments): @@ -92,14 +114,30 @@ def _build(spec, seed: int): p = torch.arange(cols * tpb).view(cols, tpb) coords = (p, (3 * p + 5 * r) % 500, (5 * p + 11 * r) % 500) position_cache[table[r, :cols].long()] = torch.stack(coords, dim=-1).to(torch.int32) - assert used <= pages and len(segments) < slots, "pool or table too small for the batch" + assert used <= pages and len(segments) <= slots, "pool or table too small for the batch" k_norm_weight = (torch.randn(head_dim, generator=gen) * 0.1).to(torch.bfloat16) logical = torch.tensor(positions, dtype=spec["logical_dtype"]) request = torch.tensor(requests, dtype=spec["request_dtype"]) - table, index_cache, position_cache = table.to(DEV), index_cache.to(DEV), position_cache.to(DEV) + table, position_cache = table.to(DEV), position_cache.to(DEV) extras = {} + if coalesced: + # `roles` layers' index pages per slot; this layer's is `ordinal`, the rest a sentinel. + roles, ordinal = spec["roles"], spec["ordinal"] + pool = torch.empty(pages, roles, tpb, 1, head_dim, dtype=torch.bfloat16, device=DEV) + pool.view(torch.int16).fill_(SENTINEL) + index_cache = pool[:, ordinal] + device_gen = torch.Generator(device=DEV).manual_seed(seed) + index_cache.copy_( + torch.randn( + pages, tpb, 1, head_dim, generator=device_gen, device=DEV, dtype=torch.bfloat16 + ) + ) + assert index_cache.stride() == (roles * tpb * head_dim, head_dim, head_dim, 1) + extras["other_layers"] = (pool, ordinal) + else: + index_cache = index_cache.to(DEV) if spec["layout"] == "shared_page": # An NVFP4 pool's page: `roles` index pages, then the position page. The other # index roles hold a sentinel and must survive, as must the position page. @@ -119,7 +157,7 @@ def _build(spec, seed: int): extras["neighbour"] = (others, others.clone()) extras["position_page"] = (position_cache, position_cache.clone()) else: - assert spec["layout"] == "plain", spec["layout"] + assert spec["layout"] in ("plain", "coalesced"), spec["layout"] args = dict( logical_positions=logical.to(DEV), @@ -144,6 +182,15 @@ def _anchors(args) -> int: return int(torch.unique(keyed[(logical + 1) % args["compress_ratio"] == 0]).numel()) +def _top_anchor_page(args) -> int: + """The highest page an anchor row writes.""" + logical = args["logical_positions"].long() + anchored = (logical + 1) % args["compress_ratio"] == 0 + request = args["request_indices"].long()[anchored] + pages = args["block_table"].long()[request, logical[anchored] // args["tokens_per_block"]] + return int(pages.max()) + + @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) def test_certified_cells(cell, certifying): # noqa: F811 certifying(op) @@ -161,12 +208,17 @@ def test_certified_cells(cell, certifying): # noqa: F811 if "position_page" in extras: view, saved = extras["position_page"] assert torch.equal(view, saved), "wrote into the position page" - - -def _scene(seed: int, **overrides): - """Two requests of 64 tokens from 0, the target's call, at the checkpoint geometry.""" - spec = {**op.CELLS[0].spec, "segments": ((0, 64),) * 2, "pages": 4, **overrides} - return _build(spec, seed=seed)[0] + if "other_layers" in extras: + pool, ordinal = extras["other_layers"] + for other in range(pool.shape[1]): + if other != ordinal: + plane = pool[:, other].view(torch.int16) + assert bool((plane == SENTINEL).all()), f"wrote into layer ordinal {other}'s plane" + if cell.spec.get("top_page"): + # The pool's last page is written, at an element offset past int32 from the base. + top = _top_anchor_page(args) + assert top == cell.spec["pages"] - 1, top + assert top * out.stride(0) >= 2**31, top * out.stride(0) #: The cell of 32 MTP7 verify blocks (256 rows, two anchors each): the segments, pool and table. @@ -174,33 +226,6 @@ def _scene(seed: int, **overrides): _MTP7_BATCH = {key: _MTP7_32[key] for key in ("segments", "pages", "slots", "rounding_flips")} -def test_cuda_graph_capture_and_replay() -> None: - """The launch grid is fixed by the row count, so it captures and replays: 32 MTP7 - verify blocks, the decode target's captured batch.""" - batch = _MTP7_BATCH - args = _scene(seed=16, **batch) - cache = args["index_cache"] - before = cache.clone() - ref = op.reference(**args) - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - for _ in range(3): - cache.copy_(before) - op(**args) - torch.cuda.current_stream().wait_stream(warmup) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - cache.copy_(before) - with torch.cuda.graph(graph): - op(**args) - for _ in range(2): - cache.copy_(before) - graph.replay() - torch.cuda.synchronize() - _compare(cache, ref, batch["rounding_flips"]) - - def _strided(tensor: torch.Tensor, dim: int) -> torch.Tensor: """The same values as a stride-2 view along `dim` of a zero-filled buffer.""" shape = list(tensor.shape) @@ -210,28 +235,3 @@ def _strided(tensor: torch.Tensor, dim: int) -> torch.Tensor: index[dim] = slice(None, None, 2) wide[tuple(index)] = tensor return wide[tuple(index)] - - -def test_guard_refuses_strided_logical_positions() -> None: - """Shown unguarded first -- accepted and wrong -- then refused: a 1-D operand the - kernel reads flat (`request_indices` and `k_norm_weight` carry the same - `is_contiguous` guard). Old receipt's max abs error against the right answer: 4.72; - the floor below states a bound, not one draw. - """ - args = _scene(seed=12) - before = args["index_cache"].clone() - ref = op.reference(**args) - bad = dict(args, logical_positions=_strided(args["logical_positions"], 0)) - assert not bad["logical_positions"].is_contiguous() - - out = op(**bad) - error = (out.float() - ref.float()).abs().max().item() - assert error > 0.25, ( - f"strided logical_positions was expected to be read wrongly; off by {error}" - ) - - args["index_cache"].copy_(before) - with pytest.raises(AssertionError, match="logical_positions must be contiguous"): - with validating(op): - op(**bad) - assert torch.equal(args["index_cache"], before), "a refused call still wrote" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_rollback_planes.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_rollback_planes.py new file mode 100644 index 000000000000..1568de7ac079 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_rollback_planes.py @@ -0,0 +1,219 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_rollback_planes variant: `qsa_rollback_gather` and +`qsa_rollback_restore`, each cell through the entry's own reference, bitwise (both are copies). + +The pool is filled with a NaN pattern `randn` never draws, then every slot a cell names is +written, so a whole-pool count of what no longer holds the fill shows that a call wrote nothing +outside those slots. The tp1 cells run over separate contiguous planes and over the packed +side-plane page of an FP8 / NVFP4 pool; the default target's cells over core's coalesced bf16 pool at a +tep4 rank's measured page count, with the requests' pages drawn across it and the pool's first +and last page among them, so the last is addressed past 2^31 elements from a plane's base. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rollback_planes import ( # noqa: E501 + ROW, + plane_table, + qsa_rollback_gather, + qsa_rollback_restore, +) + +__extra_import_path__ = ["../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "qsa_rollback_planes requires a CUDA device" + +DEV = "cuda" +TPB = 32 +#: The tp1 cells' pool pages. +SMALL_POOL = 600 +#: One plane's page: `tokens_per_block` 128-wide rows. +PLANE_PAGE = TPB * ROW +#: The QSA position page an FP8 / NVFP4 pool packs after the 12 index pages (384 bytes). +SIDE_PLANES_PAD = 192 +#: `max_seq_len` under MTP: the last position a verify row can hold is one less. +MTP_MAX_SEQ_LEN = 16384 +#: bf16 NaN with a payload: the pool's fill. +SENTINEL = 0x7FA5 + + +class _Pool: + """One rank's index planes, every element at `SENTINEL` until a slot is written. + + `contiguous`: one `[pages, 32, 1, 128]` buffer a plane. `shared_page`: the planes side by + side in one page with the position page after them. `coalesced`: core's KV cache manager V2 + pool, layer k's buffer at element 4096 k of a 49152-element page. + """ + + def __init__(self, spec: dict, layout: str): + count, pages = spec["planes"], spec.get("pool_pages", SMALL_POOL) + if layout == "contiguous": + self.buffers = [ + torch.empty(pages, TPB, 1, ROW, dtype=torch.bfloat16, device=DEV) + for _ in range(count) + ] + self.planes = list(self.buffers) + else: + pad = SIDE_PLANES_PAD if layout == "shared_page" else 0 + page = torch.empty(pages, count * PLANE_PAGE + pad, dtype=torch.bfloat16, device=DEV) + self.buffers = [page] + self.planes = [ + page[:, k * PLANE_PAGE : (k + 1) * PLANE_PAGE].view(pages, TPB, 1, ROW) + for k in range(count) + ] + for b in self.buffers: + b.view(torch.int16).fill_(SENTINEL) + self.pages = pages + + def written(self) -> int: + """Elements anywhere in the pool that no longer hold the fill.""" + total = 0 + for b in self.buffers: + for chunk in b.view(torch.int16).reshape(b.shape[0], -1).split(8192): + total += int((chunk != SENTINEL).sum()) + return total + + +def _layouts(cell): + return (cell.spec["layout"],) if "layout" in cell.spec else ("contiguous", "shared_page") + + +def _cases(op): + return [(c, layout) for c in op.CELLS for layout in _layouts(c)] + + +def _ids(cases): + return [f"{c.why[:44]}-{layout}" for c, layout in cases] + + +GATHER_CASES = _cases(qsa_rollback_gather) +RESTORE_CASES = _cases(qsa_rollback_restore) + + +def _verify_slots(spec, seed): + """`requests` x `p` consecutive positions each, every request's QSA pages its own physical + pages of the pool. Request 0 straddles a page boundary from the pool's first page into its + last; the last request ends at position 16383.""" + n, p, pool = spec["requests"], spec["p"], spec["pool_pages"] + starts = ( + [TPB - p // 2] + + [(29 + 127 * s) % (MTP_MAX_SEQ_LEN - p) for s in range(1, n - 1)] + + [MTP_MAX_SEQ_LEN - p] + ) + positions = torch.tensor(starts)[:, None] + torch.arange(p)[None, :] + logical = positions // TPB + rest = torch.randperm(pool - 2, generator=torch.Generator().manual_seed(seed)) + 1 + physical = torch.cat([torch.tensor([0, pool - 1]), rest]) + pages = torch.empty(n, p, dtype=torch.int64) + cursor = 0 + for s in range(n): + for lp in logical[s].unique().tolist(): + pages[s][logical[s] == lp] = physical[cursor] + cursor += 1 + return pages.reshape(-1).to(DEV), (positions % TPB).reshape(-1).to(DEV) + + +def _slots(spec, seed, g): + if "requests" in spec: + return _verify_slots(spec, seed) + rows = spec["rows"] + pages = torch.randperm(SMALL_POOL, generator=torch.Generator().manual_seed(rows))[:rows] + within = torch.randint(0, TPB, (rows,), generator=g, device=DEV) + return pages.to(DEV).to(torch.int64), within.to(torch.int64) + + +def _rejected(spec, g): + """The call site's mask for a verify step (bool, `ordinal >= accepted`; one request rejects + every draft, one none), or the tp1 cells' int32 one rejecting every other row, so each + cell restores some slots and keeps the rest.""" + if spec.get("rejected") != "bool": + return (torch.arange(spec["rows"], device=DEV) % 2).to(torch.int32) + n, p = spec["requests"], spec["p"] + accepted = torch.randint(1, p + 1, (n,), generator=g, device=DEV) + accepted[0], accepted[1] = 1, p + return (torch.arange(p, device=DEV)[None, :] >= accepted[:, None]).reshape(-1) + + +def _fill(planes, pages, within, g) -> None: + """Write fresh values into every slot of every plane (a step's index store).""" + for t in planes: + t[pages, within] = torch.randn(pages.shape[0], 1, ROW, generator=g, device=DEV).to(t.dtype) + + +def _slots_now(planes, pages, within) -> torch.Tensor: + return torch.stack([t[pages, within].reshape(-1, ROW) for t in planes]) + + +def _assert_far_page(spec, pool, far_pages) -> None: + """A measured-pool cell addresses the pool's last page, past 2^31 elements from the base of + every plane, where a 32-bit offset would wrap.""" + if "pool_pages" not in spec: + return + last = int(far_pages.max()) + assert last == pool.pages - 1, f"the pool's last page {pool.pages - 1} is unread ({last})" + assert last * pool.planes[-1].stride(0) >= 2**31 + + +def _setup(cell, layout): + seed = zlib.crc32(f"{cell.why}/{layout}".encode()) % 2**31 + g = torch.Generator(device=DEV).manual_seed(seed) + pool = _Pool(cell.spec, layout) + pages, within = _slots(cell.spec, seed, g) + assert pages.shape == (cell.spec["rows"],) + _fill(pool.planes, pages, within, g) + return g, pool, pages, within + + +@pytest.mark.parametrize("cell, layout", GATHER_CASES, ids=_ids(GATHER_CASES)) +def test_gather_cells(cell, layout, certifying): # noqa: F811 + certifying(qsa_rollback_gather) + _, pool, pages, within = _setup(cell, layout) + ptrs, strides = plane_table(pool.planes) + count, rows = len(pool.planes), pages.shape[0] + + want = qsa_rollback_gather.reference(pool.planes, pages, within, count) + got = qsa_rollback_gather(ptrs, strides, pages, within, count) + + assert got.shape == (count, rows, ROW) and got.dtype == torch.bfloat16 + assert torch.equal(got.view(torch.int16), want.view(torch.int16)) + assert torch.equal( + _slots_now(pool.planes, pages, within).view(torch.int16), want.view(torch.int16) + ) + assert pool.written() == count * rows * ROW, "the gather wrote into the pool" + _assert_far_page(cell.spec, pool, pages) + + +@pytest.mark.parametrize("cell, layout", RESTORE_CASES, ids=_ids(RESTORE_CASES)) +def test_restore_cells(cell, layout, certifying): # noqa: F811 + certifying(qsa_rollback_restore) + g, pool, pages, within = _setup(cell, layout) + ptrs, strides = plane_table(pool.planes) + count, rows = len(pool.planes), pages.shape[0] + # What the begin step's gather captured (its own cells hold it to this), then the step + # overwrites every slot. + old = _slots_now(pool.planes, pages, within) + _fill(pool.planes, pages, within, g) + rejected = _rejected(cell.spec, g) + assert bool(rejected.bool().any()) and not bool(rejected.bool().all()) + + want = qsa_rollback_restore.reference(pool.planes, pages, within, rejected, old) + assert qsa_rollback_restore(ptrs, strides, pages, within, rejected, old) is None + + got = _slots_now(pool.planes, pages, within) + assert torch.equal(got.view(torch.int16), want.view(torch.int16)) + assert pool.written() == count * rows * ROW, "the restore wrote outside the rows' slots" + _assert_far_page(cell.spec, pool, pages[rejected.bool()]) + + +def _small(): + spec = dict(planes=12, rows=16) + g = torch.Generator(device=DEV).manual_seed(5) + pool = _Pool(spec, "contiguous") + pages, within = _slots(spec, 5, g) + _fill(pool.planes, pages, within, g) + return g, pool, pages, within diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py similarity index 88% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py index a79e34fa0142..cc612e23c742 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py @@ -11,19 +11,18 @@ """ import math -import re import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention import ( # noqa: E501 qsa_sparse_gqa_rowtile as _module, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_sparse_gqa_rowtile import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_sparse_gqa_rowtile import ( # noqa: E501 qsa_sparse_gqa_rowtile as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 assert torch.cuda.is_available(), "qsa_sparse_gqa_rowtile requires a CUDA device" @@ -59,7 +58,9 @@ def __call__(self, selected_tokens, *args, **kwargs): _KERNELS = ( "_row_bits", - "_rowtile_plan_kernel", + "_rowtile_union_kernel", + "_rowtile_group_kernel", + "_rowtile_slots_kernel", "_rowtile_lone_kernel", "_rowtile_attention_cute", ) @@ -185,12 +186,16 @@ def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811 out = op(**args) assert [name for name, _ in launches] == list(_KERNELS) groups = math.ceil(rows / 10) - assert launches[0][1] == (rows,) and launches[1][1] == (groups,) - assert launches[2][1] == (min(rows, 256), spec["kv_heads"]) + # The plan's union and slot passes: a program per (group, 32 bitmap words of the table). + chunks = math.ceil(math.ceil(spec["pages"] * spec["tokens_per_block"] / 64) / 32) + assert launches[0][1] == (rows,) + assert launches[1][1] == launches[3][1] == (groups, chunks) + assert launches[2][1] == (groups,) + assert launches[4][1] == (min(rows, 256), spec["kv_heads"]) # The row-tile pass is persistent: one CTA per SM, never more than (group, KV head) items. items = groups * spec["kv_heads"] - assert launches[3][1] == (_module._rowtile_ctas(items, args["q"].device),) - assert launches[3][1][0] == min( + assert launches[5][1] == (_module._rowtile_ctas(items, args["q"].device),) + assert launches[5][1][0] == min( items, torch.cuda.get_device_properties(0).multi_processor_count ) assert out.shape == (rows, spec["q_heads"], spec["head_dim"]) @@ -224,14 +229,10 @@ def test_small_launches_match_the_reference(segments) -> None: _compare(op(**args), op.reference(**args)) -def test_calls_are_bitwise_deterministic() -> None: - args = _build(_with([(0, 2000), (500, 300), (7000, 1)], slots=8), seed=8800) - assert torch.equal(op(**args), op(**args)) - - -def test_rejects_a_call_without_the_prefill_pair() -> None: - """The op's own loud refusal: the row-tile pass is bounded by the prefill pair.""" - args = _build(_with([(600, 4)], slots=8), seed=9100) - args.update(query_positions=None, compress_ratio=None) - with pytest.raises(ValueError, match=re.escape("the row-tile kernel needs the prefill pair")): - op(**args) +def test_long_context_disjoint_groups_match_the_reference() -> None: + """Rows deep in a long context with independent selections: a group's list is about + ten times one row's columns (unions of up to 20480 tokens over 2690 pages, the plan's + words spanning many 32-word chunks). Every group of one request takes the row-tile + pass, however long its list.""" + args = _build(_with([(86000, 197)], slots=4, pages=2700), seed=9300) + _compare(op(**args), op.reference(**args)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_expand_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_expand_decode.py new file mode 100644 index 000000000000..28f84c8d508f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_expand_decode.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_topk_expand_decode entry. + +Cells against the entry's reference (a stable descending sort, ties to the higher block, +then the expansion entry's reference): the block list and the token selection exactly. +Columns past each row's visible count carry garbage (huge values, NaN) that must never +be read. Beyond the cells: the wide rows against the canonical composition, and a +stale-read check under graph replay at each decode-path cell's shape. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_canonical_selection import ( # noqa: E501 + qsa_canonical_selection, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_expand_block_indices import ( # noqa: E501 + qsa_expand_block_indices, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_topk_expand_decode import ( # noqa: E501 + qsa_topk_expand_decode as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_topk_select import ( # noqa: E501 + qsa_topk_select, +) + +__extra_import_path__ = ["../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "qsa_topk_expand_decode requires a CUDA device" + +DEV = "cuda" +K, CR, TOKEN_TOPK = 512, 4, 2048 + + +def _build(spec, seed): + g = torch.Generator(device=DEV).manual_seed(seed) + W = spec["width"] + if "chunk" in spec or "requests" in spec: + # Each request's consecutive rows (a context chunk, or a verify step's 1 + draft + # rows): every row's KV length is the request's last row + 1. + requests = spec["requests"] if "requests" in spec else (spec["chunk"],) + pos = torch.cat( + [torch.arange(s, s + c, dtype=torch.int64, device=DEV) for s, c in requests] + ) + seqlen = torch.cat( + [torch.full((c,), s + c, dtype=torch.int32, device=DEV) for s, c in requests] + ) + else: + pos = torch.tensor(spec["pos"], dtype=torch.int64, device=DEV) + seqlen = (pos + 1).to(torch.int32) + rows = pos.numel() + assert rows == spec["rows"], (rows, spec["rows"]) + scores = torch.randn(rows, W, generator=g, device=DEV).abs() * 3.0 + if spec.get("relu_sparse"): + # Fewer positive scores than the budget, so the zero class straddles it. + scores = torch.relu(torch.randn(rows, W, generator=g, device=DEV) - 2.2) + scores[:, ::7] = -0.0 # negative zero ties with zero + if spec.get("relu_sum"): + # The indexer's form: four heads' ReLUs summed, times a scale (zero in 1 of 16). + scores = torch.zeros(rows, W, device=DEV) + for _ in range(4): + scores += torch.relu(torch.randn(rows, W, generator=g, device=DEV)) + scores *= 0.0884 + if spec.get("equal"): + # Every score tied: the budget is the highest-index blocks. + scores = torch.full((rows, W), 0.5, device=DEV) + if spec.get("narrow"): + scores = 1.0 + 0.05 * torch.rand(rows, W, generator=g, device=DEV) + if spec.get("lognormal"): + scores = torch.exp(1.74 + 0.2 * torch.randn(rows, W, generator=g, device=DEV)) + # bf16-like resolution, so equal scores tie; a zero in every other row. + scores = scores.to(torch.bfloat16).float() + scores[::2, 5] = 0.0 + ends = ((pos + 1) // CR).to(torch.int32) + assert int(ends.max()) <= W, "a window wider than the score table" + # Garbage past the visible blocks: huge values, and a NaN three columns on. + cols = torch.arange(W, device=DEV) + scores = torch.where(cols[None, :] >= ends[:, None], 1.0e30, scores) + scores = torch.where(cols[None, :] == ends[:, None] + 3, float("nan"), scores) + topk = torch.full((rows, K), -7, dtype=torch.int32, device=DEV) + return [scores, ends, pos, seqlen, topk, CR, TOKEN_TOPK] + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + want_topk, want_sel = op.reference(*args) + margin = cell.spec.get("out_margin") + if margin is None: + sel = op(*args) + else: + # Rows [margin, margin + rows) of a step-sized selection; the rows around stay put. + rows = args[0].shape[0] + whole = torch.full( + (rows + 2 * margin, TOKEN_TOPK + CR - 1), -9, dtype=torch.int32, device=DEV + ) + sel = op(*args, out=whole[margin : margin + rows]) + assert sel.data_ptr() == whole[margin].data_ptr() + assert (whole[:margin] == -9).all() and (whole[margin + rows :] == -9).all(), ( + "rows outside out" + ) + assert sel.dtype == torch.int32 and sel.shape == want_sel.shape + assert torch.equal(args[4], want_topk), "block list" + assert torch.equal(sel, want_sel.to(torch.int32)), "token selection" + + +@pytest.mark.parametrize("width", [65536, 65568, 65600]) +@pytest.mark.parametrize("distribution", ["relu_sum", "relu_sparse", "equal"]) +def test_wide_rows_equal_the_canonical_composition(width: int, distribution: str) -> None: + """Match the fused block order and expansion to exact top-k plus canonical rewrite. + + The TP1 table includes eight pages of headroom; the default table includes four. + Ragged windows cross the block budget and slab boundaries and carry every causal + tail width. Both paths must ignore the poisoned score columns past each window. + """ + positions = (0, 1, 2, 3, 2043, 2047, 2051, 16383, 84000, 99999, 262143) + spec = dict(rows=len(positions), pos=positions, width=width, **{distribution: True}) + args = _build(spec, seed=width) + scores, ends, pos, lengths, fused_topk, cr, token_topk = args + lengths[8] += 3 # A verify row whose request continues past this query position. + starts = torch.zeros_like(ends) + composed_topk = torch.full_like(fused_topk, -1) + qsa_topk_select(scores, starts, ends, composed_topk, index_topk=K) + qsa_canonical_selection(scores, starts, ends, composed_topk) + composed = qsa_expand_block_indices( + composed_topk, pos, lengths, compress_ratio=cr, token_topk=token_topk + ) + fused = op(*args) + expected_topk, expected = op.reference(*args) + assert torch.equal(composed_topk, expected_topk), "canonical block list" + assert torch.equal(composed, expected), "canonical token expansion" + assert torch.equal(fused_topk, composed_topk), "fused block order" + assert torch.equal(fused, composed), "fused token expansion" + + +# The decode-path shapes, captured in the decode CUDA graph: tp1's one row at 2560 columns, +# the default target's relu-sparse decode and verify cells (1 and 128 rows at core's width, 1024 +# at MTP's), and the wide tables' decode cells (`graph` in their spec). +_GRAPH_SPECS = [dict(rows=1, pos=(8700,), width=2560)] + [ + c.spec + for c in op.CELLS + if c.spec.get("graph") or (c.spec.get("relu_sparse") and "out_margin" not in c.spec) +] + + +@pytest.mark.parametrize( + "spec", _GRAPH_SPECS, ids=[f"{s['rows']}x{s['width']}" for s in _GRAPH_SPECS] +) +def test_reads_after_the_wait_under_graph_replay(spec) -> None: + args = _build(spec, seed=21) + src = args[0].clone() + + def step(): + args[0].copy_(src * 1.0) + return op(*args) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + step() + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_sel = step() + for trial in range(3): + src.copy_(_build(spec, seed=100 + trial)[0]) + graph.replay() + torch.cuda.synchronize() + got_topk = args[4].clone() + sel = op(src.clone(), *args[1:]) + assert torch.equal(graph_sel, sel) and torch.equal(got_topk, args[4]), f"replay {trial}" + want_topk, want_sel = op.reference(src, *args[1:]) + assert torch.equal(got_topk, want_topk), f"replay {trial}: block list" + assert torch.equal(graph_sel, want_sel), f"replay {trial}: token selection" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_select.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_select.py similarity index 86% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_select.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_select.py index e0ee54787a3a..ecb3dfe2c753 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_select.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qsa_topk_select.py @@ -10,14 +10,14 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_select import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_topk_select import ( # noqa: E501 TopkExpectation, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_select import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_topk_select import ( # noqa: E501 qsa_topk_select as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 assert torch.cuda.is_available(), "qsa_topk_select requires a CUDA device" @@ -97,12 +97,30 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 _compare(out, op.reference(*args)) -def test_cuda_graph_capture_and_replay() -> None: - """Captured, then replayed against new scores and a new causal bound: the 256 rows a - verify step of 32 requests carries at max_draft_len 7.""" - rows = 256 +#: (rows, score columns, replay window): tp1's 256-row verify step, then the default +#: target's decode batches 1 and 128 and its 1024-row MTP7 verify step at the +#: 262144-token envelope's 65536 columns. Every capture runs full windows (the +#: histogram path); the replays move the bound, 128 rows down to the budget's +#: copy-through shortcut. +GRAPH_CASES = ((256, 2048, 1500), (1, 65536, 65535), (128, 65536, 511), (1024, 65536, 30000)) + + +def _layout_case(): + torch.manual_seed(10) + rows, cols, top_k = 4, 128, 8 + logits = torch.randn(rows, cols, device=DEV) + row_starts = torch.zeros(rows, dtype=torch.int32, device=DEV) + row_ends = torch.full((rows,), cols, dtype=torch.int32, device=DEV) + return logits, row_starts, row_ends, top_k + + +@pytest.mark.parametrize( + "rows, cols, window", GRAPH_CASES, ids=[f"{r}x{c}-to-{w}" for r, c, w in GRAPH_CASES] +) +def test_cuda_graph_capture_and_replay(rows, cols, window) -> None: + """Captured over full windows, then replayed against new scores and a new causal + bound through the captured addresses.""" gen = torch.Generator(device=DEV).manual_seed(8) - cols = 2048 logits = _logits("relu", rows, cols, gen) row_starts = torch.zeros(rows, dtype=torch.int32, device=DEV) row_ends = torch.full((rows,), cols, dtype=torch.int32, device=DEV) @@ -120,22 +138,13 @@ def test_cuda_graph_capture_and_replay() -> None: op(logits, row_starts, row_ends, indices, BLOCK_TOPK) logits.copy_(_logits("relu", rows, cols, gen)) - row_ends.fill_(1500) + row_ends.fill_(window) indices.fill_(SENTINEL) graph.replay() torch.cuda.synchronize() _compare(indices, op.reference(logits, row_starts, row_ends, indices, BLOCK_TOPK)) -def _layout_case(): - torch.manual_seed(10) - rows, cols, top_k = 4, 128, 8 - logits = torch.randn(rows, cols, device=DEV) - row_starts = torch.zeros(rows, dtype=torch.int32, device=DEV) - row_ends = torch.full((rows,), cols, dtype=torch.int32, device=DEV) - return logits, row_starts, row_ends, top_k - - def test_guard_refuses_the_token_budget_default() -> None: """The default `index_topk` is 2048, the token budget; the target's buffer is 512 wide. That is the wider-buffer fault the other way round, and the diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py similarity index 98% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py index cb7a12755720..a49258c73753 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py @@ -17,11 +17,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qwen4_exp_attn_decode_prologue import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qwen4_exp_attn_decode_prologue import ( # noqa: E501 qwen4_exp_attn_decode_prologue as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "qwen4_exp_attn_decode_prologue requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py similarity index 96% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py index d5404666e2c6..0674746fabec 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py @@ -8,14 +8,14 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qwen4_exp_qk_norm_rope import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qwen4_exp_qk_norm_rope import ( # noqa: E501 CKPT, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qwen4_exp_qk_norm_rope import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qwen4_exp_qk_norm_rope import ( # noqa: E501 qwen4_exp_qk_norm_rope as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import ( # noqa: E402,F401 — path above; fixture by name assert_within_ulp, certifying, diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_spec_rollback_commit.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_spec_rollback_commit.py new file mode 100644 index 000000000000..2631822136f5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_spec_rollback_commit.py @@ -0,0 +1,319 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for the spec_rollback_commit variant: `spec_rollback_stage` and +`spec_rollback_commit`, every cell through the entry's own reference, bitwise (both are data +movement). + +A cell is one verify step at the call site's layout: `n` generation requests of `p` rows, each +request's rows at consecutive positions on pages of their own, int32 `accepted` / `slots` / +`ids` and int64 `pages` / `within` as `SpeculativeRollback` hands them. The index-K planes are +`[pages, 32, 1, 128]` bf16 views addressed through `plane_table`: separate buffers in the tp1 +cells, the KV cache manager's coalesced pool (12 planes side by side in a 49152-element page) at +a tep4 rank's measured page count in the default target's. Request 0's rows straddle the pool's first +and last page and it accepts one row, so its rejected rows are restored on the last page, past +2^31 elements from a plane's base. The pools are filled with a NaN pattern `randn` never draws +and the trunk's writes are put at every row's slot, so a whole-pool count shows the commit +wrote nothing else; the n-gram, conv and position pools are compared whole. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rollback_planes import ( # noqa: E501 + plane_table, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.spec_rollback_commit import ( # noqa: E501 + ROW, + spec_rollback_commit, + spec_rollback_stage, +) + +__extra_import_path__ = ["../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "spec_rollback_commit requires a CUDA device" + +DEV = "cuda" +TPB = 32 +#: The tp1 cells' PLE geometry and plane pool, where a cell does not name its own. +TP1_GEOMETRY = dict(context=2, channels=640, state_len=3, layout="separate", pool_pages=600) +#: `max_seq_len` under MTP: the last position a verify row can hold is one less. +MTP_MAX_SEQ_LEN = 16384 +#: The vocabulary the ids are drawn from. +VOCAB = 248320 +#: bf16 NaN with a payload: the index pool's fill. +SENTINEL = 0x7FA5 +#: The position plane's fill; positions stay below 2^24. +POSITION_FILL = 0x7FA5A5A5 +#: One plane's page: `tokens_per_block` 128-wide rows. +PLANE_PAGE = TPB * ROW + + +def _seed(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +def _geometry(spec: dict) -> dict: + g = dict(TP1_GEOMETRY, **spec) + g.setdefault("state_slots", 2 * spec["n"] + 4) + return g + + +def _slots(n: int, state_slots: int, gen: torch.Generator) -> torch.Tensor: + """`n` distinct recurrent slots; at an engine-sized pool its first and last among them.""" + if state_slots < n + 2 or n < 2: + picked = torch.randperm(state_slots - 1, generator=gen)[:n] + 1 + else: + inner = torch.randperm(state_slots - 2, generator=gen)[: n - 2] + 1 + picked = torch.cat([torch.tensor([0, state_slots - 1]), inner]) + picked = picked[torch.randperm(n, generator=gen)] + return picked.to(DEV, torch.int32) + + +def _ple_pools(g: dict, gen: torch.Generator): + ngram_pool = torch.randint(0, 1 << 40, (g["state_slots"], g["context"]), generator=gen) + conv_pool = torch.randn(g["state_slots"], g["channels"], g["state_len"], generator=gen) + return ngram_pool.to(DEV), conv_pool.to(DEV, torch.bfloat16) + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.view(torch.int16) if t.dtype == torch.bfloat16 else t + + +def _equal(got: torch.Tensor, want: torch.Tensor, what: str) -> None: + assert got.dtype == want.dtype and got.shape == want.shape, (what, got.shape, want.shape) + differ = (_bits(got) != _bits(want)).reshape(got.shape[0], -1).any(dim=1) + assert not bool(differ.any()), f"{what}: {int(differ.sum())} of {got.shape[0]} rows differ" + + +class _Planes: + """One rank's index-K planes, every element at `SENTINEL` until a slot is written.""" + + def __init__(self, count: int, pages: int, layout: str): + if layout == "separate": + self.buffers = [ + torch.empty(pages, TPB, 1, ROW, dtype=torch.bfloat16, device=DEV) + for _ in range(count) + ] + self.views = list(self.buffers) + else: + assert layout == "coalesced", layout + page = torch.empty(pages, count * PLANE_PAGE, dtype=torch.bfloat16, device=DEV) + self.buffers = [page] + self.views = [ + page[:, k * PLANE_PAGE : (k + 1) * PLANE_PAGE].view(pages, TPB, 1, ROW) + for k in range(count) + ] + self.fill() + + def fill(self) -> None: + for b in self.buffers: + b.view(torch.int16).fill_(SENTINEL) + + def written(self) -> int: + """Elements anywhere in the pool that no longer hold the fill.""" + total = 0 + for b in self.buffers: + for chunk in b.view(torch.int16).reshape(b.shape[0], -1).split(8192): + total += int((chunk != SENTINEL).sum()) + return total + + def rows(self, pages: torch.Tensor, within: torch.Tensor) -> torch.Tensor: + return torch.stack([v[pages, within].reshape(-1, ROW) for v in self.views]) + + +class _Step: + """One verify step's staged state, the pools it commits into, and the trunk's writes.""" + + def __init__(self, spec: dict, seed: int): + g = _geometry(spec) + self.n, self.p, self.count = spec["n"], spec["p"], spec["planes"] + self.planes = _Planes(self.count, g["pool_pages"], g["layout"]) + self.position = torch.empty(g["pool_pages"], TPB, 3, dtype=torch.int32, device=DEV) + self.ptrs, self.strides = plane_table(self.planes.views) + gen = torch.Generator().manual_seed(seed) + self.ngram_pool, self.conv_pool = _ple_pools(g, gen) + self.g = g + n, p, rows = self.n, self.p, self.n * self.p + self.accepted = torch.empty(n, dtype=torch.int32, device=DEV) + self.slots = torch.empty(n, dtype=torch.int32, device=DEV) + self.pages = torch.empty(rows, dtype=torch.int64, device=DEV) + self.within = torch.empty(rows, dtype=torch.int64, device=DEV) + self.stream = torch.empty(n, g["context"] + p, dtype=torch.int64, device=DEV) + self.conv_prev = torch.empty( + n, g["channels"], g["state_len"], dtype=torch.bfloat16, device=DEV + ) + self.conv_x = torch.empty(n, p, g["channels"], dtype=torch.bfloat16, device=DEV) + self.position_old = torch.empty(rows, 3, dtype=torch.int32, device=DEV) + self.index_old = torch.empty(self.count, rows, ROW, dtype=torch.bfloat16, device=DEV) + self.draw(gen) + + def draw(self, gen: torch.Generator) -> None: + """Draw every staged input into its buffer (in place, so a captured graph re-reads it) + and put the trunk's writes at the rows' slots of a freshly filled index and position + pool.""" + g, n, p = self.g, self.n, self.p + pool_pages = g["pool_pages"] + accepted = torch.randint(1, p + 1, (n,), generator=gen) + accepted[0] = 1 + if n > 1: + accepted[1] = p + self.accepted.copy_(accepted) + self.slots.copy_(_slots(n, g["state_slots"], gen)) + + # Request 0 straddles its first two blocks, mapped to the pool's first and last page; + # every other block of every request gets a page of its own. + starts = torch.randint(0, MTP_MAX_SEQ_LEN - p + 1, (n,), generator=gen) + starts[0] = TPB - p // 2 + positions = starts[:, None] + torch.arange(p)[None, :] + blocks = positions // TPB + free = torch.randperm(pool_pages - 2, generator=gen) + 1 + page_of = {} + pages = torch.empty(n, p, dtype=torch.int64) + for i in range(n): + for t in range(p): + key = (i, int(blocks[i, t])) + if key not in page_of: + if i == 0: + page_of[key] = 0 if int(blocks[i, t]) == 0 else pool_pages - 1 + else: + page_of[key] = int(free[len(page_of)]) + pages[i, t] = page_of[key] + self.pages.copy_(pages.reshape(-1)) + self.within.copy_((positions % TPB).reshape(-1)) + + rows = n * p + self.stream.copy_(torch.randint(0, 1 << 40, tuple(self.stream.shape), generator=gen)) + self.conv_prev.copy_(torch.randn(tuple(self.conv_prev.shape), generator=gen)) + self.conv_x.copy_(torch.randn(tuple(self.conv_x.shape), generator=gen)) + self.position_old.copy_(torch.randint(0, 1 << 20, (rows, 3), generator=gen)) + self.index_old.copy_(torch.randn(self.count, rows, ROW, generator=gen)) + + # What the trunk wrote at every row: kept where the request accepted the row. + self.planes.fill() + self.position.fill_(POSITION_FILL) + trunk = torch.randn(self.count, rows, ROW, generator=gen).to(DEV, torch.bfloat16) + for k, view in enumerate(self.planes.views): + view[self.pages, self.within] = trunk[k].reshape(rows, 1, ROW) + self.position[self.pages, self.within] = torch.randint( + 0, 1 << 20, (rows, 3), generator=gen + ).to(DEV, torch.int32) + + def call_args(self, index_old=None, conv_prev=None) -> tuple: + return ( + self.accepted, + self.slots, + self.ngram_pool, + self.stream, + self.conv_pool, + self.conv_prev if conv_prev is None else conv_prev, + self.conv_x, + self.position, + self.position_old, + self.pages, + self.within, + self.ptrs, + self.strides, + self.index_old if index_old is None else index_old, + ) + + def reference(self): + return spec_rollback_commit.reference( + self.accepted, + self.slots, + self.ngram_pool, + self.stream, + self.conv_pool, + self.conv_prev, + self.conv_x, + self.position, + self.position_old, + self.pages, + self.within, + self.planes.views, + self.index_old, + ) + + def snapshot(self): + return self.ngram_pool.clone(), self.conv_pool.clone(), self.position.clone() + + def check(self, ref, before) -> None: + """The named slots hold `ref`, bit for bit, and nothing else in any pool moved.""" + ngram, conv, pos, index = ref + ngram_before, conv_before, pos_before = before + slots = self.slots.long() + want = ngram_before.clone() + want[slots] = ngram + _equal(self.ngram_pool, want, "n-gram pool") + want = conv_before.clone() + want[slots] = conv + _equal(self.conv_pool, want, "PLE conv pool") + want = pos_before.clone() + want[self.pages, self.within] = pos + _equal(self.position, want, "position plane") + _equal(self.planes.rows(self.pages, self.within), index, "index-K planes") + rows = self.n * self.p + assert self.planes.written() == self.count * rows * ROW, ( + "the commit wrote index-K pool elements outside the rows' slots" + ) + + +@pytest.mark.parametrize( + "cell", spec_rollback_stage.CELLS, ids=[c.why for c in spec_rollback_stage.CELLS] +) +def test_stage_cells(cell, certifying) -> None: # noqa: F811 + certifying(spec_rollback_stage) + spec = cell.spec + n, p = spec["n"], spec["p"] + g = _geometry(spec) + gen = torch.Generator().manual_seed(_seed(cell.why)) + ngram_pool, conv_pool = _ple_pools(g, gen) + slots = _slots(n, g["state_slots"], gen) + ids = torch.randint(0, VOCAB, (n * p,), generator=gen).to(DEV, torch.int32) + stream, conv_prev = spec_rollback_stage(ngram_pool, conv_pool, slots, ids, p) + want_stream, want_prev = spec_rollback_stage.reference(ngram_pool, conv_pool, slots, ids, p) + _equal(stream, want_stream, "stream") + _equal(conv_prev, want_prev, "conv_prev") + + +@pytest.mark.parametrize( + "cell", spec_rollback_commit.CELLS, ids=[c.why for c in spec_rollback_commit.CELLS] +) +def test_commit_cells(cell, certifying) -> None: # noqa: F811 + certifying(spec_rollback_commit) + step = _Step(cell.spec, _seed(cell.why)) + ref = step.reference() + before = step.snapshot() + spec_rollback_commit(*step.call_args()) + torch.cuda.synchronize() + step.check(ref, before) + + +def test_commit_graph_replay_rereads_its_inputs() -> None: + """One graph captured at the call site's 1024 rows (128 requests x 8, the 10240 x 9 conv + state over 1874 slots), replayed after every staged input, slot and page is redrawn in place: + what the decode target's captured verify step does each iteration.""" + spec = dict(spec_rollback_commit.CELLS[-1].spec, pool_pages=4096) + step = _Step(spec, 4242) + args = step.call_args() + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + spec_rollback_commit(*args) + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + spec_rollback_commit(*args) + for seed in (7, 8): + step.draw(torch.Generator().manual_seed(seed)) + ref = step.reference() + before = step.snapshot() + graph.replay() + torch.cuda.synchronize() + step.check(ref, before) + + +_GUARD_SPEC = dict(n=8, p=4, planes=12) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_verify_step_setup.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_verify_step_setup.py similarity index 95% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_verify_step_setup.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_verify_step_setup.py index d7824aaa72e2..ea58f90dda07 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_verify_step_setup.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/attention/test_modeling_v2_verify_step_setup.py @@ -7,14 +7,14 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_rollback_planes import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_rollback_planes import ( # noqa: E501 plane_table, qsa_rollback_gather, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.verify_step_setup import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.verify_step_setup import ( # noqa: E501 verify_step_setup, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.cu_seqlens_from_seq_lens import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.cu_seqlens_from_seq_lens import ( # noqa: E501 cu_seqlens_from_seq_lens, ) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/embedding/test_modeling_v2_ngram_hash.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/embedding/test_modeling_v2_ngram_hash.py similarity index 68% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/embedding/test_modeling_v2_ngram_hash.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/embedding/test_modeling_v2_ngram_hash.py index a402400fc12d..ce804c238fc9 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/embedding/test_modeling_v2_ngram_hash.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/embedding/test_modeling_v2_ngram_hash.py @@ -1,10 +1,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the ngram_hash entry. +"""GPU test for the default target's ngram_hash variant (models/qwen4_exp/ops/embedding/ngram_hash). Exact integer arithmetic, gated bitwise against the entry's scalar reference. Beyond the cells: the EOS reset that `ngram_token_history`'s EOS seed relies -on, the narrow-multiplier guard and graph capture. +on. """ import zlib @@ -12,7 +12,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.embedding.ngram_hash import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.embedding.ngram_hash import ( # noqa: E501 CKPT_CONTEXT_LEN, CKPT_EOS_TOKEN_ID, CKPT_HEAD_OFFSETS, @@ -20,11 +20,11 @@ CKPT_MULTIPLIERS, CKPT_VOCAB_SIZE, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.embedding.ngram_hash import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.embedding.ngram_hash import ( # noqa: E501 ngram_hash as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "ngram_hash requires a CUDA device" @@ -97,36 +97,3 @@ def test_eos_restarts_the_window() -> None: no_eos = torch.tensor([[11, 22, 33, 44, token]], device=DEV) assert torch.equal(_hash(fresh, 1), _hash(after_eos, 1)) assert not torch.equal(_hash(fresh, 1), _hash(no_eos, 1)) - - -def test_guard_refuses_narrow_multipliers() -> None: - ids = _tokens(1, CKPT_CONTEXT_LEN + 4, seed=47) - multipliers, sizes, offsets = _constants() - narrow = multipliers.to(torch.int32) - unguarded = op(ids, narrow, sizes, offsets, EOS, 4) - assert not torch.equal(unguarded, _hash(ids, 4)), "the hazard this guard covers is gone" - assert int(unguarded.min()) >= 0 and int(unguarded.max()) < PADDED_TABLE_ROWS - with pytest.raises(AssertionError, match="int64"): - with validating(op): - op(ids, narrow, sizes, offsets, EOS, 4) - - -def test_cuda_graph_capture_and_replay() -> None: - ids = _tokens(8, CKPT_CONTEXT_LEN + 1, seed=53, lead_eos=False, eos_density=0.2) - constants = _constants() - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - for _ in range(3): - op(ids, *constants, EOS, 1) - torch.cuda.current_stream().wait_stream(warmup) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured = op(ids, *constants, EOS, 1) - - fresh = _tokens(8, CKPT_CONTEXT_LEN + 1, seed=54, lead_eos=False, eos_density=0.2) - ids.copy_(fresh) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(captured, op.reference(fresh, *constants, EOS, 1)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_bf16_gemv_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_bf16_gemv_pdl.py similarity index 82% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_bf16_gemv_pdl.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_bf16_gemv_pdl.py index ba0d3a05c9df..6f2cd3278b41 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_bf16_gemv_pdl.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_bf16_gemv_pdl.py @@ -5,8 +5,7 @@ Cells against the fp32 reference within the band `low_m_bf16_mm` uses (the same products, summed in another order). Beyond the cells: a PDL chain under CUDA-graph capture and replay, where each GEMV consumes the previous one's -output and so would read a stale activation if it read before its wait; and -the layout guard. +output and so would read a stale activation if it read before its wait. """ import zlib @@ -14,11 +13,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.bf16_gemv_pdl import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.bf16_gemv_pdl import ( # noqa: E501 bf16_gemv_pdl as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import ( # noqa: E402,F401 — path above; fixture by name ULP, assert_within_ulp, @@ -53,15 +52,26 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 _compare(out, op.reference(mat_a, mat_b, s["pdl"])) -def test_dependent_chain_under_graph_replay() -> None: +@pytest.mark.parametrize( + "rows, dims, last_n", + [ + (1, [2560, 6144, 2560, 640, 2560], 16480), + # Four rows (an MTP-3 verify step) at a tp4/tep4 rank's K=1536. + (4, [1536, 1536, 1536, 1536], 2560), + ], + ids=["tp1-one-row", "tp4-four-rows"], +) +def test_dependent_chain_under_graph_replay(rows, dims, last_n) -> None: """Decode runs this under CUDA graphs, back to back with other kernels. Each link of the chain reads the output of the one before it: a GEMV that read its activation before `griddepcontrol.wait` would see the previous replay's values, and the comparison below would fail on the second replay.""" pdl = True - dims = [2560, 6144, 2560, 640, 2560] # K of each link; N is the next link's K - weights = [_build(1, n, k, seed=i)[1] for i, (k, n) in enumerate(zip(dims, dims[1:] + [16480]))] - x = torch.randn(1, dims[0], device="cuda").to(torch.bfloat16) + # K of each link; N is the next link's K. + weights = [ + _build(rows, n, k, seed=i)[1] for i, (k, n) in enumerate(zip(dims, dims[1:] + [last_n])) + ] + x = torch.randn(rows, dims[0], device="cuda").to(torch.bfloat16) def chain(): h = x @@ -90,9 +100,3 @@ def chain(): # Each link against the reference of what the previous link actually stored. _compare(out, op.reference(h, w)) h = out - - -def test_guard_refuses_a_row_major_mat_b() -> None: - mat_a, mat_b = _build(1, 512, 2560, seed=3) - with validating(op), pytest.raises(AssertionError): - op(mat_a, mat_b.contiguous()) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_fp8_gemv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_gemv.py similarity index 82% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_fp8_gemv.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_gemv.py index 2500b3fa4722..368839e703e2 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_fp8_gemv.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_gemv.py @@ -3,13 +3,13 @@ """GPU test: `fp8_gemv` / `fp8_gemv_conv1d` / `fp8_gemv_conv1d_verify` against an fp32 reference over the dequantized weight (128x128-block e4m3), per cell; the verify entry also bitwise against the two-launch path it replaces (`fp8_gemv` + `causal_conv1d_triton_update`); -strided activations; CUDA-graph replay.""" +strided activations.""" import pytest import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.fp8_gemv import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.fp8_gemv import ( # noqa: E501 dequantize_fp8_block, fp8_gemv_conv1d_op, fp8_gemv_conv1d_verify_op, @@ -17,7 +17,7 @@ pack_weight, quantize_fp8_block, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_triton_update import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.causal_conv1d_triton_update import ( # noqa: E501 causal_conv1d_triton_update, ) @@ -92,20 +92,6 @@ def test_conv1d_cells(cell) -> None: assert torch.equal(conv_state[others], before[others]), "other slots untouched" -def test_graph_replay_matches_eager() -> None: - _, _, packed = _weight(2560, 6144, 5) - x = torch.randn(4, 6144, device=DEV).to(torch.bfloat16) - fp8_gemv_op(x, packed) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - out = fp8_gemv_op(x, packed) - x.copy_(torch.randn(4, 6144, device=DEV).to(torch.bfloat16)) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(out, fp8_gemv_op(x, packed)) - - def _verify_args(spec, seed): """`fp8_gemv_conv1d_verify`'s arguments for a cell: the activation a column slice of a wider row, the pool a strided view (every 4th column skipped), slot 0's window zero (a fresh @@ -190,32 +176,3 @@ def test_conv1d_verify_cells(cell) -> None: assert not torch.equal(window_a[r, :p], window[r, :p]), r if spec["snapshot_steps"] > p: assert torch.equal(window_a[:, p:], window[:, p:]), "steps past p untouched" - - -def test_conv1d_verify_graph_replay_matches_eager() -> None: - spec = fp8_gemv_conv1d_verify_op.CELLS[2].spec - _, _, packed = _weight(spec["n"], spec["k"], 5) - x, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad = _verify_args(spec, 9) - state0, window0 = conv_state.clone(), window.clone() - fp8_gemv_conv1d_verify_op( - x, packed, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad - ) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - conv, rest = fp8_gemv_conv1d_verify_op( - x, packed, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad - ) - x.copy_(torch.randn(x.shape, device=DEV).to(torch.bfloat16)) - conv_state.copy_(state0) - window.copy_(window0) - graph.replay() - torch.cuda.synchronize() - got = (conv.clone(), rest.clone(), conv_state.clone(), window.clone()) - conv_state.copy_(state0) - window.copy_(window0) - want = fp8_gemv_conv1d_verify_op( - x, packed, conv_state, conv_weight, idx, window, rows, conv_dim, p, pad - ) - for a, b in zip(got, (*want, conv_state, window)): - assert torch.equal(a.view(torch.int16), b.view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_fp8_per_tensor_mm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_per_tensor_mm.py similarity index 82% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_fp8_per_tensor_mm.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_per_tensor_mm.py index dce1286b3460..7329eb4f2890 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_fp8_per_tensor_mm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_fp8_per_tensor_mm.py @@ -2,12 +2,12 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test: `fp8_per_tensor_mm` against its fp32 reference (the activation quantized as the op does, both operands dequantized) per cell, the activation quantization bitwise against -torch's e4m3 cast, and CUDA-graph replay with a new activation.""" +torch's e4m3 cast, and a zero activation and strided rows.""" import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.fp8_per_tensor_mm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.fp8_per_tensor_mm import ( # noqa: E501 E4M3_MAX, fp8_per_tensor_mm, fp8_per_tensor_weight, @@ -68,16 +68,3 @@ def test_zero_activation_and_strided_rows() -> None: assert torch.equal( fp8_per_tensor_mm(wide[:, :2560], codes_t, scale), fp8_per_tensor_mm(x, codes_t, scale) ) - - -def test_graph_replay_matches_eager() -> None: - x, _, codes_t, scale = _operands(64, 2560, 16480, 7) - fp8_per_tensor_mm(x, codes_t, scale) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - out = fp8_per_tensor_mm(x, codes_t, scale) - x.copy_(_operands(64, 2560, 16480, 8)[0] * 4) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(out, fp8_per_tensor_mm(x, codes_t, scale)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_low_m_bf16_mm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_low_m_bf16_mm.py similarity index 63% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_low_m_bf16_mm.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_low_m_bf16_mm.py index d3b3a17e22a0..699442c6d0f3 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_low_m_bf16_mm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_low_m_bf16_mm.py @@ -5,9 +5,6 @@ Every cell names the branch `plan` must pick, and a host-only check pins the branch of the table's other shapes and M buckets, so a table edit that moves a shape to another kernel fails here before it can change a number silently. -Beyond the cells: CUDA-graph capture and replay after warmup (the way the -decode forward runs it), and the guard against a padded `mat_b`, which the -cuBLAS branch would read as dense. """ import zlib @@ -15,14 +12,14 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.low_m_bf16_mm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.low_m_bf16_mm import ( # noqa: E501 low_m_bf16_mm as op, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.low_m_bf16_mm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.low_m_bf16_mm import ( # noqa: E501 plan, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import ( # noqa: E402,F401 — path above; fixture by name ULP, assert_within_ulp, @@ -90,6 +87,15 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 (1, 2560, 640, "cublas"), (8192, 2560, 6144, "cublas"), (1, 1024, 2048, "cublas"), + (3, 4120, 2560, "splitk"), + (9, 4120, 2560, "cublas"), + (2, 4352, 2560, "direct"), + (3, 4352, 2560, "splitk"), + (9, 4352, 2560, "cublas"), + (2, 2560, 1536, "direct"), + (5, 2560, 1536, "splitk"), + (9, 2560, 1536, "cublas"), + (33, 512, 2560, "cublas"), ) @@ -99,50 +105,9 @@ def test_plan_pins_the_branch_of_the_other_table_shapes() -> None: assert not wrong, f"(M, N, K, pinned, picked): {wrong}" -def test_graph_replay_after_warmup() -> None: - """Decode runs this under CUDA graphs, with PDL: compile in warmup, capture, - then replay with new activations copied into the captured input. One call per - branch, in a chain, so a FlashInfer kernel follows cuBLAS and vice versa.""" - pdl = True - shapes = [(1, 10240, 2560), (8, 2560, 6144), (1, 2560, 640), (8, 12288, 2560)] - assert [plan(*s).kind for s in shapes] == ["direct", "splitk", "cublas", "splitk"] - ins = [_build(m, n, k, seed=i) for i, (m, n, k) in enumerate(shapes)] - static = [a.clone() for a, _ in ins] - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - for _ in range(2): - [op(a, b, pdl) for a, (_, b) in zip(static, ins)] - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - outs = [op(a, b, pdl) for a, (_, b) in zip(static, ins)] - for trial in range(3): - g = torch.Generator(device="cuda").manual_seed(100 + trial) - for a in static: - a.copy_(torch.randn(a.shape, generator=g, device="cuda").to(a.dtype)) - graph.replay() - torch.cuda.synchronize() - for a, (_, b), out in zip(static, ins, outs): - _compare(out, op.reference(a, b)) - - def _padded(m, n, k): """A row-padded mat_a and a column-padded mat_b: dense along K, leading dimension 2K. Both are views a caller gets by slicing a wider tensor.""" a = torch.randn(m, 2 * k, dtype=torch.bfloat16, device="cuda")[:, :k] w = (torch.randn(n, 2 * k, device="cuda") * k**-0.5).to(torch.bfloat16)[:, :k] return a, w.t() - - -def test_guard_refuses_a_padded_mat_b() -> None: - """Why the guard exists: the cuBLAS branch ignores mat_b's leading dimension - and silently computes a different product.""" - a, b = _padded(4, 640, 640) - a = a.contiguous() - assert plan(4, 640, 640).kind == "cublas" and b.stride() == (1, 1280) - with pytest.raises(AssertionError): - _compare(op(a, b), op.reference(a, b)) - with pytest.raises(AssertionError, match="column-major"): - with validating(op): - op(a, b) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_mid_rows_mm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_mid_rows_mm.py similarity index 84% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_mid_rows_mm.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_mid_rows_mm.py index 959822b4b891..d6dfb627ad4f 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_mid_rows_mm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_mid_rows_mm.py @@ -2,20 +2,19 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test: `mid_rows_mm` (FP8 with the per-row activation quantization, and bf16) against an fp32 reference over the operands it multiplies, per cell; `fp8_row_quant` bitwise against -torch's e4m3 cast; strided activation rows; CUDA-graph replay with new activations.""" +torch's e4m3 cast; strided activation rows.""" import pytest import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.mid_rows_mm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.mid_rows_mm import ( # noqa: E501 AMAX_FLOOR, E4M3_MAX, MidWeight, fp8_row_quant, fp8_row_quant_op, mid_rows_mm_op, - precompile, ) assert torch.cuda.is_available(), "mid_rows_mm requires a CUDA device" @@ -101,19 +100,3 @@ def test_row_quant_is_torchs_cast(cell) -> None: assert torch.equal(s, (amax / E4M3_MAX).float()) want = (x.float() * (E4M3_MAX / amax).float()[:, None]).to(torch.float8_e4m3fn) assert torch.equal(q.view(torch.uint8), want.view(torch.uint8)) - - -def test_graph_replay_matches_eager() -> None: - for mode, n, k, m in (("tensor", 16480, 2560, 64), ("bf16", 512, 2560, 200)): - _, mw = _weight(n, k, mode, 3) - precompile(mw) - x = _activation(m, k, 1).contiguous() - mid_rows_mm_op(x, mw) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - out = mid_rows_mm_op(x, mw) - x.copy_(_activation(m, k, 2) * 4) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(out, mid_rows_mm_op(x, mw)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_nvfp4_gemv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_nvfp4_gemv.py similarity index 87% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_nvfp4_gemv.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_nvfp4_gemv.py index 7f5a273112d2..6fe44edb45fa 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_nvfp4_gemv.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/gemm/test_modeling_v2_nvfp4_gemv.py @@ -1,8 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """GPU test: `nvfp4_gemv` against an fp32 reference over the dequantized NVFP4 weight, per -cell; the quantizer against its definition; CUDA-graph replay; one compiled kernel per row -octet; the f16-range guard. The split entries (NVFP4 units on `fp8_gemv`'s kernel) per cell: +cell; the quantizer against its definition; one compiled kernel per row +octet. The split entries (NVFP4 units on `fp8_gemv`'s kernel) per cell: the projection against the fp32 reference, the conv entries also against the two-launch paths they replace (the split projection + `causal_conv1d_update` / `causal_conv1d_triton_update`).""" @@ -11,13 +11,13 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm import ( # noqa: E501 nvfp4_gemv as nv, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_triton_update import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.causal_conv1d_triton_update import ( # noqa: E501 causal_conv1d_triton_update, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_update import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.causal_conv1d_update import ( # noqa: E501 causal_conv1d_update, ) @@ -105,29 +105,6 @@ def test_one_kernel_per_row_octet() -> None: assert nv._compiled.cache_info().currsize - before <= 2 -def test_graph_replay_matches_eager() -> None: - _, _, _, _, packed = _weight(4096, 2560, 9) - x = torch.randn(4, 2560, device=DEV).to(torch.bfloat16) - nvfp4_gemv_op(x, packed) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - out = nvfp4_gemv_op(x, packed) - x.copy_(torch.randn(4, 2560, device=DEV).to(torch.bfloat16)) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(out, nvfp4_gemv_op(x, packed)) - - -def test_is_valid_rejects_activation_beyond_f16() -> None: - _, _, _, _, packed = _weight(256, 256, 1) - x = torch.randn(2, 256, device=DEV).to(torch.bfloat16) - nvfp4_gemv_op.is_valid(x, packed) - x[1, 3] = 1e5 - with pytest.raises(AssertionError): - nvfp4_gemv_op.is_valid(x, packed) - - # ------------------------------------------------------------------------------------------ # split entries: NVFP4 units on fp8_gemv's kernel # ------------------------------------------------------------------------------------------ @@ -283,24 +260,9 @@ def test_split_conv1d_verify_cells(cell) -> None: assert (got.float() - ref_.float()).abs().max() <= 2.0**-7 * ref_.float().abs().max() -def test_split_graph_replay_matches_eager() -> None: - _, codes, scales, gscale, _ = _weight(GDN_N, GDN_K, 9) - packed = nv.pack_split(codes, scales, gscale) - x = torch.randn(4, GDN_K, device=DEV).to(torch.bfloat16) - split_op(x, packed) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - out = split_op(x, packed) - x.copy_(torch.randn(4, GDN_K, device=DEV).to(torch.bfloat16)) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(out, split_op(x, packed)) - - def test_split_precompile_covers_the_decode_variants() -> None: """After `precompile_split`, no projection / conv / verify call within its rows compiles.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm import ( # noqa: E501 fp8_gemv as f8, ) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_gated_residual.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_gated_residual.py similarity index 95% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_gated_residual.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_gated_residual.py index 9d2a95df0a97..155544922e6b 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_gated_residual.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_gated_residual.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qwen4_exp_gated_residual entry. +"""GPU test for the qwen4_exp_gated_residual catalog entry. The thing under test is a formula -- which operand is gated, which axis is reduced, where each division by `hc_count` sits, the factor on the injection @@ -15,14 +15,14 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 gate_ratio, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 qwen4_exp_gated_residual as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "qwen4_exp_gated_residual requires a CUDA device" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py similarity index 94% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py index 2cd74f1c0b3d..877e8a866659 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py @@ -6,8 +6,7 @@ combine's and the finalize's own rounding at the MoE sites), `part` within `2^-16` of the reference relative to its scale, `ss` to fp32 accuracy. Beyond the cells: a stale-read check under CUDA-graph replay with PDL, where a kernel that read its -activations before `griddepcontrol.wait` would see the previous replay's; the one-row -guard. +activations before `griddepcontrol.wait` would see the previous replay's. """ import zlib @@ -15,11 +14,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_first_cute import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_decode_first_cute import ( # noqa: E501 qwen4_exp_hc_decode_first_cute as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "qwen4_exp_hc_decode_first_cute requires a CUDA device" @@ -187,10 +186,3 @@ def step(): want = op(*[*args[:10], args[10]]) for a, b in zip(graph_out, want): assert torch.equal(a, b), f"replay {trial}" - - -def test_guards() -> None: - args = list(_build(dict(), seed=6)) - two = args[0].expand(2, -1).contiguous() - with validating(op), pytest.raises(AssertionError, match="one row"): - op(two, args[1].expand(2, -1).contiguous(), args[2].expand(2, -1).contiguous(), *args[3:]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py similarity index 91% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py index 5958e069eb6a..831b6bfba981 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py @@ -4,9 +4,8 @@ (`qwen4_exp_hc_decode_mix_fp4_cute`) and the plain build (`qwen4_exp_hc_decode_mix_cute`) of the CuTe mix kernel. Each entry's cells against its float64 reference in its band (`mixed` and `inject_out` in `rtol = eps`, `atol = eps * max|ref|`; for the NVFP4 build, -the NVFP4 exact on the returned `mixed` outside the reciprocal tie window), a -stale-read check under CUDA-graph replay with PDL for the NVFP4 build, the one-row -guard.""" +the NVFP4 exact on the returned `mixed` outside the reciprocal tie window), and a +stale-read check under CUDA-graph replay with PDL for the NVFP4 build.""" import zlib @@ -17,12 +16,12 @@ Fp4Quantize, fp4_quantize, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_mix_fp4_cute import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_decode_mix_fp4_cute import ( # noqa: E501 qwen4_exp_hc_decode_mix_cute, qwen4_exp_hc_decode_mix_fp4_cute, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "qwen4_exp_hc_decode_mix_fp4_cute requires a CUDA device" @@ -151,12 +150,3 @@ def step(): torch.cuda.synchronize() for a, b in zip(graph_out, op(*args)): assert torch.equal(a, b), f"replay {trial}" - - -def test_guards() -> None: - """Both entries share the shape checks; the NVFP4 build's `is_valid` refuses two rows.""" - op = qwen4_exp_hc_decode_mix_fp4_cute - args = _build(dict(), seed=6) - two = args[3].expand(2, -1).contiguous() - with validating(op), pytest.raises(AssertionError): - op(args[0].expand(-1, 2, -1), args[1].expand(-1, 2), args[2], two, *args[4:]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_rows_tc.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_rows_tc.py similarity index 89% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_rows_tc.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_rows_tc.py index 9f4f27c1fe15..c9a8f96ab984 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_rows_tc.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_decode_rows_tc.py @@ -15,7 +15,7 @@ import pytest import torch -P = "tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc." +P = "tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc." rows_mod = importlib.import_module(P + "qwen4_exp_hc_decode_rows_tc") first_op = rows_mod.qwen4_exp_hc_decode_rows_first mix_op = rows_mod.qwen4_exp_hc_decode_rows_mix @@ -132,6 +132,21 @@ def test_mix_cells(cell) -> None: _check(cell) +@pytest.mark.parametrize("rows", [2, 7, 16]) +@pytest.mark.parametrize("site", ["combine", "fold"]) +def test_gate_column_of_the_fused_shared_gemm(rows, site) -> None: + """The target's gate logits are a column of the shared expert's fused gate|up GEMM output + (row stride 2 * 640 + 8): every output is bitwise the dense `[R, 1]` gate's, row by row, + so no row reads another's logit.""" + o = _operands(rows, 500 + rows) + gug = torch.randn(rows, 2 * 640 + 8, device=DEV).to(torch.bfloat16) + gug[:, 1280:1281] = o["gate"] + want = _site(o, site, True) + got = _site(dict(o, gate=gug[:, 1280:1281]), site, True) + for a, b in zip(got, want): + assert (a is None and b is None) or torch.equal(a, b) + + @pytest.mark.parametrize("rows", [2, 16]) @pytest.mark.parametrize("n_cols", [RANK + C, RANK]) def test_fold_bitwise_vs_runner_finalize(rows, n_cols) -> None: @@ -155,19 +170,3 @@ def test_fold_bitwise_vs_runner_finalize(rows, n_cols) -> None: want = first_op(o["hc"], routed, o["inject"], o["norm"], o["first"], C, n_cols, EPS, **kw) for a, b in zip(got, want): assert (a is None and b is None) or torch.equal(a, b) - - -def test_replay_matches_eager() -> None: - o = _operands(8, 77) - _site(o, "inject", True) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured = _site(o, "inject", True) - fresh = _operands(8, 78) - o["hc"].copy_(fresh["hc"]) - o["block"].copy_(fresh["block"]) - graph.replay() - torch.cuda.synchronize() - eager = _site(o, "inject", True) - assert torch.equal(captured[1], eager[1]) and torch.equal(captured[0], eager[0]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_first_rescale_tc.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_first_rescale_tc.py similarity index 71% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_first_rescale_tc.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_first_rescale_tc.py index 1645a331fa44..cd294ad6c6d8 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_first_rescale_tc.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_first_rescale_tc.py @@ -2,8 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the qwen4_exp_hc_first_rescale_tc entry: `low` and `inject` in the gate band (`qwen4_exp_gated_residual`'s) of the float64 projection of the -once-rounded normed stream, the stack-final form, the weight-layout guard, -graph capture. +once-rounded normed stream, and the stack-final form. """ import zlib @@ -11,14 +10,14 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 gate_ratio, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_first_rescale_tc import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_first_rescale_tc import ( # noqa: E501 qwen4_exp_hc_first_rescale_tc as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "qwen4_exp_hc_first_rescale_tc requires a CUDA device" @@ -91,32 +90,3 @@ def test_stack_final_low_is_the_injecting_form_bitwise() -> None: low_f, none = op(*args[:6], False) assert none is None and inject.shape == (300, HC_COUNT) assert torch.equal(low, low_f) - - -def test_guard_refuses_a_contiguous_first() -> None: - """The kernel streams the `[N, C*H]` weight row-major: `first` must be its - transposed view, not a `[C*H, N]` copy. - """ - args = list(_build(dict(rows=4), seed=4)) - args[3] = args[3].contiguous() - with validating(op), pytest.raises(AssertionError, match="transposed"): - op(*args) - - -def test_cuda_graph_capture_and_replay() -> None: - args = _build(dict(rows=8), seed=5) - side = torch.cuda.Stream() - side.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(side): - op(*args) - torch.cuda.current_stream().wait_stream(side) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - g_low, g_inject = op(*args) - fresh = _build(dict(rows=8), seed=6) - args[0].copy_(fresh[0]) - args[1].copy_(fresh[1]) - graph.replay() - torch.cuda.synchronize() - e_low, e_inject = op(*args) - assert torch.equal(g_low, e_low) and torch.equal(g_inject, e_inject) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py similarity index 89% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py index b69e82f5c120..f87d5115878d 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_norm.py @@ -1,11 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qwen4_exp_hc_inject_norm entry. +"""GPU test for the qwen4_exp_hc_inject_norm catalog entry. Cells are gated at one ulp of the activation dtype against the float64 reference; the MoE-combine cells widen the residual's band by the combine's own rounding (see `_compare_moe`). Beyond them: the CuTe fold kernel bitwise the -Triton kernel at the last row count before the switch, and the slot-map guard. +Triton kernel at the last row count before the switch. """ import sys @@ -14,11 +14,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_inject_norm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_inject_norm import ( # noqa: E501 qwen4_exp_hc_inject_norm as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "qwen4_exp_hc_inject_norm requires a CUDA device" @@ -33,6 +33,12 @@ # This checkpoint's grouped-norm weight statistics (layer-0 hyper-connection norm). GAMMA_RMS = 0.477 GAMMA_ABSMAX = 6.625 +# The stack-final mixer's norm weight (`hyper_connection_mixer.hc_norm.weight`), which the +# last layer's MoE site normalizes with: RMS 3.106, from -0.367 to 12.375, nearly all +# positive -- so `1 + w` scales the normed stream by up to 13.4. +FINAL_GAMMA_RMS = 3.106 +FINAL_GAMMA_MIN = -0.367 +FINAL_GAMMA_MAX = 12.375 def _ulp_ratio(out: torch.Tensor, ref: torch.Tensor) -> float: @@ -115,8 +121,13 @@ def _build(spec, seed): width = hc_count * hidden g = torch.Generator(device=DEV).manual_seed(seed) hyper_input = (torch.randn(rows, width, generator=g, device=DEV) * 3.0).to(dtype) - gamma = torch.randn(width, generator=g, device=DEV) * GAMMA_RMS - gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX + if spec.get("site", "trunk") == "stack_final": + gamma = torch.randn(width, generator=g, device=DEV).abs() * FINAL_GAMMA_RMS + gamma = gamma.clamp(max=FINAL_GAMMA_MAX) + gamma[0], gamma[width // 2] = FINAL_GAMMA_MAX, FINAL_GAMMA_MIN + else: + gamma = torch.randn(width, generator=g, device=DEV) * GAMMA_RMS + gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX gamma = gamma.to(dtype) block_out = inject = shared = gate_logits = None if spec.get("inject", True): @@ -173,13 +184,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by _compare_moe(args, (hc, normed), op.reference(*args)) -def test_guard_refuses_an_int64_slot_map() -> None: - """The fold's slot map is the runner's int32 index map; `is_valid` refuses an int64 one.""" - fold = _build(dict(rows=4, moe=True, fold=10), seed=9) - with validating(op), pytest.raises(AssertionError, match="int32"): - op(*fold[:9], fold[9].long()) - - def test_fold_cute_is_the_triton_kernel_bitwise(monkeypatch) -> None: """Below `_FEW_WARPS_MIN_ROWS` rows the fold runs the CuTe DSL kernel; its `hc_out` and `normed` are bitwise the Triton kernel's (run by refusing the diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_rstd.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_rstd.py similarity index 83% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_rstd.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_rstd.py index 20558035d0b0..5159e2281416 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_inject_rstd.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_inject_rstd.py @@ -6,8 +6,7 @@ widened by the combine's own rounding, as `qwen4_exp_hc_inject_norm`'s test), rstd within 2^-16 relative of the float64 statistic of the stored residual. Beyond them: `hc_out` and the normed stream rebuilt from rstd bitwise -`qwen4_exp_hc_inject_norm`'s at the cells' sizes, the operand guard, graph -capture. +`qwen4_exp_hc_inject_norm`'s at the cells' sizes. """ import zlib @@ -15,14 +14,14 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_inject_norm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_inject_norm import ( # noqa: E501 qwen4_exp_hc_inject_norm, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_inject_rstd import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_inject_rstd import ( # noqa: E501 qwen4_exp_hc_inject_rstd as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "qwen4_exp_hc_inject_rstd requires a CUDA device" @@ -138,28 +137,3 @@ def test_normed_rebuilt_from_rstd_is_inject_norm_bitwise(cell) -> None: normed = (x.flatten(-2) * (1.0 + gamma.float())).to(hc.dtype) diff = (normed != want_normed).sum().item() assert diff == 0, f"{diff} normed elements differ" - - -def test_guard_refuses_block_out_without_inject() -> None: - hyper_input, block_out, _, eps, c, _, _ = _build(dict(rows=4), seed=4) - with validating(op), pytest.raises(AssertionError, match="together"): - op(hyper_input, block_out, None, eps, c) - - -def test_cuda_graph_capture_and_replay() -> None: - args = list(_build(dict(rows=8, moe=True), seed=5)) - side = torch.cuda.Stream() - side.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(side): - op(*args) - torch.cuda.current_stream().wait_stream(side) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - g_hc, g_rstd = op(*args) - fresh = _build(dict(rows=8, moe=True), seed=6) - for i in (0, 1, 2, 5, 6): - args[i].copy_(fresh[i]) - graph.replay() - torch.cuda.synchronize() - e_hc, e_rstd = op(*fresh) - assert torch.equal(g_hc, e_hc) and torch.equal(g_rstd, e_rstd) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py similarity index 63% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py index 8641b2b8b7b5..07b816f9ae61 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py @@ -1,18 +1,19 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qwen4_exp_hc_lowrank_act entry: cells within one ulp of the -float64 reference, the fp32-logits guard, graph capture.""" +"""GPU test for the default target's qwen4_exp_hc_lowrank_act variant +(models/qwen4_exp/ops/hc/qwen4_exp_hc_lowrank_act): cells within one ulp of the float64 +reference.""" import zlib import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_lowrank_act import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_lowrank_act import ( # noqa: E501 qwen4_exp_hc_lowrank_act as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "qwen4_exp_hc_lowrank_act requires a CUDA device" @@ -54,26 +55,3 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by low, inject = op(*args) assert low.shape == (cell.spec["rows"], RANK) _compare((low, inject), op.reference(*args)) - - -def test_guard_refuses_bf16_logits() -> None: - logits, rank, c, wi, dt = _build(dict(rows=4), seed=4) - with validating(op), pytest.raises(AssertionError): - op(logits.to(torch.bfloat16), rank, c, wi, dt) - - -def test_cuda_graph_capture_and_replay() -> None: - logits, rank, c, wi, dt = _build(dict(rows=8), seed=5) - side = torch.cuda.Stream() - side.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(side): - op(logits, rank, c, wi, dt) - torch.cuda.current_stream().wait_stream(side) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - g_low, g_inject = op(logits, rank, c, wi, dt) - logits.copy_(_build(dict(rows=8), seed=6)[0]) - graph.replay() - torch.cuda.synchronize() - e_low, e_inject = op(logits, rank, c, wi, dt) - assert torch.equal(g_low, e_low) and torch.equal(g_inject, e_inject) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py similarity index 77% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py index 3e392c8e6f75..458a7dbb0390 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py @@ -9,14 +9,15 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 - gate_ratio, +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_up_gate_mean import ( # noqa: E501 + _BLOCK_H, + _COMPILED, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_up_gate_mean import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_up_gate_mean import ( # noqa: E501 qwen4_exp_hc_up_gate_mean as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "qwen4_exp_hc_up_gate_mean requires a CUDA device" @@ -24,6 +25,17 @@ DEV = "cuda" +def gate_ratio(out: torch.Tensor, ref: torch.Tensor) -> float: + """Worst element's distance from `ref` as a fraction of the test's band: + `rtol = eps`, `atol = eps * max|ref|`, both of `ref`'s dtype. Below 1 is + inside. (`qwen4_exp_gated_residual`'s band, restated so this test reads + no other target's code.)""" + eps = torch.finfo(ref.dtype).eps + ref64 = ref.double() + atol = eps * ref64.abs().max().item() + return ((out.double() - ref64).abs() / (atol + eps * ref64.abs())).max().item() + + def _compare(out, ref) -> None: """`rtol = eps`, `atol = eps * max|ref|` (`gate_ratio <= 1`), the band `qwen4_exp_gated_residual` derives for `mixed`.""" @@ -51,6 +63,8 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) mixed = op(*args) assert mixed.shape == (cell.spec["rows"], cell.spec.get("hidden", 2560)) + # The answer came from the CuTe DSL kernel, JIT-compiled for this device. + assert (_BLOCK_H, args[3], mixed.shape[1], args[0].shape[1], False) in _COMPILED _compare(mixed, op.reference(*args)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_rstd.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_rstd.py similarity index 77% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_rstd.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_rstd.py index 3f4c2f41f174..cee012140e7b 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_rstd.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_rstd.py @@ -4,7 +4,7 @@ `mixed` in the gate band of the float64 reference, the epilogue's fp32 operation sequence bitwise (gate at zero logits), the NVFP4 outputs against the native-torch quantization of the returned `mixed` outside the reciprocal tie window, the NVFP4 -build's `mixed` bitwise the bf16 build's, the global-scale guard, graph capture. +build's `mixed` bitwise the bf16 build's. """ import zlib @@ -16,17 +16,17 @@ Fp4Quantize, fp4_quantize, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 gate_ratio, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_up_gate_mean_rstd import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_up_gate_mean_rstd import ( # noqa: E501 qwen4_exp_hc_up_gate_mean_rstd as op, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_up_gate_mean_rstd import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_up_gate_mean_rstd import ( # noqa: E501 qwen4_exp_hc_up_gate_mean_rstd_fp4 as op4, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "qwen4_exp_hc_up_gate_mean_rstd requires a CUDA device" @@ -135,41 +135,3 @@ def test_epilogue_operation_sequence_bitwise(rows) -> None: def test_fp4_mixed_is_the_bf16_entry_bitwise(rows) -> None: args = _build(dict(rows=rows), seed=100 + rows) assert torch.equal(op4(*args)[0], op(*args[:6])) - - -def test_guards() -> None: - """The kernel loads `global_scale[0]` only; a per-row scale is refused. rstd must be - fp32. - """ - low, up, hc_out, rstd, w1, c, gs = _build(dict(rows=4), seed=4) - with validating(op4), pytest.raises(AssertionError): - op4(low, up, hc_out, rstd, w1, c, torch.cat([gs, gs])) - with validating(op), pytest.raises(AssertionError, match="rstd"): - op(low, up, hc_out, rstd.to(torch.bfloat16), w1, c) - - -@pytest.mark.parametrize("fp4", [False, True]) -def test_cuda_graph_capture_and_replay(fp4) -> None: - args = list(_build(dict(rows=8), seed=5)) - if not fp4: - args = args[:6] - entry = op4 if fp4 else op - side = torch.cuda.Stream() - side.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(side): - entry(*args) - torch.cuda.current_stream().wait_stream(side) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - g_out = entry(*args) - fresh = _build(dict(rows=8), seed=6) - for i in (0, 2, 3): - args[i].copy_(fresh[i]) - if fp4: - args[6].mul_(0.5) - graph.replay() - torch.cuda.synchronize() - want = entry(*args) - g_out, want = (g_out, want) if fp4 else ((g_out,), (want,)) - for got, exp in zip(g_out, want): - assert torch.equal(got, exp) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_post_inject.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_post_inject.py similarity index 60% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_post_inject.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_post_inject.py index e609e8632440..917c563af1b0 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_post_inject.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/hc/test_modeling_v2_qwen4_exp_post_inject.py @@ -1,9 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qwen4_exp_post_inject entry. +"""GPU test for the qwen4_exp_post_inject catalog entry. -The cells are gated bitwise (16-bit) by `_compare`. Beyond them: the -leading-dim guard and graph capture. +The cells are gated bitwise (16-bit) by `_compare`. """ import zlib @@ -11,11 +10,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_post_inject import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_post_inject import ( # noqa: E501 qwen4_exp_post_inject as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "qwen4_exp_post_inject requires a CUDA device" @@ -54,11 +53,14 @@ def _build(spec, seed): layout = spec.get("layout", "contiguous") width = hc_count * hidden if layout == "col_slice": - buf = torch.randn(hyper_input.shape[0], 2 * width, device=DEV).to(dtype) + g = torch.Generator(device=DEV).manual_seed(seed + 1) + buf = torch.randn(hyper_input.shape[0], 2 * width, generator=g, device=DEV).to(dtype) buf[:, :width] = hyper_input hyper_input = buf[:, :width] - if layout != "contiguous": - assert not hyper_input.is_contiguous() and hyper_input.stride(-1) == 1 + # Strides rather than is_contiguous(): at one row torch ignores the + # size-1 dim's stride and calls the slice contiguous. + assert hyper_input.stride() == (2 * width, 1), hyper_input.stride() + assert hyper_input.shape[0] == 1 or not hyper_input.is_contiguous() return hyper_input, block_out, inject @@ -72,35 +74,3 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res assert out.data_ptr() != hyper_input.data_ptr(), "the result is a fresh tensor" assert torch.equal(hyper_input, before), "the residual it was given is read-only" _compare(out, op.reference(hyper_input, block_out, inject)) - - -def test_guard_refuses_unequal_leading_dims() -> None: - """Unguarded, a 1-D `inject` broadcasts one row over every token.""" - hyper_input, block_out, inject = _operands((8,), seed=10) - unguarded = op(hyper_input, block_out, inject[0]) - assert unguarded.shape == hyper_input.shape, "the hazard this guard covers is gone" - assert not torch.equal(unguarded, op.reference(hyper_input, block_out, inject)) - - with pytest.raises(AssertionError, match="leading dims must match"): - with validating(op): - op(hyper_input, block_out, inject[0]) - - -def test_cuda_graph_capture_and_replay() -> None: - hyper_input, block_out, inject = _operands((64,), seed=14) - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - for _ in range(3): - op(hyper_input, block_out, inject) - torch.cuda.current_stream().wait_stream(warmup) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured = op(hyper_input, block_out, inject) - - fresh, _, _ = _operands((64,), seed=15) - hyper_input.copy_(fresh) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(captured, op.reference(hyper_input, block_out, inject)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py similarity index 88% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py index f23d2789ee81..a75ead10c13d 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py @@ -4,8 +4,8 @@ The entry owns the contract. This file builds the checkpoint-shaped operands -- 512 experts in the CuteDSL weight layout, the activation as the target quantizes -it, the sort as moe_sort returns it -- and holds the claims no cell carries: -per-pair invariance across tiles and tactics, and the contiguity guard. +it, the sort as moe_sort returns it -- and holds the claim no cell carries: +per-pair invariance across tiles and tactics. """ import functools @@ -14,7 +14,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import ( # noqa: E501 E2M1, E2M1_MAX, GATE_UP_GROUP, @@ -23,13 +23,13 @@ e4m3_value, unswizzle, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import ( # noqa: E501 cute_dsl_nvfp4_gather_grouped_gemm_act_fusion as op, ) from tensorrt_llm._torch.autotuner import AutoTuner from tensorrt_llm._torch.utils import ActivationType -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "the CuteDSL FC1 requires a CUDA device" @@ -364,36 +364,3 @@ def test_tiles_and_tactics_agree_per_pair(certifying): # noqa: F811 first = seen[(128, 128)] for key, got in seen.items(): _assert_pairs_equal(got, first, f"T={tokens} {key} vs (128, 128)") - - -# ── the guard ───────────────────────────────────────────────────────────── - - -def _strided(t: torch.Tensor) -> torch.Tensor: - """The same values as `t`, as a non-contiguous view into a wider buffer.""" - u = t.view(torch.uint8) if t.element_size() == 1 else t - wide = torch.zeros(*u.shape[:-1], 2 * u.shape[-1], dtype=u.dtype, device=DEV) - wide[..., : u.shape[-1]] = u - view = wide[..., : u.shape[-1]] - return view.view(t.dtype) if t.element_size() == 1 else view - - -def _u8(t: torch.Tensor) -> torch.Tensor: - return t.view(torch.uint8) if t.element_size() == 1 else t - - -def test_noncontiguous_operand_answers_wrongly_and_is_refused(): - """The runner hands data_ptr() to the kernel and never reads a stride: the - same values through a strided view come back finite and wrong. The activation; - the guard checks input_scale, weight, weight_scale and alpha the same way.""" - name = "input" - tokens, tile = 64, 128 - args, _ = _problem(tokens, tile, seed=21) - ref = op.reference(**args) - bad = dict(args, **{name: _strided(args[name])}) - assert not bad[name].is_contiguous() - assert torch.equal(_u8(bad[name]), _u8(args[name])) - with pytest.raises(AssertionError): - _compare(op(**bad), ref) - with validating(op), pytest.raises(AssertionError, match=f"^{name} must be contiguous"): - op(**bad) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py similarity index 90% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py index f744f13a4402..b7f9d97efb50 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py @@ -6,8 +6,7 @@ 512 experts' down_proj in the CuteDSL weight layout, the input as the CuteDSL FC1 returns it (permuted NVFP4 rows, swizzled block scales, zero padding rows, unwritten bytes past the live tiles), the sort as moe_sort returns it -- and -holds the claims no cell carries: per-pair invariance across tiles and tactics, -and the contiguity guard. +holds the claim no cell carries: per-pair invariance across tiles and tactics. """ import functools @@ -16,15 +15,15 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.cute_dsl_nvfp4_grouped_gemm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.cute_dsl_nvfp4_grouped_gemm import ( # noqa: E501 Fc2Expected, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.cute_dsl_nvfp4_grouped_gemm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.cute_dsl_nvfp4_grouped_gemm import ( # noqa: E501 cute_dsl_nvfp4_grouped_gemm as op, ) from tensorrt_llm._torch.autotuner import AutoTuner -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "the CuteDSL FC2 requires a CUDA device" @@ -342,34 +341,8 @@ def test_tiles_and_tactics_agree_per_pair(certifying): # noqa: F811 _assert_pairs_equal(got, first, f"T={tokens} {key} vs (128, 128, 1)") -# ── the guard ───────────────────────────────────────────────────────────── +# ── helpers ─────────────────────────────────────────────────────────────── def _u8(t: torch.Tensor) -> torch.Tensor: return t.view(torch.uint8) if t.element_size() == 1 else t - - -def _strided(t: torch.Tensor) -> torch.Tensor: - """The same values as `t`, as a non-contiguous view into a wider buffer.""" - u = _u8(t) - wide = torch.zeros(*u.shape[:-1], 2 * u.shape[-1], dtype=u.dtype, device=DEV) - wide[..., : u.shape[-1]] = u - view = wide[..., : u.shape[-1]] - return view.view(t.dtype) if t.element_size() == 1 else view - - -def test_noncontiguous_operand_answers_wrongly_and_is_refused(): - """The runner hands data_ptr() to the kernel and never reads a stride: the - same values through a strided view come back finite and wrong. The activation; - the guard checks the other five operands the same way.""" - name = "input" - tokens, tile = 64, 128 - args, _ = _problem(tokens, tile, seed=21) - ref = op.reference(**args) - bad = dict(args, **{name: _strided(args[name])}) - assert not bad[name].is_contiguous() - assert torch.equal(_u8(bad[name]), _u8(args[name])) - with pytest.raises(AssertionError): - _compare(op(**bad), ref) - with validating(op), pytest.raises(AssertionError, match=f"^{name} must be contiguous"): - op(**bad) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py similarity index 67% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py index 6d6cecf25444..34ce9252cc48 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py @@ -15,20 +15,20 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_block_scale_moe_renorm_permuted import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_block_scale_moe_renorm_permuted import ( # noqa: E501 finalize, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_block_scale_moe_renorm_permuted import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_block_scale_moe_renorm_permuted import ( # noqa: E501 fp4_block_scale_moe_renorm_permuted as op, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_block_scale_moe_renorm_routed import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_block_scale_moe_renorm_routed import ( # noqa: E501 fp4_block_scale_moe_renorm_routed as finalized, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.renorm_moe_routing_op import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.renorm_moe_routing_op import ( renorm_moe_routing_op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] import _fp4_moe_test_utils as H # noqa: E402 from _fp4_moe_test_utils import many_expert_layer # noqa: E402,F401 — fixture by name from _precision import true_fp32_matmul # noqa: E402 @@ -84,34 +84,3 @@ def test_certified_cells(cell, certifying, many_expert_layer): # noqa: F811 with true_fp32_matmul(): exp = op.reference(logits, x, k, ref["up"], ref["gate"], ref["down"], g2) H._compare(y, exp.to(torch.bfloat16)) - - -def test_uncertified_pair_refused(many_expert_layer): # noqa: F811 - spec = op.CELLS[0].spec - args, _, data, sf, _, logits, _, _, _, scalars = H.renorm_operands( - many_expert_layer(), spec, 13, renorm_moe_routing_op - ) - with validating(op), pytest.raises(AssertionError, match="certified Renormalize"): - op(*_args(args, data, sf, logits, scalars, spec, top_k=9)) - - -def test_cuda_graph_capture_and_replay(many_expert_layer): # noqa: F811 - spec = next(c.spec for c in op.CELLS if c.spec["tokens"] == 1) - args, _, data, sf, _, logits, _, _, _, scalars = H.renorm_operands( - many_expert_layer(), spec, 43, renorm_moe_routing_op - ) - call = list(_args(args, data, sf, logits, scalars, spec)) - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - for _ in range(3): - op(*call) - torch.cuda.current_stream().wait_stream(warmup) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - g_rows, g_w, g_idx = op(*call) - logits.copy_(torch.randn_like(logits.float()).to(logits.dtype)) - graph.replay() - torch.cuda.synchronize() - want = finalized(*call) - assert torch.equal(finalize(g_rows, g_w, g_idx), want) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py similarity index 72% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py index 249a87534aa9..78019fcf751d 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py @@ -6,8 +6,8 @@ replaces (`renorm_moe_routing_op` -> pre-routed `fp4_block_scale_moe_runner`), and it sits inside the pre-routed entry's band against this entry's own fp32 reference. The NVFP4 operand construction (layout, activations, scalars, -the band) uses ordinary shared test helpers. Beyond the cells: the refusal -of an uncertified (num_experts, top_k) pair, and per `DECODE_TACTICS` bucket +the band) uses ordinary shared test helpers. Beyond the cells: the prefill +tactic resolution bitwise the custom op's, and per `DECODE_TACTICS` bucket that the pinned tactic is a valid runner config and answers bitwise what the AutoTuner path's untuned (runner-default) tactic answers. """ @@ -17,23 +17,22 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_block_scale_moe_renorm_routed import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_block_scale_moe_renorm_routed import ( # noqa: E501 ACT_SWIGLU, - CERTIFIED_PAIRS, DECODE_TACTICS, DECODE_TACTICS_GEOMETRY, ROUTING_RENORMALIZE, _ckpt, _pinned_tactic, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_block_scale_moe_renorm_routed import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_block_scale_moe_renorm_routed import ( # noqa: E501 fp4_block_scale_moe_renorm_routed as op, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.renorm_moe_routing_op import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.renorm_moe_routing_op import ( # noqa: E501 renorm_moe_routing_op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] import _fp4_moe_test_utils as H # noqa: E402 from _fp4_moe_test_utils import many_expert_layer # noqa: E402,F401 — fixture by name from _precision import true_fp32_matmul # noqa: E402 @@ -61,7 +60,11 @@ def _routed(args, data, sf, logits, scalars, spec, top_k=None): ) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +#: The whole-layer cells (tp1); the windows have their own test file. +_WHOLE = [c for c in op.CELLS if "local_expert_offset" not in c.spec] + + +@pytest.mark.parametrize("cell", _WHOLE, ids=[c.why[:44] for c in _WHOLE]) def test_certified_cells(cell, certifying, many_expert_layer): # noqa: F811 certifying(op) spec = cell.spec @@ -95,14 +98,55 @@ def test_certified_cells(cell, certifying, many_expert_layer): # noqa: F811 H._compare(y, exp.to(torch.bfloat16)) -def test_uncertified_pair_refused(many_expert_layer): # noqa: F811 - spec = op.CELLS[0].spec - assert (spec["num_experts"], 9) not in CERTIFIED_PAIRS +def test_prefill_tactic_resolution_matches_the_custom_op_bitwise( + monkeypatch, + many_expert_layer, # noqa: F811 +): + """A non-pinned prefill bucket keeps the direct op's routing and tactic.""" + spec = _ckpt(2048) args, _, data, sf, _, logits, _, _, _, scalars = H.renorm_operands( - many_expert_layer(), spec, 13, renorm_moe_routing_op + many_expert_layer(), spec, 821, renorm_moe_routing_op ) - with validating(op), pytest.raises(AssertionError, match="certified Renormalize"): - _routed(args, data, sf, logits, scalars, spec, top_k=9) + module = __import__(op.__class__.__module__, fromlist=["run_renorm_moe"]) + calls = [] + catalog = module._pre_routed + + def recorded(*args, **kwargs): + calls.append(True) + return catalog(*args, **kwargs) + + monkeypatch.setattr(module, "_pre_routed", recorded) + got = _routed(args, data, sf, logits, scalars, spec) + expected = torch.ops.trtllm.fp4_block_scale_moe_runner( + logits, + None, + data, + sf.view(torch.float8_e4m3fn), + args["gemm1_weights"], + args["gemm1_weights_scale"], + None, + None, + None, + None, + args["gemm2_weights"], + args["gemm2_weights_scale"], + None, + *scalars, + spec["num_experts"], + spec["top_k"], + None, + None, + spec["intermediate"], + 0, + spec["num_experts"], + None, + ROUTING_RENORMALIZE, + True, + ACT_SWIGLU, + )[0] + assert calls == [True] + assert torch.equal(got, expected) + assert torch.equal(_routed(args, data, sf, logits, scalars, spec), expected) @pytest.mark.parametrize("bucket", sorted(DECODE_TACTICS)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed_windows.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed_windows.py new file mode 100644 index 000000000000..80a89853c909 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed_windows.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the fp4_block_scale_moe_renorm_routed entry's expert windows (the whole-layer cells +are in test_modeling_v2_fp4_block_scale_moe_renorm_routed.py). + +Two claims per cell: the logits-routed call over a rank's expert window is **bitwise** +the composition it replaces at the call site (`renorm_moe_routing_op` -> pre-routed +`fp4_block_scale_moe_runner` over the same window), and it sits inside the pre-routed +entry's band against this entry's own fp32 reference. The NVFP4 operand construction +(layout, the 512-expert stack, windows, activations, scalars, the band) is the pre-routed +entry's test's, loaded by path so the two tests cannot drift apart on what a valid operand +is; the tp4 zero-padded stack is built here from the same helpers. +""" + +import importlib.util +import zlib +from pathlib import Path + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_block_scale_moe_renorm_routed import ( # noqa: E501 + fp4_block_scale_moe_renorm_routed as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.renorm_moe_routing_op import ( + renorm_moe_routing_op, +) + +__extra_import_path__ = ["../../.."] +from _fp4_moe_test_utils import many_expert_layer # noqa: E402,F401 — fixture by name +from _precision import true_fp32_matmul # noqa: E402 +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "fp4_block_scale_moe_renorm_routed requires a CUDA device" + +DEV = "cuda" +G1 = 137.0 # activation global scale, as the pre-routed entry's test draws it + + +def _helpers(): + path = ( + Path(__file__).resolve().parents[3] + / "moe" + / "test_modeling_v2_fp4_block_scale_moe_runner.py" + ) + spec = importlib.util.spec_from_file_location("_fp4_moe_runner_r4_test_helpers", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +H = _helpers() + +# tp4 holds each rank's 160-wide slice of the 640 intermediate, stored zero-padded to 256. +_TP4_LOGICAL_INTER = 160 +_TP4_PADDED_INTER = 256 +_TP4 = None + + +def _pad_zero(codes: torch.Tensor, sf: torch.Tensor, dim: int, size: int): + """Extend a `(codes, scales)` stack with zero codes and zero (e4m3 +0.0) block scales + along `dim` (1 = rows, 2 = the K columns) to `size`.""" + pad_codes = list(codes.shape) + pad_codes[dim] = size - codes.shape[dim] + pad_sf = list(sf.shape) + pad_sf[dim] = pad_codes[dim] if dim == 1 else pad_codes[dim] // H.SV + codes = torch.cat([codes, torch.zeros(pad_codes, dtype=codes.dtype, device=DEV)], dim=dim) + sf_u8 = torch.cat( + [sf.view(torch.uint8), torch.zeros(pad_sf, dtype=torch.uint8, device=DEV)], dim=dim + ) + return codes.contiguous(), sf_u8.view(torch.float8_e4m3fn).contiguous() + + +def _pack(codes: torch.Tensor) -> torch.Tensor: + return (codes[..., 0::2] | (codes[..., 1::2] << 4)).contiguous() + + +def _tp4_padded(): + """The tp4 layer, built once: kernel-ready operands at I = 256 and reference operands + at I = 256 (the padded ones the kernel reads).""" + global _TP4 + if _TP4 is None: + e, h, i, ip = 512, 2560, _TP4_LOGICAL_INTER, _TP4_PADDED_INTER + gen = torch.Generator(device=DEV).manual_seed(512160) + _, up_s, up_c = H._rand_nvfp4(e, i, h, gen) + _, gt_s, gt_c = H._rand_nvfp4(e, i, h, gen) + _, dn_s, dn_c = H._rand_nvfp4(e, h, i, gen) + up_cp, up_sp = _pad_zero(up_c, up_s, 1, ip) + gt_cp, gt_sp = _pad_zero(gt_c, gt_s, 1, ip) + dn_cp, dn_sp = _pad_zero(dn_c, dn_s, 2, ip) + up_p, gt_p, dn_p = _pack(up_cp), _pack(gt_cp), _pack(dn_cp) + w1, s1 = H._prep_fc1(up_p, gt_p, up_sp, gt_sp) + w2, s2 = H._prep_fc2(dn_p, dn_sp) + args = dict( + gemm1_weights=w1, + gemm1_weights_scale=s1, + gemm2_weights=w2, + gemm2_weights_scale=s2, + intermediate_size=ip, + ) + ref = dict( + up=(up_cp, up_sp), + gate=(gt_cp, gt_sp), + down=(dn_cp, dn_sp), + raw=(up_p, gt_p, up_sp, gt_sp, dn_p, dn_sp), + hidden=h, + inter=ip, + num_experts=e, + gen=gen, + ) + _TP4 = (args, ref) + return _TP4 + + +def _stack(spec, layer): + """`(args, ref)` of the whole 512-expert stack at the cell's intermediate width.""" + if spec["intermediate"] == 256: + return _tp4_padded() + assert spec["intermediate"] == 640, spec["intermediate"] + return layer() + + +def _operands(spec, seed, layer): + gen = torch.Generator(device=DEV).manual_seed(seed) + args, ref = _stack(spec, layer) + off, n, e = spec["local_expert_offset"], spec["local_num_experts"], spec["num_experts"] + data, sf, x = H._rand_activation(spec["tokens"], spec["hidden"], G1, gen) + logits = torch.randn(spec["tokens"], e, generator=gen, device=DEV).to(spec["logits_dtype"]) + ids, wts = renorm_moe_routing_op(logits, spec["top_k"], torch.bfloat16) + g2 = H._calibrate_g2(x, ids, ref) + s1, sg, s2 = H._scalars(e, G1, g2) + win_args = args if n == e else H._window_args(args, off, n) + win_ref = ref if n == e else H._window_ref(ref, off, n) + scalars = H._window_scalars(s1, sg, s2, off, n) + return win_args, win_ref, data, sf, x, logits, ids, wts, g2, scalars + + +def _routed(win_args, data, sf, logits, scalars, spec, top_k=None, weights=None): + w = win_args if weights is None else weights + return op( + logits, + data, + sf.view(torch.float8_e4m3fn), + w["gemm1_weights"], + w["gemm1_weights_scale"], + w["gemm2_weights"], + w["gemm2_weights_scale"], + scalars["output1_scale_scalar"], + scalars["output1_scale_gate_scalar"], + scalars["output2_scale_scalar"], + spec["num_experts"], + spec["top_k"] if top_k is None else top_k, + spec["intermediate"], + spec["local_expert_offset"], + spec["local_num_experts"], + ) + + +#: The window cells; the whole-layer cells have their own test file. +_WINDOWS = [c for c in op.CELLS if "local_expert_offset" in c.spec] + + +@pytest.mark.parametrize("cell", _WINDOWS, ids=[c.why[:44] for c in _WINDOWS]) +def test_certified_cells(cell, certifying, many_expert_layer): # noqa: F811 + certifying(op) + spec = cell.spec + seed = zlib.crc32(cell.why.encode()) % 2**31 + win_args, win_ref, data, sf, x, logits, ids, wts, g2, scalars = _operands( + spec, seed, many_expert_layer + ) + off, n = spec["local_expert_offset"], spec["local_num_experts"] + + y = _routed(win_args, data, sf, logits, scalars, spec) + assert y.shape == (spec["tokens"], spec["hidden"]) and y.dtype == torch.bfloat16 + + # 1. Bitwise the composition it replaces, on identical operands. + y_pre = H._call( + data, + sf, + win_args, + spec["num_experts"], + spec["top_k"], + local_expert_offset=off, + local_num_experts=n, + topk_ids=ids, + topk_weights=wts, + **scalars, + )[0] + differing = (y.view(torch.int16) != y_pre.view(torch.int16)).any(dim=1).sum().item() + assert differing == 0, f"{differing} of {spec['tokens']} rows differ from the composition" + + # 2. Its own reference, through the pre-routed entry's gate. + with true_fp32_matmul(): + exp = op.reference( + logits, + x, + spec["top_k"], + win_ref["up"], + win_ref["gate"], + win_ref["down"], + g2, + local_expert_offset=off, + local_num_experts=n, + ) + H._compare(y, exp.to(torch.bfloat16)) + elt, rms = H._dev(y, exp.to(torch.bfloat16)) + print(f" cell {cell.why[:64]!r}: {elt:.2f} elt / {rms:.2f} rms ulp") diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py similarity index 72% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py index 4278c0822fcc..6c5e478ff738 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py @@ -4,10 +4,9 @@ The experts use the shared test helpers' drawn NVFP4 codes and power-of-two scales, so every weight is exact; this entry's layout is built -from their raw per-expert tensors by the target's own `weights.py` transforms, +from their raw per-expert tensors by the default target's `weights.py` transforms, so the test also pins those. Per cell: inside the routed entry's band against -the fp32 reference. Beyond the cells: the guard against a per-expert FC2 input -scale. +the fp32 reference. """ import zlib @@ -16,17 +15,17 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - weights as W, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp4_moe_renorm_cute_dsl import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp4_moe_renorm_cute_dsl import ( # noqa: E501 fp4_moe_renorm_cute_dsl as op, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.renorm_moe_routing_op import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.renorm_moe_routing_op import ( # noqa: E501 renorm_moe_routing_op, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501 + weights as W, +) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] import _fp4_moe_test_utils as H # noqa: E402 from _fp4_moe_test_utils import many_expert_layer # noqa: E402,F401 — fixture by name from _precision import true_fp32_matmul # noqa: E402 @@ -84,27 +83,3 @@ def test_certified_cells(cell, certifying, many_expert_layer, cute_dsl_experts): with true_fp32_matmul(): exp = op.reference(logits, x, spec["top_k"], ref["up"], ref["gate"], ref["down"], g2) H._compare(y, exp.to(torch.bfloat16)) - - -def test_guard_refuses_a_per_expert_fc2_scale(many_expert_layer, cute_dsl_experts): # noqa: F811 - spec = {**op.CELLS[0].spec, "tokens": 1024} - args, ref, data, sf, x, logits, _, _, g2, _ = H.renorm_operands( - many_expert_layer(), spec, 9, renorm_moe_routing_op - ) - fc1_w, fc1_s, fc2_w, fc2_s = cute_dsl_experts - s1, sg, s2 = H._scalars(spec["num_experts"], G1, g2) - with validating(op), pytest.raises(AssertionError): - op( - logits, - data, - sf, - fc1_w, - fc1_s, - sg, - fc2_w, - fc2_s, - s2, - s2, - spec["num_experts"], - spec["top_k"], - ) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp8_block_scale_moe_runner.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp8_block_scale_moe_runner.py similarity index 60% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp8_block_scale_moe_runner.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp8_block_scale_moe_runner.py index b2d77dcf4336..7fc469633394 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp8_block_scale_moe_runner.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fp8_block_scale_moe_runner.py @@ -1,11 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the fp8_block_scale_moe_runner catalog entry. +"""GPU test for the qwen4 default target's fp8_block_scale_moe_runner variant +(models/qwen4_exp/ops/moe/fp8_block_scale_moe_runner). The entry owns the contract: `CELLS`, `reference`, `is_valid`. This file turns a cell into operands and holds the claims no cell carries: a warm -autotuned tactic is bitwise the cold fallback the cells run, and the -contiguity guard. +autotuned tactic is bitwise the cold fallback the cells run (whole stack and +window), and the default target's four windows sum to the 512-expert layer. """ import zlib @@ -13,14 +14,14 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp8_block_scale_moe_runner import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp8_block_scale_moe_runner import ( # noqa: E501 BLOCK, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp8_block_scale_moe_runner import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp8_block_scale_moe_runner import ( # noqa: E501 fp8_block_scale_moe_runner as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "fp8_block_scale_moe_runner requires a CUDA device" @@ -41,7 +42,8 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: Measured worst over the cells: 6.65 element-wise (T=4096; the max over T*H samples grows with T -- 0.91 at T=1) and 0.306 RMS (H=I=512), the RMS sitting at 0.17-0.31 everywhere; earlier draws reached 10.03 and - 0.342. Wrong references at the checkpoint geometry, RMS: gate/up swap + 0.342. The default target's 128-expert window cells, and the four windows summed in + fp32, sit in the same range. Wrong references at the checkpoint geometry, RMS: gate/up swap 65, transposed activation scales 43, strided hidden_states 94, no FC1-output requant 3.5, neither requant 3.1, no activation requant 1.58 -- the closest. @@ -93,9 +95,18 @@ def _weights(hidden: int, inter: int, experts: int, seed: int): def _args( tokens, hidden, inter, num_experts, top_k, local_expert_offset=0, local_num_experts=None, seed=0 ): - """The op's arguments by name, pre-routed, every expert local.""" - assert local_expert_offset == 0 and local_num_experts in (None, num_experts) + """The op's arguments by name, pre-routed over all `num_experts`. + + The call answers for `[local_expert_offset, local_expert_offset + + local_num_experts)` and its weights are that window's slice of the one + stack -- what a rank of the default target holds. The routing, activations + and stack depend on `seed` and the geometry only, so the windows of one + seed split one layer. + """ + local = num_experts if local_num_experts is None else local_num_experts + window = slice(local_expert_offset, local_expert_offset + local) w1, w1s, w2, w2s = _weights(hidden, inter, num_experts, seed=hidden + inter + num_experts) + w1, w1s, w2, w2s = w1[window], w1s[window], w2[window], w2s[window] g = torch.Generator(device=DEV).manual_seed(seed) hs = (2.0 * torch.randn(tokens, hidden, device=DEV, generator=g)).to(torch.float8_e4m3fn) # [H/128, T]: the transposed layout fp8_quantize_1x128 returns here. @@ -118,8 +129,8 @@ def _args( n_group=None, topk_group=None, intermediate_size=inter, - local_expert_offset=0, - local_num_experts=num_experts, + local_expert_offset=local_expert_offset, + local_num_experts=local, routed_scaling_factor=None, routing_method_type=1, topk_weights=weights.to(torch.bfloat16).contiguous(), @@ -139,29 +150,98 @@ def _seed(cell) -> int: return zlib.crc32(cell.why.encode()) +def _window_args(cell) -> tuple[dict, torch.Tensor]: + """The cell's arguments and which slots route into its window. A window cell certifies + the window only if some slot routes into it, so a draw whose routing misses the window + entirely (likely only at a row or two) is redrawn with the next seed.""" + for attempt in range(8): + args = _args(**cell.spec, seed=_seed(cell) + attempt) + hit = _in_window(args) + if hit.any(): + return args, hit + raise AssertionError(f"no slot routes into the window: {cell.why!r} would certify nothing") + + # ── the certified cells ─────────────────────────────────────────────────── @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name certifying(op) - args = _args(**cell.spec, seed=_seed(cell)) + args, hit = _window_args(cell) out = op(**args) assert out.shape == args["hidden_states"].shape and out.dtype == torch.bfloat16 + # A row that routes nothing into the window comes back bitwise zero. + assert torch.equal(out[~hit.any(-1)], torch.zeros_like(out[~hit.any(-1)])) ref = op.reference(**args) elt, rms = _ulp(out, ref) - print(f" {cell.why[:44]}: {elt:.2f} ulp element-wise, {rms:.3f} ulp RMS") + print( + f" {cell.why[:44]}: {elt:.2f} ulp element-wise, {rms:.3f} ulp RMS; " + f"{hit.sum().item()} of {hit.numel()} slots in the window" + ) _compare(out, ref) +def _in_window(args) -> torch.Tensor: + """`[T, top_k]`: which slots the call's window answers for.""" + local = args["topk_ids"].long() - args["local_expert_offset"] + return (local >= 0) & (local < args["local_num_experts"]) + + +# ── the default target's windows ────────────────────────────────────────── + +#: The qwen4 default target's split: four 128-expert windows of the 512, one per rank. +WINDOWS = (0, 128, 256, 384) + + +@pytest.mark.parametrize("tokens", [1, 128, 1024]) +def test_windows_sum_to_the_whole_stack(tokens): + """The four windows' bf16 outputs, added in fp32 (what the target's + all-reduce or reduce-scatter does), sit inside the gate against the + 512-expert reference -- and the sum with any one window left out does not. + + Not bitwise to the 512-expert call: each window rounds its own partial to + bf16. The distance to it is printed. + """ + seed = 900 + tokens + whole = _args(tokens, **CKPT, seed=seed) + parts = [ + op(**_args(tokens, **CKPT, local_expert_offset=o, local_num_experts=128, seed=seed)) + for o in WINDOWS + ] + total = torch.stack([p.float() for p in parts]).sum(0).to(torch.bfloat16) + ref = op.reference(**whole) + elt, rms = _ulp(total, ref) + elt_call, rms_call = _ulp(total, op(**whole)) + print( + f" T={tokens}: windows summed vs reference {elt:.2f} / {rms:.3f} ulp; " + f"vs the 512-expert call {elt_call:.2f} / {rms_call:.3f} ulp" + ) + _compare(total, ref) + + # Negative control: drop each window that holds a slot. + ids = whole["topk_ids"].long() + for i, o in enumerate(WINDOWS): + if not ((ids >= o) & (ids < o + 128)).any(): + continue + short = torch.stack([p.float() for j, p in enumerate(parts) if j != i]).sum(0) + with pytest.raises(AssertionError): + _compare(short.to(torch.bfloat16), ref) + + # ── what a cell cannot carry ────────────────────────────────────────────── -def test_autotuned_tactic_matches_the_fallback_tactic(): - """A warm profiling cache changes the tactic, not the result.""" +@pytest.mark.parametrize( + "offset, local", [(0, 512), (256, 128)], ids=["whole stack", "a 128-expert window"] +) +def test_autotuned_tactic_matches_the_fallback_tactic(offset, local): + """A warm profiling cache changes the tactic, not the result -- for the + whole stack and for a window, which tunes under its own key + (`local_num_experts`) and its own config (`num_experts // local_num_experts`).""" from tensorrt_llm._torch.autotuner import AutoTuner, autotune - a = _args(256, **CKPT, seed=42) + a = _args(256, **CKPT, local_expert_offset=offset, local_num_experts=local, seed=42) AutoTuner.get().clear_cache() try: cold = op(**a) @@ -175,21 +255,3 @@ def test_autotuned_tactic_matches_the_fallback_tactic(): _compare(warm, op.reference(**a)) finally: AutoTuner.get().clear_cache() - - -# ── the guards ──────────────────────────────────────────────────────────── - - -def _strided(t: torch.Tensor) -> torch.Tensor: - wide = torch.zeros(*t.shape[:-1], 2 * t.shape[-1], device=DEV, dtype=t.dtype) - wide[..., ::2] = t - view = wide[..., ::2] - assert not view.is_contiguous() - return view - - -def test_guard_refuses_a_strided_hidden_states(): - """What `is_valid` exists for: the raw op reads a strided operand as dense.""" - a = _args(16, 256, 128, 8, 2, seed=11) - with validating(op), pytest.raises(AssertionError, match="must be contiguous"): - op(**dict(a, hidden_states=_strided(a["hidden_states"]))) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py similarity index 63% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py index 035ca9c348ce..9fb348eeb829 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py @@ -1,10 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the fused_sigmoid_gate_mul_add catalog entry. +"""GPU test for the default target's fused_sigmoid_gate_mul_add variant +(models/qwen4_exp/ops/moe/fused_sigmoid_gate_mul_add). -The entry owns the contract. This file turns a cell into operands, locks both -sides of the in-place write, and holds the claims no cell carries: the empty -batch and the loud rejection of a wrong-length gate. +The entry owns the contract. This file turns a cell into operands (a strided gate +is the call site's `gug` column), locks both sides of the in-place write, and holds +the claims no cell carries: the strided gate equal to the dense column, and the +empty batch. """ import zlib @@ -12,17 +14,21 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fused_sigmoid_gate_mul_add import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fused_sigmoid_gate_mul_add import ( # noqa: E501 fused_sigmoid_gate_mul_add as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "fused_sigmoid_gate_mul_add requires a CUDA device" DEV = "cuda" HIDDEN = 2560 +# The shared expert's fused GEMM output: gate|up (2 x 640) plus 8 gate-logit rows, the +# logit in the first of them. A strided gate is that column, as the call site reads it. +GUG_WIDTH = 2 * 640 + 8 +GATE_COL = 2 * 640 def _operands(spec, seed): @@ -37,10 +43,10 @@ def rand(*shape, dtype=dt): final, shared = rand(t, h), rand(t, h) gate = rand(t, 1, dtype=spec["gate_dtype"]) if spec["gate_layout"] == "strided": - wide = rand(t, 4, dtype=spec["gate_dtype"]) - wide[:, 2:3] = gate - gate = wide[:, 2:3] - assert not gate.is_contiguous() + gug = rand(t, GUG_WIDTH, dtype=spec["gate_dtype"]) + gug[:, GATE_COL : GATE_COL + 1] = gate + gate = gug[:, GATE_COL : GATE_COL + 1] + assert not gate.is_contiguous() or t == 1 else: assert spec["gate_layout"] == "column", spec["gate_layout"] return final, gate, shared @@ -67,17 +73,29 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th # ── what a cell cannot carry ────────────────────────────────────────────── +def _strided_spec(tokens: int) -> dict: + """The certified strided-gate cell at `tokens` rows.""" + (spec,) = [ + c.spec + for c in op.CELLS + if c.spec["tokens"] == tokens and c.spec["gate_layout"] == "strided" + ] + return dict(spec) + + +def test_strided_gate_is_the_dense_column_bit_for_bit() -> None: + """The call site's strided `gug` column and the same logits as a dense `[T, 1]`.""" + spec = _strided_spec(1024) + final, gate, shared = _operands(spec, seed=11) + dense_final = final.clone() + op(final, gate, shared) + op(dense_final, gate.contiguous(), shared) + assert torch.equal(final, dense_final) + + def test_zero_tokens_is_a_no_op(): final = torch.empty(0, HIDDEN, dtype=torch.bfloat16, device=DEV) shared = torch.empty(0, HIDDEN, dtype=torch.bfloat16, device=DEV) gate = torch.empty(0, 1, dtype=torch.bfloat16, device=DEV) returned = op(final, gate, shared) assert returned.shape == (0, HIDDEN) and returned.data_ptr() == final.data_ptr() - - -def test_a_gate_of_the_wrong_length_raises() -> None: - """Metadata violations are loud, which is why the entry has no guard.""" - bf = dict(dtype=torch.bfloat16, device=DEV) - final, shared = torch.randn(4, 128, **bf), torch.randn(4, 128, **bf) - with pytest.raises(AssertionError): - op(final, torch.randn(3, 1, **bf), shared) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_sort.py similarity index 90% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_sort.py index 8d97d4b0dd4f..f240c17b8240 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_sort.py @@ -3,8 +3,8 @@ """GPU test for the moe_sort entry of the Qwen3.8-Flash-Next target. The entry owns the contract. This file turns a cell into expert ids and holds -the claims no cell carries: CUDA-graph replay at the largest captured row -count, and the contiguity guard. +the claim no cell carries: CUDA-graph replay on new ids at the largest captured +row count. """ import zlib @@ -12,11 +12,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.moe_sort import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.moe_sort import ( # noqa: E501 moe_sort as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "moe_sort requires a CUDA device" @@ -201,15 +201,3 @@ def test_captured_rows_replay_on_new_ids(certifying): # noqa: F811 assert torch.equal(static_out[i], eager[i]), (trial, i) del graph, static_out torch.cuda.synchronize() - - -# ── the guard ───────────────────────────────────────────────────────────── - - -def test_guard_refuses_noncontiguous_ids(): - """What `is_valid` exists for: the op reads a strided view as the dense buffer - that starts at its data pointer, and answers without raising.""" - wide = _ids("router", 300, 128, seed=51, top_k=16) - sc = torch.rand(300, K, device=DEV) - with validating(op), pytest.raises(AssertionError, match="contiguous"): - op(wide[:, :K], sc, E, K, 0, E, 128) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_unpermute_inplace.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_unpermute_inplace.py similarity index 87% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_unpermute_inplace.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_unpermute_inplace.py index 1ff8171c9b9e..723fdf39ca5a 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_unpermute_inplace.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_moe_unpermute_inplace.py @@ -6,7 +6,7 @@ CuteDSL FC2 hands them over -- placed by moe_sort, every row the op must not read filled with NaN, the output pre-filled with NaN -- and holds the claims no cell carries: CUDA-graph replay at the largest captured row count, bitwise equal to -eager, and the contiguity guard. +eager. """ import zlib @@ -14,11 +14,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.moe_unpermute_inplace import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.moe_unpermute_inplace import ( # noqa: E501 moe_unpermute_inplace as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "moe_unpermute_inplace requires a CUDA device" @@ -240,29 +240,3 @@ def step(ids=ids, scales=scales, rows=rows, perm=perm, out=out, tile=tile): del graph, e2p torch.cuda.synchronize() assert stale_rows > 0, "no replay ran over a stale row: the stale-row claim is untested" - - -# ── the guards ──────────────────────────────────────────────────────────── - - -def _strided(t: torch.Tensor, pad: int) -> torch.Tensor: - """The same values in a view whose rows are `pad` elements wider apart.""" - wide = torch.zeros(t.shape[0], t.shape[1] + pad, dtype=t.dtype, device=t.device) - wide[:, : t.shape[1]] = t - return wide[:, : t.shape[1]] - - -def test_noncontiguous_operand_answers_wrongly_and_is_refused(): - """The FC2 rows the op reads; the guard checks all four operands the same way.""" - name = "permuted_input" - args = _problem(_spec(300), seed=61) - ref = op.reference(**args) - bad = dict(args, **{name: _strided(args[name], 64)}) - out = op(**bad) - with pytest.raises(AssertionError, match="differ bitwise"): - _compare(out.contiguous(), ref) - with validating(op): - with pytest.raises(AssertionError, match=f"{name} must be contiguous"): - op(**bad) - out = op(**dict(args, output=_output(300))) - _compare(out, ref) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_renorm_moe_routing_op.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_renorm_moe_routing_op.py similarity index 85% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_renorm_moe_routing_op.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_renorm_moe_routing_op.py index 5c7fe656dd28..c98408673a3e 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_renorm_moe_routing_op.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/moe/test_modeling_v2_renorm_moe_routing_op.py @@ -1,10 +1,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the renorm_moe_routing_op catalog entry. +"""GPU test for the qwen4 default target's renorm_moe_routing_op variant +(models/qwen4_exp/ops/moe/renorm_moe_routing_op). The entry owns the contract. This file turns a cell into logits and holds the -claims no cell carries: equivalence to HF's `norm_topk_prob` routing, and the -contiguity guard. +claim no cell carries: equivalence to HF's `norm_topk_prob` routing. """ import zlib @@ -12,11 +12,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.renorm_moe_routing_op import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.renorm_moe_routing_op import ( # noqa: E501 renorm_moe_routing_op as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "renorm_moe_routing_op requires a CUDA device" @@ -113,13 +113,3 @@ def test_equivalent_to_hf_norm_topk_prob(): atol=1e-7, ) assert (weights - hf_values).abs().max().item() > 1e-2 - - -# ── the guard ───────────────────────────────────────────────────────────── - - -def test_guard_refuses_a_column_slice(): - """What `is_valid` exists for: the raw op reads a slice as the dense buffer.""" - wide = _logits(8, 1024, torch.bfloat16, seed=12) - with validating(op), pytest.raises(AssertionError, match="contiguous"): - op(wide[:, :512], 10) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py similarity index 74% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py index 3d86b01b2f55..cb2e79386da1 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py @@ -7,11 +7,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.norm.flashinfer_gemma_rmsnorm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.norm.flashinfer_gemma_rmsnorm import ( # noqa: E501 flashinfer_gemma_rmsnorm as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "flashinfer_gemma_rmsnorm requires a CUDA device" @@ -47,6 +47,9 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res assert tuple(x.shape) == cell.spec["shape"] and x.stride(-1) == 1 if "row_width" in cell.spec and x.shape[0] > 1: assert x.stride(0) == cell.spec["row_width"] and not x.is_contiguous() + elif "row_width" in cell.spec: + # One stream row: torch calls it contiguous, so flashinfer takes the compact kernel. + assert x.is_contiguous() buffer_before = buffer.clone() out = op(x, weight, eps) @@ -55,16 +58,3 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res assert out.data_ptr() != x.data_ptr(), "the op must return a new tensor" assert torch.equal(buffer, buffer_before), "x (and the stream it slices) is read-only" torch.testing.assert_close(out, op.reference(x, weight, eps)) - - -def test_guard_refuses_fp8_before_the_kernel_answers() -> None: - """fp8 is in the DSL dtype table: accepted, and a tensor comes back.""" - x = torch.randn(4, 2560, device="cuda").to(torch.float8_e4m3fn) - w = torch.randn(2560, device="cuda").to(torch.float8_e4m3fn) - unguarded = torch.ops.trtllm.flashinfer_gemma_rmsnorm(x, w, 1e-6) - assert unguarded.dtype is x.dtype and unguarded.shape == x.shape, ( - "the hazard this guard covers is gone" - ) - with pytest.raises(AssertionError, match="dequantize"): - with validating(op): - op(x, w, 1e-6) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_layer_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_layer_norm.py similarity index 78% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_layer_norm.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_layer_norm.py index 0e07d7975518..b64080199d2b 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_layer_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_layer_norm.py @@ -1,10 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qwen4_exp layer_norm entry. +"""GPU test for the layer_norm catalog entry. Weight and bias are drawn over this checkpoint's own 110 vision LayerNorm tensors' ranges: the affine is wide enough (weights to 17.5, biases to -25.6) -that a wrong normalization is not masked by a near-unit scale. +that a wrong normalization is not masked by a near-unit scale. The default +bf16 gate against the entry's fp64 reference is the band. """ import zlib @@ -12,11 +13,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.norm.layer_norm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.norm.layer_norm import ( # noqa: E501 layer_norm as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 assert torch.cuda.is_available(), "layer_norm requires a CUDA device" @@ -60,12 +61,3 @@ def test_certified_cells(cell, certifying): # noqa: F811 assert all(torch.equal(a, b) for a, b in zip(before, after)), "inputs must not be mutated" assert torch.equal(out, op(x, weight, bias, eps)), "two launches must agree bitwise" torch.testing.assert_close(out, op.reference(x, weight, bias, eps)) - - -def test_rejects_a_weight_whose_length_is_not_the_normalized_width(): - """Loud: `_layer_norm_fwd` asserts the shapes, with no message.""" - gen = _gen(88) - x = torch.randn(8, WIDTH, generator=gen, device=DEVICE).to(torch.bfloat16) - weight, bias = _affine(gen) - with pytest.raises(AssertionError): - op(x, weight[:-1].contiguous(), bias, EPS) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py similarity index 58% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py index 21569aeef4e0..6db54e25f309 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/norm/test_modeling_v2_qwen4_exp_rms_norm_gated_token_major.py @@ -7,22 +7,27 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.norm.rms_norm_gated_token_major import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.norm.rms_norm_gated_token_major import ( # noqa: E501 rms_norm_gated_token_major as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "rms_norm_gated_token_major requires a CUDA device" -#: The target's concatenated GDN input projection: [q k v | z | a | b] = 10240 | 6144 | 48 | 48. -_IN_PROJ_WIDTH, _Z_OFFSET = 16480, 10240 +#: The width of the concatenated GDN input projection `[q k v | z | a | b]` a cell names no +#: `z_row_stride` for: all 48 heads, 10240 | 6144 | 48 | 48. A tp4/tep4 rank's is 4120: +#: 2560 | 1536 | 12 | 12. +_IN_PROJ_WIDTH = 16480 -#: The checkpoint's GDN epilogue geometry over a packed `z`, for the guard test. -_PACKED = dict(op.CELLS[0].spec, tokens=4, z="packed") +def _z_offset(row_stride: int, heads: int, n: int) -> int: + """First column of `z` in a `[q k v | z | a | b]` row: `a` and `b` are one column a head.""" + offset = row_stride - heads * (n + 2) + assert offset > 0, (row_stride, heads, n) + return offset def _seed(text: str) -> int: @@ -41,7 +46,9 @@ def randn(*shape): layout = spec.get("z", "packed") if layout == "in_proj": - z = randn(t, _IN_PROJ_WIDTH)[:, _Z_OFFSET : _Z_OFFSET + heads * n].view(t, heads, n) + width = spec.get("z_row_stride", _IN_PROJ_WIDTH) + offset = _z_offset(width, heads, n) + z = randn(t, width)[:, offset : offset + heads * n].view(t, heads, n) else: assert layout == "packed", layout z = randn(t, heads, n) @@ -55,7 +62,12 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res certifying(op) args = _build(cell.spec, seed=_seed(cell.why)) x, z, weight = args[:3] - assert z.stride() == (_IN_PROJ_WIDTH, cell.spec["width"], 1) and not z.is_contiguous() + assert z.stride()[1:] == (cell.spec["width"], 1) + if cell.spec["tokens"] > 1: + # A one-token view reports the dense stride for its size-1 token dim (and is then + # contiguous); the kernel never reads past token 0 there. + row_stride = cell.spec.get("z_row_stride", _IN_PROJ_WIDTH) + assert z.stride(0) == row_stride and not z.is_contiguous() before = [x.clone(), z.clone(), weight.clone()] out = op(*args) @@ -64,17 +76,3 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res for t, b in zip((x, z, weight), before): assert torch.equal(t, b), "inputs are read-only" torch.testing.assert_close(out, op.reference(*args)) - - -def test_guard_refuses_a_weight_of_the_wrong_length() -> None: - """Unguarded, a length-192 weight is silently truncated to weight[:128].""" - x, z, weight, eps, _, act = _build(_PACKED, seed=12) - long_weight = torch.cat([weight, torch.randn(64, dtype=weight.dtype, device="cuda")]) - assert torch.equal( - torch.ops.trtllm.rms_norm_gated_token_major(x, z, long_weight, eps, None, act), - torch.ops.trtllm.rms_norm_gated_token_major(x, z, weight, eps, None, act), - ), "the hazard this guard covers is gone" - - with pytest.raises(AssertionError, match="weight must be 1-D of length"): - with validating(op): - op(x, z, long_weight, eps, None, act) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ngram_token_history.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ngram_token_history.py similarity index 66% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ngram_token_history.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ngram_token_history.py index 0e16ce5bdedf..a5e97badb438 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ngram_token_history.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ngram_token_history.py @@ -6,8 +6,7 @@ against the entry's scalar reference, which builds each sequence's stream as a list and shares no structure with the masked per-shift select under test. Beyond the cells: chunked prefill against one-shot (the cross-step carry -the forward relies on), the `query_start_loc` length guard and CUDA-graph -replay. +the forward relies on). """ import zlib @@ -15,17 +14,16 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ngram_token_history import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ngram_token_history import ( # noqa: E501 CKPT_EOS_TOKEN_ID, CKPT_NGRAM_SIZE, CKPT_VOCAB_SIZE, - ENGINE_SLOTS, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ngram_token_history import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ngram_token_history import ( # noqa: E501 ngram_token_history as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "ngram_token_history requires a CUDA device" @@ -134,79 +132,3 @@ def call(pool, chunk, carry): split = torch.cat([call(pool, ids[:41], False), call(pool, ids[41:], True)]) assert torch.equal(split, one_shot) assert torch.equal(pool, one_shot_pool) - - -# ----------------------------------------------------------- guards - - -def _two_sequences(): - pool = torch.randint(0, CKPT_VOCAB_SIZE, (4, CONTEXT_LEN), device=DEV) - return ( - torch.tensor([11, 12, 13, 21], device=DEV), - pool, - torch.tensor([0, 1], dtype=torch.int32, device=DEV), - torch.tensor([True, False], device=DEV), - torch.tensor([0, 3, 4], dtype=torch.int32, device=DEV), - EOS, - NGRAM, - ) - - -def _misses(args, bad_index, bad_value): - """Run the op unguarded with one argument swapped; True if it returned - and its answer differs from the reference on the intended arguments.""" - expected = _expect(args) - bad = list(args) - bad[1] = args[1].clone() - bad[bad_index] = bad_value - window = op(*bad) - return not (torch.equal(window, expected[0]) and torch.equal(bad[1], expected[1])) - - -def test_guard_refuses_a_short_query_start_loc() -> None: - """Unguarded, a 2-entry vector for 2 sequences misattributes the tail.""" - args = _two_sequences() - short = torch.tensor([0, 3], dtype=torch.int32, device=DEV) - assert _misses(args, 4, short), "the hazard this guard covers is gone" - bad = list(args) - bad[4] = short - with pytest.raises(AssertionError, match=r"query_start_loc must be \[num_seqs"): - with validating(op): - op(*bad) - - -# ----------------------------------------------------------- capture - - -def test_cuda_graph_capture_and_replay() -> None: - """The captured decode shape, 8 sequences at the engine pool width.""" - spec = dict( - lengths=(1,) * 8, - slots=(0, ENGINE_SLOTS - 1, 4, 17, 30, 43, 56, 69), - carries=(True,) * 8, - pool_slots=ENGINE_SLOTS, - ids_dtype=torch.int32, - qsl_dtype=torch.int64, - ) - args = _build(spec, seed=11001) - start = args[1].clone() - - side = torch.cuda.Stream() - side.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(side): - for _ in range(3): - op(*args) - torch.cuda.current_stream().wait_stream(side) - - args[1].copy_(start) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured = op(*args) - - replay_ids = torch.randint(0, CKPT_VOCAB_SIZE, (8,), dtype=torch.int32, device=DEV) - args[0].copy_(replay_ids) - args[1].copy_(start) - expected = op.reference(replay_ids, start.clone(), *args[2:]) - graph.replay() - torch.cuda.synchronize() - _compare((captured, args[1]), expected) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_gate_conv_inject_norm.py similarity index 82% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_gate_conv_inject_norm.py index de9c02457d60..4581ad5e52db 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_gate_conv_inject_norm.py @@ -1,11 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the ple_gate_conv_inject_norm entry. +"""GPU test for the default target's ple_gate_conv_inject_norm variant. Cells are gated against the float64 reference that rounds where the op rounds: -the stream (`hc_out`) and the published conv state within one ulp of each -element, the normed stream within two (see `_compare`). Beyond the cells: -the operand-width guard and graph capture. +the stream (`hc_out`) and the normed stream in units of their row's scale, the +published conv state per element (see `_compare`); a verify cell's staged conv +input is checked bitwise against the state it published. Beyond the cells: the +staged conv input is what the conv consumed, and `out` writes only its rows of a +larger buffer. """ import zlib @@ -13,17 +15,17 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ple_gate_conv_inject_norm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ple_gate_conv_inject_norm import ( # noqa: E501 CKPT_CLAMP_MIN, CKPT_DILATION, CKPT_HC_COUNT, CKPT_HIDDEN, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ple_gate_conv_inject_norm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ple_gate_conv_inject_norm import ( # noqa: E501 ple_gate_conv_inject_norm as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import ( # noqa: E402,F401 — path above; fixture by name assert_within_ulp, certifying, @@ -71,6 +73,9 @@ def _stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: #: reference over three seeds at 16384 rows, the composition this replaces #: reaches 2.91 row-max ulp (normed) and the op 2.94; at 8192 rows 1.89 / 1.89. #: Cells above 1024 rows are therefore gated at 4; the RMS gate is unchanged. +#: The default target's per-rank pools (1874 / 850 slots) sit inside the same gates; a +#: reference reading each sequence's neighbouring slot falls far outside the RMS gate +#: (measured). CONTEXT_ROWS, CONTEXT_ROW_ULP = 1024, 4.0 @@ -124,24 +129,38 @@ def rand(*shape, scale=1.0): ] # fmt: skip -def _run(args): +def _run(args, conv_input=None): """`(hc_out, normed, conv_states)` from the op, on a copy of the pool.""" args = list(args) args[7] = args[7].clone() - hc_out, normed = op(*args) + hc_out, normed = op(*args, conv_input=conv_input) return hc_out, normed, args[7] +def _assert_staged_conv_input(xc, conv_states, slots, lengths) -> None: + """With `L <= state_len` rows in every sequence, the published state's last `L` + columns are those rows' conv inputs: the staged buffer must hold exactly them.""" + rows_per_seq = lengths[0] + assert set(lengths) == {rows_per_seq} and rows_per_seq <= conv_states.shape[-1] + tail = conv_states[slots.long()][:, :, -rows_per_seq:] + assert torch.equal(torch.reshape(torch.transpose(tail, 1, 2), xc.shape), xc), ( + "the staged conv input is not what the conv consumed" + ) + + @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name certifying(op) args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) hc_before = args[0].clone() ref = op.reference(*args) - out = _run(args) + xc = torch.empty_like(args[0]) if cell.spec.get("conv_input") else None + out = _run(args, conv_input=xc) assert torch.equal(args[0], hc_before), "hc is read-only" assert out[0].is_contiguous() and out[1].is_contiguous() _compare(out, ref) + if xc is not None: + _assert_staged_conv_input(xc, out[2], args[8], cell.spec["lengths"]) for name, got, want in zip(("hc_out", "normed", "conv_states"), out, ref): worst, frac = _stats(got, want) print( @@ -178,35 +197,3 @@ def test_out_writes_the_rows_of_larger_buffers() -> None: assert torch.equal(bufs[0][:rows], plain[0]) and torch.equal(bufs[1][:rows], plain[1]) assert torch.equal(args[7], plain[2]) assert bool((bufs[0][rows:] == 7.0).all()) and bool((bufs[1][rows:] == 7.0).all()) - - -def test_guard_refuses_a_short_value() -> None: - args = _build(dict(lengths=(1, 1)), seed=4) - args[2] = args[2][:, :-1] - with validating(op), pytest.raises(AssertionError): - op(*args) - - -def test_cuda_graph_capture_and_replay() -> None: - args = _build(dict(lengths=(1,) * 8, pool_slots=461), seed=8) - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - pool = args[7].clone() - with torch.cuda.stream(warmup): - for _ in range(2): - op(*args) - torch.cuda.current_stream().wait_stream(warmup) - args[7].copy_(pool) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_hc, graph_normed = op(*args) - fresh = _build(dict(lengths=(1,) * 8, pool_slots=461), seed=18) - args[0].copy_(fresh[0]) - args[1].copy_(fresh[1]) - args[7].copy_(pool) - eager = _run(args) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(graph_hc, eager[0]) - assert torch.equal(graph_normed, eager[1]) - assert torch.equal(args[7], eager[2]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_ngram_embed.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_ngram_embed.py similarity index 69% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_ngram_embed.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_ngram_embed.py index c46ea883dc07..aad5d62cc216 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_ngram_embed.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ple/test_modeling_v2_ple_ngram_embed.py @@ -1,43 +1,44 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the ple_ngram_embed entry. +"""GPU test for the ple_ngram_embed catalog entry. Integer bookkeeping plus a one-rounding dequantize, so gated bitwise -- the embedding and the whole published pool -- against the entry's reference, which is the scalar references of `ngram_token_history` and `ngram_hash` and a float64 gather-and-scale. Beyond the cells: the `query_start_loc` length guard -and graph capture. +and graph capture at the note's replay cells. """ import zlib +from types import SimpleNamespace import pytest import torch +from cuda.bindings import runtime as cudart -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.embedding.ngram_hash import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.embedding.ngram_hash import ( # noqa: E501 CKPT_HEAD_OFFSETS, CKPT_HEAD_VOCAB_SIZES, CKPT_MULTIPLIERS, ngram_hash, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ngram_token_history import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ngram_token_history import ( # noqa: E501 ngram_token_history, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ple_ngram_embed import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ple_ngram_embed import ( # noqa: E501 CKPT_EOS_TOKEN_ID, CKPT_HEAD_DIM, CKPT_SCALE, CKPT_TABLE_ROWS, ENGINE_SLOTS, + TEP4_SLOTS, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ple_ngram_embed import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ple_ngram_embed import ( # noqa: E501 ple_ngram_embed as op, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.weights import ( # noqa: E501 - _host_resident_empty, -) +from tensorrt_llm._utils import CUASSERT -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "ple_ngram_embed requires a CUDA device" @@ -50,10 +51,44 @@ _HOSTS = [] +def _host_mapped(src: torch.Tensor) -> torch.Tensor: + """A device view of `src`'s bytes in page-locked host memory. + + The stock engine's placement of the n-gram table under its PLE host offload + (`Qwen4ExpPinnedHostEmbedding` in `tensorrt_llm/_torch/modules/qwen4_exp/ple.py`): + a pinned CPU allocation reached through `cudaHostGetDevicePointer`, so the + kernel gathers rows over the host link and no byte of the table is in device + memory -- which the pointer's own attributes are asked to confirm. + """ + rows, cols = src.shape + host = torch.empty(rows * cols * src.element_size(), dtype=torch.uint8, pin_memory=True) + host.copy_(src.contiguous().view(torch.uint8).reshape(-1)) + _HOSTS.append(host) + (pointer,) = CUASSERT(cudart.cudaHostGetDevicePointer(host.data_ptr(), 0)) + # No public torch API views a pointer the caching allocator did not hand out; + # `|u1` because the array interface has no float8 typestr. + interface = SimpleNamespace( + __cuda_array_interface__={ + "shape": (rows, cols * src.element_size()), + "typestr": "|u1", + "data": (int(pointer), False), + "version": 3, + "strides": None, + } + ) + table = torch.as_tensor(interface, device=DEV).view(src.dtype) + assert table.data_ptr() == int(pointer), "torch.as_tensor copied the host buffer" + (attributes,) = CUASSERT(cudart.cudaPointerGetAttributes(table.data_ptr())) + assert attributes.type == cudart.cudaMemoryType.cudaMemoryTypeHost, attributes.type + assert table.shape == src.shape and table.dtype == src.dtype + return table + + def _scattered(n: int, pool_slots: int, g) -> list: - """`n` distinct slots over the whole pool, both ends included.""" + """`n` distinct slots over the whole pool, both ends included; a lone + sequence takes the last slot, the farthest row of the pool.""" middle = (torch.randperm(pool_slots - 2, generator=g) + 1).tolist() - slots = ([0, pool_slots - 1] + middle)[:n] + slots = ([pool_slots - 1, 0] + middle)[:n] assert len(set(slots)) == n and max(slots) < pool_slots return slots @@ -83,10 +118,7 @@ def _table(kind: str, rows: int, touched: torch.Tensor, g) -> torch.Tensor: return table src = torch.randn(rows, CKPT_HEAD_DIM, generator=g).to(FP8) if kind == "host": - table, host = _host_resident_empty(rows, CKPT_HEAD_DIM, FP8, torch.device(DEV)) - _HOSTS.append(host) - table.copy_(src.to(DEV)) - return table + return _host_mapped(src) return src.to(DEV) @@ -159,8 +191,15 @@ def test_guard_refuses_short_query_start_loc() -> None: op(*args) -def test_cuda_graph_capture_and_replay() -> None: - args = _build(dict(lengths=(1,) * 16, pool_slots=ENGINE_SLOTS, eos_density=0.1), seed=11) +@pytest.mark.parametrize( + "lengths, pool_slots", + [((1,) * 16, ENGINE_SLOTS), ((1,) * 128, TEP4_SLOTS), ((8,) * 128, TEP4_SLOTS)], + ids=["16x1-461", "128x1-1874", "128x8-1874"], +) +def test_cuda_graph_capture_and_replay(lengths, pool_slots) -> None: + """The note's replay cells: a replay reads the ids and pool written since capture.""" + spec = dict(lengths=lengths, pool_slots=pool_slots, eos_density=0.1) + args = _build(spec, seed=11) warmup = torch.cuda.Stream() warmup.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(warmup): @@ -169,7 +208,7 @@ def test_cuda_graph_capture_and_replay() -> None: graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): graph_out = op(*args) - fresh = _build(dict(lengths=(1,) * 16, pool_slots=ENGINE_SLOTS, eos_density=0.1), seed=12) + fresh = _build(spec, seed=12) args[0].copy_(fresh[0]) pool_before = args[1].clone() want, want_pool = op.reference(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_bf16_to_fp32.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_bf16_to_fp32.py similarity index 63% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_bf16_to_fp32.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_bf16_to_fp32.py index 3bf21aa13b26..37cca3ab7333 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_bf16_to_fp32.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_bf16_to_fp32.py @@ -1,18 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the bf16_to_fp32 entry: the cells bit for bit `x.float()`, graph capture, the -contiguity guard.""" +"""GPU test for the bf16_to_fp32 entry: the cells bit for bit `x.float()`.""" import zlib import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.quantization.bf16_to_fp32 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.quantization.bf16_to_fp32 import ( # noqa: E501 bf16_to_fp32 as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "bf16_to_fp32 requires a CUDA device" @@ -43,25 +42,3 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by ref = op.reference(x) assert out.dtype == torch.float32 and out.is_contiguous() assert _bitwise(out, ref), "widening is exact; a bit difference is a defect" - - -def test_cuda_graph_capture_and_replay() -> None: - x = _build(dict(shape=(128, 4099)), 1) - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - op(x) - torch.cuda.current_stream().wait_stream(warmup) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - out = op(x) - x.copy_(_build(dict(shape=(128, 4099)), 2)) - graph.replay() - torch.cuda.synchronize() - assert _bitwise(out, x.float()) - - -def test_guard_refuses_a_strided_input() -> None: - x = _build(dict(shape=(8, 64)), 3) - with validating(op), pytest.raises(AssertionError): - op(x[:, ::2]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_fp8_quantize_1x128.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_fp8_quantize_1x128.py similarity index 91% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_fp8_quantize_1x128.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_fp8_quantize_1x128.py index caa087518151..719270a41cd4 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_fp8_quantize_1x128.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/quantization/test_modeling_v2_fp8_quantize_1x128.py @@ -13,15 +13,15 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.quantization.fp8_quantize_1x128 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.quantization.fp8_quantize_1x128 import ( # noqa: E501 SCALE_ULP_GATE, block_scales, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.quantization.fp8_quantize_1x128 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.quantization.fp8_quantize_1x128 import ( # noqa: E501 fp8_quantize_1x128 as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "fp8_quantize_1x128 requires a CUDA device" @@ -122,15 +122,20 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th _compare((data, scales), ref) rel = _assert_misses_are_ties(x, data, ref[0]) if not spec["use_ue8m0"] else 0.0 exact = (data.view(torch.uint8) == ref[0].view(torch.uint8)).float().mean().item() - print(f" {cell.why[:44]}: {100 * exact:.4f}% bit-exact, worst tie {rel:.3e}") + sms = torch.cuda.get_device_properties(x.device).multi_processor_count + print( + f" {cell.why[:44]}: {100 * exact:.4f}% bit-exact, worst tie {rel:.3e}, " + f"{spec['tokens'] * nb} groups over a {sms * 8 * 8 * 2}-group grid sweep" + ) # ── what a cell cannot carry ────────────────────────────────────────────── -def test_cuda_graph_capture_and_replay(): +@pytest.mark.parametrize("rows", [64, 128, 1024]) +def test_cuda_graph_capture_and_replay(rows): """A replay reproduces a fresh eager call bitwise; so does a side stream.""" - x = _input(64, 2560, seed=17, magnitude=1.0) + x = _input(rows, 2560, seed=17, magnitude=1.0) side = torch.cuda.Stream() side.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(side): @@ -140,7 +145,7 @@ def test_cuda_graph_capture_and_replay(): graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): g_data, g_sf = op(x) - x.copy_(_input(64, 2560, seed=18, magnitude=1.0)) + x.copy_(_input(rows, 2560, seed=18, magnitude=1.0)) e_data, e_sf = op(x) graph.replay() torch.cuda.synchronize() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/__init__.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/__init__.py new file mode 100644 index 000000000000..113c2e6ba7f9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/__init__.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Package marker. + +`test_routing.py` repeats the glm_moe_dsa family's basename. Without a package, pytest's +default import mode loads both under one module name and refuses the second. +""" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qsa_fp8_pages_to_bf16__prefill.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qsa_fp8_pages_to_bf16__prefill.py new file mode 100644 index 000000000000..74f52efb90b5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qsa_fp8_pages_to_bf16__prefill.py @@ -0,0 +1,153 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for this target's qsa_fp8_pages_to_bf16 entry (`__prefill`). + +The paged state is the engine's kv_fp8 trunk: e4m3 `[pages, 24, Hkv, tpb, D]`, layer 11's K +and V two of the role planes (a straddle cell sized so the requests' pages lie on both sides +of 2^31 elements of them), and the `[slots, 8196]` table at the tune's max_seq_len, the +listed requests on seeded distinct slots, their pages a seeded permutation, `-1` past each +sequence. Scratch and table are pre-filled with sentinels, so a missing or a stray write +shows. Every cell is bitwise `reference` (a widening copy), and the scratch is what the +e4m3 sparse GQA kernel reads: e4m3 codes at unit scale. +""" + +import math +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.attention.qsa_fp8_pages_to_bf16__prefill import ( # noqa: E501 + qsa_fp8_pages_to_bf16 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "qsa_fp8_pages_to_bf16 requires a CUDA device" + +DEV = "cuda" +SCRATCH_FILL = 0x7FC1 # a bf16 NaN +TABLE_FILL = -7 +_EDGE = 2**31 + + +def _pool_pages(spec, needed: int) -> int: + if not spec["straddle"]: + return needed + 5 + role_page = spec["kv_heads"] * spec["tokens_per_block"] * spec["head_dim"] + edge_page = math.ceil( + (_EDGE - 2 * spec["layer"] * role_page) / (spec["slot_pages"] * role_page) + ) + return max(edge_page + needed // 2, needed) + 5 + + +def _build(spec, seed): + gen = torch.Generator(device=DEV).manual_seed(seed) + tpb, kv, d = spec["tokens_per_block"], spec["kv_heads"], spec["head_dim"] + counts = [math.ceil(n / tpb) for n in spec["kv_lens"]] + needed = sum(counts) + 7 + pool_pages = _pool_pages(spec, needed) + role_pages, layer = spec["slot_pages"], spec["layer"] + # Random e4m3 codes (NaN codes 0x7F / 0xFF excluded), in the layer's two planes only. + data = torch.zeros((pool_pages, role_pages, kv, tpb, d), dtype=torch.uint8, device=DEV) + first = pool_pages - needed + for p0 in range(first, pool_pages, 1024): + p1 = min(p0 + 1024, pool_pages) + codes = torch.randint(0, 254, (p1 - p0, 2, kv, tpb, d), generator=gen, device=DEV) + codes = torch.where(codes >= 0x7F, codes + 1, codes) + data[p0:p1, 2 * layer : 2 * layer + 2] = codes.to(torch.uint8) + pool = data.view(torch.float8_e4m3fn) + width, slots_total = spec["width"], spec["slots"] + table = torch.full((slots_total, width), -1, dtype=torch.int32, device=DEV) + slots = torch.randperm(slots_total, generator=torch.Generator().manual_seed(seed)) + slots = slots[: len(counts)] + page_ids = first + torch.randperm(needed, generator=gen, device=DEV).to(torch.int32) + cursor = 0 + for slot, count in zip(slots.tolist(), counts): + table[slot, :count] = page_ids[cursor : cursor + count] + cursor += count + if spec["holes"]: + head = slots[0].item() + table[head, 1] = -1 # unmapped inside the sequence + table[head, 2] = pool_pages + 3 # past the pool + counts_t = torch.tensor(counts, dtype=torch.int64, device=DEV) + listed = int(counts_t.sum()) + out_shape = (listed, kv, tpb, d) + return dict( + k_cache=pool[:, 2 * layer], + v_cache=pool[:, 2 * layer + 1], + block_table=table, + slots=slots.to(DEV, torch.int32), + counts=counts_t, + offsets=torch.cumsum(counts_t, 0) - counts_t, + table=torch.full_like(table, TABLE_FILL), + k_out=torch.full(out_shape, SCRATCH_FILL, dtype=torch.int16, device=DEV).view( + torch.bfloat16 + ), + v_out=torch.full(out_shape, SCRATCH_FILL, dtype=torch.int16, device=DEV).view( + torch.bfloat16 + ), + ) + + +def _bitwise(a: torch.Tensor, b: torch.Tensor) -> bool: + if a.dtype == torch.bfloat16: + a, b = a.view(torch.int16), b.view(torch.int16) + return a.shape == b.shape and torch.equal(a, b) + + +def _assert_straddles(args) -> None: + table = args["block_table"] + pages = table[(table >= 0) & (table < args["k_cache"].shape[0])].long() + for name in ("k_cache", "v_cache"): + plane = args[name] + ends = plane.storage_offset() + pages * plane.stride(0) + plane[0].numel() - 1 + assert bool((ends >= _EDGE).any()) and bool((ends < _EDGE).any()), name + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 -- fixture by name + certifying(op) + args = _build(cell.spec, zlib.crc32(cell.why.encode())) + if cell.spec["straddle"]: + _assert_straddles(args) + ref = op.reference(**{k: v.clone() if torch.is_tensor(v) else v for k, v in args.items()}) + out = op(**args) + for name, got, want in zip(("k_out", "v_out", "table"), out, ref): + assert _bitwise(got, want), f"{name} differs from the reference" + # Every listed page was written, and only the listed requests' table entries. + assert not bool(out[0].isnan().any()) and not bool(out[1].isnan().any()) + listed_rows = torch.zeros_like(args["block_table"], dtype=torch.bool) + for slot, count in zip(args["slots"].tolist(), args["counts"].tolist()): + listed_rows[slot, :count] = True + assert bool((out[2][~listed_rows] == TABLE_FILL).all()) + + +def test_scratch_is_the_e4m3_values_at_unit_scale() -> None: + """Every listed page's scratch holds the pool's e4m3 values exactly (e4m3 widens to bf16 + without rounding), and the table addresses it.""" + spec = op.CELLS[3].spec + args = _build(spec, 11) + k_out, v_out, table = op(**args) + for slot, count, base in zip( + args["slots"].tolist(), args["counts"].tolist(), args["offsets"].tolist() + ): + pages = args["block_table"][slot, :count].long() + for got, plane in ((k_out, args["k_cache"]), (v_out, args["v_cache"])): + want = plane[pages].float() + assert torch.equal(got[base : base + count].float(), want) + assert torch.equal( + table[slot, :count], torch.arange(base, base + count, device=DEV, dtype=torch.int32) + ) + + +def test_listing_past_the_scratch_is_dropped() -> None: + """A request whose pages run past `listed` writes only those below it.""" + args = _build(op.CELLS[3].spec, 5) + short = args["k_out"].shape[0] - 3 + args["k_out"], args["v_out"] = args["k_out"][:short], args["v_out"][:short] + ref = op.reference(**{k: v.clone() if torch.is_tensor(v) else v for k, v in args.items()}) + out = op(**args) + for got, want in zip(out, ref): + assert _bitwise(got, want) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qsa_sparse_gqa_rowtile__prefill.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qsa_sparse_gqa_rowtile__prefill.py new file mode 100644 index 000000000000..62429f6004ac --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qsa_sparse_gqa_rowtile__prefill.py @@ -0,0 +1,275 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for this target's qsa_sparse_gqa_rowtile variant (`__prefill`). + +The paged state is the target's: a trunk cell's K and V are layer 11's planes of a `[pages, +24, kv, tpb, D]` HND pool (24 role pages a slot), sized so the requests' pages straddle 2^31 +elements, every other page and plane NaN; a scratch cell's are a contiguous `[pages, kv, tpb, +D]` bf16 pool, as `qsa_fp8_pages_to_bf16` widens a kv_fp8 pool into. The slot table is the +engine's at the tune's max_seq_len (8196 columns, a row per slot, a seeded permutation, -1 +past each sequence). Selection rows are laid out the way QSA's top-k + expansion emit them +(complete groups expanded to tokens, the trailing partial group, then -1, packed at the +front). The band and the reference are `qsa_paged_sparse_gqa`'s: this op replaces it on +context launches. +""" + +import math + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.attention import ( # noqa: E501 + qsa_sparse_gqa_rowtile__prefill as _module, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.attention.qsa_sparse_gqa_rowtile__prefill import ( # noqa: E501 + qsa_sparse_gqa_rowtile as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.attention.qsa_sparse_gqa_rowtile__prefill import ( # noqa: E501 + rows_per_group, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "qsa_sparse_gqa_rowtile requires a CUDA device" + +DEV = "cuda" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """`qsa_paged_sparse_gqa`'s band: 4 ulp of each (row, head)'s magnitude, 2 ulp relative + RMS, in q's dtype, against the fp32 reference rounded once. The row-tile pass rounds P to + bf16 for the PV MMA and sums it in fp32, as the per-row kernel rounds its probabilities to + the cache dtype before P.V.""" + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + if out.numel() == 0: + return + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) + + +class _LaunchRecorder: + """Records which of the entry's kernels a call indexed with a grid; the bitmap pass + (`_row_bits`, a plain call) is recorded with its one-CTA-a-row grid.""" + + def __init__(self, inner, name, log): + self._inner, self._name, self._log = inner, name, log + + def __getitem__(self, grid): + self._log.append((self._name, tuple(grid))) + return self._inner[grid] + + def __call__(self, selected_tokens, *args, **kwargs): + self._log.append((self._name, (selected_tokens.shape[0],))) + return self._inner(selected_tokens, *args, **kwargs) + + +_KERNELS = ( + "_row_bits", + "_rowtile_plan_kernel", + "_rowtile_lone_kernel", + "_rowtile_attention_cute", +) + + +@pytest.fixture +def launches(monkeypatch): + log: list = [] + for name in _KERNELS: + inner = getattr(_module, name) + monkeypatch.setattr(_module, name, _LaunchRecorder(inner, name, log)) + return log + + +def _selected_tokens(positions, seed, compress_ratio, token_topk, shared_scores=False): + """QSA's expansion layout: a random `min(available, block_topk)`-subset of the complete + visible groups, expanded, then the trailing partial group, then `-1`, packed at the + front. With `shared_scores` every row ranks the groups by one score vector, so adjacent + rows pick nearly the same set (a sparse union, as on real selections).""" + block_topk = token_topk // compress_ratio + final_topk = token_topk + compress_ratio - 1 + gen = torch.Generator(device=DEV).manual_seed(seed) + rows = positions.numel() + visible = positions.long() + 1 + available = visible // compress_ratio + taken = available.clamp(max=block_topk) + width = max(block_topk, int(available.max().item()) if rows else 0) + out = torch.empty(rows, final_topk, dtype=torch.int32, device=DEV) + column = torch.arange(final_topk, device=DEV)[None, :] + for start in range(0, rows, 1024): + stop = min(start + 1024, rows) + avail, take, vis = available[start:stop], taken[start:stop], visible[start:stop] + if shared_scores: + scores = torch.rand(1, width, generator=gen, device=DEV).expand(stop - start, -1) + gen.manual_seed(seed) + else: + scores = torch.rand(stop - start, width, generator=gen, device=DEV) + group_ids = torch.arange(width, device=DEV)[None, :] + scores = torch.where(group_ids < avail[:, None], scores, torch.inf) + chosen = scores.argsort(dim=1)[:, :block_topk] + group_of_column = (column // compress_ratio).clamp(max=block_topk - 1) + body = torch.gather(chosen, 1, group_of_column.expand(stop - start, -1)) + body = body * compress_ratio + column % compress_ratio + in_body = column < (take * compress_ratio)[:, None] + tail_column = column - (take * compress_ratio)[:, None] + tail = (avail * compress_ratio)[:, None] + tail_column + in_tail = ~in_body & (tail_column >= 0) & (tail < vis[:, None]) + out[start:stop] = torch.where(in_body, body, torch.where(in_tail, tail, -1)).to(torch.int32) + return out + + +def _nan_pool(shape, dtype): + assert dtype == torch.bfloat16 + return torch.full(shape, 0x7FC0, dtype=torch.int16, device=DEV).view(dtype) + + +def _straddling(spec, needed: int) -> int: + """The pool page count that puts the requests -- the pool's last `needed` pages -- on + both sides of `offset_edge` elements of layer `layer`'s K plane.""" + role_page = spec["kv_heads"] * spec["tokens_per_block"] * spec["head_dim"] + edge, slot_pages, layer = spec["offset_edge"], spec["slot_pages"], spec["layer"] + edge_page = math.ceil((edge - 2 * layer * role_page) / (slot_pages * role_page)) + return max(edge_page + needed // 2, needed) + + +def _build(spec, seed): + """Real paged state for one spec, keyed by `__call__`'s parameter names.""" + gen = torch.Generator(device=DEV).manual_seed(seed) + tpb, kv, d, q_heads = ( + spec["tokens_per_block"], + spec["kv_heads"], + spec["head_dim"], + spec["q_heads"], + ) + segments = list(spec["segments"]) + positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments]) + req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)]) + rows, slots = positions.numel(), spec["slots"] + assert len(segments) <= slots + slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(segments)] + mapped = [math.ceil((s + n) / tpb) for s, n in segments] + assert max(mapped) <= spec["pages"] + needed = sum(mapped) + 7 + trunk = spec["cache_layout"] == "trunk" + pool_pages = _straddling(spec, needed) if trunk else needed + first = pool_pages - needed + page_ids = first + torch.randperm(needed, generator=gen, device=DEV).to(torch.int32) + block_table = torch.full((slots, spec["pages"]), -1, dtype=torch.int32, device=DEV) + cursor = 0 + for slot, count in zip(slot_of.tolist(), mapped): + block_table[slot, :count] = page_ids[cursor : cursor + count] + cursor += count + + dtype = spec["cache_dtype"] + if trunk: + # Layer `layer`'s K and V planes of the trunk; every other page and plane is NaN. + layer = spec["layer"] + pool = _nan_pool((pool_pages, spec["slot_pages"], kv, tpb, d), dtype) + for p0 in range(first, pool_pages, 1024): + p1 = min(p0 + 1024, pool_pages) + pool[p0:p1, 2 * layer : 2 * layer + 2] = torch.randn( + (p1 - p0, 2, kv, tpb, d), generator=gen, device=DEV + ).to(dtype) + k_cache, v_cache = pool[:, 2 * layer], pool[:, 2 * layer + 1] + else: + # The widened scratch: two contiguous bf16 pools. + assert spec["cache_layout"] == "scratch" + k_cache = torch.randn((pool_pages, kv, tpb, d), generator=gen, device=DEV).to(dtype) + v_cache = torch.randn((pool_pages, kv, tpb, d), generator=gen, device=DEV).to(dtype) + assert spec["q_layout"] == "qkv_slice" + qkv = torch.randn(rows, q_heads + 2 * kv, d, generator=gen, device=DEV).to(spec["q_dtype"]) + q = qkv[:, :q_heads] + + selected = _selected_tokens(positions, seed + 3, spec["compress_ratio"], spec["token_topk"], + spec.get("shared_scores", False)) # fmt: skip + request_indices = slot_of.to(DEV)[req].to(spec["request_dtype"]) + return dict( + q=q, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + selected_tokens=selected, + request_indices=request_indices, + tokens_per_block=tpb, + softmax_scale=d**-0.5, + query_positions=positions.to(spec["request_dtype"]), + compress_ratio=spec["compress_ratio"], + ) + + +def _bound(positions, spec): + cr = spec["compress_ratio"] + visible = positions.long() + 1 + return (visible // cr).clamp(max=spec["token_topk"] // cr) * cr + visible % cr + + +def _is_prefix_packed(table): + valid = table >= 0 + columns = torch.arange(table.shape[1], device=table.device)[None, :] + return (valid == (columns < valid.sum(1)[:, None])).all(1) + + +def _assert_straddles(args, edge: int) -> None: + """The mapped pages put the K and V planes' element offsets on both sides of `edge`.""" + table = args["block_table"] + pages = table[table >= 0].long() + for name in ("k_cache", "v_cache"): + plane = args[name] + ends = plane.storage_offset() + pages * plane.stride(0) + plane[0].numel() - 1 + assert bool((ends >= edge).any()), f"{name}: no mapped page lies past {edge}" + assert bool((ends < edge).any()), f"{name}: every mapped page lies past {edge}" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811 + certifying(op) + spec = cell.spec + args = _build(spec, seed=sum(map(ord, cell.why))) + if spec["cache_layout"] == "trunk": + _assert_straddles(args, spec["offset_edge"]) + rows = args["q"].shape[0] + # The per-row kernel's precondition, which this op shares. + selected = args["selected_tokens"] + assert bool(_is_prefix_packed(selected).all()) + assert torch.equal((selected >= 0).sum(1), _bound(args["query_positions"], spec)) + out = op(**args) + assert [name for name, _ in launches] == list(_KERNELS) + group_rows = rows_per_group(spec["q_heads"] // spec["kv_heads"]) + assert group_rows == {12: 10, 6: 20}[spec["q_heads"] // spec["kv_heads"]] + groups = math.ceil(rows / group_rows) + assert launches[0][1] == (rows,) and launches[1][1] == (groups,) + assert launches[2][1] == (min(rows, 256), spec["kv_heads"]) + # The row-tile pass is persistent: one CTA per SM, never more than (group, KV head) items. + items = groups * spec["kv_heads"] + assert launches[3][1] == ( + min(items, torch.cuda.get_device_properties(0).multi_processor_count), + ) + assert out.shape == (rows, spec["q_heads"], spec["head_dim"]) + assert out.dtype == args["q"].dtype and out.is_contiguous() + assert not bool(out.isnan().any()) + _compare(out, op.reference(**args)) + + +TP4 = op.CELLS[0].spec +assert (TP4["q_heads"], TP4["kv_heads"], TP4["head_dim"]) == (6, 1, 256) + + +def _with(segments, **overrides): + return {**TP4, "segments": tuple(segments), **overrides} + + +@pytest.mark.parametrize( + "segments", + [ + # Fewer items (2 groups x 1 KV head) than SMs: a 2-CTA grid, one item a CTA. + [(5000, 37)], + # A lone group (rows of two requests) between row-tile groups of each request. + [(3000, 50), (7000, 50)], + ], + ids=["fewer items than SMs", "a lone group between row-tile groups"], +) +def test_small_launches_match_the_reference(segments) -> None: + """The persistent pass's edges at 20-row groups: CTAs left without a second item, and + an item the pass skips (a lone group) between items it serves.""" + args = _build(_with(segments, slots=8), seed=7700 + len(segments)) + _compare(op(**args), op.reference(**args)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue__decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue__decode.py new file mode 100644 index 000000000000..db073f810319 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue__decode.py @@ -0,0 +1,261 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the default target's qwen4_exp_attn_decode_prologue entry: every cell against +the entry's fp32 `reference` on every allocation either writes -- the packed QKV, the index +query, the K/V pool, the index plane, the position cache -- and a stale-read check under +graph replay with PDL. + +The rotated values (the packed q and k, the index query, the K pool's k and the compressed +keys) are judged in units of their row's scale: the kernel and the reference evaluate the +same fp32 formulas and differ only in the rms's reduction order, which moves a value across +a bf16 rounding boundary at most once -- one bf16 spacing of an element no larger than the +row's largest, at most 2^-7 of the row scale (2 ulp in `assert_within_ulp`'s 2^-8 units). +Everything else -- v in the packed buffer and the V pool, the raw index keys, the +coordinates and every byte neither side writes -- is compared bitwise, and the K pool's +rows are bitwise the kernel's own packed k converted to the pool's dtype. + +The pools are laid out as the engine serves them: K/V the two role planes of a +`[pages, 2, Hkv, tpb, D]` pool, requests on random pages.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.attention.qwen4_exp_attn_decode_prologue__decode import ( # noqa: E501 + qwen4_exp_attn_decode_prologue as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "qwen4_exp_attn_decode_prologue requires a CUDA device" + +DEV = "cuda" +D, IH, ID, TPB, CR, ROT, WIDTH, SLOTS = 256, 4, 128, 32, 4, 64, 8196, 8 +#: Per-rank geometry: (q heads, KV heads); the fused row is +#: `[q | k | v | gate | index q 4 x 128, k 128 | v slot 128]`. +GEOMS = {"tp4": (6, 1), "dep4": (24, 2)} +#: One bf16 spacing of an element no larger than its row's largest, in 2^-8 units of the +#: row scale; the relative RMS bound is measured, see the module docstring. +ELEMENT_ULP, RMS_ULP = 2.0, 0.05 + + +def _geometry(geom): + hq, hk = GEOMS[geom] + qkv = (hq + 2 * hk) * D + iq_col = qkv + hq * D + return hq, hk, qkv, iq_col, iq_col + (IH + 2) * ID + + +def _edge_rows(proj, hq, hk) -> None: + """Zero, -0.0 and outlier blocks in the k head (past the rotary prefix, where the norm + keeps zeros zero) and in the v head, over the first rows.""" + k0, v0 = hq * D, (hq + hk) * D + edges = ( + (0, slice(k0 + 128, k0 + 144), 0.0), + (1, slice(k0 + 160, k0 + 176), -0.0), + (2, slice(k0 + 200, k0 + 201), 900.0), + (0, slice(v0 + 16, v0 + 32), 0.0), + (1, slice(v0 + 48, v0 + 64), torch.tensor([-0.0] * 8 + [0.0] * 8)), + (3, slice(v0 + 40, v0 + 41), -900.0), + (2, slice(v0 + 64, v0 + 80), 1e-30), + ) + for r, cols, value in edges: + if r < proj.shape[0]: + proj[r, cols] = value + + +def _rows(spec): + """(request, position) per row: one row a request, or `rows_per_request` consecutive + positions of each request from its first.""" + per = spec.get("rows_per_request", 1) + return [(r, p + i) for r, p in enumerate(spec["pos"]) for i in range(per)] + + +def _build(spec, seed): + g = torch.Generator(device=DEV).manual_seed(seed) + hq, hk, qkv_width, iq_col, row_width = _geometry(spec["geom"]) + layout = _rows(spec) + pos = [p for _, p in layout] + rows = len(pos) + last = {} + for req, p in layout: + last[req] = max(last.get(req, 0), p) + slots = max(SLOTS, len(last) + 3) + pages = sum(p // TPB + 1 for p in last.values()) + 64 + perm = torch.randperm(pages, generator=g, device=DEV).to(torch.int32) + block_table = torch.full((slots, WIDTH), -1, dtype=torch.int32, device=DEV) + used = 0 + for req, p in last.items(): + n = p // TPB + 1 + block_table[req, :n] = perm[used : used + n] + used += n + request = torch.tensor([req for req, _ in layout], dtype=torch.int32, device=DEV) + if spec.get("unmapped") is not None: + request[spec["unmapped"]] = slots + 5 + logical = torch.tensor(pos, dtype=torch.int64, device=DEV) + ids = torch.stack([logical, logical + 3, logical + 7]).to(torch.int32).contiguous() + buf = torch.empty(3, 1, rows + 5, dtype=torch.int32, device=DEV) + buf[:, 0, :rows] = ids + coords = torch.transpose(buf[:, 0, :rows], 0, 1) + table_rows = max(pos) + 8 + angles = torch.rand(table_rows, ROT // 2, generator=g, device=DEV) * 6.28 + cos_sin = torch.cat([torch.cos(angles), torch.sin(angles)], dim=1).contiguous() + proj = (torch.randn(rows, row_width, generator=g, device=DEV) * 2.0).to(torch.bfloat16) + if spec.get("edge"): + _edge_rows(proj, hq, hk) + out = torch.randn(rows, qkv_width + 64, generator=g, device=DEV).to(torch.bfloat16) + index_scratch = torch.randn(rows, IH + 2, ID, generator=g, device=DEV).to(torch.bfloat16) + weights = [ + (0.3 * torch.randn(n, generator=g, device=DEV)).to(torch.bfloat16) for n in (D, D, ID, ID) + ] + kv_dtype = torch.float8_e4m3fn if spec.get("fp8") else torch.bfloat16 + pool = torch.randn(pages, 2, hk, TPB, D, generator=g, device=DEV).to(kv_dtype) + index_cache = torch.randn(pages, TPB, 1, ID, generator=g, device=DEV).to(torch.bfloat16) + position_cache = torch.randint( + 0, table_rows, (pages, TPB, 3), generator=g, device=DEV, dtype=torch.int32 + ) + return dict( + proj=proj, out=out[:, :qkv_width], index_q=index_scratch[:, :IH], position_ids=ids, + coords=coords, cos_sin=cos_sin, q_weight=weights[0], k_weight=weights[1], + index_q_weight=weights[2], index_k_weight=weights[3], k_cache=pool[:, 0], + v_cache=pool[:, 1], index_cache=index_cache, position_cache=position_cache, + request_indices=request, logical_positions=logical, block_table=block_table, + tokens_per_block=TPB, compress_ratio=CR, num_heads_q=hq, num_heads_k=hk, head_dim=D, + index_heads=IH, index_head_dim=ID, index_q_col=iq_col, rotary_dim=ROT, eps=1e-6, + mrope_section1=11, mrope_section2=10, rows_per_request=spec.get("rows_per_request", 1), + ) # fmt: skip + + +def _bytes(storage, device) -> torch.Tensor: + return torch.empty(0, dtype=torch.uint8, device=device).set_(storage) + + +def _clone(args): + """The arguments over fresh copies of every allocation they view, each view rebuilt at + its offset and strides; returns them and the copies, keyed by the first argument + viewing each allocation.""" + out, copies, keys = dict(args), {}, {} + for name, t in args.items(): + if not isinstance(t, torch.Tensor): + continue + key = keys.setdefault(t.untyped_storage().data_ptr(), name) + if key not in copies: + copies[key] = _bytes(t.untyped_storage(), t.device).clone() + view = torch.empty(0, dtype=t.dtype, device=t.device) + view.set_(copies[key].untyped_storage(), t.storage_offset(), t.shape, t.stride()) + out[name] = view + return out, copies + + +def _written_slots(args): + """(row, page, slot) of every row whose page lookup holds, from the host's view.""" + table = args["block_table"].cpu() + pages = args["k_cache"].shape[0] + found = [] + for t, (req, p) in enumerate( + zip(args["request_indices"].tolist(), args["logical_positions"].tolist()) + ): + column = p // TPB + if 0 <= req < table.shape[0] and column < table.shape[1]: + page = int(table[req, column]) + if 0 <= page < pages: + found.append((t, page, p % TPB)) + return found + + +def _check(got, want, where="") -> dict: + """Judge the rotated values within the band, then neutralize them in `got` and require + every byte of every allocation bitwise. Returns the measured worst errors.""" + g, w = got[0], want[0] + hq, hk = g["num_heads_q"], g["num_heads_k"] + width = (hq + hk) * D + measured = {} + + def banded(name, a, b): + a, b = a.float().reshape(-1, a.shape[-1]), b.float().reshape(-1, b.shape[-1]) + scale = b.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + measured[name] = ((a - b).abs() / scale).max().item() / 2**-8 + assert_within_ulp(a.to(torch.bfloat16), b.to(torch.bfloat16), ELEMENT_ULP, RMS_ULP) + + banded("packed q/k", g["out"][:, :width], w["out"][:, :width]) + banded("index query", g["index_q"], w["index_q"]) + slots = _written_slots(g) + k_got = g["out"][:, hq * D : width].reshape(-1, hk, D) + for t, page, slot in slots: + own = k_got[t].to(g["k_cache"].dtype) + assert torch.equal(g["k_cache"][page, :, slot].view(torch.uint8), own.view(torch.uint8)), ( + f"{where}row {t}: the K pool's row is not the packed k converted" + ) + g["k_cache"][page, :, slot] = w["k_cache"][page, :, slot] + ends = [(t, page, slot) for t, page, slot in slots if (slot + 1) % CR == 0] + if ends: + a = torch.stack([g["index_cache"][page, slot, 0] for _, page, slot in ends]) + b = torch.stack([w["index_cache"][page, slot, 0] for _, page, slot in ends]) + banded("compressed keys", a, b) + for _, page, slot in ends: + g["index_cache"][page, slot, 0] = w["index_cache"][page, slot, 0] + g["out"][:, :width] = w["out"][:, :width] + g["index_q"].copy_(w["index_q"]) + assert got[1].keys() == want[1].keys() + for key, copy in got[1].items(): + differ = (copy != want[1][key]).nonzero() + assert differ.numel() == 0, f"{where}{key}: {differ.shape[0]} bytes differ" + return measured + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + want = _clone(args) + op.reference(**want[0]) + got = _clone(args) + op(**got[0]) + torch.cuda.synchronize() + measured = _check(got, want) + print(f"{cell.why[:48]}: worst element / row scale in 2^-8 units {measured}") + + +@pytest.mark.parametrize( + "spec", + [ + dict(geom="tp4", pos=(8699, 8700)), + dict(geom="tp4", pos=(84000,), rows_per_request=4), + dict(geom="dep4", pos=(8697, 8700), rows_per_request=4, fp8=True), + ], + ids=["tp4_decode", "tp4_verify", "dep4_fp8_verify"], +) +def test_reads_after_the_wait_under_graph_replay(spec) -> None: + """The projection is rewritten inside the graph right before the op; with PDL a read of + it before the wait would see the previous replay's row (in a verify step, also the raw + index keys of a group's earlier rows).""" + args = _build(spec, seed=23) + src = args["proj"].clone() + work, work_copies = _clone(args) + + def step(): + work["proj"].copy_(src * 1.0) + op(**work) + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + step() + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + step() + # The replays' projections from a seeded generator: the global RNG's state depends on the + # tests run before this one in the process. + g = torch.Generator(device=DEV).manual_seed(29) + for trial in range(3): + src.copy_((torch.randn(src.shape, generator=g, device=DEV) * 2.0).to(src.dtype)) + before = _clone(work) + graph.replay() + torch.cuda.synchronize() + want = _clone(before[0]) + want[0]["proj"].copy_(src) + op.reference(**want[0]) + _check((work, work_copies), want, where=f"replay {trial}: ") diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_vision_rope_2d_packed_qk.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_vision_rope_2d_packed_qk.py new file mode 100644 index 000000000000..f47dc0ec1c37 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/attention/test_modeling_v2_vision_rope_2d_packed_qk.py @@ -0,0 +1,196 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the vision_rope_2d_packed_qk variant: the vision tower's 2-D rotary, +one Triton launch in place on the q|k prefix of the packed qkv projection. + +Each cell builds a packed `[T, 3 * 16, 72]` bf16 qkv projection and the caller's angle +table from the cell's grid (spatial-merge positions, fp64 angles rounded once to fp32, +fp32 cos / sin), rotates `qkv[:, :32]` in place, and gates: + 1. bitwise against the entry's `reference` on the same table; + 2. bitwise against the whole-op reference of `vision_rope_2d__r52` (copied below: + positions by nested loops, fp64 angles, `polar`, complex multiply), q and k + separately, with that entry's distribution gates beside it; + 3. the v heads of every row untouched. +""" + +import zlib +from typing import List, Sequence, Tuple + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.attention.vision_rope_2d_packed_qk import ( # noqa: E501 + HEAD_DIM, + NUM_HEADS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.attention.vision_rope_2d_packed_qk import ( # noqa: E501 + vision_rope_2d_packed_qk as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "vision_rope_2d_packed_qk requires a CUDA device" + +MERGE = 2 +THETA = 10000.0 +HALF = HEAD_DIM // 2 + + +def _num_tokens(grid) -> int: + return sum(t * h * w for t, h, w in grid) + + +def _spatial_merge_positions(grid: Sequence[Sequence[int]], merge: int) -> List[Tuple[int, int]]: + """Per-token `(row, col)` by four nested loops, in spatial-merge order.""" + positions: List[Tuple[int, int]] = [] + for frames, height, width in grid: + frame = [ + (block_row * merge + r, block_col * merge + c) + for block_row in range(height // merge) + for block_col in range(width // merge) + for r in range(merge) + for c in range(merge) + ] + positions.extend(frame * frames) + return positions + + +def _angle_table(grid_thw: torch.Tensor, merge: int, theta: float, head_dim: int): + """The caller's glue: per-token fp32 `cos` / `sin` `[T, head_dim // 2]`, height angles + then width, from fp64 angles rounded once to fp32.""" + rope_dim = head_dim // 2 + inv_freq = 1.0 / (theta ** (torch.arange(0, rope_dim, 2, dtype=torch.float64) / rope_dim)) + pos = torch.tensor(_spatial_merge_positions(grid_thw.tolist(), merge), dtype=torch.float64) + angles = (pos[:, :, None] * inv_freq).flatten(1).to(device="cuda", dtype=torch.float32) + return angles.cos().contiguous(), angles.sin().contiguous() + + +def _r52_reference(q, k, grid_thw, spatial_merge_size, theta=THETA): + """`vision_rope_2d__r52.reference`, verbatim: the whole op from the grid, as a fresh + `(q_out, k_out)`.""" + rope_dim = q.shape[-1] // 2 + axis_dim = rope_dim // 2 + positions = _spatial_merge_positions(grid_thw.tolist(), spatial_merge_size) + inverse = torch.tensor( + [theta ** (-(2.0 * j) / rope_dim) for j in range(axis_dim)], dtype=torch.float64 + ) + pos = torch.tensor(positions, dtype=torch.float64).reshape(-1, 2) + angles = (pos[:, :, None] * inverse).flatten(1).to(device=q.device, dtype=torch.float32) + rotor = torch.polar(torch.ones_like(angles), angles).unsqueeze(-2) + + def rotate(x: torch.Tensor) -> torch.Tensor: + fp32 = x.float() + rotated = torch.complex(fp32[..., :rope_dim], fp32[..., rope_dim:]) * rotor + return torch.cat((rotated.real, rotated.imag), dim=-1).to(x.dtype) + + return rotate(q), rotate(k) + + +def _packed_qkv(tokens: int, dtype: torch.dtype, seed: int) -> torch.Tensor: + g = torch.Generator(device="cuda").manual_seed(seed) + return torch.randn(tokens, 3 * NUM_HEADS, HEAD_DIM, generator=g, device="cuda").to(dtype) + + +def _qk(qkv: torch.Tensor) -> torch.Tensor: + qk = qkv[:, : 2 * NUM_HEADS] + assert qk.stride() == (3 * NUM_HEADS * HEAD_DIM, HEAD_DIM, 1) and not qk.is_contiguous() + return qk + + +def _distribution(got: torch.Tensor, want: torch.Tensor) -> Tuple[float, float, float]: + """(cosine, relative L2, p99.9 absolute error) of `got` against `want`, flattened.""" + g, w = got.float().reshape(-1), want.float().reshape(-1) + cos = torch.nn.functional.cosine_similarity(g, w, dim=0).item() + rel_l2 = ((g - w).norm() / w.norm().clamp_min(1e-30)).item() + err = (g - w).abs() + # torch.quantile caps its input size; a strided sample of 2^24 is plenty for p99.9. + if err.numel() > 2**24: + err = err[:: -(-err.numel() // 2**24)] + return cos, rel_l2, torch.quantile(err, 0.999).item() + + +#: `vision_rope_2d__r52`'s distribution gates: both sides round the same fp32 rotation to bf16 +#: once, so they differ by at most one bf16 ulp per element, at most 2^-7 of its value: the +#: relative L2 is bounded by that, and 1 - cosine by half its square. +_REL_L2_BOUND = 2.0**-7 +_COS_BOUND = 1.0 - 2.0**-15 + + +def _mismatch(got: torch.Tensor, want: torch.Tensor) -> Tuple[int, float]: + """(elements differing in their bits, largest absolute difference).""" + differing = (got.reshape(-1).view(torch.int16) != want.reshape(-1).view(torch.int16)).sum() + return int(differing.item()), (got.float() - want.float()).abs().max().item() + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + spec = cell.spec + grid_thw = torch.tensor([list(row) for row in spec["grid"]], dtype=spec["grid_dtype"]) + tokens = _num_tokens(spec["grid"]) + qkv = _packed_qkv(tokens, spec["dtype"], seed=zlib.crc32(cell.why.encode()) % 2**31) + before = qkv.clone() + cos, sin = _angle_table(grid_thw, MERGE, THETA, HEAD_DIM) + qk = _qk(qkv) + + want = op.reference(qk, cos, sin) + assert torch.equal(qkv, before), "reference must not mutate its input" + out = op(qk, cos, sin) + torch.cuda.synchronize() + + assert out is qk, "the op returns the view it rotated" + assert torch.equal(qkv[:, 2 * NUM_HEADS :], before[:, 2 * NUM_HEADS :]), "v was written" + assert not torch.equal(qk, before[:, : 2 * NUM_HEADS]), "q|k was not rotated" + + # 1. The entry's own reference on the same table: bitwise. + differing, max_dev = _mismatch(qk, want) + print(f"[{cell.why}] tokens={tokens} vs reference: {differing} differing, max {max_dev:.3e}") + assert differing == 0, f"{differing} elements differ from reference (max {max_dev:.3e})" + + # 2. vision_rope_2d__r52's whole-op reference from the grid: bitwise, plus its + # distribution gates. + q_r52, k_r52 = _r52_reference( + before[:, :NUM_HEADS], before[:, NUM_HEADS : 2 * NUM_HEADS], grid_thw, MERGE, THETA + ) + for name, got, ref in ( + ("q", qkv[:, :NUM_HEADS], q_r52), + ("k", qkv[:, NUM_HEADS : 2 * NUM_HEADS], k_r52), + ): + differing, max_dev = _mismatch(got, ref) + cos_sim, rel_l2, p999 = _distribution(got, ref) + print( + f"[{cell.why}] {name} vs r52: {differing} differing, max {max_dev:.3e}, " + f"cos={cos_sim:.9f} rel_l2={rel_l2:.3e} p99.9={p999:.3e}" + ) + assert cos_sim >= _COS_BOUND and rel_l2 <= _REL_L2_BOUND, (name, cos_sim, rel_l2) + assert differing == 0, f"{name}: {differing} elements differ from r52 (max {max_dev:.3e})" + + +def _small(): + grid_thw = torch.tensor([(1, 8, 12)], dtype=torch.int64) + qkv = _packed_qkv(96, torch.bfloat16, seed=17) + cos, sin = _angle_table(grid_thw, MERGE, THETA, HEAD_DIM) + return qkv, cos, sin + + +#: The silent faults the op's note lists, each a table the kernel accepts: (cos, sin) -> table. +_FAULTS = { + "bf16 table": lambda t: t.to(torch.bfloat16), + "short table": lambda t: t[:-8], + "narrow table": lambda t: t[:, : HALF - 4].contiguous(), + "non-contiguous table": lambda t: torch.cat((t, t), dim=-1)[:, :HALF], +} + + +@pytest.mark.parametrize("fault", list(_FAULTS)) +def test_is_valid_refuses_each_silent_table_fault(fault): + """Unguarded, the kernel returns a wrong rotation for each fault; `is_valid` refuses it.""" + qkv, cos, sin = _small() + want = op.reference(_qk(qkv), cos, sin) + bad_cos, bad_sin = _FAULTS[fault](cos), _FAULTS[fault](sin) + got = op(_qk(qkv.clone()), bad_cos, bad_sin) + torch.cuda.synchronize() + assert not torch.equal(got, want), f"{fault}: the kernel got it right after all" + with validating(op), pytest.raises(AssertionError): + op(_qk(qkv.clone()), bad_cos, bad_sin) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first__decode.py similarity index 70% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first__decode.py index bf2d62d45c13..4550b39ed729 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_decode_tc_first__decode.py @@ -1,14 +1,15 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qwen4_exp_hc_decode_tc_first entry: cells in the band against -the float64 reference, the row-count guard, graph capture.""" +"""GPU test for this target's qwen4_exp_hc_decode_tc_first variant +(ops/hc/qwen4_exp_hc_decode_tc_first__decode): cells in the band against the float64 +reference at the decode rows the members run.""" import zlib import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_decode_tc_first import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.hc.qwen4_exp_hc_decode_tc_first__decode import ( # noqa: E501 qwen4_exp_hc_decode_tc_first as op, ) @@ -66,28 +67,3 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by low, inject = op(*args) assert low.shape == (cell.spec["rows"], RANK) _compare((low, inject), op.reference(*args)) - - -def test_guard_refuses_more_than_two_tiles_of_rows() -> None: - normed, first, rank, c, wi = _build(dict(rows=257), seed=4) - with validating(op), pytest.raises(AssertionError): - op(normed, first, rank, c, wi) - - -@pytest.mark.parametrize("rows", [8, 256]) -def test_cuda_graph_capture_and_replay(rows) -> None: - normed, first, rank, c, wi = _build(dict(rows=rows), seed=5) - side = torch.cuda.Stream() - side.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(side): - op(normed, first, rank, c, wi) - torch.cuda.current_stream().wait_stream(side) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - g_low, g_inject = op(normed, first, rank, c, wi) - fresh = _build(dict(rows=rows), seed=6)[0] - normed.copy_(fresh) - graph.replay() - torch.cuda.synchronize() - low, inject = op(fresh, first, rank, c, wi) - assert torch.equal(g_low, low) and torch.equal(g_inject, inject) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4__moe.py similarity index 69% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4__moe.py index eecc9d836bef..1d9bcf987b1d 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4__moe.py @@ -1,9 +1,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qwen4_exp_hc_up_gate_mean_fp4 entry: `mixed` in the band -`qwen4_exp_hc_up_gate_mean`'s test uses, the NVFP4 outputs against the -native-torch quantization of the returned `mixed` outside the reciprocal tie -window, the global-scale guard, graph capture.""" +"""GPU test for the qwen4 default target's qwen4_exp_hc_up_gate_mean_fp4__moe +entry: `mixed` in the band `qwen4_exp_hc_up_gate_mean`'s test uses and bitwise +that entry's own output, the NVFP4 outputs against the native-torch +quantization of the returned `mixed` outside the reciprocal tie window, the +global-scale guard, graph capture.""" import zlib @@ -14,10 +15,10 @@ Fp4Quantize, fp4_quantize, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_gated_residual import ( # noqa: E501 - gate_ratio, +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_up_gate_mean import ( + qwen4_exp_hc_up_gate_mean, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_up_gate_mean_fp4 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.hc.qwen4_exp_hc_up_gate_mean_fp4__moe import ( # noqa: E501 qwen4_exp_hc_up_gate_mean_fp4 as op, ) @@ -29,16 +30,14 @@ DEV = "cuda" -def nvfp4_reference(x, gs): - """`fp4_quantize`'s native-torch NVFP4 of `x` `[T, K]` with linear scales: - `(packed [T, K/2] uint8, sf [T*K/16] uint8)`.""" - return fp4_quantize.reference(x, gs, 16, False, False) - - -def nvfp4_near_tie(x, gs): - """`[T, K]` bool: elements within `fp4_quantize`'s tie window of an e2m1 midpoint, - where the kernels' approximate reciprocals may round either way.""" - return Fp4Quantize.near_tie(fp4_quantize.scaled(x, gs, 16)) +def gate_ratio(out: torch.Tensor, ref: torch.Tensor) -> float: + """Worst element's distance from `ref` as a fraction of the band `rtol = eps`, + `atol = eps * max|ref|` (both of `ref`'s dtype), the band + `qwen4_exp_hc_up_gate_mean`'s test certifies `mixed` in. Below 1 is inside.""" + eps = torch.finfo(ref.dtype).eps + ref64 = ref.double() + atol = eps * ref64.abs().max().item() + return ((out.double() - ref64).abs() / (atol + eps * ref64.abs())).max().item() def _unpack(xq: torch.Tensor) -> torch.Tensor: @@ -49,25 +48,16 @@ def _check_quant(mixed, xq, xsf, gs) -> None: """NVFP4 outputs against the native-torch quantization of the `mixed` this call returned: scale bytes equal, e2m1 codes equal except inside the reciprocal tie window.""" - want_q, want_sf = nvfp4_reference(mixed, gs) + want_q, want_sf = fp4_quantize.reference(mixed, gs, 16, False, False) assert xq.dtype == torch.uint8 and xq.shape == want_q.shape and xq.is_contiguous() assert xsf.dtype == torch.uint8 and xsf.shape == want_sf.shape assert torch.equal(xsf, want_sf), f"{(xsf != want_sf).sum().item()} scale bytes differ" diff = _unpack(xq) != _unpack(want_q) - stray = diff & ~nvfp4_near_tie(mixed, gs) + near = Fp4Quantize.near_tie(fp4_quantize.scaled(mixed, gs, 16)) + stray = diff & ~near assert not stray.any(), f"{stray.sum().item()} e2m1 codes differ outside the tie window" -def _compare(out, ref, gs) -> None: - """`mixed`: gate_ratio <= 1, the band `qwen4_exp_hc_up_gate_mean` is - certified in. Quantization: exact on the returned `mixed`.""" - (mixed, xq, xsf), (ref_mixed, _, _) = out, ref - assert mixed.dtype == ref_mixed.dtype and mixed.shape == ref_mixed.shape - assert mixed.is_contiguous() - assert gate_ratio(mixed, ref_mixed) <= 1.0, f"{gate_ratio(mixed, ref_mixed):.4g}x the band" - _check_quant(mixed, xq, xsf, gs) - - def _build(spec, seed): hc, hidden, rank = spec.get("hc_count", 4), spec.get("hidden", 2560), spec.get("rank", 320) g = torch.Generator(device=DEV).manual_seed(seed) @@ -92,9 +82,14 @@ def _seed(cell) -> int: def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name certifying(op) args = _build(cell.spec, seed=_seed(cell)) - out = op(*args) - assert out[0].shape == (cell.spec["rows"], cell.spec.get("hidden", 2560)) - _compare(out, op.reference(*args), args[4]) + mixed, xq, xsf = op(*args) + assert mixed.shape == (cell.spec["rows"], cell.spec.get("hidden", 2560)) + ref_mixed = op.reference(*args)[0] + assert mixed.dtype == ref_mixed.dtype and mixed.is_contiguous() + assert gate_ratio(mixed, ref_mixed) <= 1.0, f"{gate_ratio(mixed, ref_mixed):.4g}x the band" + # The same kernel and epilogue arithmetic as the bf16-only build the target runs elsewhere. + assert torch.equal(mixed, qwen4_exp_hc_up_gate_mean(*args[:4])) + _check_quant(mixed, xq, xsf, args[4]) def test_guards() -> None: diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ssm/test_modeling_v2_gdn_cached_replay_cute__verify.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ssm/test_modeling_v2_gdn_cached_replay_cute__verify.py new file mode 100644 index 000000000000..fabebf6190ed --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/ssm/test_modeling_v2_gdn_cached_replay_cute__verify.py @@ -0,0 +1,276 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the CuTe DSL cached-replay entry (`gdn_cached_replay_cute__verify`) at a +tp4/tep4 rank's 4 key / 12 value heads. + +Each cell drives consecutive verify steps the way the engine does, through the Triton +cached-replay entry's test harness (`_Bench`: the stock manager's bf16 state page at the tep4 engine's slot +count, work items, the all-layer commit, `_advance_replay_state` with random acceptance), and +holds every call to this entry's `reference` on that call's own inputs (output and all four +state buffers, the Triton entry's bands) and every output to the exact twin (the fp32 token-by-token +recurrence over every accepted token). With `strided_gates` the raw a / b reach the kernel as +column views of an input-projection-shaped buffer, as the target hands them over. +""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +import pytest +import torch + +import tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.ssm.gdn_cached_replay_cute__verify as cr # noqa: E501 + +_MV2_TESTS = Path(__file__).resolve().parents[4] +if str(_MV2_TESTS) not in sys.path: + sys.path.insert(0, str(_MV2_TESTS)) +from _validating import certifying # noqa: E402,F401 — fixture by name + +_spec = importlib.util.spec_from_file_location( + "cached_replay_harness", + # The Triton entry's test, in the family tree (tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm). + _MV2_TESTS + / "models/qwen4_exp/ssm" + / "test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py", +) +trep = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(trep) + +assert torch.cuda.is_available(), "gdn_cached_replay_cute requires a CUDA device" + +op = cr.gdn_cached_replay_cute +DEV = "cuda" +K, V = cr.K, cr.V +#: The input projection's column split at a tp4/tep4 rank: conv (q|k|v), z, a, b. +_CONV, _Z = 2 * cr.RANK_K_HEADS * K + cr.RANK_V_HEADS * V, cr.RANK_V_HEADS * V + + +def _gates(bench, act, strided): + """The raw a / b `[N, T, HV]`: the harness's contiguous tensors, or the same values as + column views of a `[N*T, conv + z + 2 HV]` projection (row stride = its width).""" + if not strided: + return act["a"], act["b"] + n, t, hv = bench.n, bench.t, bench.hv + proj = torch.randn(n * t, _CONV + _Z + 2 * hv, device=DEV, dtype=torch.bfloat16) + a = proj[:, _CONV + _Z : _CONV + _Z + hv] + b = proj[:, _CONV + _Z + hv :] + a.copy_(act["a"].reshape(n * t, hv)) + b.copy_(act["b"].reshape(n * t, hv)) + a, b = a.reshape(n, t, hv), b.reshape(n, t, hv) + assert a.stride() == (t * proj.shape[1], proj.shape[1], 1), a.stride() + return a, b + + +def _reference(bench, kw): + """This entry's `reference` on a call's arguments, every per-slot buffer through the + harness's `compact` (the touched slots' rows of an engine page).""" + if bench.idx is None: + return op.reference(**kw) + per_slot = ( + "ssm_states", + "old_u", + "old_k", + "old_G", + "old_beta", + "cache_buf_idx", + "prev_num_accepted_tokens", + ) + compact = {name: bench.compact(kw[name]) for name in per_slot} + return op.reference( + **dict(kw, **compact, state_indices=bench.compact_indices(kw["state_indices"])) + ) + + +def _drive(cell, bench) -> None: + strided = cell.spec.get("strided_gates", False) + wraps = appends = 0 + for _ in range(bench.steps): + mask, wrapped, _ = bench.writes() + wraps += wrapped.numel() + appends += bench.live.numel() - wrapped.numel() + bench.prepare() + after_commit = [] + for layer in range(bench.layers): + act = bench.inputs() + kw = bench.kwargs(layer, act) + kw["g"], kw["beta"] = _gates(bench, act, strided) + _, u, k, G = bench.state(layer) + before = dict(u=u.clone(), k=k.clone(), G=G.clone(), dA=bench.old_dA[layer].clone()) + ref = _reference(bench, kw) + exact = op.reference(**bench.exact_kwargs(layer, kw)) + committed = None + if bench.path == "all_layer": + committed = _reference( + bench, + dict(kw, replay_work_items=None, n_writes=None, use_all_layer_commit=False), + ) + got = op(**kw) + trep._compare((got, *bench.observed(layer)), ref) + _, u, k, G = bench.state(layer) + assert torch.equal(bench.old_dA[layer], before["dA"]), "old_beta is never written" + assert torch.equal(u[~mask], before["u"][~mask]), "old_u written outside the append" + assert torch.equal(k[~mask], before["k"][~mask]), "old_k written outside the append" + Gt, Gt0 = G.transpose(-1, -2), before["G"].transpose(-1, -2) + assert torch.equal(Gt[~mask], Gt0[~mask]), "old_G written outside the append" + bench.adopt_exact(layer, exact) + after_commit.append((got, exact, committed)) + bench.commit() + for layer, (got, exact, committed) in enumerate(after_commit): + if committed is not None: + # The partitioned call plus the manager's commit is the call that commits + # in-kernel. + trep._compare((got, *bench.observed(layer)), committed) + trep._compare((got, None, None, None, None), (exact[0], None, None, None, None)) + if bench.page is not None: + bench.page.assert_rest_unchanged() + bench.accept(bench.draw_acceptance()) + assert wraps > 0 and appends > 0, f"the cell never wrapped or appended ({wraps}, {appends})" + + +def _drive_padding(cell) -> None: + """Trailing rows sharing one slot, as CUDA-graph padding repeats the manager's dummy + slot: the rows on unique slots, and every slot but the shared one, are held to the + reference run on the unique rows alone (a dense pool: the shared slot's rows are + indexed out).""" + spec = {k_: v_ for k_, v_ in cell.spec.items() if k_ not in ("page_layers",)} + spec["slots"] = 24 + bench = trep._Bench(spec, seed=trep._seed(cell.why), exact=False) + dup = spec["dup"] + unique = bench.n - len(dup) + shared = int(bench.state_indices[unique]) + others = torch.ones(bench.slots, dtype=torch.bool, device=DEV) + others[shared] = False + for _ in range(bench.steps): + for layer in range(bench.layers): + act = bench.inputs() + kw = bench.kwargs(layer, act) + kw["g"], kw["beta"] = _gates(bench, act, spec.get("strided_gates", False)) + alone = dict( + kw, + q=kw["q"][:unique], + k=kw["k"][:unique], + v=kw["v"][:unique], + g=kw["g"][:unique], + beta=kw["beta"][:unique], + state_indices=kw["state_indices"][:unique], + packed_qkv=act["packed"][: unique * bench.t], + ) + ref = op.reference(**alone) + got = op(**kw) + pool, u, k, G = bench.state(layer) + trep._compare( + (got[:unique], pool[others], u[others], k[others], G[others]), + (ref[0], *(x[others] for x in ref[1:])), + ) + bench.accept(bench.draw_acceptance()) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + trep._MEASURED.clear() + if "dup" in cell.spec: + _drive_padding(cell) + else: + bench = trep._Bench(cell.spec, seed=trep._seed(cell.why)) + try: + _drive(cell, bench) + finally: + bench.close() + print( + f"\n[{cell.why}] " + + ", ".join(f"{k_} {v_:.3f}" for k_, v_ in sorted(trep._MEASURED.items())) + ) + + +def _rank_spec(**spec): + return dict(state_dtype=torch.bfloat16, H=cr.RANK_K_HEADS, HV=cr.RANK_V_HEADS, **spec) + + +@pytest.mark.parametrize("pdl", [False, True]) +@pytest.mark.parametrize("n,path", [(1, "small"), (4, "small"), (16, "all_layer")]) +def test_cuda_graph_replay_on_new_data_is_bitwise_eager(n, path, pdl) -> None: + """Two layers captured once (strided a / b, the all-layer commit after them when + partitioned), replayed three times on new activations with acceptance in between, + bitwise against the same steps run eagerly on a twin state; with `pdl` the replay + launches with programmatic dependent launch, as the target's verify step does.""" + t = 4 + spec = _rank_spec(path=path, n=n, t=t, layers=2, slots=40) + seed = 9100 + n + eager, graph, warm = (trep._Bench(spec, seed, exact=False) for _ in range(3)) + acts = [warm.inputs() for _ in range(2)] + gates = [_gates(warm, act, True) for act in acts] + + def step(b): + b.prepare() + outs = [] + for layer in range(2): + kw = b.kwargs(layer, acts[layer]) + kw["g"], kw["beta"] = gates[layer] + kw["launch_with_pdl"] = pdl + outs.append(op(**kw)) + b.commit() + return outs + + step(warm) + torch.cuda.synchronize() + captured = torch.cuda.CUDAGraph() + with torch.cuda.graph(captured): + outs_graph = step(graph) + refill = torch.Generator(device=DEV).manual_seed(seed + 1) + for _ in range(3): + for act, (a, b) in zip(acts, gates): + for x in (act["packed"], a, b): + x.copy_(torch.randn(x.shape, generator=refill, device=DEV, dtype=x.dtype)) + outs_eager = step(eager) + captured.replay() + torch.cuda.synchronize() + for e, g in zip(outs_eager, outs_graph): + assert torch.equal(e, g), "replayed output differs from eager" + for name in ("backing", "old_x", "old_B", "old_dt", "pnat", "buf"): + assert torch.equal(getattr(eager, name), getattr(graph, name)), f"{name} differs" + accepted = eager.draw_acceptance() + eager.accept(accepted) + graph.accept(accepted) + + +#: Launch mappings the 4 / 12 geometry can take, each checked against the tp1 default +#: small-batch mapping: every one computes the same instructions per state row. +_MAPPINGS = ( + dict(tpc=1, warps=1, pf=0, minb=4, early_acc=True), + dict(tpc=2, warps=2, pf=0, minb=4, early_acc=True), + dict(tpc=4, warps=4, pf=0, minb=4, early_acc=True), + dict(tpc=8, warps=4, pf=1, minb=4), + dict(tpc=24, warps=8, pf=1, minb=2), +) + + +@pytest.mark.parametrize("pdl", [False, True]) +@pytest.mark.parametrize("n,t", [(1, 4), (4, 4), (2, 2), (1, 8)]) +def test_launch_mappings_agree(n, t, pdl) -> None: + """One step bitwise equal across launch mappings (output, checkpoint with in-kernel + commits, u / k / G history), with wrapping and appending rows and, from 4 rows, a + negative slot.""" + spec = _rank_spec(path="small", n=n, t=t, steps=1, layers=1, slots=24) + bench = trep._Bench(spec, 4242 + n * 10 + t, exact=False) + sidx = bench.state_indices.clone() + if n > 2: + sidx[n // 2] = -1 + act = bench.inputs() + a, b = _gates(bench, act, True) + snap = [x.clone() for x in bench.state(0)] + results = [] + for cfg in _MAPPINGS: + for x, y in zip(bench.state(0), snap): + x.copy_(y) + kw = bench.kwargs(0, act) + sentinel = torch.full((n, t, bench.hv, V), 3.0, dtype=torch.bfloat16, device=DEV) + kw.update(g=a, beta=b, state_indices=sidx, output=sentinel, config=cfg, launch_with_pdl=pdl) + out = op(**kw) + torch.cuda.synchronize() + results.append([out.clone()] + [x.clone() for x in bench.state(0)]) + for cfg, got in zip(_MAPPINGS[1:], results[1:]): + for name, x, y in zip(("out", "state", "u", "k", "G"), got, results[0]): + assert torch.equal(x, y), f"{name}: {cfg} vs {_MAPPINGS[0]}" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_attention_plumbing.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_attention_plumbing.py new file mode 100644 index 000000000000..2daa23f1e54d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_attention_plumbing.py @@ -0,0 +1,386 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The target's attention plumbing around its kernels. + +* **The FP8 KV range probe** (`TRTLLM_MODELING_V2_VALIDATE` over an fp8 pool). The pool stores + K/V at the unit scale, so a magnitude past the e4m3 maximum saturates. The probe reports the + largest magnitude of the first eager forward's attention layers, once, and never reads a + tensor under CUDA graph capture. +* **The context pass over an NVFP4 pool, a group of requests at a time** (one rank). The + context rows read a bf16 scratch of their requests' pages (`_nvfp4_context_plan`), filled and + attended one group of consecutive requests at a time, so it holds the largest group's pages + rather than the step's. Under a budget of a few pages, with one request longer than the + budget, a prefill step's rows match one group over the whole step bitwise: its requests' + rows come in multiples of the row-tile pass's 10-row groups, so a launch per group splits + none of them and every row takes the same pass as in the one launch. +* **MTP IndexShare.** The draft selection (`_selection`) in its fused and its two-launch form, + the capture of each request's draft-extend row into the shared buffers (decode and + mixed-batch forms), and the reuse entry's order over canonical captures. The targets are + built without weights: the capture methods read only their arguments, the metadata and the + core's draft binding. +""" + +import math +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_mtp_reuse_topk import ( # noqa: E501 + qsa_mtp_reuse_topk as reuse, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501 + modeling as M, +) + +_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="runs on a CUDA device") +DEV = "cuda" + + +# ---- the FP8 KV range probe --------------------------------------------------------------- + +LAST = 11 + + +def _kv(peak: float = 1.0): + """K and V of one layer, 5 rows x 2 heads x 8 dims, with `peak` as V's largest value.""" + k = torch.linspace(-1.0, 1.0, 80, device=DEV).reshape(5, 2, 8).to(torch.bfloat16) + v = torch.linspace(-0.5, 0.5, 80, device=DEV).reshape(5, 2, 8).to(torch.bfloat16) + v[3, 1, 4] = peak + return k, v + + +def _run_forward(probe, peaks): + """One forward's attention layers 3, 7 and LAST; returns each `observe` result.""" + return [probe.observe(layer, *_kv(peak)) for layer, peak in zip((3, 7, LAST), peaks)] + + +@_cuda +def test_a_forward_within_the_e4m3_range_is_silent(monkeypatch): + logger = Mock() + monkeypatch.setattr(M, "logger", logger) + + assert _run_forward(M._KvRangeProbe(LAST), (1.0, 300.0, 448.0)) == [True, True, False] + + logger.warning.assert_not_called() + + +@_cuda +@pytest.mark.parametrize("peak", [-600.0, float("nan")]) +def test_a_magnitude_past_the_e4m3_maximum_is_reported_once(monkeypatch, peak): + """A negative value counts by its magnitude; a NaN counts as out of range.""" + logger = Mock() + monkeypatch.setattr(M, "logger", logger) + + # The oversized value is in a middle layer: the report waits for the forward's last one. + assert _run_forward(M._KvRangeProbe(LAST), (1.0, peak, 1.0)) == [True, True, False] + + logger.warning.assert_called_once() + assert "e4m3 maximum 448" in logger.warning.call_args.args[0] + + +@pytest.mark.cpu_only +def test_the_probe_runs_only_when_asked_over_an_fp8_pool(): + core = SimpleNamespace(quant_mode=M.QM_FP8_KV_CACHE, attn_layers=[3, 7, LAST]) + assert M._KvRangeProbe.for_target(core, True)._last == LAST + assert M._KvRangeProbe.for_target(core, False) is None + for quant_mode in (0, M.QM_NVFP4_KV_CACHE): + other = SimpleNamespace(quant_mode=quant_mode, attn_layers=core.attn_layers) + assert M._KvRangeProbe.for_target(other, True) is None + + +@_cuda +def test_observing_under_capture_reads_nothing(monkeypatch): + logger = Mock() + monkeypatch.setattr(M, "logger", logger) + probe = M._KvRangeProbe(LAST) + k, v = _kv(1000.0) + scratch = torch.zeros(1, device=DEV) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + scratch.add_(1) # warm the allocator off the capture stream + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + kept = [] + with torch.cuda.graph(graph): + scratch.add_(1) + # The last layer, with an out-of-range value: a read here would fail the capture + # (a host sync) and a report would fire. + kept.append(probe.observe(LAST, k, v)) + + assert kept == [True], "a captured forward does not end the probe" + logger.warning.assert_not_called() + # The first eager forward after capture is what reports. + assert _run_forward(probe, (1.0, 1.0, 1000.0)) == [True, True, False] + logger.warning.assert_called_once() + + +# ---- the context pass over an NVFP4 pool --------------------------------------------------- + +TPB, HKV, D, QH, RATIO, TOPK = 32, 2, 256, 24, 4, 2048 +# (cached, new) tokens a context request: 10, 30, 1, 47 and 4 pages. +SEGMENTS = ((0, 300), (900, 40), (20, 10), (1400, 100), (64, 60)) +BUDGET_PAGES = 32 # request 3's 47 pages form a group of their own + + +def _context_step(seed: int = 5): + gen = torch.Generator(device=DEV).manual_seed(seed) + kv_lens = [s + n for s, n in SEGMENTS] + pages = [math.ceil(k / TPB) for k in kv_lens] + pool_pages = sum(pages) + 3 + shape = (pool_pages, 2, HKV, TPB) + data = torch.randint(0, 256, (*shape, D // 2), generator=gen, device=DEV).to(torch.uint8) + scales = torch.randint(0x08, 0x31, (*shape, D // 16), generator=gen, device=DEV) + scales = scales.to(torch.uint8) + slots = 2 * len(SEGMENTS) + slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed)) + slot_of = slot_of[: len(SEGMENTS)].to(DEV, torch.int32) + page_ids = torch.randperm(pool_pages, generator=gen, device=DEV).to(torch.int32) + block_table = torch.full((slots, max(pages)), -1, dtype=torch.int32, device=DEV) + cursor = 0 + for slot, count in zip(slot_of.tolist(), pages): + block_table[slot, :count] = page_ids[cursor : cursor + count] + cursor += count + new = torch.tensor([n for _, n in SEGMENTS], dtype=torch.int32) + md = SimpleNamespace( + num_contexts=len(SEGMENTS), + kv_lens_runtime=torch.tensor(kv_lens, dtype=torch.int32), + kv_lens_cuda_runtime=torch.tensor(kv_lens, dtype=torch.int32, device=DEV), + seq_lens=new, + seq_lens_cuda=new.to(DEV), + ) + logical = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in SEGMENTS]) + rows = logical.numel() + # Every visible token, as the expansion lays out a selection under its top-k. + column = torch.arange(TOPK + RATIO - 1, device=DEV)[None, :] + selected = torch.where(column <= logical[:, None], column, -1).to(torch.int32) + core = SimpleNamespace( + heads_kv=HKV, + head_dim=D, + q_width=QH * D, + kv_scale_quant_orig=torch.ones(3, dtype=torch.float32, device=DEV), + softmax_scale=D**-0.5, + compress_ratio=RATIO, + ) + kv = SimpleNamespace(k=data[:, 0], v=data[:, 1], k_scale=scales[:, 0], v_scale=scales[:, 1]) + return dict( + core=core, + kv=kv, + md=md, + block_table=block_table, + req_idx=slot_of.repeat_interleave(new.to(DEV).long()), + q=torch.randn(rows, QH, D, generator=gen, device=DEV).to(torch.bfloat16), + selected=selected, + logical=logical.to(torch.int32), + rows=rows, + ) + + +def _context_pass(step): + plan = M._nvfp4_context_plan( + step["core"], step["md"], step["req_idx"], step["block_table"], TPB + ) + out = torch.empty(step["rows"], QH * D, dtype=torch.bfloat16, device=DEV) + M._sparse_gqa_nvfp4_context( + step["core"], + step["kv"], + plan, + step["q"], + step["block_table"], + step["selected"], + step["req_idx"], + step["logical"], + TPB, + step["rows"], + out, + 1, + ) + return plan, out + + +@_cuda +def test_grouped_context_pass_equals_one_group(monkeypatch) -> None: + step = _context_step() + whole, want = _context_pass(step) + assert len(whole[0]) == 1 and whole[2].shape[0] == 92 + + page_bytes = 2 * HKV * TPB * D * 2 + monkeypatch.setattr(M, "_NVFP4_CONTEXT_SCRATCH_BYTES", BUDGET_PAGES * page_bytes) + (groups, _, k_scratch, _), got = _context_pass(step) + # Requests [0], [1, 2], [3] (over the budget, alone), [4]. + assert [(g[0], g[1]) for g in groups] == [(0, 300), (300, 350), (350, 450), (450, 510)] + assert k_scratch.shape[0] == 47 + assert torch.equal(got.view(torch.int16), want.view(torch.int16)) + + +# ---- MTP IndexShare ----------------------------------------------------------------------- + +K = 512 +SLOT = 1 +COLS = 2080 + + +def _target(cls): + target = object.__new__(cls) + target.core = SimpleNamespace(_draft_bound=SimpleNamespace(share_slot=SLOT)) + return target + + +def _metadata(batch: int, accepted, seq_lens=None, num_contexts: int = 0): + return SimpleNamespace( + num_seqs=batch, + num_contexts=num_contexts, + seq_lens_cuda=None + if seq_lens is None + else torch.tensor(seq_lens, dtype=torch.int32, device=DEV), + qsa_mtp_num_accepted=torch.tensor(accepted, dtype=torch.int32, device=DEV), + qsa_shared_topk_indices=torch.full((3, 8, K), -7, dtype=torch.int32, device=DEV), + qsa_shared_topk_visible_blocks=torch.full((3, 8), -7, dtype=torch.int32, device=DEV), + qsa_shared_topk_captured_slots=set(), + ) + + +def _scores_and_windows(rows: int, gen: torch.Generator): + """Relu-shaped scores (half zeros: ties at the budget's last score whenever fewer than + K are positive) over windows of 0 to 3K groups; the columns past a window are NaN.""" + ends = torch.randint(0, COLS + 1, (rows,), generator=gen).to(torch.int32).to(DEV) + scores = torch.relu(torch.randn(rows, COLS, generator=gen)).to(DEV) + scores = scores * (torch.rand(rows, COLS, generator=gen).to(DEV) < 0.3) + outside = torch.arange(COLS, device=DEV)[None, :] >= ends[:, None] + scores.masked_fill_(outside, float("nan")) + return scores, torch.zeros(rows, dtype=torch.int32, device=DEV), ends + + +def _select(scores, starts, ends) -> torch.Tensor: + out = torch.empty((scores.shape[0], K), dtype=torch.int32, device=DEV) + M._select(SimpleNamespace(block_topk=K), scores, starts, ends, out) + return out + + +@_cuda +@pytest.mark.parametrize("width", [2080, 65600]) +@pytest.mark.parametrize("fused", [False, True]) +def test_selection_matches_the_canonical_composition(width, fused, monkeypatch): + """The draft selection keeps the canonical tie order and expansion in the fused launch + and in the selection and the expansion apart.""" + positions = torch.tensor([2043, 2047, 84000, 99999], dtype=torch.int64, device=DEV) + visible = ((positions + 1) // 4).to(torch.int32) + visible.clamp_(max=width) + positions = torch.minimum(positions, (visible.to(torch.int64) * 4 + 2)) + lengths = (positions + 3).to(torch.int32) + starts = torch.zeros_like(visible) + generator = torch.Generator(device=DEV).manual_seed(85) + scores = torch.relu(torch.randn((4, width), generator=generator, device=DEV) - 1.8) + cols = torch.arange(width, device=DEV) + scores.masked_fill_(cols[None, :] >= visible[:, None], float("nan")) + expected = torch.empty((4, K), dtype=torch.int32, device=DEV) + core = SimpleNamespace(block_topk=K, compress_ratio=4, token_topk=2048) + M._select(core, scores, starts, visible, expected) + expected_tokens = M.qsa_expand_block_indices( + expected, positions, lengths, compress_ratio=4, token_topk=2048 + ) + if not fused: + monkeypatch.setattr(M, "_topk_expand_supports", lambda scores: False) + result = torch.full_like(expected, -7) + selected = torch.full_like(expected_tokens, -7) + M._selection(core, scores, visible, positions, lengths, result, starts, selected) + assert torch.equal(result, expected) + assert torch.equal(selected, expected_tokens) + + +@_cuda +def test_decode_capture_at_draft_seven_takes_row_s_times_8_plus_a_minus_1(): + """max_draft_len 7: eight rows per request, accepted 1..8, so the captured row is + `s * 8 + a - 1` -- the first and the last row of a block included.""" + p, batch = 8, 4 + accepted = [1, 8, 5, 3] + topk = torch.arange(batch * p * K, dtype=torch.int32, device=DEV).reshape(batch * p, K) + visible = torch.arange(300, 300 + batch * p, dtype=torch.int32, device=DEV) + md = _metadata(batch, accepted) + + _target(M.DecodeTarget)._capture_selection(md, topk, visible, batch, p) + + rows = [s * p + a - 1 for s, a in enumerate(accepted)] + assert rows == [0, 15, 20, 26] + torch.testing.assert_close(md.qsa_shared_topk_indices[SLOT, :batch], topk[rows]) + torch.testing.assert_close(md.qsa_shared_topk_visible_blocks[SLOT, :batch], visible[rows]) + # Other slots and rows past the batch are untouched. + assert torch.all(md.qsa_shared_topk_indices[SLOT, batch:] == -7) + assert torch.all(md.qsa_shared_topk_indices[0] == -7) + assert md.qsa_shared_topk_captured_slots == {SLOT} + + +@_cuda +@pytest.mark.parametrize("compact", [False, True]) +def test_prefill_capture_at_draft_seven(compact): + """A mixed batch at max_draft_len 7: two context chunks of 5 and 3 rows, then three + generation requests of 8 rows with 8, 1 and 6 accepted. A context request captures its + chunk's last row whatever accepted count the worker left there. A `compact` selection + (the draft-extend step's kept rows under attention TP) holds each context request's last + row, then every generation row.""" + seq_lens = [5, 3, 8, 8, 8] + accepted = [2, 3, 8, 1, 6] + rows = [4, 7, 8 + 7, 16 + 0, 24 + 5] + if compact: + seq_lens_rows = [1, 1, 8, 8, 8] + rows = [0, 1, 2 + 7, 10 + 0, 18 + 5] + else: + seq_lens_rows = seq_lens + rows_total = sum(seq_lens_rows) + topk = torch.arange(rows_total * K, dtype=torch.int32, device=DEV).reshape(rows_total, K) + visible = torch.arange(400, 400 + rows_total, dtype=torch.int32, device=DEV) + md = _metadata(5, accepted, seq_lens=seq_lens, num_contexts=2) + md.num_ctx_tokens = 8 + + _target(M.PrefillTarget)._capture_selection(md, topk, visible, compact=compact) + + torch.testing.assert_close(md.qsa_shared_topk_indices[SLOT, :5], topk[rows]) + torch.testing.assert_close(md.qsa_shared_topk_visible_blocks[SLOT, :5], visible[rows]) + assert md.qsa_shared_topk_captured_slots == {SLOT} + + +def _reuse_case(fresh: int): + """A canonical capture of 16 requests and the reuse call's arguments for it.""" + gen = torch.Generator().manual_seed(11 + fresh) + batch = 16 + scores, starts, ends = _scores_and_windows(batch, gen) + captured_visible = ends.cpu() + captured = _select(scores, starts, ends) + shared = torch.full((8, 32, K), -1, dtype=torch.int32, device=DEV) + shared[SLOT, :batch] = captured + shared_visible = torch.zeros((8, 32), dtype=torch.int32, device=DEV) + shared_visible[SLOT, :batch] = captured_visible.to(DEV) + args = dict( + shared_indices=shared[SLOT], + shared_visible_blocks=shared_visible[SLOT], + request_indices=torch.arange(batch, dtype=torch.int32, device=DEV), + visible_blocks=(captured_visible + fresh).to(DEV), + output=torch.empty((batch, K), dtype=torch.int32, device=DEV), + ) + return args, captured_visible + + +@_cuda +def test_reuse_of_a_canonical_capture_stays_ascending(): + """A canonical capture plus the groups completed since is still ascending: appended ids + are newer than every captured one, and a full row gives up its oldest ids to them, so + the newest captured groups -- the ones next to the new ones -- are kept.""" + for fresh in (0, 2): + args, captured_visible = _reuse_case(fresh) + + out = reuse(**args) + + assert torch.equal(out, reuse.reference(**args)) + for row, shared_row, vis in zip( + out.cpu().tolist(), args["shared_indices"].cpu().tolist(), captured_visible.tolist() + ): + ids = [i for i in row if i >= 0] + assert row == ids + [-1] * (K - len(ids)), "padding must stay last" + assert ids == sorted(ids) and len(set(ids)) == len(ids) + assert len(ids) == min(vis + fresh, K) + assert all(i < vis + fresh for i in ids) + captured = [i for i in shared_row if i >= 0] + kept = captured[max(0, len(captured) + fresh - K) :] + assert ids == kept + list(range(vis, vis + fresh)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_core_layout.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_core_layout.py new file mode 100644 index 000000000000..57b3e5f537e5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_core_layout.py @@ -0,0 +1,769 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The default target's core at tp4, tep4, dep4 and dep2: each rank's forward widths and its +experts. + +The core is built on the meta device from the checkpoint's text and vision config values at +every rank of a tp4, tep4, dep4 or dep2 mapping, so nothing here touches a device. What is pinned is +what the forward addresses: the fused attention row and the split the norm + rope reads it +through, the Gated-DeltaNet projection's column split, the routed experts a rank holds (an +expert window, or tp4's padded intermediate slice of every expert) and the per-expert scalars +cut to them, and the attention-DP round trip's padding and slicing. Under attention DP the +draft layer holds every attention head and runs its FP8 expert window over the rows the group +gathers, padded by the row counts the spec worker hands each draft step. +""" + +from __future__ import annotations + +import ast +import inspect +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501 + modeling as M, +) +from tensorrt_llm._torch.configs import Qwen4ExpTextConfig, Qwen4ExpVisionConfig +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm.llmapi import QSASparseAttentionConfig +from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +pytestmark = pytest.mark.cpu_only + +NUM_LAYERS = 48 +WORLD = 4 + +# Qwen3.8-Flash-Next's text config. +_TEXT = dict( + architectures=["Qwen4ExpForCausalLM"], + attention_bias=False, + attention_dropout=0.0, + bos_token_id=248044, + dtype="bfloat16", + eos_token_id=248044, + full_attention_interval=4, + hc_count=4, + hc_lowrank=320, + head_dim=256, + heads_per_ngram=8, + hidden_act="silu", + hidden_size=2560, + indexer_budget=2048, + indexer_compress_ratio=4, + indexer_head_dim=128, + indexer_kv_heads=1, + indexer_n_heads=4, + layer_types=["full_attention" if i % 4 == 3 else "linear_attention" for i in range(NUM_LAYERS)], + linear_conv_kernel_dim=4, + linear_key_head_dim=128, + linear_num_key_heads=16, + linear_num_value_heads=48, + linear_value_head_dim=128, + make_ngram_vocab_size_divisible_by=128, + mamba_ssm_dtype="float32", + max_position_embeddings=262144, + moe_intermediate_size=640, + mtp=dict( + hybrid=True, + layer_types=["full_attention"], + mtp_use_hidden_state_from_layer=None, + num_hidden_layers=1, + rope_theta=10000000, + ), + mtp_num_hidden_layers=1, + mtp_use_dedicated_embeddings=False, + ngram_size=3, + ngram_vocab_size_base=20000000, + norm_topk_prob=True, + num_attention_heads=24, + num_experts=512, + num_experts_per_tok=10, + num_hidden_layers=NUM_LAYERS, + num_key_value_heads=2, + number_of_conv_states=3, + output_gate_type="sigmoid", + partial_rotary_factor=0.25, + ple_conv_kernel_size=4, + ple_embed_dim=2560, + ple_layer_ids=[2], + rms_norm_eps=1e-06, + rope_parameters=dict( + mrope_interleaved=True, + mrope_section=[11, 11, 10], + partial_rotary_factor=0.25, + rope_theta=10000000, + rope_type="default", + ), + shared_expert_intermediate_size=640, + split_ngram_parts=128, + tie_word_embeddings=False, + vocab_size=248320, +) +_VISION = dict( + deepstack_visual_indexes=[], + depth=27, + dtype="bfloat16", + hidden_act="gelu_pytorch_tanh", + hidden_size=1152, + in_channels=3, + intermediate_size=4304, + num_heads=16, + num_position_embeddings=2304, + out_hidden_size=2560, + patch_size=16, + spatial_merge_size=2, + temporal_patch_size=2, +) + +# (attention DP, MoE TP, MoE EP) of the three members. +_TEP4 = dict(enable_attention_dp=False, moe_tp_size=1, moe_ep_size=4) +_DEP4 = dict(enable_attention_dp=True, moe_tp_size=1, moe_ep_size=4) +_DEP2 = dict(enable_attention_dp=True, moe_tp_size=1, moe_ep_size=2) +_TP4 = dict(enable_attention_dp=False, moe_tp_size=4, moe_ep_size=1) + + +def _core(rank: int, layout: dict, spec_config=None, world: int = WORLD): + """The core of rank `rank` of a `world`-rank engine at `layout`, every parameter on meta.""" + model_config = ModelConfig( + pretrained_config=Qwen4ExpTextConfig(**_TEXT), + mapping=Mapping(world_size=world, tp_size=world, rank=rank, **layout), + quant_config=QuantConfig(quant_algo=QuantAlgo.MIXED_PRECISION), + moe_backend="TRTLLM", + moe_backend_requested="TRTLLM", + sparse_attention_config=QSASparseAttentionConfig(), + spec_config=spec_config, + ) + composite = SimpleNamespace( + vision_config=Qwen4ExpVisionConfig(**_VISION), image_token_id=248056, video_token_id=248057 + ) + # The core asserts the device it runs on; the weights it declares stay on meta. + with ( + mock.patch.object(torch.cuda, "get_device_capability", return_value=M._SM), + torch.device("meta"), + ): + return M.Qwen4ExpDefaultModelingV2(model_config, composite) + + +@pytest.mark.parametrize("rank", range(WORLD)) +def test_a_tep4_rank_runs_its_share_of_the_heads(rank) -> None: + """Six q heads and one of the two KV heads (repeated on rank pairs) over a 4352-wide + fused row: `[q 1536 | k 256 | v 256 | gate 1536 | index q 512, k 128 | v slot 128]`. + The main norm + rope reads it as 6 q, 1 k and 10 "v" heads of 256 (the row stride + alone), which is the cell `fused_qk_norm_rope` certifies.""" + core = _core(rank, _TEP4) + assert (core.heads_q, core.heads_kv, core.kv_head_start) == (6, 1, rank * 2 // WORLD) + assert (core.q_width, core.kv_width, core.qkv_width) == (1536, 256, 2048) + assert (core.index_width, core.index_scratch_width) == (640, 768) + assert core.attn_in_width == 4352 + assert core.attn_in_tail_heads == 10 + assert core.index_col == 3584 + assert tuple(core.w["l3_attn_in_w"].shape) == (core.attn_in_width, 2560) + assert tuple(core.w["l3_o_w"].shape) == (2560, core.q_width) + + +@pytest.mark.parametrize("rank", range(WORLD)) +def test_a_tep4_rank_runs_its_share_of_the_recurrent_heads(rank) -> None: + """Four key and twelve value heads: the conv runs over `[q | k | v]` = 512 + 512 + 1536 + channels, and the one input projection's output splits `[q|k|v 2560 | z 1536 | a 12 | + b 12]`.""" + core = _core(rank, _TEP4) + assert (core.num_k_heads, core.num_v_heads) == (4, 12) + assert (core.key_dim, core.value_dim, core.conv_dim) == (512, 1536, 2560) + assert tuple(core.in_proj_split) == (2560, 1536, 12, 12) + assert tuple(core.w["l0_in_proj_w"].shape) == (sum(core.in_proj_split), 2560) + assert tuple(core.w["l0_out_w"].shape) == (2560, core.value_dim) + assert tuple(core.w["l0_conv_w"].shape) == (core.conv_dim, 1, 4) + + +@pytest.mark.parametrize("rank", range(WORLD)) +def test_a_tep4_rank_holds_one_128_expert_window(rank) -> None: + """The routed experts' window is `[128r, 128r + 128)`; everything the router, the shared + expert and the hyper-connections read is whole on every rank.""" + core = _core(rank, _TEP4) + assert (core.local_experts, core.expert_offset) == (128, 128 * rank) + assert core.w["l0_fc1_w"].shape[0] == 128 + assert core.w["l0_fc2_w"].shape[0] == 128 + assert tuple(core.w["l0_router"].shape) == (512, 2560) + assert tuple(core.w["embed"].shape) == (248320, 2560) + + +@pytest.mark.parametrize("rank", range(WORLD)) +def test_a_tp4_rank_holds_every_expert_at_its_padded_intermediate_slice(rank) -> None: + """tp4 cuts attention and Gated-DeltaNet as tep4 does, and holds all 512 experts at a + 160-row intermediate slice stored 256 wide: the runner is handed the whole routing space + (`local_expert_offset` 0, `local_num_experts` 512) at `intermediate_size` 256, the + stack `fp4_block_scale_moe_runner` certifies, and the scalars span every expert.""" + core = _core(rank, _TP4) + assert (core.heads_q, core.heads_kv, core.kv_head_start) == (6, 1, rank * 2 // WORLD) + assert (core.num_k_heads, core.num_v_heads, core.attn_in_width) == (4, 12, 4352) + assert (core.dp_size, core.attn_tp_size) == (1, WORLD) + assert (core.local_experts, core.expert_offset, core.expert_inter) == (512, 0, 256) + assert core.moe_inter == 640 + assert tuple(core.w["l0_fc1_w"].shape) == (512, 2 * 256, 2560 // 2) + assert tuple(core.w["l0_fc1_s"].shape) == (512, 2 * 256, 2560 // 16) + assert tuple(core.w["l0_fc2_w"].shape) == (512, 2560, 256 // 2) + assert tuple(core.w["l0_fc2_s"].shape) == (512, 2560, 256 // 16) + assert tuple(core.w["l0_e_isc2"].shape) == (512,) + + +@pytest.mark.parametrize("layout", [_TEP4, _DEP4], ids=["tep4", "dep4"]) +def test_an_expert_window_runs_at_the_whole_intermediate(layout) -> None: + core = _core(1, layout) + assert (core.local_experts, core.expert_offset, core.expert_inter) == (128, 128, 640) + + +def test_the_expert_scalars_are_cut_to_the_window() -> None: + """The runner takes one scalar per expert of its window: each is the whole routing + space's product, sliced, and the FC1 activation scale is the layer's one number.""" + num_experts, start, count = 512, 256, 128 + gen = torch.Generator().manual_seed(0) + isc1 = torch.full((num_experts,), 0.5) + ws2_1 = torch.rand(num_experts, generator=gen) + 0.5 + isc2 = torch.rand(num_experts, generator=gen) + 0.5 + ws2_2 = torch.rand(num_experts, generator=gen) + 0.5 + w = { + "l0_e_isc1": isc1, + "l0_e_isc1_up": isc1.clone(), + "l0_e_ws2_1": ws2_1, + "l0_e_ws2_1_up": ws2_1.clone(), + "l0_e_isc2": isc2, + "l0_e_ws2_2": ws2_2, + } + s = M._moe_layer_scalars(w, 0, start, count) + window = slice(start, start + count) + gate1 = isc1 * ws2_1 + torch.testing.assert_close(s["global_scale"], torch.tensor([2.0]), rtol=0, atol=0) + torch.testing.assert_close(s["output1_scale_gate_scalar"], gate1[window], rtol=0, atol=0) + torch.testing.assert_close(s["output1_scale_scalar"], (gate1 / isc2)[window], rtol=0, atol=0) + torch.testing.assert_close(s["output2_scale_scalar"], (isc2 * ws2_2)[window], rtol=0, atol=0) + assert all(t.shape == (count,) for k, t in s.items() if k != "global_scale") + + +def test_a_spread_fc1_input_scale_is_refused() -> None: + """One activation quantization feeds every expert, so the layer must have one + `gate_proj.input_scale`.""" + isc1 = torch.full((512,), 0.5) + isc1[7] = 0.25 + ones = torch.ones(512) + w = { + "l0_e_isc1": isc1, + "l0_e_isc1_up": isc1.clone(), + "l0_e_ws2_1": ones, + "l0_e_ws2_1_up": ones.clone(), + "l0_e_isc2": ones, + "l0_e_ws2_2": ones, + } + with pytest.raises(AssertionError, match="one gate_proj.input_scale"): + M._moe_layer_scalars(w, 0, 0, 128) + + +@pytest.mark.parametrize("rank", range(WORLD)) +def test_a_dep4_rank_runs_every_head(rank) -> None: + """Attention DP keeps the single-rank geometry: 24 q and both KV heads over the 14080-wide + fused row (55 heads of 256: 24 q, 2 k, 29 "v"), and all 16 / 48 GDN heads over 10240 conv + channels. Nothing is summed after the mixers.""" + core = _core(rank, _DEP4) + assert (core.heads_q, core.heads_kv, core.kv_head_start) == (24, 2, 0) + assert (core.q_width, core.kv_width, core.qkv_width) == (6144, 512, 7168) + assert core.attn_in_width == 14080 + assert core.attn_in_tail_heads == 29 + assert core.index_col == 13312 + assert tuple(core.w["l3_attn_in_w"].shape) == (14080, 2560) + assert tuple(core.w["l3_o_w"].shape) == (2560, 6144) + assert (core.num_k_heads, core.num_v_heads, core.conv_dim) == (16, 48, 10240) + assert tuple(core.in_proj_split) == (10240, 6144, 48, 48) + assert core.attn_tp_size == 1 + + +@pytest.mark.parametrize("rank", range(WORLD)) +def test_a_dep4_rank_runs_one_128_expert_window_for_the_group(rank) -> None: + """The MoE group is the TP group, which spans the DP ranks: rank `r` is DP rank `r` and + holds the window `[128r, 128r + 128)`, over the rows the whole group gathers.""" + core = _core(rank, _DEP4) + assert (core.dp_size, core.dp_rank, core.dp_group) == (WORLD, rank, list(range(WORLD))) + assert (core.local_experts, core.expert_offset) == (128, 128 * rank) + assert core.w["l0_fc1_w"].shape[0] == 128 + assert _core(rank, _TEP4).dp_size == 1 + + +@pytest.mark.parametrize("rank", range(2)) +def test_a_dep2_rank_runs_every_head_and_one_256_expert_window(rank) -> None: + """dep2, the disaggregated-serving engines' layout, is dep4 at half the group: the same + single-rank attention and Gated-DeltaNet geometry, and rank `r` holds the 256-expert + window `[256r, 256r + 256)` over the rows both ranks gather.""" + core, dep4 = _core(rank, _DEP2, world=2), _core(rank, _DEP4) + for name in ( + "heads_q", + "heads_kv", + "kv_head_start", + "qkv_width", + "attn_in_width", + "attn_in_tail_heads", + "index_col", + "num_k_heads", + "num_v_heads", + "conv_dim", + "attn_tp_size", + ): + assert getattr(core, name) == getattr(dep4, name), name + assert tuple(core.in_proj_split) == tuple(dep4.in_proj_split) + for key in ("l3_attn_in_w", "l3_o_w", "l0_in_proj_w", "l0_router", "l0_sh_gu"): + assert core.w[key].shape == dep4.w[key].shape, key + assert (core.dp_size, core.dp_rank, core.dp_group) == (2, rank, [0, 1]) + assert (core.local_experts, core.expert_offset) == (256, 256 * rank) + assert core.w["l0_fc1_w"].shape[0] == 256 + assert core.w["l0_fc1_w"].shape[1:] == dep4.w["l0_fc1_w"].shape[1:] + + +def test_attention_dp_pads_to_the_groups_largest_row_count() -> None: + """Each rank pads to the largest of `all_rank_num_tokens`; a count that disagrees with + the rows this rank holds would gather the wrong split, and is refused.""" + core = SimpleNamespace(dp_size=WORLD, dp_rank=2) + md = SimpleNamespace(all_rank_num_tokens=[5, 17, 9, 1]) + assert M.Qwen4ExpDefaultModelingV2._dp_pad_rows(core, md, 9) == 8 + with pytest.raises(AssertionError, match="wrong split"): + M.Qwen4ExpDefaultModelingV2._dp_pad_rows(core, md, 10) + # Without attention DP the metadata is not read. + assert M.Qwen4ExpDefaultModelingV2._dp_pad_rows(SimpleNamespace(dp_size=1), None, 3) == 0 + + +@pytest.mark.parametrize("rank", range(WORLD)) +def test_the_attention_dp_round_trip_returns_this_ranks_rows_through_every_window( + rank, monkeypatch +) -> None: + """The round trip pads this rank's rows, gathers the group's in rank order, runs this + rank's window over all of them, reduce-scatters and drops the padding: with each window + a function of the gathered rows, this rank gets back the sum over windows of its own rows. + The collectives are stood in for by the group's arithmetic.""" + counts = [3, 7, 1, 4] + hidden = 8 + gen = torch.Generator().manual_seed(rank) + xs = [torch.randn(n, hidden, generator=gen) for n in counts] + rows = max(counts) + gathered = torch.cat([nn_pad(x, rows) for x in xs], 0) + + def window(w): + return lambda t: t * (w + 1) + w + + full = sum(window(w)(gathered) for w in range(WORLD)) + seen = {} + + def fake_allgather(t, sizes, group, process_group): + assert sizes is None and group == list(range(WORLD)) and process_group is None + assert tuple(t.shape) == (rows, hidden) + torch.testing.assert_close(t, nn_pad(xs[rank], rows), rtol=0, atol=0) + return gathered + + def fake_reducescatter(t, sizes, group, process_group): + assert sizes is None and group == list(range(WORLD)) and process_group is None + seen["partial"] = t + return full[rank * rows : (rank + 1) * rows] + + monkeypatch.setattr(M, "allgather", fake_allgather) + monkeypatch.setattr(M, "reducescatter", fake_reducescatter) + monkeypatch.setattr(M, "mpi_disabled", lambda: False) + core = SimpleNamespace(dp_size=WORLD, dp_group=list(range(WORLD)), mapping=None) + expert_sum = M.Qwen4ExpDefaultModelingV2._expert_sum_op(core) + out = expert_sum(window(rank), xs[rank], rows - counts[rank]) + torch.testing.assert_close(seen["partial"], window(rank)(gathered), rtol=0, atol=0) + torch.testing.assert_close(out, sum(window(w)(xs[rank]) for w in range(WORLD))) + assert tuple(out.shape) == (counts[rank], hidden) + + +def nn_pad(x: torch.Tensor, rows: int) -> torch.Tensor: + """`x` with zero rows appended up to `rows`.""" + return torch.cat([x, x.new_zeros(rows - x.shape[0], x.shape[1])], 0) + + +@pytest.mark.parametrize("max_draft_len", [1, 3, 7]) +def test_mtp_widens_the_verify_step_and_keys_the_draft_layer_after_the_trunk( + max_draft_len, +) -> None: + """A verify step carries `max_draft_len + 1` rows a generation request; the draft layer is + global layer 48, the index every KV manager maps it at and the weight table keys it at. + Without speculation the MTP module is not declared.""" + core = _core(0, _TEP4, MTPDecodingConfig(max_draft_len=max_draft_len)) + assert core.mtp_enabled and core.max_verify_rows == max_draft_len + 1 + assert core.mtp_layer_idx == NUM_LAYERS + assert core.mtp_theta == core.theta == 1e7 + assert core.spec_rollback is not None + assert any(k.startswith(f"l{NUM_LAYERS}_mtp_") for k in core.w) + plain = _core(0, _TEP4) + assert not plain.mtp_enabled and plain.max_verify_rows == 1 and plain.spec_rollback is None + assert not any("_mtp_" in k for k in plain.w) + + +@pytest.mark.parametrize("rank", range(WORLD)) +def test_a_tep4_rank_holds_its_share_of_the_draft_layer(rank) -> None: + """The draft layer's attention is cut like a trunk attention layer's: six q heads (their + q and interleaved gate rows), one KV head and the o_proj columns of the six heads; its + QSA indexer, front end, router, shared expert and mixers are whole. Its FP8 experts are + the rank's window `[128r, 128r + 128)` over the TP group, FP8 128x128 block scales.""" + core = _core(rank, _TEP4, MTPDecodingConfig(max_draft_len=3)) + p = f"l{NUM_LAYERS}_mtp_" + w = core.w + assert tuple(w[f"{p}q_w"].shape) == (2 * 1536, 2560) + assert tuple(w[f"{p}k_w"].shape) == tuple(w[f"{p}v_w"].shape) == (256, 2560) + assert tuple(w[f"{p}o_w"].shape) == (2560, 1536) + assert tuple(w[f"{p}idx_w"].shape) == (640, 2560) + assert tuple(w[f"{p}fc_e"].shape) == tuple(w[f"{p}fc_h"].shape) == (2560, 2560) + assert tuple(w[f"{p}pre_h_norm"].shape) == (4 * 2560,) + assert tuple(w[f"{p}router"].shape) == (512, 2560) + assert tuple(w[f"{p}sh_gu"].shape) == (2 * 640 + 8, 2560) + assert (core.mtp_local_experts, core.mtp_expert_offset) == (128, 128 * rank) + assert tuple(w[f"{p}fc1_w"].shape) == (128, 2 * 640, 2560) + assert w[f"{p}fc1_w"].dtype == torch.float8_e4m3fn + assert tuple(w[f"{p}fc1_s"].shape) == (128, 2 * 640 // 128, 2560 // 128) + assert tuple(w[f"{p}fc2_w"].shape) == (128, 2560, 640) + assert tuple(w[f"{p}fc2_s"].shape) == (128, 2560 // 128, 640 // 128) + + +@pytest.mark.parametrize("rank", range(WORLD)) +def test_a_tp4_rank_holds_whole_expert_draft_windows_beside_its_intermediate_slice(rank) -> None: + """At tp4 the trunk's routed experts are every expert at a 160-row intermediate slice + stored 256 wide, but the draft layer's FP8 experts cannot be cut that way (their 128x128 + weight blocks would split mid-block). They are the rank's whole-expert window + `[128r, 128r + 128)` over the TP group, exactly tep4's, and the four windows partition + the routing space. The draft attention is cut as at tep4.""" + max_draft_len = 3 + core = _core(rank, _TP4, MTPDecodingConfig(max_draft_len=max_draft_len)) + assert core.mtp_enabled and core.max_verify_rows == max_draft_len + 1 + assert (core.local_experts, core.expert_offset, core.expert_inter) == (512, 0, 256) + assert (core.mtp_local_experts, core.mtp_expert_offset, core.moe_inter) == ( + 128, + 128 * rank, + 640, + ) + p = f"l{NUM_LAYERS}_mtp_" + w = core.w + assert tuple(w[f"{p}fc1_w"].shape) == (128, 2 * 640, 2560) + assert w[f"{p}fc1_w"].dtype == torch.float8_e4m3fn + assert tuple(w[f"{p}fc1_s"].shape) == (128, 2 * 640 // 128, 2560 // 128) + assert tuple(w[f"{p}fc2_w"].shape) == (128, 2560, 640) + assert tuple(w[f"{p}fc2_s"].shape) == (128, 2560 // 128, 640 // 128) + assert tuple(w[f"{p}q_w"].shape) == (2 * 1536, 2560) + assert tuple(w[f"{p}k_w"].shape) == tuple(w[f"{p}v_w"].shape) == (256, 2560) + assert tuple(w[f"{p}o_w"].shape) == (2560, 1536) + assert tuple(w[f"{p}router"].shape) == (512, 2560) + windows = [_core(r, _TP4, MTPDecodingConfig(max_draft_len=1)) for r in range(WORLD)] + covered = sorted( + e + for c in windows + for e in range(c.mtp_expert_offset, c.mtp_expert_offset + c.mtp_local_experts) + ) + assert covered == list(range(512)) + + +def test_the_tp4_draft_windows_sum_to_the_whole_draft_moe(monkeypatch) -> None: + """`_mtp_moe_block` at tp4 hands the FP8 runner this rank's whole-expert window -- the + draft layer's `local_expert_offset` / `local_num_experts` and the whole 640-row + intermediate, never the trunk's 256-wide padded slice -- and completes the window's + partial with the TP sum once. The entries are stood in for by arithmetic that names the + expert a slot routes to, so the four ranks' partials sum to the all-expert output.""" + tokens, hidden, num_experts, topk = 5, 2560, 512, 10 + gen = torch.Generator().manual_seed(0) + x = torch.randn(tokens, hidden, generator=gen) + router = torch.randn(hidden, num_experts, generator=gen) + logits = x @ router + top_logits, top_ids = torch.topk(logits, topk, dim=-1) + top_w = torch.softmax(top_logits, dim=-1) + + def expert_out(ids, wts, start, count): + keep = ((ids >= start) & (ids < start + count)).to(x.dtype) + return (wts * keep * (ids + 1).to(x.dtype)).sum(-1, keepdim=True) * x + + full = expert_out(top_ids, top_w, 0, num_experts) + + def fake_cublas_mm(a, b, *args, **kwargs): + return a @ b + + def fake_routing(lg, topk, output_dtype): + assert topk == 10 and lg.shape == (tokens, num_experts) + return top_ids, top_w + + def fake_runner(**kw): + core = seen["core"] + assert kw["intermediate_size"] == core.moe_inter == 640 + assert kw["local_expert_offset"] == core.mtp_expert_offset + assert kw["local_num_experts"] == core.mtp_local_experts == 128 + assert tuple(kw["gemm1_weights"].shape) == (128, 2 * 640, hidden) + assert tuple(kw["gemm2_weights"].shape) == (128, hidden, 640) + assert kw["num_experts"] == num_experts and kw["top_k"] == topk + return expert_out( + kw["topk_ids"], kw["topk_weights"], kw["local_expert_offset"], kw["local_num_experts"] + ) + + monkeypatch.setattr(M, "cublas_mm", fake_cublas_mm) + monkeypatch.setattr(M, "renorm_moe_routing_op", fake_routing) + monkeypatch.setattr(M, "fp8_quantize_1x128", lambda t, use_ue8m0: (t, None)) + monkeypatch.setattr(M, "fp8_block_scale_moe_runner", fake_runner) + monkeypatch.setattr(M, "silu_and_mul_rows", lambda gug, inter: gug[:, :inter]) + # The shared expert and its gate are replicated and checked elsewhere: keep the routed sum. + monkeypatch.setattr(M, "fused_sigmoid_gate_mul_add", lambda routed, gate, shared: routed) + + seen = {} + partials = [] + for rank in range(WORLD): + core = _core(rank, _TP4, MTPDecodingConfig(max_draft_len=3)) + seen["core"] = core + summed = [] + core.tp_sum = lambda t: summed.append(t) or t + p = f"l{NUM_LAYERS}_mtp_" + views = SimpleNamespace( + router=router, + fc1_w=core.w[f"{p}fc1_w"], + fc1_s=core.w[f"{p}fc1_s"], + fc2_w=core.w[f"{p}fc2_w"], + fc2_s=core.w[f"{p}fc2_s"], + sh_gug=torch.zeros(hidden, 2 * 640 + 8), + sh_dn=torch.zeros(640, hidden), + # The FP8 experts (no NVFP4 copy away from one rank). + moe_gscale=None, + ) + out = M._mtp_moe_block(core, views, x) + assert len(summed) == 1, "the window's partial is summed over the TP group once" + partials.append(out) + torch.testing.assert_close(sum(partials), full) + + +def test_the_draft_layer_views_are_the_trunks_gemm_operands() -> None: + """`_mtp_views` hands the GEMMs column-major views: the shared expert's one + `[gate | up | gate logit | pad]` operand (the trunk layers' 2560 -> 1288 shape), the q + projection with its interleaved gate, and each hyper-connection site in the trunk's + form -- its whole grouped-norm weight, its `[down; block_inject; 0]` first projection + padded to 384 rows (the checkpoint's parameters rebound as its row views) and its up + projection -- plus the unit injection weight of the front end's site.""" + core = _core(0, _TEP4, MTPDecodingConfig(max_draft_len=3)) + with torch.device("cpu"): + for key, t in core.w.items(): + core.w[key] = torch.nn.Parameter( + torch.empty(t.shape, dtype=t.dtype, device="cpu"), requires_grad=False + ) + core._fuse_hc_first_projections() + M_ = M._mtp_views(core) + assert tuple(M_.sh_gug.shape) == (2560, 1288) and M_.sh_gug.stride() == (1, 2560) + assert tuple(M_.sh_dn.shape) == (640, 2560) and M_.sh_dn.stride() == (1, 640) + assert tuple(M_.o.shape) == (1536, 2560) and M_.o.stride() == (1, 1536) + assert tuple(M_.q.shape) == (2560, 3072) and M_.q.stride() == (1, 2560) + for norm in (M_.pre_h_norm, M_.ahc_norm, M_.mhc_norm, M_.fhc_norm): + assert tuple(norm.shape) == (4 * 2560,) and norm.is_contiguous() + p = f"l{NUM_LAYERS}_mtp_" + for tag, first, up, inject in ( + ("ahc", M_.ahc_first, M_.ahc_up, True), + ("mhc", M_.mhc_first, M_.mhc_up, True), + ("fhc", M_.fhc_first, M_.fhc_up, False), + ): + assert tuple(first.shape) == (4 * 2560, 384) and first.stride() == (1, 4 * 2560) + assert tuple(up.shape) == (320, 4 * 2560) and up.stride() == (1, 320) + assert core.w[f"{p}{tag}_down"].data_ptr() == first.data_ptr() + if inject: + assert core.w[f"{p}{tag}_inject"].data_ptr() == first[:, 320:].data_ptr() + assert tuple(M_.unit_inject.shape) == (core.max_num_tokens, 4) + assert bool((M_.unit_inject == 1).all()) + + +# ---- MTP under attention DP --------------------------------------------------------------- + +HIDDEN, NUM_EXPERTS, TOPK = 2560, 512, 10 + + +def _mtp(max_draft_len: int = 3) -> MTPDecodingConfig: + return MTPDecodingConfig(max_draft_len=max_draft_len) + + +@pytest.mark.parametrize("rank", range(WORLD)) +def test_a_dep4_rank_holds_the_whole_draft_attention_and_one_fp8_window(rank) -> None: + """The constructor builds MTP under attention DP. The draft layer's attention is unsharded, + as the trunk's dep4 attention is: 24 gated q heads, both KV heads and the whole o_proj; + nothing sums its output. Its FP8 experts are this rank's 128-expert window at the whole + 640-row intermediate, which the group's four windows complete.""" + core = _core(rank, _DEP4, _mtp()) + assert core.mtp_enabled and core.max_verify_rows == 4 + assert (core.dp_size, core.dp_rank, core.attn_tp_size) == (WORLD, rank, 1) + assert core.spec_rollback is not None + p = f"l{NUM_LAYERS}_mtp_" + w = core.w + assert tuple(w[f"{p}q_w"].shape) == (2 * 6144, HIDDEN) + assert tuple(w[f"{p}k_w"].shape) == tuple(w[f"{p}v_w"].shape) == (512, HIDDEN) + assert tuple(w[f"{p}o_w"].shape) == (HIDDEN, 6144) + assert tuple(w[f"{p}idx_w"].shape) == (640, HIDDEN) + assert tuple(w[f"{p}router"].shape) == (NUM_EXPERTS, HIDDEN) + assert (core.mtp_local_experts, core.mtp_expert_offset, core.moe_inter) == ( + 128, + 128 * rank, + 640, + ) + assert tuple(w[f"{p}fc1_w"].shape) == (128, 2 * 640, HIDDEN) + assert w[f"{p}fc1_w"].dtype == torch.float8_e4m3fn + assert tuple(w[f"{p}fc2_w"].shape) == (128, HIDDEN, 640) + windows = [_core(r, _DEP4, _mtp(1)) for r in range(WORLD)] + covered = sorted( + e + for c in windows + for e in range(c.mtp_expert_offset, c.mtp_expert_offset + c.mtp_local_experts) + ) + assert covered == list(range(NUM_EXPERTS)) + + +def test_a_draft_step_pads_by_the_counts_the_worker_hands_it() -> None: + """The spec worker hands each draft step the group's row counts -- every rank's trunk + tokens on the loop's first step, its request count on the later ones -- and does not write + them into the metadata. A count that disagrees with this rank's rows would gather the wrong + split, and a missing one cannot pad at all; without attention DP nothing is padded.""" + core = _core(2, _DEP4, _mtp()) + # Draft-extend: rank 2 holds 9 trunk rows, the group's largest is 17. + assert M._draft_pad_rows(core, [5, 17, 9, 1], 9) == 8 + # A later draft step: one row a request, attention-DP dummy requests included. + assert M._draft_pad_rows(core, [2, 4, 3, 1], 3) == 1 + # A captured decode step pads every rank to one batch. + assert M._draft_pad_rows(core, [4, 4, 4, 4], 4) == 0 + with pytest.raises(AssertionError, match="wrong split"): + M._draft_pad_rows(core, [5, 17, 9, 1], 10) + with pytest.raises(AssertionError, match="without the group's row counts"): + M._draft_pad_rows(core, None, 9) + assert M._draft_pad_rows(_core(0, _TEP4, _mtp()), None, 7) == 0 + + +def _expert_out(x, ids, wts, start, count): + """A stand-in MoE whose output names the experts a row routes to inside `[start, start + + count)`: windows of whole experts sum to the all-expert output.""" + keep = ((ids >= start) & (ids < start + count)).to(x.dtype) + return (wts * keep * (ids + 1).to(x.dtype)).sum(-1, keepdim=True) * x + + +@pytest.mark.parametrize( + "counts", + [[3, 7, 1, 4], [1, 2, 1, 3], [4, 4, 4, 4]], + ids=["draft-extend", "draft-decode-with-a-dummy-rank", "captured-batch"], +) +def test_every_rank_completes_its_draft_rows_through_the_round_trip(counts, monkeypatch) -> None: + """At every draft step every rank pads its rows to the group's largest count, gathers the + group's rows (one allgather), routes, quantizes and runs its window over all of them in + chunks of at most `_MTP_MOE_MAX_ROWS`, and reduce-scatters (one reduce-scatter): the same + two collectives at the same shape on every rank, a rank holding one dummy request included. + Each rank gets back the all-expert MoE of its own rows; nothing goes through the TP sum.""" + chunk_rows = 5 + monkeypatch.setattr(M, "_MTP_MOE_MAX_ROWS", chunk_rows) + gen = torch.Generator().manual_seed(0) + dt = torch.float64 + xs = [torch.randn(n, HIDDEN, generator=gen, dtype=dt) for n in counts] + router = torch.randn(HIDDEN, NUM_EXPERTS, generator=gen, dtype=dt) + rows = max(counts) + gathered = torch.cat([nn_pad(x, rows) for x in xs], 0) + + def route(t): + top_logits, top_ids = torch.topk(t @ router, TOPK, dim=-1) + return top_ids, torch.softmax(top_logits, dim=-1) + + ids, wts = route(gathered) + full = _expert_out(gathered, ids, wts, 0, NUM_EXPERTS) + + state = {} + + def fake_cublas_mm(a, b, *args, **kwargs): + return a @ b + + def fake_routing(lg, topk, output_dtype): + assert topk == TOPK and lg.shape[1] == NUM_EXPERTS and output_dtype == dt + top_logits, top_ids = torch.topk(lg, topk, dim=-1) + return top_ids, torch.softmax(top_logits, dim=-1) + + def fake_runner(**kw): + core = state["core"] + assert kw["intermediate_size"] == core.moe_inter == 640 + assert kw["local_expert_offset"] == 128 * core.dp_rank + assert kw["local_num_experts"] == core.mtp_local_experts == 128 + x = kw["hidden_states"] + assert 0 < x.shape[0] <= chunk_rows + state["runner_rows"] += x.shape[0] + return _expert_out(x, kw["topk_ids"], kw["topk_weights"], kw["local_expert_offset"], 128) + + def fake_allgather(t, sizes, group, process_group): + assert sizes is None and group == list(range(WORLD)) and process_group is None + state["collectives"].append(("allgather", tuple(t.shape))) + torch.testing.assert_close(t, nn_pad(xs[state["rank"]], rows), rtol=0, atol=0) + return gathered + + def fake_reducescatter(t, sizes, group, process_group): + assert sizes is None and group == list(range(WORLD)) and process_group is None + state["collectives"].append(("reducescatter", tuple(t.shape))) + state["partial"] = t + r = state["rank"] + return full[r * rows : (r + 1) * rows] + + monkeypatch.setattr(M, "cublas_mm", fake_cublas_mm) + monkeypatch.setattr(M, "renorm_moe_routing_op", fake_routing) + monkeypatch.setattr(M, "fp8_quantize_1x128", lambda t, use_ue8m0: (t, None)) + monkeypatch.setattr(M, "fp8_block_scale_moe_runner", fake_runner) + monkeypatch.setattr(M, "silu_and_mul_rows", lambda gug, inter: gug[:, :inter]) + # The shared expert and its gate are replicated and checked elsewhere: keep the routed sum. + monkeypatch.setattr(M, "fused_sigmoid_gate_mul_add", lambda routed, gate, shared: routed) + monkeypatch.setattr(M, "allgather", fake_allgather) + monkeypatch.setattr(M, "reducescatter", fake_reducescatter) + monkeypatch.setattr(M, "mpi_disabled", lambda: False) + + per_rank = [] + for rank in range(WORLD): + core = _core(rank, _DEP4, _mtp()) + core.expert_sum = core._expert_sum_op() + + def no_tp_sum(t): + raise AssertionError("the draft MoE under attention DP went through the TP sum") + + core.tp_sum = no_tp_sum + state.update(core=core, rank=rank, collectives=[], runner_rows=0) + p = f"l{NUM_LAYERS}_mtp_" + views = SimpleNamespace( + router=router, + fc1_w=core.w[f"{p}fc1_w"], + fc1_s=core.w[f"{p}fc1_s"], + fc2_w=core.w[f"{p}fc2_w"], + fc2_s=core.w[f"{p}fc2_s"], + sh_gug=torch.zeros(HIDDEN, 2 * 640 + 8, dtype=dt), + sh_dn=torch.zeros(640, HIDDEN, dtype=dt), + # The FP8 experts (no NVFP4 copy away from one rank). + moe_gscale=None, + ) + pad_rows = M._draft_pad_rows(core, counts, counts[rank]) + out = M._mtp_moe_block(core, views, xs[rank], pad_rows) + # The window ran over every gathered row once, chunk by chunk. + assert state["runner_rows"] == WORLD * rows + torch.testing.assert_close( + state["partial"], _expert_out(gathered, ids, wts, 128 * rank, 128) + ) + assert tuple(out.shape) == (counts[rank], HIDDEN) + torch.testing.assert_close(out, _expert_out(xs[rank], *route(xs[rank]), 0, NUM_EXPERTS)) + per_rank.append(state["collectives"]) + assert all(c == per_rank[0] for c in per_rank), per_rank + assert per_rank[0] == [ + ("allgather", (rows, HIDDEN)), + ("reducescatter", (WORLD * rows, HIDDEN)), + ] + + +def _calls(fn: ast.AST, name: str) -> list[ast.Call]: + return [ + node + for node in ast.walk(fn) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == name + ] + + +@pytest.mark.parametrize("target", [M.PrefillTarget, M.DecodeTarget]) +def test_both_draft_targets_pad_the_draft_moe_by_the_workers_counts(target) -> None: + """A draft step runs on the prefill target while the metadata still carries context rows + and on the decode target otherwise; ranks of one attention-DP group can be on either at the + same draft step. Both pad the draft MoE by the `all_rank_num_tokens` the worker passes them + for that step, never by the metadata's, which still holds the trunk step's counts.""" + fn = ast.parse(inspect.getsource(target.mtp_forward).lstrip()) + (pad,) = _calls(fn, "_draft_pad_rows") + assert [ast.unparse(a) for a in pad.args] == ["core", "all_rank_num_tokens", "rows"] + (tail,) = _calls(fn, "_mtp_tail") + assert ast.unparse(tail.args[-1]) == "pad_rows" + assert "_dp_pad_rows" not in ast.unparse(fn) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_decode_paths.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_decode_paths.py new file mode 100644 index 000000000000..6c88e8525cb2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_decode_paths.py @@ -0,0 +1,295 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The target's decode-side paths at one rank (tp1): the draft layer's front end, the heads +and their logits, and the Gated-DeltaNet verify step's fused conv. + +* **The draft front end**: the token embedding under rejection sampling, and the dense + projections' dispatch between `fp8_gemv` and the bf16 projection. +* **The logits**: the target's logits processor returns the shell's logits bit for bit (the + same row gather, the same fp32 values, its head's bf16 logits widened by one vectorized + kernel); a consumer of the row argmax alone may take the bf16 logits; the dense FP8 verify + head reads the NVFP4 copy at decode rows. +* **The verify step's conv** (`DecodeTarget._gdn_block`) under the dense FP8 mode and over the + in_proj's NVFP4 copy (`TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4`, up to + `_INPROJ_NVFP4_MAX_ROWS`). A captured verify step runs the in_proj with the multi-token conv + step as its epilogue (`fp8_gemv_conv1d_verify` / `nvfp4_split_gemv_conv1d_verify`); an eager + one runs the projection and then `causal_conv1d_triton_update`. Both must hand the cached + replay the same thing -- the conv output (contiguous, as the replay reads it), `a`, `b` and + `z` bitwise -- and leave the conv pool and the per-step snapshot buffer bitwise the same. The + replay is stubbed to record its inputs (it is the same call on both sides), the epilogue to + record `z`. Over a bf16 pool (the CuTe replay) `a` and `b` reach it as views of the + projection, not copies. +""" + +import types +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.fp8_gemv import ( # noqa: E501 + pack_from_bf16, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.gemm.nvfp4_gemv import ( # noqa: E501 + pack_split_from_bf16, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501 + modeling as M, +) +from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor +from tensorrt_llm._torch.speculative.interface import rejection_sampling_one_model + +_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="runs on a CUDA device") +DEV, BF16 = "cuda", torch.bfloat16 + + +# ---- the draft front end ------------------------------------------------------------------ + + +@_cuda +def test_draft_embedding_reads_row_0_for_rejected_slots() -> None: + """Rejection sampling pads the slots past a request's accepted count with -1 (here every + draft is rejected at slot 0); the front end embeds those as row 0 and every valid id as + the plain lookup does.""" + batch, draft, vocab, hidden = 8, 3, 4096, 256 + draft_tokens = torch.arange(batch * draft, dtype=torch.int32, device=DEV).reshape(batch, draft) + draft_probs = torch.zeros(batch, draft, vocab, device=DEV) + draft_probs.scatter_(2, draft_tokens.long().unsqueeze(-1), 1.0) + target_probs = torch.zeros(batch, draft + 1, vocab, device=DEV) + target_probs[..., vocab - 1] = 1.0 + seed = torch.zeros(batch, dtype=torch.int64, device=DEV) + offset = torch.arange(batch, dtype=torch.int64, device=DEV) + accepted, _ = rejection_sampling_one_model( + draft_probs, draft_tokens, target_probs, seed, offset + ) + assert bool((accepted[:, 1:] == -1).all()) and bool((accepted[:, 0] == vocab - 1).all()) + + g = torch.Generator(device=DEV).manual_seed(0) + table = torch.randn(vocab, hidden, generator=g, device=DEV).to(BF16) + valid_ids = torch.randint(0, vocab, (257,), generator=g, device=DEV, dtype=torch.int32) + for ids in (accepted.flatten(), valid_ids, valid_ids.long()): + e = M._draft_embedding(ids, table) + valid = ids >= 0 + assert torch.equal(e[valid], table[ids[valid].long()]) + assert torch.equal(e[~valid], table[0].expand(int((~valid).sum()), -1)) + + +@_cuda +@pytest.mark.parametrize("rows", [1, 16, 17]) +def test_draft_projection_dispatch(rows: int) -> None: + """`_draft_mm`: `fp8_gemv` over the packed copy up to `_MTP_DENSE_FP8_MAX_ROWS` rows, the + bf16 projection past them and without a copy.""" + g = torch.Generator(device=DEV).manual_seed(rows) + w = (torch.randn(2560, 2560, generator=g, device=DEV) * 0.02).to(BF16) + x = torch.randn(rows, 2560, generator=g, device=DEV).to(BF16) + packed = pack_from_bf16(w) + got = M._draft_mm(x, w.t(), packed, cublas_mm) + if rows <= M._MTP_DENSE_FP8_MAX_ROWS: + assert torch.equal(got, M.fp8_gemv_op(x, packed)) + ref = x.float() @ w.float().t() + assert ((got.float() - ref).norm() / ref.norm()).item() < 0.05 + else: + assert torch.equal(got, cublas_mm(x, w.t())) + assert torch.equal(M._draft_mm(x, w.t(), None, cublas_mm), cublas_mm(x, w.t())) + + +# ---- the logits --------------------------------------------------------------------------- + +VOCAB, HIDDEN = 248320, 2560 + + +@pytest.fixture(scope="module") +def head(): + g = torch.Generator(device=DEV).manual_seed(0) + lin = torch.nn.Linear(HIDDEN, VOCAB, bias=False, device=DEV, dtype=BF16) + with torch.no_grad(): + lin.weight.copy_(torch.randn(VOCAB, HIDDEN, generator=g, device=DEV) * 0.02) + return lin + + +@_cuda +@pytest.mark.parametrize("seq_lens", [[1] * 128, [5000, 3000, 1] + [1] * 121]) +def test_logits_match_the_shell_processor_bitwise(head, seq_lens) -> None: + g = torch.Generator(device=DEV).manual_seed(len(seq_lens)) + lens = torch.tensor(seq_lens, dtype=torch.int32, device=DEV) + hidden = torch.randn(int(lens.sum()), HIDDEN, generator=g, device=DEV).to(BF16) + md = SimpleNamespace(seq_lens_cuda=lens) + with torch.no_grad(): + want = LogitsProcessor().forward(hidden, head, md) + got = M._WideningLogitsProcessor().forward(hidden, head, md) + assert got.dtype == torch.float32 and got.shape == (len(seq_lens), VOCAB) + assert torch.equal(got.view(torch.int32), want.view(torch.int32)) + + +@_cuda +def test_context_logits_and_no_metadata(head) -> None: + g = torch.Generator(device=DEV).manual_seed(9) + hidden = torch.randn(33, HIDDEN, generator=g, device=DEV).to(BF16) + with torch.no_grad(): + for md, ctx in ((None, False), (SimpleNamespace(seq_lens_cuda=None), True)): + want = LogitsProcessor().forward(hidden, head, md, ctx) + got = M._WideningLogitsProcessor().forward(hidden, head, md, ctx) + assert torch.equal(got.view(torch.int32), want.view(torch.int32)) + + +@_cuda +@pytest.mark.parametrize("rows", [1, 4, 256]) +def test_unwidened_logits_give_the_same_greedy_tokens(head, rows) -> None: + """`widen=False` returns the head's own bf16 logits, whose greedy samplers (the verify + sampler's `greedy_argmax`, the draft sampler's int32 one) pick the widened logits' tokens.""" + from tensorrt_llm._torch.speculative.greedy_argmax import greedy_argmax + + g = torch.Generator(device=DEV).manual_seed(rows) + hidden = torch.randn(rows, HIDDEN, generator=g, device=DEV).to(BF16) + proc = M._WideningLogitsProcessor() + with torch.no_grad(): + wide = proc.forward(hidden, head, None, True) + narrow = proc.forward(hidden, head, None, True, widen=False) + # Without return_context_logits the shell's gather and fp32 contract stay. + gathered = proc.forward(hidden, head, None, False, widen=False) + assert narrow.dtype == BF16 and wide.dtype == torch.float32 + assert gathered.dtype == torch.float32 + assert torch.equal(narrow.float().view(torch.int32), wide.view(torch.int32)) + narrow[0, 17] = narrow[0].max() # an exact tie, broken to the lower index either way + wide = narrow.float() + assert torch.equal(greedy_argmax(narrow), greedy_argmax(wide)) + assert torch.equal(greedy_argmax(narrow, torch.int32), greedy_argmax(wide, torch.int32)) + assert torch.equal(greedy_argmax(narrow), torch.argmax(wide, dim=-1)) + + +@_cuda +@pytest.mark.parametrize("rows", [1, 3, 16, 17]) +def test_fp8_verify_head_reads_the_nvfp4_copy_at_decode_rows(head, rows) -> None: + """With an NVFP4 copy (`_TARGET_HEAD_NVFP4_ENV`) the dense FP8 verify head's decode rows + are `nvfp4_gemv` over it; past `_FP8_GEMV_HEAD_ROWS` they take the per-row FP8 copy as + the head without it does.""" + rowwise = M._fp8_rowwise(head.weight.data) + nv = M.nvfp4_pack_from_bf16(head.weight.data) + g = torch.Generator(device=DEV).manual_seed(100 + rows) + hidden = torch.randn(rows, HIDDEN, generator=g, device=DEV).to(BF16) + with torch.no_grad(): + got = M._Fp8Head(*rowwise, None, nv)(hidden) + if rows <= M._FP8_GEMV_HEAD_ROWS: + want = M.nvfp4_gemv_op(hidden, nv) + else: + want = M._Fp8Head(*rowwise, None, None)(hidden) + assert got.shape == (rows, VOCAB) and torch.equal(got, want) + + +def _gate(guided=None, greedy=True, rejection=False, penalty=False, relaxed=False): + worker = SimpleNamespace( + guided_decoder=guided, + spec_config=SimpleNamespace(use_relaxed_acceptance_for_thinking=relaxed), + ) + md = SimpleNamespace( + is_all_greedy_sample=greedy, + wants_advanced_draft_sampling=rejection and not greedy, + enable_penalty=penalty, + ) + return M._spec_argmax_only(worker, md) + + +@pytest.mark.cpu_only +def test_argmax_only_gate() -> None: + assert _gate() == (True, True) + # A non-greedy batch: the verify sampler samples; the draft stays argmax unless + # rejection sampling needs its distribution. + assert _gate(greedy=False) == (False, True) + assert _gate(greedy=False, rejection=True) == (False, False) + # The occurrence penalties and relaxed acceptance read the verify logits only. + assert _gate(penalty=True) == (False, True) + assert _gate(relaxed=True) == (False, True) + # A guided decoder edits both. + assert _gate(guided=object()) == (False, False) + + +# ---- the Gated-DeltaNet verify step's conv ------------------------------------------------ + +KEY, HV, V, CONV, VD, SLOTS, ROWS = 16, 48, 128, 10240, 6144, 461, 32 + + +@pytest.fixture(scope="module", params=["fp8", "nvfp4"]) +def gdn_layer(request): + g = torch.Generator(device=DEV).manual_seed(0) + w = (torch.randn(CONV + VD + 2 * HV, HIDDEN, generator=g, device=DEV) * HIDDEN**-0.5).to(BF16) + nv = request.param == "nvfp4" + return types.SimpleNamespace( + in_proj=w.t(), + in_proj8=None, + in_proj_m=None, + in_proj_g=None if nv else pack_from_bf16(w), + in_proj_nv=pack_split_from_bf16(w) if nv else None, + conv_w=(torch.randn(CONV, 4, generator=g, device=DEV) * 0.5).to(BF16), + a_log=None, + dt_bias=None, + ) + + +def _gdn_block(layer, x, si, int_idx, conv0, win0, p, captured, monkeypatch): + """One verify `_gdn_block` on copies of the pools; returns what the cached replay and the + epilogue were handed, and the pools.""" + seen = {} + + def replay(**kw): + mq_out = kw["packed_qkv"] + n, p_ = kw["g"].shape[:2] + seen.update( + mq_out=mq_out.clone(), + a=kw["g"].clone(), + b=kw["beta"].clone(), + pdl=kw["launch_with_pdl"], + ) + assert mq_out.is_contiguous() and mq_out.shape == (n * p_, CONV) + # The CuTe replay reads a / b through their row stride: views, not copies. + assert kw["g"].stride(1) > HV and kw["beta"].stride(1) > HV + return torch.zeros(n, p_, HV, V, dtype=BF16, device=DEV) + + monkeypatch.setattr(M, "gdn_cached_replay_cute", replay) + history = torch.zeros(1, device=DEV) + core = types.SimpleNamespace( + conv_dim=CONV, + key_dim=KEY * 128, + value_dim=VD, + num_k_heads=KEY, + head_k_dim=128, + num_v_heads=HV, + head_v_dim=V, + in_proj_split=(CONV, VD, HV, HV), + decode_in_mm=cublas_mm, + _bound=types.SimpleNamespace( + pools=types.SimpleNamespace( + conv={0: conv0.clone()}, ssm={0: torch.zeros(1, dtype=BF16, device=DEV)} + ), + int_conv={0: win0.clone()}, + replay=types.SimpleNamespace( + old={0: (history,) * 4}, cache_buf_idx=None, prev_num_accepted=None, history=16 + ), + ), + ) + target = object.__new__(M.DecodeTarget) + target.core = core + target._gdn_epilogue = lambda out, z, L, num_tokens: z.clone() + verify = (int_idx, None, None, False) + z = target._gdn_block(x, 0, layer, si, verify, p, captured=captured) + pools = core._bound.pools + return dict(seen, z=z, conv=pools.conv[0], win=core._bound.int_conv[0]) + + +@_cuda +@pytest.mark.parametrize("n,p", [(1, 4), (2, 4), (8, 2), (2, 8)]) +def test_fused_verify_conv_matches_the_two_launches(gdn_layer, n, p, monkeypatch) -> None: + if gdn_layer.in_proj_nv is not None and n * p > M._INPROJ_NVFP4_MAX_ROWS: + pytest.skip("past the NVFP4 copy's rows") + g = torch.Generator(device=DEV).manual_seed(10 * n + p) + conv0 = torch.randn(SLOTS, CONV, 3, generator=g, device=DEV).to(BF16) + win0 = torch.randn(ROWS, p, CONV, 3, generator=g, device=DEV).to(BF16) + x = torch.randn(n * p, HIDDEN, generator=g, device=DEV).to(BF16) + si = torch.randperm(SLOTS, generator=g, device=DEV)[:n].to(torch.int32) + int_idx = torch.arange(ROWS, dtype=torch.int32, device=DEV)[:n] + eager = _gdn_block(gdn_layer, x, si, int_idx, conv0, win0, p, False, monkeypatch) + fused = _gdn_block(gdn_layer, x, si, int_idx, conv0, win0, p, True, monkeypatch) + for key in ("mq_out", "a", "b", "z", "conv", "win"): + assert torch.equal(eager[key].view(torch.int16), fused[key].view(torch.int16)), key + assert fused["pdl"] == eager["pdl"] == M._GDN_REPLAY_PDL + assert not torch.equal(fused["conv"], conv0) and not torch.equal(fused["win"], win0) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_modeling_v2_spec_rollback.py similarity index 72% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_modeling_v2_spec_rollback.py index fc1099ea6537..1936894dbafa 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_modeling_v2_spec_rollback.py @@ -1,35 +1,37 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test for the target's speculative rollback handler (`SpeculativeRollback`). - -A verify step is simulated with the target's own state-writing ops -- the n-gram id -history, the dilated short conv and the QSA index store, plus an anchor-slot overwrite -standing in for the compression -- run over every draft row the way the forward runs -them, with the handler staged at the forward's call sites and committed afterwards. - -The reference shares no code with the commit: it runs the same ops from the same -pre-step state over each generation request's accepted prefix only, which is by -definition the state a verify step must leave behind. Every pool and plane is compared -bitwise in full, so a write to a slot the step does not own fails too. The step is also -run sync-free and under CUDA-graph capture and replay with the accepted counts changed -in place. Each check runs over the three handler forms: per-layer plane staging, the -batched plane gather / restore, and the fused one-launch staging and commit. +"""GPU test for the default target's speculative rollback handler (`SpeculativeRollback`). + +A verify step is simulated with the state writes the forward makes over every draft row: +the n-gram id history (`ngram_token_history`), the PLE dilated-conv state (a plain-torch +stand-in for the conv inside `ple_gate_conv_inject_norm`: the state after a run of rows is +the last `state_len` columns of `[carried state | the rows' conv input]`), the QSA index +store, and an anchor-slot overwrite standing in for the compression. The handler is staged +at the forward's call sites and committed afterwards, as the spec worker commits it. + +The reference shares no code with the commit: it runs the same writes from the same +pre-step state over each generation request's accepted prefix only, which is by definition +the state a verify step must leave behind. Every pool and plane is compared bitwise in full, +so a write to a slot the step does not own fails too. The step also runs sync-free and under +CUDA-graph capture and replay with the accepted counts changed in place. Each check runs over +both handler forms: the batched plane gather / restore with torch PLE staging (the prefill +target's and an eager decode step's), and the one-launch staging and commit (a captured +decode step's). """ +from __future__ import annotations + import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.modeling import ( # noqa: E501 - SpeculativeRollback, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_index_store import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_paged_index_store import ( # noqa: E501 qsa_paged_index_store, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ngram_token_history import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ple.ngram_token_history import ( # noqa: E501 ngram_token_history, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.dilated_short_conv_paged import ( # noqa: E501 - dilated_short_conv_paged, +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.modeling import ( # noqa: E501 + SpeculativeRollback, ) assert torch.cuda.is_available(), "the rollback handler test requires a CUDA device" @@ -52,12 +54,12 @@ class _Batch: - """One verify step: `ctx_lens` context requests (fresh) followed by - generation requests of `p` rows each, at `starts` cached tokens, in the - recurrent `slots` given (contexts first). Generation request `s` reuses the - block-table row of request `dup_of[s]` when given -- the CUDA-graph padding - dummy, repeated. With `shared_page` the live index and position planes are strided - views of one per-slot buffer, as an NVFP4 KV cache's manager serves them.""" + """One verify step: `ctx_lens` context requests (fresh) followed by generation requests + of `p` rows each, at `starts` cached tokens, in the recurrent `slots` given (contexts + first). Generation request `s` reuses the block-table row of request `dup_of[s]` when + given -- the CUDA-graph padding dummy, repeated. With `shared_page` the live index and + position planes are strided views of one per-slot buffer, as an NVFP4 KV cache's manager + serves them.""" def __init__(self, seed, ctx_lens, gen_starts, p, slots, dup_of=None, shared_page=False): g = torch.Generator().manual_seed(seed) @@ -99,7 +101,6 @@ def __init__(self, seed, ctx_lens, gen_starts, p, slots, dup_of=None, shared_pag self.position0 = torch.randint(0, 1 << 20, (PAGES, TPB, 3), generator=g).to( DEV, torch.int32 ) - self.weight = (0.3 * torch.randn(CHANNELS, 1, KERNEL, generator=g)).to(DEV, torch.bfloat16) # Step inputs. self.ids = torch.randint(0, VOCAB, (self.rows,), generator=g).to(DEV, torch.int32) @@ -166,9 +167,22 @@ def _qsl(lengths): return torch.tensor(starts, dtype=torch.int64, device=DEV) +def _conv_state(conv, x, lengths, slots, carry): + """The PLE conv pool after each sequence's rows: the last `STATE_LEN` columns of its + carried state (zeros for a fresh one) followed by its rows' conv input, channel-major. + Host-sized slices only, so it captures.""" + start = 0 + for s, length in enumerate(lengths): + slot = slots[s : s + 1].long() + prev = conv[slot][0] if carry[s] else torch.zeros_like(conv[0]) + window = torch.cat([prev, x[start : start + length].t()], 1) + conv[slot] = window[:, -STATE_LEN:].unsqueeze(0) + start += length + + def _write_planes(b, layer_index, position, req_idx, logical, keys, anchor_keys, coords): - """The index store, then a compression stand-in that overwrites each anchor - row's own slot -- the two writers of an attention layer's planes.""" + """The index store, then a compression stand-in that overwrites each anchor row's own + slot -- the two writers of an attention layer's planes.""" qsa_paged_index_store(keys, coords, layer_index, position, req_idx, logical, b.block_table, TPB) pages = b.block_table[req_idx.long(), logical // TPB].long() within = logical % TPB @@ -176,13 +190,12 @@ def _write_planes(b, layer_index, position, req_idx, logical, keys, anchor_keys, layer_index[pages, within] = torch.where(anchor, anchor_keys, layer_index[pages, within]) -def _verify_step(b, rollback, fused=False): - """The forward's write sequence over all `p` rows of every generation - request, staged at the forward's call sites, then the worker's commit. The - PLE pools stage as the fused PLE launches do: `stage_ple_verify` before the - id history publishes, `stage_ple_conv_input` from the conv input.""" +def _verify_step(b, rollback, fused): + """The forward's write sequence over all `p` rows of every generation request, staged + at the forward's call sites, then the worker's commit: `stage_ple_verify` before the id + history publishes, `stage_ple_conv_input` from the conv input, `stage_index` before + each attention layer's store.""" rollback.discard() - qsl = b.qsl rollback.begin( num_contexts=b.num_contexts, prefix_tokens=b.prefix, @@ -196,11 +209,9 @@ def _verify_step(b, rollback, fused=False): fused=fused, ) rollback.stage_ple_verify(b.ngram, b.conv, b.ids) - ngram_token_history(b.ids, b.ngram, b.slots, b.carry, qsl, EOS, NGRAM) + ngram_token_history(b.ids, b.ngram, b.slots, b.carry, b.qsl, EOS, NGRAM) rollback.stage_ple_conv_input(b.x) - dilated_short_conv_paged( - b.x.t(), b.weight, None, b.conv, qsl, b.slots, b.carry, DILATION, True, -1 - ) + _conv_state(b.conv, b.x, b.lengths, b.slots, [False] * b.num_contexts + [True] * b.n) for layer in range(LAYERS): rollback.stage_index(b.index[layer], b.position) _write_planes( @@ -217,8 +228,8 @@ def _verify_step(b, rollback, fused=False): def _reference(b): - """The state after running the same ops from the pre-step state over the - context rows and each generation request's accepted prefix only.""" + """The state after running the same writes from the pre-step state over the context + rows and each generation request's accepted prefix only.""" keep = list(range(b.prefix)) lengths = list(b.ctx_lens) for s, a in enumerate(b.gen_accepted): @@ -226,14 +237,11 @@ def _reference(b): keep += range(first, first + a) lengths.append(a) keep = torch.tensor(keep, dtype=torch.long, device=DEV) - qsl = _qsl(lengths) ngram, conv = b.ngram0.clone(), b.conv0.clone() index = [t.clone() for t in b.index0] position = b.position0.clone() - ngram_token_history(b.ids[keep], ngram, b.slots, b.carry, qsl, EOS, NGRAM) - dilated_short_conv_paged( - b.x[keep].t(), b.weight, None, conv, qsl, b.slots, b.carry, DILATION, True, -1 - ) + ngram_token_history(b.ids[keep], ngram, b.slots, b.carry, _qsl(lengths), EOS, NGRAM) + _conv_state(conv, b.x[keep], lengths, b.slots, [False] * b.num_contexts + [True] * b.n) for layer in range(LAYERS): _write_planes( b, @@ -261,12 +269,13 @@ def _assert_state(b, expected, what): assert torch.equal(b.position, position), f"{what}: position plane" -# (ctx_lens, gen_starts, p, gen_accepted). Positions straddle pages and land on -# compression anchors; one request accepts everything. All three are the accepted-prefix check; -# `decode_p8` also drives the shared-page and capture tests. +# (ctx_lens, gen_starts, p, gen_accepted). Positions straddle pages and land on compression +# anchors; one request accepts everything. `decode_p8` also drives the shared-page test. CASES = { "decode": ((), (5, 30, 61, 100), 4, (1, 2, 4, 3)), "mixed": ((7, 3), (12, 29, 95), 4, (2, 1, 3)), + # max_draft_len 1: the shortest verify step. + "decode_p2": ((), (3, 31, 64), 2, (1, 2, 1)), # max_draft_len 7: every acceptance 1..8 once; blocks at 27, 60 and 90 straddle a page, # and an 8-row block completes two compression groups. "decode_p8": ((), (0, 27, 60, 90, 101, 5, 44, 119), 8, (1, 2, 3, 4, 5, 6, 7, 8)), @@ -277,26 +286,25 @@ def _slots(num): return list(range(3, 3 + 2 * num, 2)) -#: Handler forms: per-layer plane staging; the planes bound (`begin` gathers them all in one -#: launch and the commit restores them in one, `qsa_rollback_gather` / `qsa_rollback_restore`); -#: the planes bound and the step opened `fused` (`spec_rollback_stage` / `spec_rollback_commit`). -MODES = ["per_layer", "batched", "fused"] +#: Handler forms: torch PLE staging with the planes gathered / restored in one launch each +#: (`qsa_rollback_gather` / `qsa_rollback_restore`), and the step opened `fused` +#: (`spec_rollback_stage` / `spec_rollback_commit`). +MODES = ["batched", "fused"] -def _rollback(b, mode): +def _rollback(b): rollback = SpeculativeRollback() - if mode != "per_layer": - rollback.set_index_planes(b.index) + rollback.set_index_planes(b.index) return rollback @pytest.mark.parametrize("mode", MODES) -@pytest.mark.parametrize("case", ["decode", "mixed", "decode_p8"]) +@pytest.mark.parametrize("case", sorted(CASES)) def test_commit_matches_the_accepted_prefix(case, mode): ctx_lens, gen_starts, p, gen_accepted = CASES[case] b = _Batch(11, ctx_lens, gen_starts, p, _slots(len(ctx_lens) + len(gen_starts))) b.set_accepted(gen_accepted) - _verify_step(b, _rollback(b, mode), mode == "fused") + _verify_step(b, _rollback(b), mode == "fused") _assert_state(b, _reference(b), case) @@ -312,7 +320,7 @@ def test_commit_over_the_strided_planes_of_a_shared_page(mode): 13, ctx_lens, gen_starts, p, _slots(len(ctx_lens) + len(gen_starts)), shared_page=True ) b.set_accepted(gen_accepted) - _verify_step(b, _rollback(b, mode), mode == "fused") + _verify_step(b, _rollback(b), mode == "fused") _assert_state(b, _reference(b), case) @@ -320,8 +328,7 @@ def test_commit_without_staging_is_a_no_op(): ctx_lens, gen_starts, p, gen_accepted = CASES["decode"] b = _Batch(14, ctx_lens, gen_starts, p, _slots(len(gen_starts))) b.set_accepted(gen_accepted) - rollback = SpeculativeRollback() - rollback.commit_speculative_states(b.accepted, b.slots, 0) + _rollback(b).commit_speculative_states(b.accepted, b.slots, 0) _assert_state(b, (b.ngram0, b.conv0, b.index0, b.position0), "unstaged") @@ -330,7 +337,7 @@ def test_discard_drops_a_stale_staging(): ctx_lens, gen_starts, p, gen_accepted = CASES["decode"] b = _Batch(15, ctx_lens, gen_starts, p, _slots(len(gen_starts))) b.set_accepted(gen_accepted) - rollback = SpeculativeRollback() + rollback = _rollback(b) rollback.begin( num_contexts=0, prefix_tokens=0, @@ -347,16 +354,34 @@ def test_discard_drops_a_stale_staging(): _assert_state(b, (b.ngram0, b.conv0, b.index0, b.position0), "discarded") +def test_a_step_without_bound_planes_is_refused(): + """The trunk's bind gives the handler every attention layer's index plane before any + verify step; a step opened without them would leave the planes unrolled.""" + ctx_lens, gen_starts, p, _ = CASES["decode"] + b = _Batch(19, ctx_lens, gen_starts, p, _slots(len(gen_starts))) + with pytest.raises(AssertionError, match="index planes are not bound"): + SpeculativeRollback().begin( + num_contexts=0, + prefix_tokens=0, + num_gens=b.n, + rows_per_gen=p, + slots=b.slots, + request_indices=b.req_idx, + logical_positions=b.logical, + block_table=b.block_table, + tokens_per_block=TPB, + ) + + @pytest.mark.parametrize("mode", MODES) def test_commit_is_sync_free(mode): ctx_lens, gen_starts, p, gen_accepted = CASES["mixed"] b = _Batch(16, ctx_lens, gen_starts, p, _slots(len(ctx_lens) + len(gen_starts))) b.set_accepted(gen_accepted) - rollback = _rollback(b, mode) - if mode == "fused": - # Triton compiles the launches on their first call, outside the checked region. - _verify_step(b, rollback, True) - b.reset() + rollback = _rollback(b) + # Triton compiles the launches on their first call, outside the checked region. + _verify_step(b, rollback, mode == "fused") + b.reset() torch.cuda.synchronize() torch.cuda.set_sync_debug_mode("error") try: @@ -368,15 +393,14 @@ def test_commit_is_sync_free(mode): @pytest.mark.parametrize("mode", MODES) def test_padding_duplicates_leave_the_unique_request_exact(mode): - # Requests 2 and 3 are the same padding dummy as request 1: one slot, one - # block-table row, written concurrently with no defined winner, so only the - # unique request is compared. + # Requests 2 and 3 are the same padding dummy as request 1: one slot, one block-table + # row, written concurrently with no defined winner, so only the unique request is + # compared. gen_starts, p = (5, 40, 40, 40), 4 slots = [3, 7, 7, 7] b = _Batch(17, (), gen_starts, p, slots, dup_of={2: 1, 3: 1}) b.set_accepted((2, 3, 3, 3)) - _verify_step(b, _rollback(b, mode), mode == "fused") - # The unique request's slot is exactly its accepted prefix. + _verify_step(b, _rollback(b), mode == "fused") pages = b.block_table[0].long() expected = _reference(_single(b, 0)) assert torch.equal(b.ngram[3], expected[0][3]) @@ -409,8 +433,8 @@ def _single(b, s): return one -#: Per case, the acceptance counts one captured graph is replayed with (the accepted count -#: of every generation request, bonus token included). +#: Per case, the acceptance counts one captured graph is replayed with (the accepted count of +#: every generation request, bonus token included). _REPLAYS = { "decode_p8": [ (1,) * 8, @@ -424,11 +448,11 @@ def _single(b, s): @pytest.mark.parametrize("mode", MODES) @pytest.mark.parametrize("case", sorted(_REPLAYS)) def test_capture_and_replay_follow_the_accepted_counts(case, mode): - """The whole staged step and commit captured once, replayed with the - accepted counts and the pre-step state rewritten in place.""" + """The whole staged step and commit captured once, replayed with the accepted counts and + the pre-step state rewritten in place.""" ctx_lens, gen_starts, p, _ = CASES[case] b = _Batch(18, ctx_lens, gen_starts, p, _slots(len(gen_starts))) - rollback = _rollback(b, mode) + rollback = _rollback(b) fused = mode == "fused" b.set_accepted((p,) * len(gen_starts)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_mtp_draft_hc.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_mtp_draft_hc.py new file mode 100644 index 000000000000..79aa96a99447 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_mtp_draft_hc.py @@ -0,0 +1,382 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MTP draft layer: its hyper-connection sites on the trunk's fused entries, against the +HuggingFace math in float64 and against the composition they replace, and its token embedding +under rejection sampling. + +`_mtp_front`, `_mtp_tail` and `MTPLayer.shared_head` run over a stand-in core that holds the +draft layer's hyper-connection weights (synthetic, at the checkpoint's element statistics): the +attention-side site with the front end's `x_h + x_e` as its injection at unit weight, the +MoE-side site after the output projection, and the head's stack-final site. `_mtp_tail`'s MoE +is the identity on its mixed rows and its TP sum the identity: the sites around them are under +test, at the rows a draft step runs -- 1 (the later draft steps and every shared head), 4 (the +first draft step of an MTP-3 verify of one request), 32 and 1024 (larger batches -- an +attention-DP rank's decode batch -- and a prefill). + +The float64 reference rounds to bf16 where the runtime's activations are bf16 between modules +(the norm outputs, the projections, the residual) and once at each mixer output; the mixers of +the fused path keep their logits in fp32 (the trunk's numerics). Each output must be within +the band of that reference that the replaced composition (per-stream `flashinfer_gemma_rmsnorm` +and `qwen4_exp_gated_residual`) reaches, widened by half. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import ( + cublas_mm as mtp_cublas_mm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_gated_residual import ( + qwen4_exp_gated_residual, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_hc_inject_norm import ( + qwen4_exp_hc_inject_norm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.hc.qwen4_exp_post_inject import ( + qwen4_exp_post_inject, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.norm.flashinfer_gemma_rmsnorm import ( + flashinfer_gemma_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501 + modeling as MOD, +) +from tensorrt_llm._torch.speculative.interface import rejection_sampling_one_model + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3), + reason="the fused hyper-connection entries are certified on sm_103", +) + +DEV = "cuda" +BF16 = torch.bfloat16 +HC = 4 +HIDDEN = 2560 +WIDTH = HC * HIDDEN +RANK = 320 +EPS = 1e-6 +#: A TP / TEP rank's gated attention output: 6 of the 24 query heads of 256. +Q_WIDTH = 1536 +VOCAB = 4096 +P = "l0_mtp_" + +# Element statistics of the checkpoint's hyper-connection tensors (the gated_residual +# entry's test): the norm weight's tail puts the normed stream at ~+-30. +DOWN_RMS, UP_RMS, INJECT_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0296, 0.0681, 0.0107, 0.477, 6.625 + + +class _Weights(dict): + """The draft layer's weights; the operands of its attention and MoE, which no site here + reads, are placeholders.""" + + def __missing__(self, key): + return torch.zeros(1, 1, dtype=BF16, device=DEV) + + +def _raw(seed: int) -> dict: + g = torch.Generator(device=DEV).manual_seed(seed) + + def normal(*shape, scale): + return (torch.randn(*shape, generator=g, device=DEV) * scale).to(BF16) + + def gamma(n): + t = normal(n, scale=GAMMA_RMS) + t[0], t[n // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX + return t + + w = { + "embed": normal(VOCAB, HIDDEN, scale=1.0), + "fhc_down": normal(RANK, WIDTH, scale=DOWN_RMS), + f"{P}pre_e_norm": gamma(HIDDEN), + f"{P}pre_h_norm": gamma(WIDTH), + f"{P}fc_e": normal(HIDDEN, HIDDEN, scale=HIDDEN**-0.5), + f"{P}fc_h": normal(HIDDEN, HIDDEN, scale=HIDDEN**-0.5), + f"{P}o_w": normal(HIDDEN, Q_WIDTH, scale=Q_WIDTH**-0.5), + } + for tag in ("ahc", "mhc", "fhc"): + w[f"{P}{tag}_norm"] = gamma(WIDTH) + w[f"{P}{tag}_down"] = normal(RANK, WIDTH, scale=DOWN_RMS) + w[f"{P}{tag}_up"] = normal(WIDTH, RANK, scale=UP_RMS) + if tag != "fhc": + w[f"{P}{tag}_inject"] = normal(HC, WIDTH, scale=INJECT_RMS) + return w + + +def _core(raw: dict): + """A stand-in core over copies of `raw`, loaded as `post_load_weights` loads the draft + layer: the first projections fused, then the views.""" + w = _Weights({k: v.clone() for k, v in raw.items()}) + core = SimpleNamespace( + w=w, + num_layers=0, + mtp_layer_idx=0, + mtp_enabled=True, + hidden=HIDDEN, + hc_count=HC, + stream=WIDTH, + eps=EPS, + hc_lowrank=RANK, + hc_first_cols=-(-(RANK + HC) // 32) * 32, + max_num_tokens=8192, + # The shared expert's operands are placeholders (`_Weights`). + shared_inter=1, + tp_sum=lambda t: t, + _hc_first=None, + ) + MOD.Qwen4ExpDefaultModelingV2._fuse_hc_first_projections(core) + core._mtp = MOD._mtp_views(core) + return core, core._mtp + + +@pytest.fixture +def moe_identity(monkeypatch): + monkeypatch.setattr(MOD, "_mtp_moe_block", lambda core, M, x, pad_rows=0: x) + + +# ------------------------------------------------------------------- references --- + + +def _bf(x: torch.Tensor) -> torch.Tensor: + return x.to(BF16) + + +def _norm64(x: torch.Tensor, gamma: torch.Tensor, groups: int) -> torch.Tensor: + """Gemma RMS norm of each of `groups` equal column groups, float64, rounded once.""" + xs = x.double().unflatten(-1, (groups, -1)) + xs = xs * torch.rsqrt(xs.pow(2).mean(-1, keepdim=True) + EPS) + return _bf(xs.flatten(-2) * (1.0 + gamma.double())) + + +def _mm64(a: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + """`a @ w^T` for an `nn.Linear` weight, float64, rounded once.""" + return _bf(a.double() @ w.double().t()) + + +def _mixer64(raw, tag, normed, inject=True): + return qwen4_exp_gated_residual.reference( + normed, + raw[f"{P}{tag}_down"], + raw[f"{P}{tag}_up"], + raw[f"{P}{tag}_inject"] if inject else None, + HC, + ) + + +def _front64(raw, ids, hs): + rows = hs.shape[0] + x_e = _mm64(_norm64(raw["embed"][ids], raw[f"{P}pre_e_norm"], 1), raw[f"{P}fc_e"]) + h_n = _norm64(hs, raw[f"{P}pre_h_norm"], HC) + x_h = _mm64(h_n.reshape(rows * HC, HIDDEN), raw[f"{P}fc_h"]).reshape(rows, WIDTH) + x = _bf(x_h.double() + x_e.double().repeat(1, HC)) + return (x, *_mixer64(raw, "ahc", _norm64(x, raw[f"{P}ahc_norm"], HC))) + + +def _inject64(x, block, inject): + gated = block.double().unsqueeze(-2) * inject.double().unsqueeze(-1) + return _bf((x.double().unflatten(-1, (HC, HIDDEN)) + gated).flatten(-2)) + + +def _tail64(raw, x, attn, inject): + x = _inject64(x, _mm64(attn, raw[f"{P}o_w"]), inject) + mixed, inject = _mixer64(raw, "mhc", _norm64(x, raw[f"{P}mhc_norm"], HC)) + return _inject64(x, mixed, inject) + + +def _head64(raw, hs): + return _mixer64(raw, "fhc", _norm64(hs, raw[f"{P}fhc_norm"], HC), inject=False)[0] + + +# The composition the fused sites replace: one `flashinfer_gemma_rmsnorm` a stream, then +# `qwen4_exp_gated_residual` (bf16 GEMM outputs), and the separate accumulate. + + +def _grouped(x, gamma): + parts = torch.split(x, [HIDDEN] * HC, 1) + return torch.cat( + [flashinfer_gemma_rmsnorm(p, g, EPS) for g, p in zip(torch.split(gamma, HIDDEN), parts)], + 1, + ) + + +def _mixer_old(raw, tag, normed, inject=True): + return qwen4_exp_gated_residual( + normed, + raw[f"{P}{tag}_down"], + raw[f"{P}{tag}_up"], + raw[f"{P}{tag}_inject"] if inject else None, + HC, + ) + + +def _front_old(raw, ids, hs): + rows = hs.shape[0] + e = torch.nn.functional.embedding(ids, raw["embed"]) + x_e = mtp_cublas_mm( + flashinfer_gemma_rmsnorm(e, raw[f"{P}pre_e_norm"], EPS), raw[f"{P}fc_e"].t() + ) + h_n = _grouped(hs, raw[f"{P}pre_h_norm"]) + x_h = mtp_cublas_mm(h_n.reshape(rows * HC, HIDDEN), raw[f"{P}fc_h"].t()).reshape(rows, WIDTH) + x = torch.add(x_h, torch.cat([x_e] * HC, -1)) + return (x, *_mixer_old(raw, "ahc", _grouped(x, raw[f"{P}ahc_norm"]))) + + +def _tail_old(raw, x, attn, inject): + x = qwen4_exp_post_inject(x, mtp_cublas_mm(attn, raw[f"{P}o_w"].t()), inject) + mixed, inject = _mixer_old(raw, "mhc", _grouped(x, raw[f"{P}mhc_norm"])) + return qwen4_exp_post_inject(x, mixed, inject) + + +def _head_old(raw, hs): + return _mixer_old(raw, "fhc", _grouped(hs, raw[f"{P}fhc_norm"]), inject=False)[0] + + +def _err(got: torch.Tensor, ref: torch.Tensor) -> float: + """Worst element's distance in units of `eps * max|ref|` (bf16 eps).""" + ref64 = ref.double() + unit = torch.finfo(BF16).eps * ref64.abs().max().item() + return (got.double() - ref64).abs().max().item() / unit + + +def _rel_l2(got: torch.Tensor, ref: torch.Tensor) -> float: + ref64 = ref.double() + return ((got.double() - ref64).norm() / ref64.norm()).item() + + +def _within(name, got, ref, old) -> None: + """`got` as close to `ref` as the replaced composition `old` is, widened by half (and + never tighter than one eps of max|ref| / 1e-3 relative L2, both below the noise of one + bf16 rounding of a mixer output).""" + assert got.shape == ref.shape and got.dtype == ref.dtype, (name, got.shape, ref.shape) + assert bool(torch.isfinite(got).all()), name + e_new, e_old = _err(got, ref), _err(old, ref) + l_new, l_old = _rel_l2(got, ref), _rel_l2(old, ref) + print( + f"{name}: max {e_new:.3f} (old {e_old:.3f}) eps*max|ref|, relL2 {l_new:.3g} (old {l_old:.3g})" + ) + assert e_new <= max(1.5 * e_old, 1.0), (name, e_new, e_old) + assert l_new <= max(1.5 * l_old, 1e-3), (name, l_new, l_old) + + +# ------------------------------------------------------------------------ tests --- + + +def _inputs(rows: int, seed: int, strided: bool = False): + """Token ids, the trunk's `[rows, stream]` residual (with `strided`, the column view of + the trunk's packed `[final | hc]` output the first draft step reads), the attention + output and an attention-site injection in the checkpoint's range.""" + g = torch.Generator(device=DEV).manual_seed(seed) + ids = torch.randint(0, VOCAB, (rows,), generator=g, device=DEV) + if strided: + packed = (torch.randn(rows, HIDDEN + WIDTH, generator=g, device=DEV) * 3).to(BF16) + hs = packed[:, HIDDEN:] + else: + hs = (torch.randn(rows, WIDTH, generator=g, device=DEV) * 3).to(BF16) + attn = torch.randn(rows, Q_WIDTH, generator=g, device=DEV).to(BF16) + inject = (0.43 + 1.15 * torch.rand(rows, HC, generator=g, device=DEV)).to(BF16) + return ids, hs, attn, inject + + +ROWS = (1, 4, 32, 1024) + + +@pytest.mark.parametrize( + "rows,strided", [(r, False) for r in ROWS] + [(4, True)], ids=lambda v: str(v) +) +def test_the_front_site(rows, strided) -> None: + raw = _raw(11) + core, M = _core(raw) + ids, hs, _, _ = _inputs(rows, 101 + rows, strided) + x, mixed, inject = MOD._mtp_front(core, M, raw["embed"], ids, hs) + ref, old = _front64(raw, ids, hs), _front_old(raw, ids, hs) + for name, got, r, o in zip(("x", "mixed", "inject"), (x, mixed, inject), ref, old): + _within(f"front[{rows}].{name}", got, r, o) + # The unit injection is the bf16 add: the residual is bitwise the replaced one's, whose + # projections are the same calls. + _, h_n = qwen4_exp_hc_inject_norm(hs, None, None, M.pre_h_norm, EPS, HC) + x_h = mtp_cublas_mm(h_n.reshape(rows * HC, HIDDEN), M.fc_h).reshape(rows, WIDTH) + e = torch.nn.functional.embedding(ids, raw["embed"]) + x_e = mtp_cublas_mm(flashinfer_gemma_rmsnorm(e, M.pre_e_norm, EPS), M.fc_e) + assert torch.equal(x, torch.add(x_h, torch.cat([x_e] * HC, -1))) + + +@pytest.mark.parametrize("rows", ROWS) +def test_the_tail_sites(rows, moe_identity) -> None: + raw = _raw(12) + core, M = _core(raw) + _, x, attn, inject = _inputs(rows, 201 + rows) + out = MOD._mtp_tail(core, M, x, attn, inject, 0) + _within(f"tail[{rows}]", out, _tail64(raw, x, attn, inject), _tail_old(raw, x, attn, inject)) + + +# 32: a shared head over one attention-DP rank's decode batch (ht_dep4 drafts under MTP). +@pytest.mark.parametrize("rows", (1, 4, 32)) +def test_the_head_site(rows) -> None: + raw = _raw(13) + core, _ = _core(raw) + _, hs, _, _ = _inputs(rows, 301 + rows) + head = SimpleNamespace( + core=core, logits_processor=SimpleNamespace(forward=lambda final, *args: final) + ) + final = MOD.MTPLayer.shared_head(head, hs, None, None, False) + _within(f"head[{rows}]", final, _head64(raw, hs), _head_old(raw, hs)) + + +def test_a_captured_draft_step_replays_the_eager_bits(moe_identity) -> None: + """Front, tail and head of a 4-row draft step captured in one CUDA graph: the replay is + bitwise the eager result.""" + raw = _raw(14) + core, M = _core(raw) + ids, hs, attn, _ = _inputs(4, 401, strided=True) + head = SimpleNamespace( + core=core, logits_processor=SimpleNamespace(forward=lambda final, *args: final) + ) + + def step(): + x, _, inject = MOD._mtp_front(core, M, raw["embed"], ids, hs) + out = MOD._mtp_tail(core, M, x, attn, inject, 0) + return out, MOD.MTPLayer.shared_head(head, out, None, None, False) + + eager = step() + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + step() + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = step() + graph.replay() + torch.cuda.synchronize() + for got, want in zip(captured, eager): + assert torch.equal(got, want) + + +def test_draft_embedding_reads_row_0_for_rejected_slots() -> None: + """Rejection sampling pads the slots past a request's accepted count with -1 (here every + draft is rejected at slot 0); the front end embeds those as row 0 and every valid id as + the plain lookup does.""" + batch, draft, vocab, hidden = 8, 3, 4096, 256 + draft_tokens = torch.arange(batch * draft, dtype=torch.int32, device="cuda").reshape( + batch, draft + ) + draft_probs = torch.zeros(batch, draft, vocab, device="cuda") + draft_probs.scatter_(2, draft_tokens.long().unsqueeze(-1), 1.0) + target_probs = torch.zeros(batch, draft + 1, vocab, device="cuda") + target_probs[..., vocab - 1] = 1.0 + seed = torch.zeros(batch, dtype=torch.int64, device="cuda") + offset = torch.arange(batch, dtype=torch.int64, device="cuda") + accepted, _ = rejection_sampling_one_model( + draft_probs, draft_tokens, target_probs, seed, offset + ) + assert bool((accepted[:, 1:] == -1).all()) and bool((accepted[:, 0] == vocab - 1).all()) + + g = torch.Generator(device="cuda").manual_seed(0) + table = torch.randn(vocab, hidden, generator=g, device="cuda").to(torch.bfloat16) + valid_ids = torch.randint(0, vocab, (257,), generator=g, device="cuda", dtype=torch.int32) + for ids in (accepted.flatten(), valid_ids, valid_ids.long()): + e = MOD._draft_embedding(ids, table) + valid = ids >= 0 + assert torch.equal(e[valid], table[ids[valid].long()]) + assert torch.equal(e[~valid], table[0].expand(int((~valid).sum()), -1)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_pool_contract.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_pool_contract.py similarity index 55% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_pool_contract.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_pool_contract.py index 2cbe7c9046bf..edf3dcd623a8 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_pool_contract.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_pool_contract.py @@ -1,10 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""qwen4_exp's KV cache pool contract rejects the mismatches the kernels would not. +"""The qwen4_exp target's KV cache pool contract (`pool.py`) rejects the mismatches the +kernels would not. The checks read only `.shape` / `.dtype`, so these run on meta tensors with no GPU and no engine. Each case is one wrong pool that the kernel consuming it would accept and -mis-address. +mis-address. The geometry is one rank's at tp1, which holds every head. """ from __future__ import annotations @@ -14,27 +15,25 @@ import pytest import torch -__extra_import_path__ = ["../../.."] # noqa: F841 -- repo's file-scoped import hook - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp import routing -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - modeling as M, +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501 + pool as P, ) -from tensorrt_llm._torch.modules.fla.cache_manager import GDNIntermediateState, GDNReplayState -from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManager pytestmark = pytest.mark.cpu_only -G = M._PoolGeometry( +G = P.PoolGeometry( tokens_per_block=32, heads_kv=2, + kv_head_start=0, head_dim=256, index_kv_heads=1, index_head_dim=128, conv_dim=8192, conv_kernel=4, num_k_heads=16, + k_head_start=0, num_v_heads=48, + v_head_start=0, head_v_dim=128, head_k_dim=128, stream=10240, @@ -74,103 +73,42 @@ def _history(slots=SLOTS, history=HISTORY): def test_valid_pools_pass() -> None: - M.check_position_plane(G, _t(9, 32, 3, dtype=torch.int32)) - M.check_attention_pools( + P.check_position_plane(G, _t(9, 32, 3, dtype=torch.int32)) + P.check_attention_pools( G, 3, _t(9, 32, 1, 128), _t(9, 2, G.heads_kv, G.tokens_per_block, G.head_dim) ) - M.check_attention_pools( + P.check_attention_pools( G_FP8, 3, _t(9, 32, 1, 128), _t(9, 2, G.heads_kv, G.tokens_per_block, G.head_dim, dtype=torch.float8_e4m3fn), ) nvfp4_kv = _t(9, 2, G.heads_kv, G.tokens_per_block, 128, dtype=torch.uint8) - M.check_attention_pools(G_NVFP4, 3, _t(9, 32, 1, 128), nvfp4_kv) + P.check_attention_pools(G_NVFP4, 3, _t(9, 32, 1, 128), nvfp4_kv) for dtype in (torch.float32, torch.bfloat16): - M.check_recurrent_pools(G, _conv(), _ssm(dtype)) - M.check_ple_pools( + P.check_recurrent_pools(G, _conv(), _ssm(dtype)) + P.check_ple_pools( G, _t(SLOTS, G.stream, G.ple_state_len), _t(SLOTS, G.ngram_context_len, dtype=torch.int64), _conv(), ) - M.check_step_buffers( + P.check_step_buffers( G, _t(SLOTS, 2, G.conv_dim, G.conv_kernel - 1), _conv(), _t(SLOTS, dtype=torch.int32) ) for dtype in (torch.float32, torch.bfloat16): - M.check_legacy_state( - G, _t(SLOTS, 2, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=dtype), _ssm(dtype) - ) - for dtype in (torch.float32, torch.bfloat16): - for width in (2, 5, M._MAX_VERIFY_WIDTH): - M.check_replay_state( + for width in (2, 5, P.MAX_VERIFY_WIDTH): + P.check_replay_state( HISTORY, width, _t(SLOTS, dtype=torch.int32), _t(SLOTS, dtype=torch.int32), _ssm(dtype), ) - M.check_replay_history(G, 3, HISTORY, SLOTS, *_history()) - for rows in range(1, M._MAX_VERIFY_WIDTH + 1): - M.check_verify_width(rows) - M.check_shared_selection(_t(2, 8, 512, dtype=torch.int32), {3: 1}, 3, 512) - - -@pytest.mark.parametrize("replay", [False, True]) -def test_bind_step_buffers_uses_manager_layer_views(replay: bool) -> None: - state = GDNReplayState(tokens_per_gen_step=2) if replay else GDNIntermediateState() - state.intermediate_conv = _t(2, SLOTS, 2, G.conv_dim, G.conv_kernel - 1) - if replay: - state.prev_num_accepted_tokens = _t(SLOTS, dtype=torch.int32) - state.cache_buf_idx = _t(SLOTS, dtype=torch.int32) - state.old_x = _t(2, SLOTS, 2, HISTORY, G.num_v_heads, G.head_v_dim) - state.old_B = _t(2, SLOTS, 2, HISTORY, G.num_k_heads, G.head_k_dim) - state.old_dt = _t(2, SLOTS, 2, G.num_v_heads, HISTORY, dtype=torch.float32) - state.old_dA_cumsum = torch.empty_like(state.old_dt) - state._state_strides = (2, 0, 1) - else: - state.intermediate_ssm = _t( - 2, SLOTS, 2, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=torch.float32 - ) - state.intermediate_indices = _t(SLOTS, dtype=torch.int32) - manager = Qwen4ExpHybridCacheManager.__new__(Qwen4ExpHybridCacheManager) - manager.spec_config = SimpleNamespace(tokens_per_gen_step=2) - manager._speculative_state = state - manager.mamba_layer_offsets = {3: 1, 7: 0} - manager.all_conv_states = [_conv(), _conv()] - manager.all_ssm_states = [_ssm(), _ssm()] - - core = M.Qwen4ExpTp1ModelingV2.__new__(M.Qwen4ExpTp1ModelingV2) - torch.nn.Module.__init__(core) - core.gdn_layers = [3, 7] - metadata = SimpleNamespace( - replay_work_items=_t(SLOTS, 4, dtype=torch.int32), - replay_n_writes=_t(1, dtype=torch.int32), - ) - bound = core._bind_step_buffers(manager, G, metadata, _conv(), _ssm()) - - assert bound["int_idx"] is state.intermediate_indices - for layer, offset in manager.mamba_layer_offsets.items(): - views = [(bound["int_conv"][layer], state.intermediate_conv)] - if replay: - assert "int_ssm" not in bound - replay_bound = bound["replay"] - assert replay_bound.all_layer_commit - assert replay_bound.cache_buf_idx is state.cache_buf_idx - assert replay_bound.prev_num_accepted is state.prev_num_accepted_tokens - views.extend( - zip( - replay_bound.old[layer], - (state.old_x, state.old_B, state.old_dt, state.old_dA_cumsum), - ) - ) - else: - assert "replay" not in bound - views.append((bound["int_ssm"][layer], state.intermediate_ssm)) - for view, storage in views: - assert view._base is storage - assert view.storage_offset() == offset * storage.stride(0) - assert view.shape == storage.shape[1:] + P.check_replay_history(G, 3, HISTORY, SLOTS, *_history()) + for rows in range(1, P.MAX_VERIFY_WIDTH + 1): + P.check_verify_width(rows) + P.check_shared_selection(_t(2, 8, 512, dtype=torch.int32), {3: 1}, 3, 512) @pytest.mark.parametrize( @@ -184,23 +122,23 @@ def test_bind_step_buffers_uses_manager_layer_views(replay: bool) -> None: ) def test_position_plane_rejects(plane, match) -> None: with pytest.raises(AssertionError, match=match): - M.check_position_plane(G, plane) + P.check_position_plane(G, plane) def test_attention_pools_reject_wrong_index_plane_and_kv_layout() -> None: kv = _t(9, 2, G.heads_kv, G.tokens_per_block, G.head_dim) with pytest.raises(AssertionError, match="QSA index plane"): - M.check_attention_pools(G, 3, _t(9, 32, 1, 64), kv) + P.check_attention_pools(G, 3, _t(9, 32, 1, 64), kv) with pytest.raises(AssertionError, match="QSA index plane"): - M.check_attention_pools(G, 3, _t(9, 32, 1, 128, dtype=torch.float16), kv) + P.check_attention_pools(G, 3, _t(9, 32, 1, 128, dtype=torch.float16), kv) with pytest.raises(AssertionError, match="paged K/V"): - M.check_attention_pools(G, 3, _t(9, 32, 1, 128), _t(9, 2, 4, 32, 256)) + P.check_attention_pools(G, 3, _t(9, 32, 1, 128), _t(9, 2, 4, 32, 256)) # NHD instead of HND: heads and block swapped. with pytest.raises(AssertionError, match="paged K/V"): - M.check_attention_pools(G, 3, _t(9, 32, 1, 128), _t(9, 2, 32, 2, 256)) + P.check_attention_pools(G, 3, _t(9, 32, 1, 128), _t(9, 2, 32, 2, 256)) # K and V not adjacent role planes. with pytest.raises(AssertionError, match="role planes"): - M.check_attention_pools(G, 3, _t(9, 32, 1, 128), _t(9, 3, 2, 32, 256)) + P.check_attention_pools(G, 3, _t(9, 32, 1, 128), _t(9, 3, 2, 32, 256)) def test_attention_pools_reject_the_other_storage_dtype() -> None: @@ -210,21 +148,21 @@ def test_attention_pools_reject_the_other_storage_dtype() -> None: bf16_kv = _t(9, 2, G.heads_kv, G.tokens_per_block, G.head_dim) fp8_kv = bf16_kv.to(torch.float8_e4m3fn) with pytest.raises(AssertionError, match="paged K/V"): - M.check_attention_pools(G, 3, index, fp8_kv) + P.check_attention_pools(G, 3, index, fp8_kv) with pytest.raises(AssertionError, match="paged K/V"): - M.check_attention_pools(G, 3, index, bf16_kv.to(torch.float16)) + P.check_attention_pools(G, 3, index, bf16_kv.to(torch.float16)) with pytest.raises(AssertionError, match="paged K/V"): - M.check_attention_pools(G_FP8, 3, index, bf16_kv) + P.check_attention_pools(G_FP8, 3, index, bf16_kv) # The index plane is a kernel contract, not KV storage: bf16 under either algorithm. with pytest.raises(AssertionError, match="QSA index plane"): - M.check_attention_pools(G_FP8, 3, index.to(torch.float8_e4m3fn), fp8_kv) + P.check_attention_pools(G_FP8, 3, index.to(torch.float8_e4m3fn), fp8_kv) # An NVFP4 pool is half a head wide: a full-width uint8 pool is the other storage. with pytest.raises(AssertionError, match="paged K/V"): - M.check_attention_pools( + P.check_attention_pools( G_NVFP4, 3, index, _t(9, 2, G.heads_kv, G.tokens_per_block, 256, dtype=torch.uint8) ) with pytest.raises(AssertionError, match="paged K/V"): - M.check_attention_pools( + P.check_attention_pools( G, 3, index, _t(9, 2, G.heads_kv, G.tokens_per_block, 128, dtype=torch.uint8) ) @@ -244,128 +182,111 @@ def test_kv_scale_planes_must_hold_the_data_pools_pages_per_slot() -> None: against 24576 of scales.""" data = _strided((9, 2, G.heads_kv, G.tokens_per_block, 128), 24 * 8192) scales = _strided((9, 2, G.heads_kv, G.tokens_per_block, 16), 24 * 1024) - M.check_kv_scale_planes(G_NVFP4, 3, data, scales) + P.check_kv_scale_planes(G_NVFP4, 3, data, scales) # A side page coalesced into the data pool: 36 pages per slot against 24. coalesced = _strided((9, 2, G.heads_kv, G.tokens_per_block, 128), 36 * 8192) with pytest.raises(AssertionError, match="36 pages per slot.*24"): - M.check_kv_scale_planes(G_NVFP4, 3, coalesced, scales) + P.check_kv_scale_planes(G_NVFP4, 3, coalesced, scales) with pytest.raises(AssertionError, match="whole number"): - M.check_kv_scale_planes(G_NVFP4, 3, data, _strided((9, 2, 2, 32, 16), 24 * 1024 + 1)) + P.check_kv_scale_planes(G_NVFP4, 3, data, _strided((9, 2, 2, 32, 16), 24 * 1024 + 1)) def test_kv_scale_planes_reject_the_wrong_block_scale_width() -> None: data = _t(9, 2, G.heads_kv, G.tokens_per_block, 128, dtype=torch.uint8) with pytest.raises(AssertionError, match="NVFP4 block scales"): - M.check_kv_scale_planes(G_NVFP4, 3, data, _t(9, 2, 2, 32, 8, dtype=torch.uint8)) + P.check_kv_scale_planes(G_NVFP4, 3, data, _t(9, 2, 2, 32, 8, dtype=torch.uint8)) def test_recurrent_pools_reject_wrong_key_dim_order_and_dtype() -> None: with pytest.raises(AssertionError, match="recurrent state pool"): - M.check_recurrent_pools( + P.check_recurrent_pools( G, _conv(), _t(SLOTS, G.num_v_heads, G.head_k_dim, G.head_v_dim + 1) ) with pytest.raises(AssertionError, match="conv state pool"): - M.check_recurrent_pools(G, _conv(dtype=torch.float32), _ssm()) + P.check_recurrent_pools(G, _conv(dtype=torch.float32), _ssm()) # The engine accepts an fp16 pool; no recurrence entry the target binds is certified at it. with pytest.raises(AssertionError, match="recurrent state pool"): - M.check_recurrent_pools(G, _conv(), _ssm(torch.float16)) + P.check_recurrent_pools(G, _conv(), _ssm(torch.float16)) def test_ple_pools_reject_slot_mismatch_and_int32_history() -> None: ple_conv = _t(SLOTS, G.stream, G.ple_state_len) ngram = _t(SLOTS, G.ngram_context_len, dtype=torch.int64) with pytest.raises(AssertionError, match="one slot space"): - M.check_ple_pools( + P.check_ple_pools( G, ple_conv, _t(SLOTS + 1, G.ngram_context_len, dtype=torch.int64), _conv() ) with pytest.raises(AssertionError, match="n-gram context pool"): - M.check_ple_pools(G, ple_conv, _t(SLOTS, G.ngram_context_len, dtype=torch.int32), _conv()) + P.check_ple_pools(G, ple_conv, _t(SLOTS, G.ngram_context_len, dtype=torch.int32), _conv()) with pytest.raises(AssertionError, match="did not allocate"): - M.check_ple_pools(G, None, ngram, _conv()) + P.check_ple_pools(G, None, ngram, _conv()) def test_step_buffers_reject_missing_conv_rank_dtype_and_index() -> None: int_conv = _t(SLOTS, 2, G.conv_dim, G.conv_kernel - 1) idx = _t(SLOTS, dtype=torch.int32) with pytest.raises(AssertionError, match="no per-step conv"): - M.check_step_buffers(G, None, _conv(), idx) + P.check_step_buffers(G, None, _conv(), idx) # Rank without the steps dim. with pytest.raises(AssertionError, match="per-step conv buffer"): - M.check_step_buffers(G, _conv(), _conv(), idx) + P.check_step_buffers(G, _conv(), _conv(), idx) # The promotion kernel copies without converting, so the dtype must match the pool. with pytest.raises(AssertionError, match="per-step conv buffer"): - M.check_step_buffers(G, int_conv.to(torch.float32), _conv(), idx) + P.check_step_buffers(G, int_conv.to(torch.float32), _conv(), idx) with pytest.raises(AssertionError, match="intermediate_state_indices"): - M.check_step_buffers(G, int_conv, _conv(), idx.to(torch.int64)) + P.check_step_buffers(G, int_conv, _conv(), idx.to(torch.int64)) with pytest.raises(AssertionError, match="intermediate_state_indices"): - M.check_step_buffers(G, int_conv, _conv(), None) - - -def test_legacy_state_rejects_missing_rank_dtype_drift_and_too_many_rows() -> None: - def per_step(steps): - return _t(SLOTS, steps, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=torch.float32) - - int_ssm = per_step(2) - with pytest.raises(AssertionError, match="no per-step recurrent"): - M.check_legacy_state(G, None, _ssm()) - with pytest.raises(AssertionError, match="per-step recurrent buffer"): - M.check_legacy_state(G, _ssm(), _ssm()) - with pytest.raises(AssertionError, match="per-step recurrent buffer"): - M.check_legacy_state(G, int_ssm.to(torch.bfloat16), _ssm()) - # More rows per request than the legacy recurrence is certified at. - M.check_legacy_state(G, per_step(M._MAX_LEGACY_VERIFY_WIDTH), _ssm()) - with pytest.raises(AssertionError, match="TRTLLM_USE_GDN_REPLAY"): - M.check_legacy_state(G, per_step(M._MAX_LEGACY_VERIFY_WIDTH + 1), _ssm()) + P.check_step_buffers(G, int_conv, _conv(), None) def test_replay_state_rejects_pool_dtype_geometry_and_bookkeeping() -> None: pnat = _t(SLOTS, dtype=torch.int32) buf = _t(SLOTS, dtype=torch.int32) with pytest.raises(AssertionError, match="fp32 or a bf16 checkpoint"): - M.check_replay_state(HISTORY, 3, pnat, buf, _ssm(torch.float16)) + P.check_replay_state(HISTORY, 3, pnat, buf, _ssm(torch.float16)) with pytest.raises(AssertionError, match="replay_history_size"): - M.check_replay_state(12, 3, pnat, buf, _ssm()) - for width in (1, M._MAX_VERIFY_WIDTH + 1): + P.check_replay_state(12, 3, pnat, buf, _ssm()) + for width in (1, P.MAX_VERIFY_WIDTH + 1): with pytest.raises(AssertionError, match="replay_step_width"): - M.check_replay_state(HISTORY, width, pnat, buf, _ssm()) + P.check_replay_state(HISTORY, width, pnat, buf, _ssm()) # The bookkeeping is one int32 per pool slot, so it must cover the whole pool. with pytest.raises(AssertionError, match="prev_num_accepted_tokens"): - M.check_replay_state(HISTORY, 3, _t(SLOTS - 1, dtype=torch.int32), buf, _ssm()) + P.check_replay_state(HISTORY, 3, _t(SLOTS - 1, dtype=torch.int32), buf, _ssm()) with pytest.raises(AssertionError, match="cache_buf_idx"): - M.check_replay_state(HISTORY, 3, pnat, buf.to(torch.int64), _ssm()) + P.check_replay_state(HISTORY, 3, pnat, buf.to(torch.int64), _ssm()) def test_verify_width_admits_one_to_eight_rows_per_request() -> None: - for rows in (0, M._MAX_VERIFY_WIDTH + 1): + for rows in (0, P.MAX_VERIFY_WIDTH + 1): with pytest.raises(AssertionError, match="predicted_tokens_per_seq"): - M.check_verify_width(rows) + P.check_verify_width(rows) def test_replay_width_admits_only_the_managers_step_width() -> None: replay = SimpleNamespace(width=3) - M.check_replay_width(None, 1) # the legacy recurrence takes any width - M.check_replay_width(replay, 3) + P.check_replay_width(None, 1) # without the cached replay any width + P.check_replay_width(replay, 3) for rows in (1, 2, 4): with pytest.raises(AssertionError, match="step width 3"): - M.check_replay_width(replay, rows) + P.check_replay_width(replay, rows) def test_replay_history_rejects_missing_short_pool_and_wrong_geometry() -> None: old_x, old_B, old_dt = _history() with pytest.raises(AssertionError, match="no old_x"): - M.check_replay_history(G, 3, HISTORY, SLOTS, None, old_B, old_dt) + P.check_replay_history(G, 3, HISTORY, SLOTS, None, old_B, old_dt) with pytest.raises(AssertionError, match="slots for a"): - M.check_replay_history(G, 3, HISTORY, SLOTS, *_history(slots=SLOTS - 1)) + P.check_replay_history(G, 3, HISTORY, SLOTS, *_history(slots=SLOTS - 1)) # The kernel addresses the history by `history_size`, so a buffer of another depth # would be read at the wrong offsets. with pytest.raises(AssertionError, match="old_x"): - M.check_replay_history(G, 3, HISTORY, SLOTS, *_history(history=12)) + P.check_replay_history(G, 3, HISTORY, SLOTS, *_history(history=12)) with pytest.raises(AssertionError, match="old_B"): - M.check_replay_history(G, 3, HISTORY, SLOTS, old_x, old_B.to(torch.float32), old_dt) + P.check_replay_history(G, 3, HISTORY, SLOTS, old_x, old_B.to(torch.float32), old_dt) with pytest.raises(AssertionError, match="old_dt"): - M.check_replay_history(G, 3, HISTORY, SLOTS, old_x, old_B, old_dt.to(torch.bfloat16)) + P.check_replay_history(G, 3, HISTORY, SLOTS, old_x, old_B, old_dt.to(torch.bfloat16)) with pytest.raises(AssertionError, match="double-buffered"): - M.check_replay_history( + P.check_replay_history( G, 3, HISTORY, @@ -379,22 +300,10 @@ def test_replay_history_rejects_missing_short_pool_and_wrong_geometry() -> None: def test_shared_selection_rejects_missing_row_and_width() -> None: shared = _t(2, 8, 512, dtype=torch.int32) with pytest.raises(AssertionError, match="no shared buffers"): - M.check_shared_selection(None, {3: 0}, 3, 512) + P.check_shared_selection(None, {3: 0}, 3, 512) with pytest.raises(AssertionError, match="no row"): - M.check_shared_selection(shared, {4: 0}, 3, 512) + P.check_shared_selection(shared, {4: 0}, 3, 512) with pytest.raises(AssertionError, match="no row"): - M.check_shared_selection(shared, {3: 2}, 3, 512) + P.check_shared_selection(shared, {3: 2}, 3, 512) with pytest.raises(AssertionError, match="blocks per row"): - M.check_shared_selection(shared, {3: 1}, 3, 256) - - -def test_routing_admits_exactly_what_the_target_implements() -> None: - """`route` declines a MoE backend, a draft length or a recurrent-state dtype the - constructor or the first bind would assert on. `test_modeling_v2_routing.py` pins - routing's behavior at today's values; this keeps routing's constants and the target's - from drifting apart.""" - assert set(routing._MOE_BACKENDS) == set(M._TRUNK_MOE_BACKENDS) - assert routing._MAX_DRAFT_LEN == M._MAX_VERIFY_WIDTH - 1 - assert routing._MAX_LEGACY_DRAFT_LEN == M._MAX_LEGACY_VERIFY_WIDTH - 1 - assert routing._MAX_CUTEDSL_DRAFT_LEN == M._CUTEDSL_MAX_DRAFT_LEN - assert routing._SSM_STATE_DTYPES == M._SSM_STATE_DTYPES + P.check_shared_selection(shared, {3: 1}, 3, 256) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_pool_layout.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_pool_layout.py new file mode 100644 index 000000000000..e84867c58dc6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_pool_layout.py @@ -0,0 +1,345 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The per-rank pool layout this target binds at tp1, tp4, tep4, dep4 and dep2, against the +stock manager. + +`pool.pool_geometry` derives one rank's pool shapes from the text config and the rank's `Mapping`. +These tests hold it to `LAYOUT` (what `core` binds at each member), then build the stock +`Qwen4ExpHybridCacheManager` the way the engine does -- `get_kv_cache_manager_cls` and +`_create_kv_cache_manager` over a `ModelConfig`, with the core runtime's cache configuration +(bf16 recurrent state, block reuse off, KV cache manager V2, which the engine resolves `auto` to +for this model) -- on the first and the last rank of each member, one GPU, and take every pool +view through `pool.take_pools`, which checks each against the geometry. Ranks 0 and 3 hold +different KV and GDN heads at tp4 and tep4; the shapes are the same, the heads are not. dep2, the +layout of both disaggregated-serving engines, binds dep4's geometry on both of its ranks, so the +transceiver moves pools of one layout between identical layouts. tp1 holds every head, the +geometry dep4 and dep2 bind on each rank. +""" + +import contextlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501 + pool as P, +) +from tensorrt_llm._torch.configs import Qwen4ExpTextConfig +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_utils import duplicate_kv_weight +from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManager +from tensorrt_llm._torch.pyexecutor._util import _create_kv_cache_manager, get_kv_cache_manager_cls +from tensorrt_llm._torch.pyexecutor.model_loader import validate_and_set_mamba_ssm_cache_dtype +from tensorrt_llm.bindings import DataType +from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MTPDecodingConfig, QSASparseAttentionConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +NUM_LAYERS = 48 + +# Qwen3.8-Flash-Next's text config: the fields the manager and the geometry read, and what the +# QSA lowering and the config class need beside them. +_TEXT = dict( + architectures=["Qwen4ExpForCausalLM"], + dtype="bfloat16", + full_attention_interval=4, + hc_count=4, + hc_lowrank=320, + head_dim=256, + heads_per_ngram=8, + hidden_size=2560, + indexer_budget=2048, + indexer_compress_ratio=4, + indexer_head_dim=128, + indexer_kv_heads=1, + indexer_n_heads=4, + layer_types=["full_attention" if i % 4 == 3 else "linear_attention" for i in range(NUM_LAYERS)], + linear_conv_kernel_dim=4, + linear_key_head_dim=128, + linear_num_key_heads=16, + linear_num_value_heads=48, + linear_value_head_dim=128, + make_ngram_vocab_size_divisible_by=128, + mamba_ssm_dtype="float32", + max_position_embeddings=262144, + moe_intermediate_size=640, + mtp=dict( + hybrid=True, + layer_types=["full_attention"], + mtp_use_hidden_state_from_layer=None, + num_hidden_layers=1, + rope_theta=10000000, + ), + mtp_num_hidden_layers=1, + mtp_use_dedicated_embeddings=False, + ngram_size=3, + ngram_vocab_size_base=20000000, + num_attention_heads=24, + num_experts=512, + num_experts_per_tok=10, + num_hidden_layers=NUM_LAYERS, + num_key_value_heads=2, + partial_rotary_factor=0.25, + ple_conv_kernel_size=4, + ple_embed_dim=2560, + ple_layer_ids=[2], + rms_norm_eps=1e-06, + rope_parameters=dict( + mrope_interleaved=True, + mrope_section=[11, 11, 10], + partial_rotary_factor=0.25, + rope_theta=10000000, + rope_type="default", + ), + shared_expert_intermediate_size=640, + split_ngram_parts=128, + tie_word_embeddings=False, + vocab_size=248320, +) + +ATTN_LAYERS = [i for i in range(NUM_LAYERS) if i % 4 == 3] +GDN_LAYERS = [i for i in range(NUM_LAYERS) if i % 4 != 3] +PLE_LAYER = 1 # `ple_layer_ids: [2]` is one-based +MTP_LAYER = NUM_LAYERS +TOKENS_PER_BLOCK = 32 # `kv_cache_config.tokens_per_block`'s default, which the runtime keeps + +#: Each member's mapping: one rank (tp1), attention TP4 with the experts split by intermediate +#: (tp4) or by expert (tep4), and attention DP with expert parallelism (dep4, dep2). The expert +#: split moves no pool. +MEMBERS = { + "tp1": dict(world_size=1, tp_size=1, moe_tp_size=1, moe_ep_size=1, enable_attention_dp=False), + "tp4": dict(world_size=4, tp_size=4, moe_tp_size=4, moe_ep_size=1, enable_attention_dp=False), + "tep4": dict(world_size=4, tp_size=4, moe_tp_size=1, moe_ep_size=4, enable_attention_dp=False), + "dep4": dict(world_size=4, tp_size=4, moe_tp_size=1, moe_ep_size=4, enable_attention_dp=True), + "dep2": dict(world_size=2, tp_size=2, moe_tp_size=1, moe_ep_size=2, enable_attention_dp=True), +} + +#: What `core` binds on each member and rank: KV heads and the first one's global index, GDN key +#: heads and the first, GDN value heads and the first, and the conv channels ([q | k | v]). +LAYOUT = { + ("tp1", 0): dict(heads_kv=2, kv_head_start=0, num_k_heads=16, k_head_start=0, num_v_heads=48, + v_head_start=0, conv_dim=10240), + ("tp4", 0): dict(heads_kv=1, kv_head_start=0, num_k_heads=4, k_head_start=0, num_v_heads=12, + v_head_start=0, conv_dim=2560), + ("tp4", 3): dict(heads_kv=1, kv_head_start=1, num_k_heads=4, k_head_start=12, num_v_heads=12, + v_head_start=36, conv_dim=2560), + ("tep4", 0): dict(heads_kv=1, kv_head_start=0, num_k_heads=4, k_head_start=0, num_v_heads=12, + v_head_start=0, conv_dim=2560), + ("tep4", 3): dict(heads_kv=1, kv_head_start=1, num_k_heads=4, k_head_start=12, num_v_heads=12, + v_head_start=36, conv_dim=2560), + ("dep4", 0): dict(heads_kv=2, kv_head_start=0, num_k_heads=16, k_head_start=0, num_v_heads=48, + v_head_start=0, conv_dim=10240), + ("dep4", 3): dict(heads_kv=2, kv_head_start=0, num_k_heads=16, k_head_start=0, num_v_heads=48, + v_head_start=0, conv_dim=10240), + ("dep2", 0): dict(heads_kv=2, kv_head_start=0, num_k_heads=16, k_head_start=0, num_v_heads=48, + v_head_start=0, conv_dim=10240), + ("dep2", 1): dict(heads_kv=2, kv_head_start=0, num_k_heads=16, k_head_start=0, num_v_heads=48, + v_head_start=0, conv_dim=10240), +} # fmt: skip + +#: The same on every member and rank: the page size, the head dims, the replicated QSA index and +#: PLE pools, and the activation dtype the conv, index and PLE pools hold. +COMMON = dict( + tokens_per_block=TOKENS_PER_BLOCK, + head_dim=256, + index_kv_heads=1, + index_head_dim=128, + conv_kernel=4, + head_k_dim=128, + head_v_dim=128, + stream=4 * 2560, + ple_state_len=(4 - 1) * 3, + ngram_context_len=3 - 1, + dtype=torch.bfloat16, +) + +#: The K/V pool's storage per resolved KV-cache algorithm: (pool dtype, innermost extent, the +#: manager's binding dtype). +KV_STORAGE = { + None: (torch.bfloat16, 256, DataType.BF16), + QuantAlgo.FP8: (torch.float8_e4m3fn, 256, DataType.FP8), + QuantAlgo.NVFP4: (torch.uint8, 128, DataType.NVFP4), +} +KV_IDS = {None: "kv_bf16", QuantAlgo.FP8: "kv_fp8", QuantAlgo.NVFP4: "kv_nvfp4"} + +_needs_gpu = pytest.mark.skipif( + not torch.cuda.is_available(), reason="the manager allocates its pools on the GPU" +) + + +def _mapping(member: str, rank: int) -> Mapping: + return Mapping(rank=rank, gpus_per_node=4, pp_size=1, **MEMBERS[member]) + + +def _geometry(member: str, rank: int, kv_cache_quant_algo=None) -> P.PoolGeometry: + return P.pool_geometry( + Qwen4ExpTextConfig(**_TEXT), + _mapping(member, rank), + tokens_per_block=TOKENS_PER_BLOCK, + kv_cache_quant_algo=kv_cache_quant_algo, + dtype=torch.bfloat16, + ) + + +# The first and the last rank of each member. +_MEMBER_RANKS = [ + (member, rank) for member, m in MEMBERS.items() for rank in sorted({0, m["world_size"] - 1}) +] +_MEMBER_RANK_IDS = [f"{member}-rank{rank}" for member, rank in _MEMBER_RANKS] +# dep2 serves without speculation (routing declines a `speculative_config` there). +_MTP_MEMBER_RANKS = [(member, rank) for member, rank in _MEMBER_RANKS if member != "dep2"] +_MTP_MEMBER_RANK_IDS = [f"{member}-rank{rank}" for member, rank in _MTP_MEMBER_RANKS] + + +@pytest.mark.parametrize("member, rank", _MEMBER_RANKS, ids=_MEMBER_RANK_IDS) +def test_the_geometry_is_the_layout_core_binds(member, rank): + g = _geometry(member, rank) + expected = dict(COMMON, **LAYOUT[member, rank], kv_dtype=torch.bfloat16, kv_width=256) + assert g._asdict() == expected + assert g.conv_sections == ( + g.num_k_heads * 128, + g.num_k_heads * 128, + g.num_v_heads * 128, + ) + + +@pytest.mark.parametrize("algo", list(KV_STORAGE), ids=list(KV_IDS.values())) +def test_the_kv_storage_follows_the_resolved_algorithm(algo): + g = _geometry("dep4", 0, algo) + assert (g.kv_dtype, g.kv_width) == KV_STORAGE[algo][:2] + assert g.kv_block_scales == (algo == QuantAlgo.NVFP4) + + +def test_another_kv_algorithm_is_refused(): + with pytest.raises(ValueError, match="kv_cache_quant_algo"): + _geometry("dep4", 0, QuantAlgo.INT8) + + +@pytest.mark.parametrize("member, rank", _MEMBER_RANKS, ids=_MEMBER_RANK_IDS) +def test_a_rank_holds_the_kv_head_the_weight_split_gives_it(member, rank): + """With more attention-TP ranks than KV heads the stock loader repeats each head + (`duplicate_kv_weight`) before the even row split, so the K/V a rank writes into its pages + is that of global head `kv_head_start`; under attention DP a rank holds all of them.""" + g = _geometry(member, rank) + heads, rows_per_head = 2, 4 + weight = torch.arange(heads * rows_per_head, dtype=torch.float32).reshape(-1, 1) + tp = P.attention_tp_size(_mapping(member, rank)) + local = duplicate_kv_weight(weight, heads, tp).chunk(tp)[rank if tp > 1 else 0] + start, stop = g.kv_head_start * rows_per_head, (g.kv_head_start + g.heads_kv) * rows_per_head + assert torch.equal(local, weight[start:stop]) + + +@contextlib.contextmanager +def _manager(member: str, rank: int, kv_cache_quant_algo=None, spec_config=None): + """The trunk manager the engine builds on `rank` of `member`, through the engine's own class + choice and factory.""" + mapping = _mapping(member, rank) + model_config = ModelConfig( + pretrained_config=Qwen4ExpTextConfig(**_TEXT), + mapping=mapping, + quant_config=QuantConfig( + quant_algo=QuantAlgo.MIXED_PRECISION, kv_cache_quant_algo=kv_cache_quant_algo + ), + sparse_attention_config=QSASparseAttentionConfig(), + spec_config=spec_config, + ) + # What the loader does with `kv_cache_config.mamba_ssm_cache_dtype: bfloat16`. + validate_and_set_mamba_ssm_cache_dtype(model_config, "bfloat16") + kv_cache_config = KvCacheConfig( + enable_block_reuse=False, + mamba_ssm_cache_dtype="bfloat16", + max_tokens=4096, + use_kv_cache_manager_v2=True, + ) + cls = get_kv_cache_manager_cls(model_config, kv_cache_config) + cm = _create_kv_cache_manager( + None, + cls, + mapping, + kv_cache_config, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=2048, + max_batch_size=4, + spec_config=spec_config, + sparse_attention_config=model_config.sparse_attention_config, + max_num_tokens=2048, + max_beam_width=1, + kv_connector_manager=None, + model_config=model_config, + dtype=torch.bfloat16, + max_cuda_graph_batch_size=4, + ) + try: + yield cm + finally: + cm.shutdown() + + +@_needs_gpu +@pytest.mark.parametrize("algo", list(KV_STORAGE), ids=list(KV_IDS.values())) +@pytest.mark.parametrize("member, rank", _MEMBER_RANKS, ids=_MEMBER_RANK_IDS) +def test_the_stock_manager_serves_the_layout(member, rank, algo): + """Every trunk pool view of the stock manager passes the contract at this rank's geometry, + and the manager's own per-rank numbers are the geometry's.""" + g = _geometry(member, rank, algo) + with _manager(member, rank, algo) as cm: + assert type(cm) is Qwen4ExpHybridCacheManager + assert cm.dtype == KV_STORAGE[algo][2] + assert cm.ssm_state_dtype == torch.bfloat16 + pools = P.take_pools(cm, g, ATTN_LAYERS, GDN_LAYERS, PLE_LAYER) + + assert {cm.num_kv_heads_per_layer[cm.layer_offsets[i]] for i in ATTN_LAYERS} == {g.heads_kv} + assert (cm.qsa_index_kv_heads, cm.qsa_index_dim) == (g.index_kv_heads, g.index_head_dim) + assert list(cm.conv_state_shape) == [g.conv_dim, g.conv_kernel - 1] + assert list(cm.ssm_state_shape) == [g.num_v_heads, g.head_v_dim, g.head_k_dim] + # Each attention layer has a pool offset of its own. + assert len(set(pools.pool_idx.values())) == len(ATTN_LAYERS) + # One slot space for every QSA view of the layers, and one for the recurrent and PLE + # state. + slots = {view.shape[0] for view in (pools.position_cache, *pools.index_cache.values())} + slots |= {planes.k.shape[0] for planes in pools.kv.values()} + assert len(slots) == 1 + assert {t.shape[0] for t in (*pools.conv.values(), *pools.ssm.values())} == { + pools.ple_state[0].shape[0] + } + + +@_needs_gpu +@pytest.mark.parametrize("member, rank", _MTP_MEMBER_RANKS, ids=_MTP_MEMBER_RANK_IDS) +def test_the_mtp_verify_buffers_hold_this_ranks_heads(monkeypatch, member, rank): + """With MTP (three drafts, the GDN cached replay) the verify buffers -- the per-step conv + windows and every GDN layer's replay history -- are this rank's heads. Under attention TP the + draft layer has a separate draft KV cache; under attention DP it is in the trunk's manager, + with the trunk's per-rank pool shapes.""" + monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY", raising=False) + g = _geometry(member, rank) + spec = MTPDecodingConfig(max_draft_len=3, num_nextn_predict_layers=1) + with _manager(member, rank, spec_config=spec) as cm: + pools = P.take_pools(cm, g, ATTN_LAYERS, GDN_LAYERS, PLE_LAYER) + assert cm.is_speculative() + ssm0 = pools.ssm[GDN_LAYERS[0]] + caches = {i: cm.mamba_layer_cache(i) for i in GDN_LAYERS} + for i, cache in caches.items(): + P.check_step_buffers( + g, cache.intermediate_conv_window, pools.conv[i], cm.intermediate_state_indices + ) + assert cache.intermediate_ssm is None, "the cached replay keeps no per-step state" + meta = cm.get_replay_state_update_metadata() + assert meta is not None + P.check_replay_state( + meta.replay_history_size, + meta.replay_step_width, + meta.prev_num_accepted_tokens, + meta.cache_buf_idx, + ssm0, + ) + assert meta.replay_step_width == spec.max_draft_len + 1 + history, slots = meta.replay_history_size, ssm0.shape[0] + for i, cache in caches.items(): + P.check_replay_history(g, i, history, slots, cache.old_x, cache.old_B, cache.old_dt) + + draft_in_trunk = MTP_LAYER in cm.layer_offsets + assert draft_in_trunk == MEMBERS[member]["enable_attention_dp"] + if draft_in_trunk: + P.kv_planes(cm, g, MTP_LAYER, cm.get_index_k_buffer(MTP_LAYER)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routed_experts.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routed_experts.py new file mode 100644 index 000000000000..2eb055a0bb2c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routed_experts.py @@ -0,0 +1,671 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The routed experts at one rank (tp1): the CuteDSL experts and the NVFP4 draft experts. + +**CuteDSL.** The CuteDSL op tests build their operands with layout code of their own. Here the +checkpoint's expert tensors (512 experts: packed NVFP4 `up_proj` / `gate_proj` / `down_proj` +codes, linear e4m3 block scales, one `weight_scale_2` per expert, per-expert `input_scale`s) +go through the target's own load-time transforms (`weights._expert_fc1_w` / `_expert_fc1_s` / +`_expert_fc2_w` / `_expert_fc2_s`) and scalar derivation (`modeling._moe_layer_scalars`, +`_moe_layer_scalars_cd`), once per routed backend, then through the call sequence of the +target's routed experts -- CuteDSL (`_cute_dsl_routed`): sort, FC1 with SwiGLU, FC2, +unpermute; trtllm-gen: the one-call block-scale runner -- with explicit routing. + +Two routings hit all 512 experts. In the first every token's only non-zero routing weight is +on its own expert, so each expert's output is checked alone: a layout error confined to one +expert (a permutation, a scale swizzle, an alpha) cannot hide in a sum over ten. On it both +backends are held to a float64 reference over the dequantized checkpoint tensors, with +`down_proj.input_scale` spread across the experts as a calibrated checkpoint's is. The second +is a weighted top-10: with a uniform `down_proj.input_scale` the two backends requantize the +SwiGLU output identically, so on it they are held to each other. The reference models the one +rounding step that is not exact: the SwiGLU output is requantized to NVFP4 before FC2, at the +backend's own global scale (the layer's single `global_sf` on CuteDSL, each expert's +`1 / down_proj.input_scale` on trtllm-gen). The scalars, the CuteDSL weight row orders and the +prefill copy's decision are also checked on the host. + +**NVFP4 draft experts** (`TRTLLM_MODELING_V2_MTP_MOE_NVFP4`). `weights.mtp_experts_to_nvfp4` +requantizes FP8 128x128 block-scaled experts to NVFP4 in the trunk's trtllm-gen layout, +`modeling._mtp_moe_nvfp4` adds the runner's scalars and `modeling._mtp_routed_nvfp4` runs +them. Against the FP8 block-scale path the draft layer runs by default, on the same random +experts, activations and router, the NVFP4 output must stay inside a relative-error band of +the FP8 one, keep its magnitude (a wrong runner scalar moves the norm, which unbiased +quantization noise does not), and leave no FP8 copy behind. +""" + +import functools +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( + fp4_block_scale_moe_runner, + quantize_nvfp4, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.fp8_block_scale_moe_runner import ( # noqa: E501 + fp8_block_scale_moe_runner, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.moe.renorm_moe_routing_op import ( # noqa: E501 + renorm_moe_routing_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.quantization.fp8_quantize_1x128 import ( # noqa: E501 + fp8_quantize_1x128, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501 + modeling as M, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501 + weights as W, +) + +_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="the experts run on a CUDA device") +DEV = "cuda" + + +# ---- the scalars and the weight layout, on the host --------------------------------------- + +SCALAR_EXPERTS = 8 + + +def _scalar_rows(isc2: torch.Tensor, seed: int = 0) -> dict: + gen = torch.Generator().manual_seed(seed) + isc1 = torch.full((SCALAR_EXPERTS,), 0.0123) + ws2_1 = torch.rand(SCALAR_EXPERTS, generator=gen) * 1e-3 + 1e-4 + return { + "l0_e_isc1": isc1, + "l0_e_isc1_up": isc1.clone(), + "l0_e_ws2_1": ws2_1, + "l0_e_ws2_1_up": ws2_1.clone(), + "l0_e_isc2": isc2, + "l0_e_ws2_2": torch.rand(SCALAR_EXPERTS, generator=gen) * 1e-3 + 1e-4, + } + + +@pytest.mark.cpu_only +def test_uniform_isc2_keeps_the_per_expert_products_bitwise() -> None: + w = _scalar_rows(torch.full((SCALAR_EXPERTS,), 0.0371)) + got = M._moe_layer_scalars_cd(w, 0) + assert torch.equal(got["alpha2"], w["l0_e_isc2"] * w["l0_e_ws2_2"]) + assert torch.equal(got["global_sf"], 1.0 / w["l0_e_isc2"][:1]) + assert got["global_sf"].shape == (1,) + assert torch.equal(got["alpha1"], w["l0_e_isc1"] * w["l0_e_ws2_1"]) + + +@pytest.mark.cpu_only +def test_isc2_spread_is_folded_into_alpha2() -> None: + """CuteDSL requantizes FC1's output at ONE global scale per layer, where trtllm-gen folds + each expert's `down_proj.input_scale` into its own FC1 scalar. Whatever the spread (here + 40x), `global_sf * alpha2[e]` is expert e's own `ws2_2[e]`, the product trtllm-gen forms + from `(1 / isc2[e]) * (isc2[e] * ws2_2[e])`.""" + isc2 = 0.01 * torch.linspace(1.0, 40.0, SCALAR_EXPERTS) + w = _scalar_rows(isc2.flip(0)) + cute = M._moe_layer_scalars_cd(w, 0) + gen = M._moe_layer_scalars(w, 0, 0, SCALAR_EXPERTS) + isc2 = w["l0_e_isc2"] + assert torch.equal(cute["global_sf"], 1.0 / isc2.max().reshape(1)) + # No expert's FC2 input is scaled above what the checkpoint calibrated for it. + assert (cute["global_sf"] <= 1.0 / isc2).all() + assert torch.allclose( + cute["global_sf"] * cute["alpha2"], (1.0 / isc2) * gen["output2_scale_scalar"], rtol=1e-6 + ) + assert torch.allclose(cute["global_sf"] * cute["alpha2"], w["l0_e_ws2_2"], rtol=1e-6) + + +@pytest.mark.cpu_only +def test_gate_up_input_scale_mismatch_is_refused() -> None: + w = _scalar_rows(torch.full((SCALAR_EXPERTS,), 0.0371)) + w["l0_e_isc1_up"] = w["l0_e_isc1_up"] * 1.5 + with pytest.raises(AssertionError, match="gate/up input_scale differ"): + M._moe_layer_scalars(w, 0, 0, SCALAR_EXPERTS) + + +@pytest.mark.cpu_only +def test_cutedsl_fc1_pairs_up_and_gate_rows_in_64_row_groups() -> None: + """Destination rows `[128k, 128k + 64)` hold up rows `[64k, 64k + 64)` and the next 64 + hold the matching gate rows, which sit `I` rows later in the `[up ; gate]` stack.""" + inter = 640 + perm = W._fc1_perm(2 * inter, "cpu", SimpleNamespace(trunk_moe_backend="CUTEDSL")) + assert sorted(perm.tolist()) == list(range(2 * inter)) + for k in range(inter // 64): + assert perm[128 * k : 128 * k + 64].tolist() == list(range(64 * k, 64 * k + 64)) + assert perm[128 * k + 64 : 128 * k + 128].tolist() == list( + range(inter + 64 * k, inter + 64 * k + 64) + ) + + +@pytest.mark.cpu_only +def test_cutedsl_fc2_keeps_the_checkpoint_row_order_and_trtllm_gen_shuffles_it() -> None: + """trtllm-gen's 32-row block shuffle moves source row `4u + v` to destination row + `8v + u`; CuteDSL reads `down_proj` as the checkpoint stores it.""" + rows = torch.arange(64).unsqueeze(1).expand(64, 8).contiguous() + kept = W._expert_fc2_w(rows, SimpleNamespace(trunk_moe_backend="CUTEDSL")) + assert torch.equal(kept, rows) + shuffled = W._expert_fc2_w(rows, SimpleNamespace(trunk_moe_backend="TRTLLM")) + for u in range(8): + for v in range(4): + assert int(shuffled[8 * v + u, 0]) == 4 * u + v + assert int(shuffled[32 + 8 * v + u, 0]) == 32 + 4 * u + v + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "one_rank, backend, max_num_tokens, env, kept", + [ + (True, "TRTLLM", 4096, None, True), + (True, "TRTLLM", 4095, None, False), + (True, "TRTLLM", 16384, "0", False), + (True, "CUTEDSL", 16384, None, False), + (False, "TRTLLM", 16384, None, False), + ], +) +def test_second_expert_copy_kept_only_where_its_launches_can_happen( + monkeypatch, one_rank, backend, max_num_tokens, env, kept +): + """The copy costs about 1.3 GiB a layer: it is kept only at one rank, under trtllm-gen, + when a step's rows can reach the CuteDSL dispatch threshold, and not when the env var is + off.""" + if env is None: + monkeypatch.delenv(M._MOE_PREFILL_COPY_ENV, raising=False) + else: + monkeypatch.setenv(M._MOE_PREFILL_COPY_ENV, env) + assert M._moe_prefill_copy(one_rank, backend, max_num_tokens) is kept + + +# ---- CuteDSL: the experts' geometry and gates --------------------------------------------- + +#: The trunk's routed experts: count, top-k, hidden, intermediate, NVFP4 block. +E, K, H, INTER, VEC = 512, 10, 2560, 640, 16 +#: The activation's input scale, one number for all experts: |x| = 5 maps to the e4m3 maximum. +ISC1 = 5.0 / (448 * 6) +#: The SwiGLU output maximum `down_proj.input_scale` is calibrated to. With `spread` each +#: expert's own value ranges over 2^-1/2 .. 2^1/2 of it, so the layer's `max` is not every +#: expert's. +ISC2 = 64.0 / (448 * 6) +E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +#: The routed runners read activation scales as a linear buffer, not the 128x4-swizzled +#: one of a dense NVFP4 GEMM; byte counts match at multiples of 128 rows, so the wrong +#: one passes silently. +SF_LINEAR = False + +#: Tokens of the two routings: one per expert, and a weighted top-10. +ONE_HOT_TOKENS = E +WEIGHTED_TOKENS = 1024 + +BACKENDS = ("CUTEDSL", "TRTLLM") +#: The float64 reference's gate, relative per token. Both backends sit at 2.5e-3 or below: +#: the bf16 rounding of the output is 1e-3 of it, the rest is the reference's fp32 / tie +#: differences. A wrong layout is ~1; a wrong per-expert scalar is its own relative error. +REFERENCE_TOLERANCE = 5e-3 +#: The same, between the backends under a uniform `down_proj.input_scale`. Each backend is +#: within `REFERENCE_TOLERANCE` of the reference, so the two are within twice that of each +#: other. Both requantize the SwiGLU output to NVFP4, and an fp32 difference in the last +#: place moves an element at an e2m1 rounding boundary by a whole code: the token holding +#: it differs by ~5e-3 while the mean stays near 5e-6. A wrong layout is still ~1. +BACKEND_TOLERANCE = 2 * REFERENCE_TOLERANCE + + +# ---- CuteDSL: the checkpoint ------------------------------------------------- + + +def _quantize(dense: torch.Tensor): + """A bf16 `[R, C]` matrix as the checkpoint stores it: packed e2m1 codes `[R, C/2]` + uint8, e4m3 block scales `[R, C/16]`, and the `weight_scale_2` that undoes the global + scale the block scales were derived under.""" + rows, cols = dense.shape + gs = (448 * 6 / dense.abs().amax().float()).reshape(1) + codes, sf = torch.ops.trtllm.fp4_quantize(dense, gs, VEC, False, False) + return ( + codes.view(torch.uint8).view(rows, cols // 2), + sf.view(torch.uint8).view(rows, cols // VEC).view(torch.float8_e4m3fn), + (1.0 / gs).squeeze(0), + ) + + +@functools.lru_cache(maxsize=None) +def _checkpoint(spread: bool) -> SimpleNamespace: + """All 512 experts of one layer, with the checkpoint's tensor names and dtypes. + + `up_proj` and `gate_proj` are quantized together, so they share one `weight_scale_2`, as + the fused FC1 requires. Expert magnitudes spread over 2^-1 .. 2^1/2 (FC1) and 2^-1 .. 2^1 + (FC2), so a wrong per-expert scalar shows.""" + g = torch.Generator(device=DEV).manual_seed(20261014) + spread1 = torch.exp2(torch.rand(E, generator=g, device=DEV) * 1.5 - 1) + spread2 = torch.exp2(torch.rand(E, generator=g, device=DEV) * 2 - 1) + fields = {k: [] for k in ("up_w", "up_s", "gate_w", "gate_s", "down_w", "down_s")} + ws2_gu, ws2_dn = [], [] + for e in range(E): + gate_up = torch.randn(2 * INTER, H, generator=g, device=DEV) * (spread1[e] / H**0.5) + down = torch.randn(H, INTER, generator=g, device=DEV) * (spread2[e] / INTER**0.5) + codes, sf, ws2 = _quantize(gate_up.to(torch.bfloat16)) + fields["up_w"].append(codes[:INTER]) + fields["gate_w"].append(codes[INTER:]) + fields["up_s"].append(sf[:INTER]) + fields["gate_s"].append(sf[INTER:]) + ws2_gu.append(ws2) + codes, sf, ws2 = _quantize(down.to(torch.bfloat16)) + fields["down_w"].append(codes) + fields["down_s"].append(sf) + ws2_dn.append(ws2) + ws2_gu, ws2_dn = torch.stack(ws2_gu), torch.stack(ws2_dn) + isc1 = torch.full((E,), ISC1, dtype=torch.float32, device=DEV) + isc2 = torch.full((E,), ISC2, dtype=torch.float32, device=DEV) + if spread: + isc2 = isc2 * torch.exp2(torch.rand(E, generator=g, device=DEV) - 0.5) + return SimpleNamespace( + **fields, + ws2_gu=ws2_gu, + ws2_dn=ws2_dn, + isc2=isc2, + # The rows `modeling._moe_layer_scalars` / `_moe_layer_scalars_cd` read, under layer 0. + scalars={ + "l0_e_isc1": isc1, + "l0_e_isc1_up": isc1.clone(), + "l0_e_ws2_1": ws2_gu, + "l0_e_ws2_1_up": ws2_gu.clone(), + "l0_e_isc2": isc2, + "l0_e_ws2_2": ws2_dn, + }, + ) + + +# ---- CuteDSL: what the target builds from it --------------------------------- + + +@functools.lru_cache(maxsize=None) +def _layer(backend: str, spread: bool) -> SimpleNamespace: + """The layer's expert operands for `backend`, built as the weight table builds them: + one transform call per expert and tensor, stacked, then viewed as `_layer_views` + views them.""" + ck = _checkpoint(spread) + core = SimpleNamespace(trunk_moe_backend=backend) + + def stack(fn, *tensors): + return torch.stack([fn(tuple(t[e] for t in tensors), core) for e in range(E)]) + + def stack1(fn, tensor): + return torch.stack([fn(tensor[e], core) for e in range(E)]) + + fc1_w = stack(W._expert_fc1_w, ck.up_w, ck.gate_w) + fc1_s = stack(W._expert_fc1_s, ck.up_s, ck.gate_s) + fc2_w = stack1(W._expert_fc2_w, ck.down_w) + fc2_s = stack1(W._expert_fc2_s, ck.down_s) + if backend == "CUTEDSL": + w_dtype, s_dtype = torch.float4_e2m1fn_x2, torch.uint8 + else: + w_dtype, s_dtype = torch.uint8, torch.float8_e4m3fn + return SimpleNamespace( + fc1_w=fc1_w.view(w_dtype), + fc1_s=fc1_s.view(s_dtype), + fc2_w=fc2_w.view(w_dtype), + fc2_s=fc2_s.view(s_dtype), + scalars=_expert_scalars(ck.scalars, backend), + ) + + +def _expert_scalars(rows: dict, backend: str) -> dict: + """Layer 0's routed-expert scalars as the target derives them for `backend`: the + trtllm-gen runner's (with the FC1 activation's global scale both backends quantize at), + plus on CuteDSL its FC1 / FC2 alphas and FC2 global scale.""" + scalars = M._moe_layer_scalars(rows, 0, 0, E) + if backend == "CUTEDSL": + scalars.update(M._moe_layer_scalars_cd(rows, 0)) + return scalars + + +def _cutedsl_routed(layer, x, ids, weights) -> torch.Tensor: + """The CuteDSL branch of `_moe_block`, from the quantization to the unpermute.""" + sc = layer.scalars + tokens = x.shape[0] + xq, xsf = fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, SF_LINEAR) + tile_expert, tile_limit, expanded_to_permuted, permuted_to_expanded, _, live_tiles = M.moe_sort( + token_selected_experts=ids, + token_final_scales=weights, + num_experts=E, + top_k=K, + local_expert_offset=0, + local_num_experts=E, + tile_tokens_dim=M._CUTEDSL_TILE_SIZE, + ) + h, h_sf = M.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion( + input=xq.view(torch.float4_e2m1fn_x2), + weight=layer.fc1_w, + input_scale=xsf.view(tokens, H // M._SF_VEC), + weight_scale=layer.fc1_s, + alpha=sc["alpha1"], + tile_idx_to_group_idx=tile_expert, + tile_idx_to_mn_limit=tile_limit, + permuted_idx_to_expanded_idx=permuted_to_expanded, + num_non_exiting_tiles=live_tiles, + global_sf=sc["global_sf"], + num_experts=E, + top_k=K, + num_local_experts=E, + local_expert_offset=0, + tile_size=M._CUTEDSL_TILE_SIZE, + scaling_vector_size=M._SF_VEC, + activation_type=int(M.ActivationType.Swiglu), + ) + y = M.cute_dsl_nvfp4_grouped_gemm( + input=h, + weight=layer.fc2_w, + input_scale=h_sf, + weight_scale=layer.fc2_s, + alpha=sc["alpha2"], + tile_idx_to_group_idx=tile_expert, + num_non_exiting_tiles=live_tiles, + num_experts=E, + top_k=K, + num_local_experts=E, + local_expert_offset=0, + tile_size=M._CUTEDSL_TILE_SIZE, + output_dtype=torch.bfloat16, + scaling_vector_size=M._SF_VEC, + ) + return M.moe_unpermute_inplace( + permuted_input=y, + output=torch.empty(tokens, H, dtype=torch.bfloat16, device=x.device), + expanded_idx_to_permuted_idx=expanded_to_permuted, + topk_scales=weights, + ) + + +def _trtllm_gen_routed(layer, x, ids, weights) -> torch.Tensor: + """The trtllm-gen branch of `_moe_block`.""" + sc = layer.scalars + xq, xsf = fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, SF_LINEAR) + return fp4_block_scale_moe_runner( + routing_logits=None, + routing_bias=None, + hidden_states=xq, + hidden_states_scale=xsf.view(torch.float8_e4m3fn), + gemm1_weights=layer.fc1_w, + gemm1_weights_scale=layer.fc1_s, + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_weights=layer.fc2_w, + gemm2_weights_scale=layer.fc2_s, + gemm2_bias=None, + output1_scale_scalar=sc["output1_scale_scalar"], + output1_scale_gate_scalar=sc["output1_scale_gate_scalar"], + output2_scale_scalar=sc["output2_scale_scalar"], + num_experts=E, + top_k=K, + n_group=None, + topk_group=None, + intermediate_size=INTER, + local_expert_offset=0, + local_num_experts=E, + routed_scaling_factor=None, + routing_method_type=M._ROUTING_METHOD_RENORM, + do_finalize=True, + act_type=M._ACT_TYPE_SWIGLU, + topk_weights=weights.to(torch.bfloat16), + topk_ids=ids, + )[0] + + +# ---- CuteDSL: the inputs and the reference ----------------------------------- + + +@functools.lru_cache(maxsize=None) +def _problem(case: str): + """`(x, ids, weights)`: bf16 activations, expert ids `[T, K]` int32 and routing weights + `[T, K]` fp32 (exact in bf16, which is the dtype trtllm-gen takes them in).""" + tokens = ONE_HOT_TOKENS if case == "one_hot" else WEIGHTED_TOKENS + g = torch.Generator(device=DEV).manual_seed(7 if case == "one_hot" else 8) + x = torch.randn(tokens, H, generator=g, device=DEV).to(torch.bfloat16) + scores = torch.rand(tokens, E, generator=g, device=DEV) + if case == "one_hot": + # The forced expert has the top score, so it sits in slot 0. + scores[torch.arange(tokens), torch.arange(tokens) % E] = 2.0 + weights = torch.zeros(tokens, K, device=DEV) + weights[:, 0] = 1.0 + else: + weights = torch.softmax(torch.randn(tokens, K, generator=g, device=DEV), dim=-1) + weights = weights.to(torch.bfloat16).float() + ids = scores.topk(K, dim=-1).indices.to(torch.int32) + return x, ids, weights + + +def _dequant(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """`[R, C/2]` packed e2m1 (element 2j in the low nibble) and `[R, C/16]` e4m3 scales + -> float64 `[R, C]`.""" + magnitude = torch.tensor(E2M1, dtype=torch.float64, device=codes.device) + table = torch.cat([magnitude, -magnitude]) + lo, hi = table[(codes & 0xF).long()], table[(codes >> 4).long()] + values = torch.stack([lo, hi], dim=-1).reshape(codes.shape[0], -1) + return values * scales.view(torch.float8_e4m3fn).double().repeat_interleave(VEC, dim=1) + + +def _reference(backend: str, case: str, spread: bool) -> torch.Tensor: + """The routed output in float64, `[T, H]`. + + The activation is what the kernels see, the dequantized NVFP4 of the target's own + quantization. Each expert's weights are dequantized from the checkpoint tensors, times + its `weight_scale_2`. The SwiGLU output is requantized at `backend`'s global scale.""" + ck, layer = _checkpoint(spread), _layer(backend, spread) + x, ids, weights = _problem(case) + sc = layer.scalars + xq, xsf = fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, SF_LINEAR) + x_dq = _dequant(xq.view(torch.uint8), xsf.view(x.shape[0], H // VEC)) * ISC1 + if backend == "CUTEDSL": + g2 = [float(sc["global_sf"])] * E + else: + g2 = (1.0 / ck.isc2).tolist() + out = torch.zeros(x.shape[0], H, dtype=torch.float64, device=DEV) + for e in range(E): + token, slot = (ids == e).nonzero(as_tuple=True) + w = weights[token, slot].double() + token, w = token[w != 0], w[w != 0] + if token.numel() == 0: + continue + ws2 = ck.ws2_gu[e].double() + xe = x_dq[token] + up = xe @ (_dequant(ck.up_w[e], ck.up_s[e]) * ws2).T + gate = xe @ (_dequant(ck.gate_w[e], ck.gate_s[e]) * ws2).T + act = up * gate * torch.sigmoid(gate) + act = quantize_nvfp4(act.float(), g2[e]).double() / g2[e] + y = act @ (_dequant(ck.down_w[e], ck.down_s[e]) * ck.ws2_dn[e].double()).T + out.index_add_(0, token, y * w[:, None]) + return out + + +@functools.lru_cache(maxsize=None) +def _routed(backend: str, case: str, spread: bool) -> torch.Tensor: + x, ids, weights = _problem(case) + run = _cutedsl_routed if backend == "CUTEDSL" else _trtllm_gen_routed + return run(_layer(backend, spread), x, ids, weights).double() + + +def _token_error(out: torch.Tensor, ref: torch.Tensor) -> torch.Tensor: + """Each token's `|out - ref| / |ref|` (L2 over the hidden width).""" + return (out - ref).norm(dim=1) / ref.norm(dim=1).clamp_min(1e-12) + + +@_cuda +@pytest.mark.parametrize("backend", BACKENDS) +def test_routed_output_matches_the_float64_reference(backend: str) -> None: + case = "one_hot" + _, ids, weights = _problem(case) + assert torch.unique(ids[weights != 0]).numel() == E, "the routing misses an expert" + err = _token_error(_routed(backend, case, True), _reference(backend, case, True)) + print(f"[err] {backend} {case}: mean {err.mean():.4e} max {err.max():.4e}") + assert err.max() <= REFERENCE_TOLERANCE, (err.mean().item(), err.max().item()) + + +@_cuda +def test_cutedsl_matches_trtllm_gen_under_a_uniform_input_scale() -> None: + case = "weighted" + _, ids, weights = _problem(case) + assert torch.unique(ids[weights != 0]).numel() == E, "the routing misses an expert" + err = _token_error(_routed("CUTEDSL", case, False), _routed("TRTLLM", case, False)) + print(f"[err] CUTEDSL vs TRTLLM {case}: mean {err.mean():.4e} max {err.max():.4e}") + assert err.max() <= BACKEND_TOLERANCE, (err.mean().item(), err.max().item()) + + +# ---- the NVFP4 draft experts -------------------------------------------------------------- + +#: The draft experts' FP8 weight block. +BLK = 128 +#: A trunk calibration to derive the static draft scales from: the checkpoint's last trunk layer +#: (FC1 input, FC2 input amax), the layer that feeds the draft. +TRUNK_AMAX = (11.1875, 57.75) +#: Relative RMS of NVFP4 - FP8 over the FP8 output, the floor on each row's cosine, and the +#: band on the output norm ratio. Measured here: 0.239-0.251 RMS, row cosine >= 0.961, at every +#: row count; on the checkpoint's real draft experts 0.222-0.231 and >= 0.964 (bench). That is +#: the NVFP4 chain's own noise -- four e2m1 roundings (activation, FC1 weight, FC1 output, FC2 +#: weight) at ~0.1 each, compounded by the SwiGLU product -- not a bias: it is flat across +#: activation scales 0.5-4x, and FP8 sits at 0.07 from the fp32 reference of the same weights. +REL_RMS_BAND = 0.3 +COS_FLOOR = 0.95 +NORM_RATIO_BAND = 0.1 + + +def _fp8_blocks(rows: int, cols: int, gen: torch.Generator): + """`[E, rows, cols]` e4m3 codes and `[E, rows/128, cols/128]` fp32 scales of a random + expert stack whose 128x128 blocks span 16x in magnitude, as trained experts' do.""" + w = torch.randn(E, rows, cols, generator=gen, device=DEV) * cols**-0.5 + spread = torch.exp2( + torch.randint(-2, 3, (E, rows // BLK, 1, cols // BLK, 1), generator=gen, device=DEV) + ).float() + blocks = w.view(E, rows // BLK, BLK, cols // BLK, BLK) * spread + scale = blocks.abs().amax(dim=(2, 4), keepdim=True) / 448.0 + codes = (blocks / scale).to(torch.float8_e4m3fn).view(E, rows, cols) + return codes, scale.view(E, rows // BLK, cols // BLK).contiguous() + + +@pytest.fixture(scope="module") +def draft_experts(): + gen = torch.Generator(device=DEV).manual_seed(20261022) + fc1_w, fc1_s = _fp8_blocks(2 * INTER, H, gen) + fc2_w, fc2_s = _fp8_blocks(H, INTER, gen) + router = (torch.randn(E, H, generator=gen, device=DEV) * H**-0.5).to(torch.bfloat16) + fp8 = SimpleNamespace(fc1_w=fc1_w, fc1_s=fc1_s, fc2_w=fc2_w, fc2_s=fc2_s, router=router.t()) + p = "l1_mtp_" + w = {p + k: getattr(fp8, k).clone() for k in ("fc1_w", "fc1_s", "fc2_w", "fc2_s")} + w["l0_e_isc1"] = torch.full((E,), TRUNK_AMAX[0] / 2688.0, device=DEV) + w["l0_e_isc2"] = torch.full((E,), TRUNK_AMAX[1] / 2688.0, device=DEV) + # No packed FP8 draft projections: the shared expert runs its bf16 GEMMs. + core = SimpleNamespace( + w=w, + mtp_layer_idx=1, + num_layers=1, + num_experts=E, + topk=K, + moe_inter=INTER, + _mtp_gemv8={}, + ) + fields = M._mtp_moe_nvfp4(core) + nvfp4 = SimpleNamespace(router=fp8.router, **fields) + return core, fp8, nvfp4 + + +def _fp8_routed(fp8, x): + """The FP8 block-scale draft path, call for call as `_mtp_moe_block` makes it.""" + logits = M.cublas_mm(x, fp8.router) + ids, wts = renorm_moe_routing_op(logits, topk=K, output_dtype=torch.bfloat16) + xq, xsf = fp8_quantize_1x128(x, use_ue8m0=False) + return fp8_block_scale_moe_runner( + routing_logits=None, + routing_bias=None, + hidden_states=xq, + hidden_states_scale=xsf, + gemm1_weights=fp8.fc1_w, + gemm1_weights_scale=fp8.fc1_s, + gemm2_weights=fp8.fc2_w, + gemm2_weights_scale=fp8.fc2_s, + num_experts=E, + top_k=K, + num_fused_shared_experts=None, + n_group=None, + topk_group=None, + intermediate_size=INTER, + local_expert_offset=0, + local_num_experts=E, + routed_scaling_factor=None, + routing_method_type=M._ROUTING_METHOD_RENORM, + topk_weights=wts, + topk_ids=ids, + act_type=M._ACT_TYPE_SWIGLU, + gemm1_clamp_limit=None, + ) + + +@_cuda +def test_draft_fp8_copies_freed_and_layout(draft_experts): + core, _, nvfp4 = draft_experts + for k in ("fc1_w", "fc1_s", "fc2_w", "fc2_s"): + assert core.w[f"l1_mtp_{k}"].numel() == 0, k + assert nvfp4.fc1_w.shape == (E, 2 * INTER, H // 2) and nvfp4.fc1_w.dtype == torch.uint8 + assert nvfp4.fc1_s.shape == (E, 2 * INTER, H // 16) and nvfp4.fc1_s.dtype == torch.float8_e4m3fn + assert nvfp4.fc2_w.shape == (E, H, INTER // 2) and nvfp4.fc2_s.shape == (E, H, INTER // 16) + for s in (nvfp4.out1_scale, nvfp4.out1_gate_scale, nvfp4.out2_scale): + assert s.shape == (E,) and s.dtype == torch.float32 and torch.isfinite(s).all() + assert nvfp4.moe_gscale.shape == (1,) + + +@_cuda +@pytest.mark.parametrize("rows", [1, 4, 32, 128, 1024]) +def test_draft_nvfp4_tracks_fp8(draft_experts, rows): + core, fp8, nvfp4 = draft_experts + gen = torch.Generator(device=DEV).manual_seed(rows) + x = torch.randn(rows, H, generator=gen, device=DEV).to(torch.bfloat16) + ref = _fp8_routed(fp8, x).float() + out = M._mtp_routed_nvfp4(core, nvfp4, x) + assert out.shape == (rows, H) and out.dtype == torch.bfloat16 + out = out.float() + rel = ((out - ref).norm() / ref.norm()).item() + cos = torch.nn.functional.cosine_similarity(out, ref, dim=1).min().item() + ratio = (out.norm() / ref.norm()).item() + print(f"rows={rows}: rel RMS {rel:.4f}, min row cosine {cos:.4f}, norm ratio {ratio:.4f}") + assert rel < REL_RMS_BAND, rel + assert cos > COS_FLOOR, cos + assert abs(ratio - 1.0) < NORM_RATIO_BAND, ratio + + +@_cuda +@pytest.mark.parametrize("rows", [1, 64, 512]) +def test_draft_shared_expert_overlap_preserves_serial_output(draft_experts, rows): + """`DecodeTarget._mtp_moe` over the NVFP4 experts runs the shared expert on the side stream + beside them: its output is the serial composition's bitwise, eager and replayed from a + CUDA graph over new inputs.""" + core, _, nvfp4 = draft_experts + gen = torch.Generator(device=DEV).manual_seed(42) + # The checkpoint's fused shared expert: gate|up rows, then the one gate-logit row. + sh_gu = (torch.randn(2 * INTER + 1, H, generator=gen, device=DEV) / H**0.5).to(torch.bfloat16) + sh_dn = (torch.randn(H, INTER, generator=gen, device=DEV) / INTER**0.5).to(torch.bfloat16) + weights = SimpleNamespace( + **vars(nvfp4), + sh_gu=sh_gu[: 2 * INTER].t(), + sh_dn=sh_dn.t(), + sh_gate=sh_gu[2 * INTER :].t(), + ) + target = object.__new__(M.DecodeTarget) + target.core = SimpleNamespace( + **vars(core), _mtp=weights, _side_stream=torch.cuda.Stream(), shared_inter=INTER + ) + x = torch.randn(rows, H, generator=gen, device=DEV).to(torch.bfloat16) + + def serial(): + routed = M._mtp_routed_nvfp4(target.core, weights, x) + up = M.silu_and_mul_rows(M.cublas_mm(x, weights.sh_gu), INTER) + shared = M.cublas_mm(up, weights.sh_dn) + gate = M.cublas_mm(x, weights.sh_gate) + return M.fused_sigmoid_gate_mul_add(routed, gate, shared) + + expected = serial().clone() + torch.testing.assert_close(target._mtp_moe(x), expected, rtol=0, atol=0) + if rows == 64: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = target._mtp_moe(x) + for _ in range(3): + x.normal_(generator=gen) + expected = serial().clone() + graph.replay() + torch.testing.assert_close(captured, expected, rtol=0, atol=0) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py new file mode 100644 index 000000000000..5a6960713b27 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_routing.py @@ -0,0 +1,396 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Qwen3.8-Flash-Next-NVFP4 routing at its layouts, driven by synthetic configs on CPU. + +`qwen4_exp` routes to its one target one rank (tp1), tp4 and tep4 with or without MTP, and +the attention-DP layouts dep4 (with or without MTP) and dep2 (the disaggregated-serving +engines' layout, without speculation), neither with a breakable prefill graph. The layouts +the target does not serve fall back under `auto`. The tp1 engine choices are exercised in +`tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py`. +""" + +from dataclasses import replace + +import pytest +import torch +from transformers import PretrainedConfig + +from tensorrt_llm._torch._experimental.modeling_v2._router_index import modeling_v2_resolve +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp import routing +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501 + modeling as M, +) +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm.llmapi.llm_args import ( + KvCacheConfig, + MTPDecodingConfig, + NGramDecodingConfig, + TorchLlmArgs, +) +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +pytestmark = pytest.mark.cpu_only + +_DUMMY_MODEL = "/tmp/dummy_model" +_ARCH = "Qwen4ExpForConditionalGeneration" +_TARGET = ( + "models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.modeling." + "ModelingV2Qwen38FlashNextNvfp4Sm103Default" +) + +# The four-rank layouts, as the LLM API's parallel arguments build them. tp4 leaves the +# expert split to Mapping's default, which is the TP size. +_LAYOUTS = { + "tp4": dict(world_size=4, tp_size=4), + "tep4": dict(world_size=4, tp_size=4, moe_tp_size=1, moe_ep_size=4), + "dep4": dict(world_size=4, tp_size=4, moe_tp_size=1, moe_ep_size=4, enable_attention_dp=True), +} + +# The two-rank layout both engines of the disaggregated-serving test run: attention DP, +# the routed experts in two 256-expert windows. +_DEP2 = dict(world_size=2, tp_size=2, moe_tp_size=1, moe_ep_size=2, enable_attention_dp=True) + +# The layouts with attention TP, which serve MTP and the breakable prefill graph, and the +# attention-DP ones, which serve neither. +_ATTENTION_TP = ["tp4", "tep4"] +_ATTENTION_DP = {"dep4": _LAYOUTS["dep4"], "dep2": _DEP2} + +# Every draft length the target runs. +_DRAFT_LENS = list(range(1, 8)) + + +def _qualified(target) -> str: + """``models...modeling.`` for a resolved target class.""" + return f"{target.__module__.split('.modeling_v2.')[1]}.{target.__name__}" + + +@pytest.fixture(autouse=True) +def _on_sm103(monkeypatch): + """Route as if this were a GB300, wherever the test actually runs, with the GDN + cached replay at its default (on).""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (10, 3)) + monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY", raising=False) + + +def _llm_args(mode="auto", **kwargs) -> TorchLlmArgs: + return TorchLlmArgs(model=_DUMMY_MODEL, modeling_v2=mode, **kwargs) + + +def _qwen4_exp_config(text_overrides=None, **vision_overrides): + """The shape Qwen3.8-Flash-Next-NVFP4's normalized composite config declares.""" + text = dict( + num_hidden_layers=48, + hidden_size=2560, + num_experts=512, + num_attention_heads=24, + linear_num_value_heads=48, + vocab_size=248320, + ) + text.update(text_overrides or {}) + vision = dict( + depth=27, + hidden_size=1152, + num_heads=16, + spatial_merge_size=2, + out_hidden_size=2560, + num_position_embeddings=2304, + deepstack_visual_indexes=[], + hidden_act="gelu_pytorch_tanh", + ) + vision.update(vision_overrides) + return PretrainedConfig( + architectures=[_ARCH], + text_config=PretrainedConfig(**text), + vision_config=PretrainedConfig(**vision), + ) + + +def _model_config(spec_config=None, pretrained_config=None, **mapping_kwargs) -> ModelConfig: + """The loaded config; the checkpoint's own KV is bf16 (no KV-cache algorithm).""" + return ModelConfig( + pretrained_config=pretrained_config or _qwen4_exp_config(), + mapping=Mapping(**mapping_kwargs), + spec_config=spec_config, + quant_config=QuantConfig(kv_cache_quant_algo=None), + ) + + +def _claimed(config, **llm_args_kwargs) -> None: + for mode in ("auto", "require"): + target = modeling_v2_resolve(config, _llm_args(mode, **llm_args_kwargs)) + assert target is not None and _qualified(target) == _TARGET, mode + + +def _declined(config, match=None, **llm_args_kwargs) -> None: + """``auto`` falls back to the built-in model and ``require`` says why.""" + assert modeling_v2_resolve(config, _llm_args("auto", **llm_args_kwargs)) is None + with pytest.raises(ValueError, match=match): + modeling_v2_resolve(config, _llm_args("require", **llm_args_kwargs)) + + +def test_routing_admits_exactly_what_the_target_implements(): + """`route` declines a MoE backend, a draft length or a recurrent-state dtype the + target does not run; this keeps routing's constants and the target's together.""" + # At tp1 the CuteDSL backend and the fp32 state too; at the other layouts neither. + assert set(routing._ONE_RANK_MOE_BACKENDS) == set(M._TRUNK_MOE_BACKENDS) + assert set(routing._MOE_BACKENDS) == set(M._TRUNK_MOE_BACKENDS) - {"CUTEDSL"} + assert routing._MAX_CUTEDSL_DRAFT_LEN == M._CUTEDSL_MAX_DRAFT_LEN + assert routing._MAX_DRAFT_LEN == M._MAX_VERIFY_WIDTH - 1 + assert routing._SSM_STATE_DTYPES == M._SSM_STATE_DTYPES + assert routing._ONE_RANK_SSM_STATE_DTYPES == M._ONE_RANK_SSM_STATE_DTYPES + assert routing._SM == M._SM + + +# --- each layout lands on the default target ------------------------------------ + + +@pytest.mark.parametrize("member", sorted(_LAYOUTS)) +def test_each_member_routes_to_the_default_target(member): + config = _model_config(**_LAYOUTS[member]) + assert routing.layout(config.mapping) == member + for dtype in ("bfloat16", "auto"): + _claimed(config, kv_cache_config=KvCacheConfig(mamba_ssm_cache_dtype=dtype)) + + +@pytest.mark.parametrize("max_draft_len", _DRAFT_LENS) +@pytest.mark.parametrize("member", _ATTENTION_TP) +def test_mtp_at_every_draft_length_routes_under_attention_tp(member, max_draft_len): + """The tune runs MTP at tp4 and tep4 over draft lengths 1 to 7.""" + config = _model_config(MTPDecodingConfig(max_draft_len=max_draft_len), **_LAYOUTS[member]) + assert routing.layout(config.mapping) == member + for dtype in ("bfloat16", "auto"): + _claimed(config, kv_cache_config=KvCacheConfig(mamba_ssm_cache_dtype=dtype)) + + +@pytest.mark.parametrize("dtype", ["fp8", "nvfp4", "auto"]) +def test_dep2_routes_to_the_default_target(dtype): + config = _model_config(**_DEP2) + assert routing.layout(config.mapping) == "dep2" + _claimed(config, kv_cache_config=KvCacheConfig(dtype=dtype, mamba_ssm_cache_dtype="bfloat16")) + + +@pytest.mark.parametrize("max_draft_len", _DRAFT_LENS) +def test_mtp_at_every_draft_length_routes_at_dep4(max_draft_len): + """Each dep4 rank drafts for its own requests and the draft layer's FP8 expert window runs + the trunk's attention-DP round trip: the tune's ht_dep4 config may serve MTP.""" + config = _model_config(MTPDecodingConfig(max_draft_len=max_draft_len), **_LAYOUTS["dep4"]) + assert routing.layout(config.mapping) == "dep4" + for dtype in ("bfloat16", "auto"): + _claimed(config, kv_cache_config=KvCacheConfig(mamba_ssm_cache_dtype=dtype)) + _claimed(config, kv_cache_config=KvCacheConfig(dtype="fp8", mamba_ssm_cache_dtype="bfloat16")) + + +# --- what the target does not run under attention DP ------------------------------ + + +@pytest.mark.parametrize("max_draft_len", [1, 3, 7]) +def test_dep2_declines_speculation(max_draft_len): + """MTP under attention DP is admitted only where its draft acceptance was measured (dep4); + neither family claims dep2 with speculation.""" + config = _model_config(MTPDecodingConfig(max_draft_len=max_draft_len), **_DEP2) + assert routing.layout(config.mapping) == "dep2" + _declined(config, match=r"speculative decoding with attention DP \(dep2\)") + + +@pytest.mark.parametrize("member", sorted(_ATTENTION_DP)) +def test_attention_dp_declines_the_breakable_prefill_graph(member): + """The engine pads every attention-DP rank to one prefill bucket, which a rank holding + only generation requests cannot take; neither family claims it.""" + config = _model_config(**_ATTENTION_DP[member]) + _declined( + config, + match=rf"prefill_cuda_graph_backend='breakable' with attention DP \({member}\)", + prefill_cuda_graph_backend="breakable", + ) + _claimed(config, prefill_cuda_graph_backend="disabled") + + +def test_attention_dp_declines_speculation_before_the_engine_choices(monkeypatch): + """MTP without the cached replay is declined at dep2 for attention DP: that is the + structural reason, and no engine choice would lift it. At dep4 the same draft is declined + for the missing cached replay alone.""" + monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "0") + config = _model_config(MTPDecodingConfig(max_draft_len=7), **_DEP2) + _declined(config, match=r"speculative decoding with attention DP \(dep2\)") + config = _model_config(MTPDecodingConfig(max_draft_len=7), **_LAYOUTS["dep4"]) + _declined(config, match="TRTLLM_USE_GDN_REPLAY=0") + + +def test_dep4_declines_mtp_with_the_breakable_prefill_graph(): + """Admitting MTP at dep4 keeps the breakable prefill graph's refusal under attention DP.""" + config = _model_config(MTPDecodingConfig(max_draft_len=3), **_LAYOUTS["dep4"]) + _declined( + config, + match=r"prefill_cuda_graph_backend='breakable' with attention DP \(dep4\)", + prefill_cuda_graph_backend="breakable", + ) + + +@pytest.mark.parametrize("spec_config", [None, MTPDecodingConfig(max_draft_len=3)]) +def test_the_one_rank_deployment_routes_to_the_default_target_at_tp1(spec_config): + config = _model_config(spec_config) + assert routing.layout(config.mapping) == "tp1" + _claimed(config) + # The one-rank engine choices: the CuteDSL experts and the fp32 recurrent state. + _claimed( + replace(config, moe_backend_requested="CUTEDSL"), + kv_cache_config=KvCacheConfig(mamba_ssm_cache_dtype="float32"), + ) + + +def test_attention_dp_at_one_rank_declines(): + config = _model_config(world_size=1, tp_size=1, enable_attention_dp=True) + assert routing.layout(config.mapping) is None + _declined(config) + + +@pytest.mark.parametrize( + "mapping_kwargs", + [ + dict(world_size=2, tp_size=2), + dict(world_size=2, tp_size=2, moe_tp_size=1, moe_ep_size=2), + dict(_DEP2, moe_tp_size=2, moe_ep_size=1), + dict(_DEP2, enable_lm_head_tp_in_adp=True), + dict(world_size=8, tp_size=8), + dict(world_size=8, tp_size=8, moe_tp_size=1, moe_ep_size=8), + dict(world_size=8, tp_size=8, moe_tp_size=1, moe_ep_size=8, enable_attention_dp=True), + ], + ids=["tp2", "tep2", "dep2-moetp2", "dep2-lm-head-tp", "tp8", "tep8", "dep8"], +) +def test_other_two_or_eight_rank_layouts_decline(mapping_kwargs): + """The target does not claim it, and ``require`` names the layout and what the target + serves.""" + config = _model_config(**mapping_kwargs) + assert modeling_v2_resolve(config, _llm_args("auto")) is None + with pytest.raises(ValueError) as raised: + modeling_v2_resolve(config, _llm_args("require")) + message = str(raised.value) + assert "the qwen4_exp target is" in message, message + assert f"world_size={mapping_kwargs['world_size']}" in message, message + + +# --- near misses at four ranks ------------------------------------------------------ + + +@pytest.mark.parametrize( + "mapping_kwargs", + [ + # attention DP with the experts split by intermediate + dict(_LAYOUTS["dep4"], moe_tp_size=4, moe_ep_size=1), + # experts split two ways by intermediate and two by expert + dict(world_size=4, tp_size=4, moe_tp_size=2, moe_ep_size=2), + # pipeline parallelism + dict(world_size=4, tp_size=2, pp_size=2), + # lm_head TP under attention DP + dict(_LAYOUTS["dep4"], enable_lm_head_tp_in_adp=True), + # context parallelism + dict(world_size=4, tp_size=2, cp_size=2), + # four ranks over two nodes + dict(_LAYOUTS["tep4"], gpus_per_node=2), + ], + ids=["dep4-moetp4", "tp4-moetp2-ep2", "tp2-pp2", "dep4-lm-head-tp", "tp2-cp2", "two-nodes"], +) +def test_other_layouts_at_four_ranks_decline(mapping_kwargs): + config = _model_config(**mapping_kwargs) + assert routing.layout(config.mapping) is None + _declined(config) + + +@pytest.mark.parametrize( + "spec_config", + [ + NGramDecodingConfig(max_draft_len=3, max_matching_ngram_size=2), + MTPDecodingConfig(max_draft_len=8), + MTPDecodingConfig(max_draft_len=3, use_mtp_vanilla=True), + ], + ids=["ngram", "draft8", "vanilla-mtp"], +) +def test_unsupported_speculation_declines(spec_config): + _declined(_model_config(spec_config, **_LAYOUTS["tep4"])) + + +def test_another_checkpoint_shape_declines(): + config = _model_config( + pretrained_config=_qwen4_exp_config(dict(num_hidden_layers=60)), **_LAYOUTS["tep4"] + ) + _declined(config, match="text shape") + + +def test_a_deepstack_vision_tower_declines(): + config = _model_config( + pretrained_config=_qwen4_exp_config(deepstack_visual_indexes=[5, 11, 17]), + **_LAYOUTS["tep4"], + ) + _declined(config, match="vision shape") + + +def test_sm100_declines(monkeypatch): + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (10, 0)) + _declined(_model_config(**_LAYOUTS["tep4"]), match=r"compute capability \(10, 0\)") + + +# --- engine choices the target does not implement --------------------------------- + + +@pytest.mark.parametrize("dtype", ["fp8", "nvfp4"]) +@pytest.mark.parametrize("member", sorted(_LAYOUTS)) +def test_quantized_kv_cache_dtypes_match(member, dtype): + _claimed(_model_config(**_LAYOUTS[member]), kv_cache_config=KvCacheConfig(dtype=dtype)) + + +def test_other_kv_cache_algorithms_decline(): + choices = routing.engine_choices(_model_config(**_LAYOUTS["tep4"]), _llm_args()) + for algo in (QuantAlgo.INT8, QuantAlgo.FP8_BLOCK_SCALES): + with pytest.raises(routing.NoTarget, match="KV cache"): + routing._check_engine(None, algo, choices, "tep4") + _declined( + _model_config(**_LAYOUTS["tep4"]), + match="KV cache", + kv_cache_config=KvCacheConfig(dtype="fp8_ds_mla"), + ) + + +@pytest.mark.parametrize("backend", ["AUTO", "TRTLLM"]) +def test_the_trtllm_moe_backend_matches(backend): + _claimed(replace(_model_config(**_LAYOUTS["tep4"]), moe_backend_requested=backend)) + + +@pytest.mark.parametrize("backend", ["CUTEDSL", "CUTLASS", "TRITON", "DEEPGEMM"]) +def test_other_moe_backends_decline(backend): + """CUTEDSL included: its expert pipeline runs every expert on one rank (tp1).""" + config = replace(_model_config(**_LAYOUTS["tep4"]), moe_backend_requested=backend) + _declined(config, match="moe_config.backend") + + +def test_a_float32_state_cache_declines(): + _declined( + _model_config(**_LAYOUTS["dep4"]), + match="recurrent-state cache", + kv_cache_config=KvCacheConfig(mamba_ssm_cache_dtype="float32"), + ) + + +@pytest.mark.parametrize("backend", ["disabled", "breakable"]) +@pytest.mark.parametrize("member", _ATTENTION_TP) +def test_no_or_the_breakable_prefill_cuda_graph_matches(member, backend): + _claimed(_model_config(**_LAYOUTS[member]), prefill_cuda_graph_backend=backend) + + +def test_the_piecewise_prefill_cuda_graph_declines(): + _declined( + _model_config(**_LAYOUTS["tep4"]), match="piecewise", prefill_cuda_graph_backend="piecewise" + ) + + +@pytest.mark.parametrize("max_draft_len", [1, 3, 7]) +def test_mtp_needs_the_cached_replay(monkeypatch, max_draft_len): + """The target verifies on the GDN cached replay only; without speculation it is unused.""" + config = _model_config(MTPDecodingConfig(max_draft_len=max_draft_len), **_LAYOUTS["tep4"]) + monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "0") + _declined(config, match="TRTLLM_USE_GDN_REPLAY=0") + _claimed(_model_config(**_LAYOUTS["tep4"])) + monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY") + _claimed(config) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_select_canonical.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_select_canonical.py new file mode 100644 index 000000000000..6ab7a2d73a3f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_select_canonical.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the default target's block selection (`_select`). + +`_select` picks each row's canonical top-`block_topk` groups of its window: the exact top-k +and then the canonical rewrite, two launches, which serve a draft step's rows (the rejection +kernel's ragged windows) at every table width the engine's graph buckets use. + +Both checks run over windows short of the budget, at it and past it, with a coarse score +grid so the k-th value ties: the rows equal the canonical definition +(`qsa_canonical_selection.reference`: LARGE tie-break, ascending, `-1` past the window), and +bounding the rewrite to the live columns returns the same rows. +""" + +from __future__ import annotations + +import types + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.attention.qsa_canonical_selection import ( # noqa: E501 + qsa_canonical_selection, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501 + modeling as U, +) + +assert torch.cuda.is_available(), "the block selection runs on a CUDA device" + +K = 512 + + +def _case(rows: int, cols: int, seed: int) -> tuple: + """Ragged windows over a coarse score grid: every row's window ends somewhere in + `[1, cols]`, so rows sit below, at and past the budget, and the k-th value ties.""" + g = torch.Generator(device="cuda").manual_seed(seed) + scores = torch.randint(0, 5, (rows, cols), generator=g, device="cuda").float() + row_starts = torch.zeros(rows, dtype=torch.int32, device="cuda") + ends = torch.randint(1, cols + 1, (rows,), generator=g, device="cuda").to(torch.int32) + ends[0] = cols + ends[1] = K + ends[2] = 1 + indices = torch.full((rows, K), -1, dtype=torch.int32, device="cuda") + return scores, row_starts, ends, indices + + +@pytest.mark.parametrize("rows,cols", [(24, 2304), (24, 4096), (8, 1024)]) +def test_rows_equal_the_canonical_definition(rows: int, cols: int) -> None: + scores, row_starts, row_ends, indices = _case(rows, cols, seed=rows * cols) + U._select(types.SimpleNamespace(block_topk=K), scores, row_starts, row_ends, indices) + + expected = qsa_canonical_selection.reference(scores, row_starts, row_ends, indices.clone()) + assert torch.equal(indices, expected) + + +def test_live_cols_returns_the_same_rows() -> None: + scores, row_starts, row_ends, indices = _case(16, 2304, seed=7) + bounded = indices.clone() + U._select(types.SimpleNamespace(block_topk=K), scores, row_starts, row_ends, indices) + U._select( + types.SimpleNamespace(block_topk=K), + scores, + row_starts, + row_ends, + bounded, + int(row_ends.max()), + ) + assert torch.equal(indices, bounded) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_thop_attention_calls.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_thop_attention_calls.py similarity index 75% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_thop_attention_calls.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_thop_attention_calls.py index 87584e317806..1b11587e562e 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_thop_attention_calls.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_thop_attention_calls.py @@ -5,11 +5,11 @@ The attention op reads the pool through `quant_mode` and the two KV scales, and the three have to agree with the pool's dtype. A call that leaves one at its default reads an fp8 pool as bf16 (or the reverse) and still returns finite values, so no run notices. The check reads -`modeling.py` as text: it needs no device and no import of the target. +`modeling.py` and `vision.py` as text: it needs no device and no import of the target. -The vision tower's call is the one exception: its attention is the encoder-only mode of the -op, which has no KV cache, so it passes `kv_scale_*` None and `quant_mode` 0 through -`_VISION_CALL_CONSTANTS`. +The vision tower's call (`vision.py`) is the one exception: its attention is the encoder-only +mode of the op, which has no KV cache, so it passes `kv_scale_*` None and `quant_mode` 0 +through `_VISION_CALL_CONSTANTS`. """ from __future__ import annotations @@ -23,33 +23,34 @@ pytestmark = pytest.mark.cpu_only -_MODELING = ( +_TARGET = ( Path(_modeling_v2.__file__).resolve().parent - / "models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py" + / "models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default" ) _KV_STORAGE_ARGS = ("quant_mode", "kv_scale_orig_quant", "kv_scale_quant_orig") -_TREE = ast.parse(_MODELING.read_text()) +_TREE = ast.parse((_TARGET / "modeling.py").read_text()) +_VISION_TREE = ast.parse((_TARGET / "vision.py").read_text()) -def _thop_attention_calls() -> list[ast.Call]: +def _thop_attention_calls(tree: ast.Module) -> list[ast.Call]: return [ node - for node in ast.walk(_TREE) + for node in ast.walk(tree) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "thop_attention" ] -def _module_dict(name: str) -> ast.Call: - """The `dict(...)` a module-level `name = dict(...)` assigns.""" - for node in _TREE.body: +def _module_dict(tree: ast.Module, name: str) -> ast.Call: + """The `dict(...)` a module-level `name = dict(...)` of `tree` assigns.""" + for node in tree.body: if isinstance(node, ast.Assign) and any( isinstance(t, ast.Name) and t.id == name for t in node.targets ): assert isinstance(node.value, ast.Call), name return node.value - raise AssertionError(f"modeling.py has no module-level {name}") + raise AssertionError(f"no module-level {name}") def _keywords(call: ast.Call) -> dict[str, ast.expr]: @@ -67,13 +68,14 @@ def _is_vision(call: ast.Call) -> bool: return "_VISION_CALL_CONSTANTS" in _unpacked(call) -_KV_CACHE_CALLS = [call for call in _thop_attention_calls() if not _is_vision(call)] +_KV_CACHE_CALLS = _thop_attention_calls(_TREE) def test_every_kv_cache_call_passes_the_pool_storage() -> None: # Against going vacuous: the trunk and the draft layer each call it from the prefill # and the decode target. assert len(_KV_CACHE_CALLS) >= 4 + assert not any(_is_vision(call) for call in _KV_CACHE_CALLS) for call in _KV_CACHE_CALLS: keywords = _keywords(call) for name in _KV_STORAGE_ARGS: @@ -89,7 +91,7 @@ def test_the_unpacked_arguments_carry_no_pool_storage() -> None: """`**_CALL_CONSTANTS` and `**step` (each target's `step_args`) are unpacked into every KV cache call; a storage argument in either would raise as a duplicate keyword or, if the call dropped its own, silently pin the pool to one storage.""" - constants = {kw.arg for kw in _module_dict("_CALL_CONSTANTS").keywords} + constants = {kw.arg for kw in _module_dict(_TREE, "_CALL_CONSTANTS").keywords} assert constants assert not constants & set(_KV_STORAGE_ARGS) for call in _KV_CACHE_CALLS: @@ -110,11 +112,13 @@ def test_the_unpacked_arguments_carry_no_pool_storage() -> None: def test_the_vision_tower_runs_without_kv_storage() -> None: """The encoder-only call, made once, states a bf16, unscaled, cache-less configuration.""" - constants = {kw.arg: kw.value for kw in _module_dict("_VISION_CALL_CONSTANTS").keywords} + constants = { + kw.arg: kw.value for kw in _module_dict(_VISION_TREE, "_VISION_CALL_CONSTANTS").keywords + } assert ast.literal_eval(constants["quant_mode"]) == 0 assert ast.literal_eval(constants["kv_scale_orig_quant"]) is None assert ast.literal_eval(constants["kv_scale_quant_orig"]) is None - vision = [call for call in _thop_attention_calls() if _is_vision(call)] - assert len(vision) == 1 + vision = _thop_attention_calls(_VISION_TREE) + assert len(vision) == 1 and _is_vision(vision[0]) for call in vision: assert not set(_keywords(call)) & set(_KV_STORAGE_ARGS) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision.py new file mode 100644 index 000000000000..6490c6a81cba --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision.py @@ -0,0 +1,568 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The default target's vision path on the CPU: the tower's weights, its calls and angle table, +the splice's bookkeeping and the admission check. + +The tower is built over the weight table's own `vis_*` declarations on the meta device, so +what it reads is what `weights.py` declares. Its per-call body is a stub where a test needs one +(the vision ops need a GPU), and so is the whole tower where the shell's bookkeeping is under +test: every row of a stub tower names the item it came from, so the modality dispatch, item +metadata, chunk slice, frame grouping and the order of concatenated calls run with no engine. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501 + modeling as M, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501 + vision as V, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501 + weights as W, +) +from tensorrt_llm._torch.configs import Qwen4ExpVisionConfig +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalItemMetadata, + MultimodalParams, +) + +pytestmark = pytest.mark.cpu_only + +Shell = M.ModelingV2Qwen38FlashNextNvfp4Sm103Default + +# Qwen3.8-Flash-Next's vision config. +_VISION = dict( + deepstack_visual_indexes=[], + depth=27, + dtype="bfloat16", + hidden_act="gelu_pytorch_tanh", + hidden_size=1152, + in_channels=3, + intermediate_size=4304, + num_heads=16, + num_position_embeddings=2304, + out_hidden_size=2560, + patch_size=16, + spatial_merge_size=2, + temporal_patch_size=2, +) +HIDDEN_OUT = 2560 + + +def _tower_over_the_weight_table() -> V.VisionTower: + """The tower over every `vis_*` weight the table declares, at the checkpoint's shapes, + on meta; `embed` stands in for the activation dtype the tower reads off it.""" + vcfg = Qwen4ExpVisionConfig(**_VISION) + d = SimpleNamespace( + hidden=HIDDEN_OUT, + dtype=torch.bfloat16, + v_depth=vcfg.depth, + v_hidden=vcfg.hidden_size, + v_inter=vcfg.intermediate_size, + v_channels=vcfg.in_channels, + v_tpatch=vcfg.temporal_patch_size, + v_patch=vcfg.patch_size, + v_pos_rows=vcfg.num_position_embeddings, + v_merge_dim=vcfg.hidden_size * vcfg.spatial_merge_size**2, + ) + w = {"embed": torch.empty(8, HIDDEN_OUT, dtype=torch.bfloat16, device="meta")} + for entry in W.MODEL_WEIGHTS.WEIGHTS: + if not entry.name.startswith("vis_"): + continue + for key, _, _ in W.MODEL_WEIGHTS._targets(entry, d): + w[key] = torch.empty(*entry.shape(d), dtype=entry.dtype or d.dtype, device="meta") + return V.VisionTower(w, vcfg) + + +def test_the_tower_reads_the_weight_tables_vision_rows_at_the_ops_layouts(): + """Every view the tower builds exists in the table, and each GEMM operand is the + column-major `[K, N]` view `cublas_mm` reads: the packed qkv row is `3 * 1152` wide and + the merger reads `2 x 2` patches of 1152.""" + tower = _tower_over_the_weight_table() + assert (tower.depth, tower.heads, tower.head_dim, tower.merge) == (27, 16, 72, 2) + assert (tower.pixel_width, tower.merge_dim, tower.grid_side) == (1536, 4608, 48) + assert tower.dtype == torch.bfloat16 + views = tower._w + assert tuple(views.patch_w.shape) == (1152, 3, 2, 16, 16) + assert tuple(views.pos.shape) == (2304, 1152) + assert len(views.blocks) == 27 + for block in views.blocks: + for gemm, (k, n) in ( + (block.qkv, (1152, 3456)), + (block.proj, (1152, 1152)), + (block.fc1, (1152, 4304)), + (block.fc2, (4304, 1152)), + ): + assert tuple(gemm.shape) == (k, n) and gemm.stride() == (1, k) + assert tuple(block.qkv_b.shape) == (3456,) and tuple(block.fc1_b.shape) == (4304,) + assert tuple(views.m_fc1.shape) == (4608, 4608) and views.m_fc1.stride() == (1, 4608) + assert tuple(views.m_fc2.shape) == (4608, HIDDEN_OUT) and views.m_fc2.stride() == (1, 4608) + + +# --------------------------------------------------------------- angle table --- + + +def _nested_loop_positions(grid, merge): + """Per-row `(row, col)` by four nested loops, in spatial-merge order: the template the + rotary entry's test builds its certified tables from.""" + positions = [] + for frames, height, width in grid: + frame = [ + (block_row * merge + r, block_col * merge + c) + for block_row in range(height // merge) + for block_col in range(width // merge) + for r in range(merge) + for c in range(merge) + ] + positions.extend(frame * frames) + return positions + + +@pytest.mark.parametrize( + "grid", [[(1, 62, 90)], [(3, 8, 8), (1, 16, 24)], [(2, 4, 6)]], ids=["62x90", "mixed", "t2"] +) +def test_the_angle_table_is_the_rotary_entrys_certified_table(grid): + """`_rope_tables` builds, bit for bit, the table `vision_rope_2d_packed_qk` is + certified with: positions in spatial-merge order, fp64 angles rounded once to fp32, + height angles then width angles, `[rows, 36]` contiguous.""" + assert V._merge_order_positions(grid, 2).tolist() == [ + list(p) for p in _nested_loop_positions(grid, 2) + ] + cos, sin = V._rope_tables(grid, 2, V.VISION_ROPE_THETA, 72, "cpu") + rope_dim = 36 + inv_freq = 1.0 / ( + V.VISION_ROPE_THETA ** (torch.arange(0, rope_dim, 2, dtype=torch.float64) / rope_dim) + ) + pos = torch.tensor(_nested_loop_positions(grid, 2), dtype=torch.float64) + angles = (pos[:, :, None] * inv_freq).flatten(1).to(torch.float32) + rows = sum(t * h * w for t, h, w in grid) + for got, want in ((cos, angles.cos()), (sin, angles.sin())): + assert got.dtype == torch.float32 and got.is_contiguous() + assert tuple(got.shape) == (rows, rope_dim) + assert torch.equal(got, want) + + +# ---------------------------------------------------------------- tower calls --- + +HIDDEN = 3 +MERGE = 2 +PIXEL_WIDTH = 4 + + +def _frames(grid): + """The item-by-item frame list a grid describes: one `(h, w)` per frame.""" + return [(h, w) for t, h, w in grid for _ in range(t)] + + +def _random_grid(gen): + return [ + (int(torch.randint(1, 6, (1,), generator=gen)), 2 * h, 2 * w) + for h, w in torch.randint( + 1, 5, (int(torch.randint(1, 6, (1,), generator=gen)), 2), generator=gen + ).tolist() + ] + + +def test_frames_are_grouped_whole_and_in_order(): + # 16-row frames, 32 rows to a call: one item's frames split across calls. + assert V._frame_calls([(5, 4, 4)], 32) == [[(2, 4, 4)], [(2, 4, 4)], [(1, 4, 4)]] + # A call takes the end of one item and the start of the next. + grid = [(1, 8, 4), (3, 4, 4), (1, 2, 2)] + assert V._frame_calls(grid, 48) == [[(1, 8, 4), (1, 4, 4)], [(2, 4, 4), (1, 2, 2)]] + assert V._frame_calls(grid, 1 << 20) == [grid] + + +def test_a_call_is_filled_to_the_ceiling_and_never_past_it(): + gen = torch.Generator().manual_seed(0) + for _ in range(200): + grid = _random_grid(gen) + cap = int(torch.randint(64, 300, (1,), generator=gen)) + calls = V._frame_calls(grid, cap) + assert [f for call in calls for f in _frames(call)] == _frames(grid) + sizes = [sum(t * h * w for t, h, w in call) for call in calls] + assert all(0 < size <= cap for size in sizes) + # Greedy: the next call's first frame would not have fit. + for size, call in zip(sizes, calls[1:]): + h, w = _frames(call)[0] + assert size + h * w > cap + + +def test_one_frame_past_the_ceiling_is_refused_and_one_at_it_is_not(): + assert V._frame_calls([(2, 8, 8)], 64) == [[(1, 8, 8)], [(1, 8, 8)]] + with pytest.raises(AssertionError, match="one frame is 8x10 = 80 patch rows"): + V._frame_calls([(1, 4, 4), (3, 8, 10)], 64) + + +def test_the_tower_call_ceiling_is_the_vision_entries_largest_cell(): + """32768 patch rows (8192 merged rows) per call: the ceiling the patch embedding, + position interpolation, layer norm, rotary, GEMM and attention cells reach.""" + assert V.VISION_CALL_MAX_PATCHES == 32768 + + +def _stub_tower(monkeypatch, cap): + """A `VisionTower` whose per-call body is a stub: merged row `j` of a call carries the + pixel value of the call's patch row `4 * j`, which `_pixels` set to the patch's index in + the request.""" + monkeypatch.setattr(V, "VISION_CALL_MAX_PATCHES", cap) + tower = object.__new__(V.VisionTower) + tower.dtype = torch.float32 + tower.pixel_width = PIXEL_WIDTH + tower.merge = MERGE + tower.calls = [] + + def run(pixel_values, grid): + tower.calls.append((pixel_values, grid)) + return pixel_values[:: MERGE * MERGE, :1].expand(-1, HIDDEN).clone() + + tower._run = run + return tower + + +def _pixels(grid): + patches = sum(t * h * w for t, h, w in grid) + return torch.arange(patches, dtype=torch.float32)[:, None].expand(-1, PIXEL_WIDTH).clone() + + +def test_the_tower_calls_whole_frames_and_concatenates_their_rows_in_order(monkeypatch): + grid = [(1, 8, 4), (3, 4, 4), (2, 2, 6), (1, 2, 2)] + tower = _stub_tower(monkeypatch, 48) + out = tower(_pixels(grid), torch.tensor(grid)) + patches = sum(t * h * w for t, h, w in grid) + assert out[:, 0].tolist() == [float(4 * j) for j in range(patches // 4)] + assert [call[1] for call in tower.calls] == V._frame_calls(grid, 48) + start = 0 + for pixel_values, call in tower.calls: + rows = sum(t * h * w for t, h, w in call) + assert pixel_values[:, 0].tolist() == [float(i) for i in range(start, start + rows)] + assert rows <= 48 + start += rows + assert start == patches + + +def test_a_request_past_one_calls_ceiling_is_served_in_two(monkeypatch): + # 33 temporal groups of 32x32 patches (66 frames of 512x512 pixels): 33792 rows, 1024 + # rows past what a single call may carry. + ceiling = V.VISION_CALL_MAX_PATCHES + grid = [(33, 32, 32)] + tower = _stub_tower(monkeypatch, ceiling) + out = tower(_pixels(grid), torch.tensor(grid)) + assert [call[1] for call in tower.calls] == [[(32, 32, 32)], [(1, 32, 32)]] + assert out.shape[0] == 33 * 32 * 32 // 4 + assert out[:, 0].tolist() == [float(4 * j) for j in range(out.shape[0])] + + +def test_the_tower_refuses_pixels_its_grid_does_not_describe(monkeypatch): + tower = _stub_tower(monkeypatch, 64) + with pytest.raises(AssertionError, match="describes 16 patch rows"): + tower(_pixels([(1, 4, 4)])[:12], torch.tensor([(1, 4, 4)])) + with pytest.raises(AssertionError, match="divide by spatial_merge_size 2"): + tower(_pixels([(1, 4, 3)]), torch.tensor([(1, 4, 3)])) + with pytest.raises(AssertionError, match="activation dtype"): + tower(_pixels([(1, 4, 4)]).to(torch.bfloat16), torch.tensor([(1, 4, 4)])) + + +# ------------------------------------------------------------------- splice --- + +IMAGE, VIDEO = 1, 2 + + +class _Tower: + """Stands in for `VisionTower`: every row of item `i` of modality `m` holds `100 * m + + i`, so a spliced row says where it came from.""" + + merge = MERGE + + def __init__(self): + self.calls = 0 + + def __call__(self, pixel_values, grid_thw): + self.calls += 1 + modality = int(pixel_values[0, 0]) + return torch.cat( + [ + torch.full((t * h * w // (MERGE * MERGE), HIDDEN), float(100 * modality + item)) + for item, (t, h, w) in enumerate(grid_thw.tolist()) + ] + ) + + +def _media(modality, grids): + """One `multimodal_data` entry; `pixel_values` only carries the stub's modality.""" + patches = sum(t * h * w for t, h, w in grids) + key = "pixel_values" if modality == IMAGE else "pixel_values_videos" + grid_key = "image_grid_thw" if modality == IMAGE else "video_grid_thw" + return {key: torch.full((patches, 1), float(modality)), grid_key: torch.tensor(grids)} + + +def _params(image=None, video=None, order=None, runtime=None, mrope=True): + data = {} + if mrope: + data["mrope_config"] = {"mrope_position_ids": torch.zeros(3, 1, 8, dtype=torch.int32)} + if image is not None: + data["image"] = _media(IMAGE, image) + if video is not None: + data["video"] = _media(VIDEO, video) + if order is not None: + data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata(item_refs=order) + return MultimodalParams(multimodal_data=data, multimodal_runtime=runtime) + + +def _splice(tower, params, num_rows, padding=0): + """The tower's rows spliced, as the prefill target's break does, over a prompt whose + tokens 1..num_rows are the placeholders and which `padding` rows follow; returns the rows + that landed there, one value per row, or None when the step carries no media.""" + kwargs = { + "multimodal_params": [params], + "mm_token_indices": torch.arange(1, 1 + num_rows, dtype=torch.int64), + } + vision = V.multimodal_embeds(tower, kwargs) + if vision is None: + return None + total = num_rows + 2 + h = torch.full((total + padding, HIDDEN), -1.0) + target = SimpleNamespace(_step=SimpleNamespace(vision=vision, num_tokens=total)) + M.PrefillTarget._splice(target, h) + assert torch.count_nonzero(h[total:]) == 0, "padding rows must be zeroed" + assert torch.all(h[0] == -1.0) and torch.all(h[total - 1] == -1.0) + return h[1 : 1 + num_rows, 0].tolist() + + +def _item(modality, index): + return modality, index + + +# Two videos of 4 and 8 rows, two images of 4 and 1. +VIDEOS = [(2, 2, 4), (2, 4, 4)] +IMAGES = [(1, 4, 4), (1, 2, 2)] +MIXED_ORDER = [_item("video", 0), _item("image", 0), _item("video", 1), _item("image", 1)] +MIXED_ROWS = [200.0] * 4 + [100.0] * 4 + [201.0] * 8 + [101.0] + + +def test_a_text_step_runs_no_tower(): + assert V.multimodal_embeds(None, {}) is None + assert V.multimodal_embeds(None, {"multimodal_params": []}) is None + # Position ids alone (a text prompt through a multimodal processor) are not media. + assert V.multimodal_embeds(None, {"multimodal_params": [_params()]}) is None + + +def test_the_splice_fills_the_placeholders_and_zeroes_the_padding(): + """An image request's rows land on its placeholders in one tower call and the padding + rows after the step are zeroed (`_splice` checks); a text step writes only the padding.""" + tower = _Tower() + assert _splice(tower, _params(image=IMAGES), 5, padding=3) == [100.0] * 4 + [101.0] + assert tower.calls == 1 + target = SimpleNamespace(_step=SimpleNamespace(vision=None, num_tokens=2)) + h = torch.full((4, HIDDEN), -1.0) + M.PrefillTarget._splice(target, h) + assert torch.all(h[:2] == -1.0) and torch.count_nonzero(h[2:]) == 0 + + +def test_video_rows_reach_the_placeholders_item_by_item(): + tower = _Tower() + assert _splice(tower, _params(video=VIDEOS), 12) == [200.0] * 4 + [201.0] * 8 + assert tower.calls == 1 + + +def test_a_mixed_request_is_reassembled_in_prompt_order(): + tower = _Tower() + params = _params(image=IMAGES, video=VIDEOS, order=MIXED_ORDER) + assert _splice(tower, params, len(MIXED_ROWS)) == MIXED_ROWS + assert tower.calls == 2 + + +def test_a_mixed_request_needs_its_manifest(): + with pytest.raises(AssertionError, match="needs item metadata"): + _splice(_Tower(), _params(image=IMAGES, video=VIDEOS), len(MIXED_ROWS)) + + +def test_a_chunk_takes_its_slice_of_the_prompt_order_rows(): + runtime = SimpleNamespace( + num_cached_mm_tokens=3, num_mm_tokens_in_chunk=6, total_embeds_in_request=len(MIXED_ROWS) + ) + params = _params(image=IMAGES, video=VIDEOS, order=MIXED_ORDER, runtime=runtime) + assert _splice(_Tower(), params, 6) == MIXED_ROWS[3:9] + + +def test_a_chunk_without_placeholders_skips_the_tower(): + runtime = SimpleNamespace( + num_cached_mm_tokens=len(MIXED_ROWS), + num_mm_tokens_in_chunk=0, + total_embeds_in_request=len(MIXED_ROWS), + ) + tower = _Tower() + params = _params(image=IMAGES, video=VIDEOS, order=MIXED_ORDER, runtime=runtime) + assert _splice(tower, params, 0) is None + assert tower.calls == 0 + + +def test_a_media_request_needs_the_processors_position_ids(): + with pytest.raises(AssertionError, match="mrope_position_ids"): + _splice(_Tower(), _params(video=VIDEOS, mrope=False), 12) + + +def test_the_placeholder_rows_must_match_the_towers_rows(): + kwargs = { + "multimodal_params": [_params(image=IMAGES)], + "mm_token_indices": torch.arange(4, dtype=torch.int64), + } + with pytest.raises(AssertionError, match="4 placeholder rows and the tower produced 5"): + V.multimodal_embeds(_Tower(), kwargs) + + +def test_the_grids_stay_on_the_host(): + """The executor moves only the pixels and the position ids; the grids the tower and the + admission check read on the host stay there.""" + paths = Shell.multimodal_data_device_paths.fget(None) + assert set(paths) == { + "image.pixel_values", + "video.pixel_values_videos", + "mrope_config.mrope_position_ids", + "mrope_config.mrope_position_deltas", + } + assert not any("grid_thw" in path for path in paths) + + +# ---------------------------------------------------------------- admission --- + +PATCH = 16 +TEXT_TOKENS = 5 + +_GEOMETRY = SimpleNamespace(merge=MERGE, patch=PATCH, pixel_width=PIXEL_WIDTH) + + +def _request_media(modality, grids): + """A payload whose pixel tensor has the right shape and no storage.""" + patches = sum(t * h * w for t, h, w in grids) + key = "pixel_values" if modality == "image" else "pixel_values_videos" + grid_key = "image_grid_thw" if modality == "image" else "video_grid_thw" + return { + key: torch.empty(patches, PIXEL_WIDTH, device="meta"), + grid_key: torch.tensor(grids), + } + + +def _request(image=None, video=None, order=None, **overrides): + """A request whose prompt is `TEXT_TOKENS` text tokens around the placeholders the media + produce, with the position ids and placeholder cumsum the processor emits.""" + rows = sum(t * h * w for t, h, w in (image or []) + (video or [])) // (MERGE * MERGE) + prompt_len = rows + TEXT_TOKENS + data = { + "mrope_config": {"mrope_position_ids": torch.zeros(3, 1, prompt_len, dtype=torch.int32)}, + "multimodal_embed_mask_cumsum": torch.cat( + [torch.arange(1, rows + 1), torch.full((TEXT_TOKENS,), rows)] + ), + } + if image is not None: + data["image"] = _request_media("image", image) + if video is not None: + data["video"] = _request_media("video", video) + if order is not None: + data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata(item_refs=order) + fields = {"py_multimodal_data": data, "prompt_len": prompt_len} + fields.update(overrides) + return SimpleNamespace(**fields) + + +def _admit(request): + """Through the shell's own hook, as the executor calls it.""" + Shell.validate_request(SimpleNamespace(vision_tower=_GEOMETRY), request) + + +def test_clean_requests_are_admitted(): + # A single-modality request needs no manifest. + _admit(_request(image=IMAGES)) + _admit(_request(video=VIDEOS)) + _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER)) + + +def test_a_request_without_media_costs_nothing(): + # Not even the other request attributes are read. + _admit(SimpleNamespace(py_multimodal_data=None)) + _admit(SimpleNamespace(py_multimodal_data={})) + # A warmup request carries position ids and no media. + _admit( + SimpleNamespace( + py_multimodal_data={"mrope_config": {"mrope_position_ids": torch.zeros(3, 1, 8)}} + ) + ) + + +def test_the_longest_admitted_frame_is_the_calls_ceiling(monkeypatch): + monkeypatch.setattr(V, "VISION_CALL_MAX_PATCHES", 64) + _admit(_request(image=[(1, 8, 8)])) + _admit(_request(video=[(3, 4, 16)])) + with pytest.raises(ValueError, match="frames of 8x10 = 80 patches"): + _admit(_request(image=[(1, 8, 8), (1, 8, 10)])) + + +def test_a_frame_past_the_ceiling_fails_the_request(): + # 3072x3072 pixels: a 192x192 grid, 36864 patches in one frame. + request = _request(video=[(1, 192, 192)]) + with pytest.raises(ValueError, match=r"video 0 has frames of 192x192 = 36864 patches"): + _admit(request) + # The message names the pixels and the bound. + with pytest.raises(ValueError, match=r"\(3072x3072 pixels\).*at most 32768 patches"): + _admit(request) + + +def test_the_frame_bound_is_per_frame_not_per_item(): + # 33 temporal groups of 32x32 patches: 33792 rows in all, 1024 past one call. + _admit(_request(video=[(33, 32, 32)])) + with pytest.raises(ValueError, match="video 1 has frames of 192x192"): + _admit(_request(video=[(4, 128, 128), (1, 192, 192)])) + + +def test_a_malformed_grid_is_refused(): + request = _request(image=[(1, 4, 4)]) + request.py_multimodal_data["image"]["image_grid_thw"] = torch.tensor([(1, 4, 5)]) + with pytest.raises(ValueError, match="spatial_merge_size 2"): + _admit(request) + request.py_multimodal_data["image"]["image_grid_thw"] = torch.tensor([1, 4, 4]) + with pytest.raises(ValueError, match="one \\(t, h, w\\) row per image"): + _admit(request) + + +def test_pixels_that_disagree_with_the_grid_are_refused(): + request = _request(video=VIDEOS) + video = request.py_multimodal_data["video"] + video["pixel_values_videos"] = torch.empty(5, PIXEL_WIDTH, device="meta") + with pytest.raises(ValueError, match="pixel_values_videos is .*describe 48 patch rows"): + _admit(request) + del video["video_grid_thw"] + with pytest.raises(ValueError, match="this target reads 'pixel_values_videos'"): + _admit(request) + + +def test_missing_or_mis_sized_position_ids_fail_the_request(): + request = _request(video=VIDEOS) + request.py_multimodal_data["mrope_config"]["mrope_position_ids"] = torch.zeros(3, 1, 40) + with pytest.raises(ValueError, match="one position per prompt token"): + _admit(request) + del request.py_multimodal_data["mrope_config"] + with pytest.raises(ValueError, match="mrope_position_ids.*got None"): + _admit(request) + + +def test_a_mixed_request_lists_its_items(): + with pytest.raises(ValueError, match="needs item metadata"): + _admit(_request(image=IMAGES, video=VIDEOS)) + with pytest.raises(ValueError, match="lists image items \\[0\\]; the payload carries 2"): + _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER[:3])) + with pytest.raises(ValueError, match="is not an item of this request"): + _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER + [_item("audio", 0)])) + + +def test_the_placeholder_count_is_the_sum_over_both_modalities(): + request = _request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER) + request.py_multimodal_data["multimodal_embed_mask_cumsum"] = torch.arange(1, 14) + with pytest.raises(ValueError, match="holds 13 image/video placeholder tokens.*17"): + _admit(request) + del request.py_multimodal_data["multimodal_embed_mask_cumsum"] + _admit(request) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_weight_shards.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_weight_shards.py new file mode 100644 index 000000000000..a25b790cd6f0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_weight_shards.py @@ -0,0 +1,799 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The default target's per-rank weight table, checked against the checkpoint on the CPU. + +From `model.safetensors.index.json` and the safetensors headers alone -- no tensor data is +read -- at each member (tp4, tep4, dep4; dep2, the disaggregated engines' layout, at the end): + +* **coverage**: at every rank, with MTP off and on, every checkpoint tensor is either read by + the rank's table or predicted unread by `expected_unconsumed`, never both; across a member's + ranks every tensor is read by some rank, and each routed or draft expert's tensors by + exactly one -- except tp4's routed experts, which every rank reads at its intermediate + slice; +* **shapes**: at ranks 0 and 3, every source read through its cut lies inside the stored + tensor, and the cut run through the table's own transform (on meta tensors) has exactly + the shape and dtype of the parameter it fills. A rank reads a split tensor at its share + only, and every other tensor whole, once; +* **the modules' shapes**: the per-rank widths the forward is written against, and the head + ranges the same rank's K/V and recurrent pools hold (`pool.pool_geometry`). + +And, on small synthetic tensors whose values name their rows, that each cut selects the +rank's own heads in the order the fused operands lay them out, and that tp4's routed +experts hold the rank's intermediate slice with zero padding at the end of each half. +""" + +from __future__ import annotations + +import collections +import json +import re +import struct +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from utils.llm_data import llm_models_root + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501 + pool as P, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501 + weights as WT, +) +from tensorrt_llm._torch.configs import Qwen4ExpTextConfig, Qwen4ExpVisionConfig +from tensorrt_llm.mapping import Mapping + +pytestmark = pytest.mark.cpu_only + +CHECKPOINT = "Qwen3.8-Flash-Next-NVFP4" + +#: Each member's mapping: attention TP4 with the experts split by intermediate (tp4) or by +#: expert (tep4), and attention DP with expert parallelism (dep4). +MEMBERS = { + "tp4": dict(world_size=4, tp_size=4, moe_tp_size=4, moe_ep_size=1, enable_attention_dp=False), + "tep4": dict(world_size=4, tp_size=4, moe_tp_size=1, moe_ep_size=4, enable_attention_dp=False), + "dep4": dict(world_size=4, tp_size=4, moe_tp_size=1, moe_ep_size=4, enable_attention_dp=True), +} +#: Both engines of the disaggregated-serving test: attention DP over two ranks, the routed +#: experts in two 256-expert windows. It serves without speculation. +DEP2 = dict(world_size=2, tp_size=2, moe_tp_size=1, moe_ep_size=2, enable_attention_dp=True) +_LAYOUTS = {**MEMBERS, "dep2": DEP2} +RANKS = (0, 1, 2, 3) +SHAPE_RANKS = (0, 3) + +_STORED_DTYPES = { + "BF16": torch.bfloat16, + "F32": torch.float32, + "U8": torch.uint8, + "F8_E4M3": torch.float8_e4m3fn, + "I64": torch.int64, +} + +#: A routed or draft expert's tensors, as opposed to its per-tensor scalars (which every +#: rank reads, for the whole routing space). +_EXPERT_TENSOR = re.compile( + r"\.mlp\.experts\.\d+\.\w+_proj\.(weight|weight_scale|weight_scale_inv)$" +) +_TRUNK_EXPERT_TENSOR = re.compile( + r"^model\.language_model\.layers\.\d+\.mlp\.experts\.\d+\.\w+_proj\.(weight|weight_scale)$" +) + + +# ------------------------------------------------------------------ fixtures -- + + +@pytest.fixture(scope="module") +def checkpoint() -> Path: + root = llm_models_root() + path = None if root is None else root / CHECKPOINT + if path is None or not (path / "model.safetensors.index.json").is_file(): + pytest.skip(f"{CHECKPOINT} is not under LLM_MODELS_ROOT") + return path + + +@pytest.fixture(scope="module") +def stored(checkpoint) -> dict[str, tuple[tuple[int, ...], torch.dtype]]: + """Every checkpoint tensor's stored shape and dtype, from the index and the headers.""" + index = json.loads((checkpoint / "model.safetensors.index.json").read_text())["weight_map"] + specs = {} + for file in sorted(set(index.values())): + with open(checkpoint / file, "rb") as fh: + (length,) = struct.unpack(" Mapping: + return Mapping(rank=rank, gpus_per_node=4, pp_size=1, **_LAYOUTS[member]) + + +def _dims(configs, member: str, rank: int, mtp: bool = True, prefill_copy: bool = False): + text, vision = configs + return WT.rank_dims( + text, + vision, + _mapping(member, rank), + mtp_enabled=mtp, + trunk_moe_backend="TRTLLM", + moe_prefill_copy=prefill_copy, + dtype=torch.bfloat16, + ) + + +def _attention_tp(member: str) -> int: + layout = _LAYOUTS[member] + return 1 if layout["enable_attention_dp"] else layout["tp_size"] + + +def _lm_head_rows(member: str, vocab: int) -> int: + """The rows of the lm_head the shell builds: vocab-parallel under attention TP, whole + under attention DP (`DecoderModelForCausalLM`).""" + return vocab // _attention_tp(member) + + +def _rows_per_shard(stored, d) -> int: + shapes = {stored[key][0] for key in WT.ngram_shard_keys(d)} + assert len(shapes) == 1, shapes + return next(iter(shapes))[0] + + +def _rank_manifest(stored, d, member: str) -> dict[str, list[tuple]]: + """The rank's manifest as `load` builds it, with the n-gram table and the shell's lm_head.""" + manifest = WT.MODEL_WEIGHTS.manifest(d) + manifest["ple_table"] = WT.ngram_table_sources(d, _rows_per_shard(stored, d)) + cut, pad = WT.lm_head_share(_lm_head_rows(member, d.vocab), d.vocab, d.tp_rank) + manifest["lm_head.weight"] = [("lm_head.weight", None, cut, pad)] + return manifest + + +def _keys_and_cuts(ckpt_key, cut): + if isinstance(ckpt_key, tuple): + return list(zip(ckpt_key, cut if cut is not None else (None,) * len(ckpt_key))) + return [(ckpt_key, cut)] + + +def _claimed(manifest) -> set: + return { + key + for sources in manifest.values() + for ckpt_key, _, cut, _ in sources + for key, _ in _keys_and_cuts(ckpt_key, cut) + } + + +def _cut_numel(shape: tuple[int, ...], cut) -> int: + """Elements of a `shape` tensor `cut` reads, after checking each slice lies inside it.""" + if cut is None: + return int(torch.Size(shape).numel()) + extents = list(shape) + for axis, s in enumerate(cut): + assert isinstance(s, slice) and s.step in (None, 1), (shape, cut) + start = 0 if s.start is None else s.start + stop = shape[axis] if s.stop is None else s.stop + assert 0 <= start < stop <= shape[axis], f"cut {cut} leaves the stored {shape}" + extents[axis] = stop - start + return int(torch.Size(extents).numel()) + + +class _HeaderEntry: + """A checkpoint tensor known by its header: indexing it returns a meta tensor of the + indexed shape, and nothing is ever read.""" + + def __init__(self, shape, dtype): + self.shape, self.dtype = torch.Size(shape), dtype + self.ndim = len(shape) + + def __getitem__(self, index): + return torch.empty(self.shape, dtype=self.dtype, device="meta")[index] + + +def _share(member: str, key: str) -> float: + """The fraction of `key` a rank of `member` holds, worked out here from the member alone: + attention q/o by q head and k/v by KV head (two heads over the attention TP ranks), the + Gated-DeltaNet tensors by head over the same ranks, the shell's lm_head by vocabulary, + and at tp4 each routed expert's tensors by intermediate. Everything else is read whole.""" + tp = _attention_tp(member) + if member == "tp4" and _TRUNK_EXPERT_TENSOR.match(key): + return 1 / MEMBERS[member]["moe_tp_size"] + if re.search(r"\.self_attn\.(q_proj|o_proj)\.weight$", key) or key == "lm_head.weight": + return 1 / tp + if re.search(r"\.self_attn\.(k_proj|v_proj)\.weight$", key): + return 1 / min(tp, 2) + if re.search( + r"\.linear_attn\.((in_proj_qkv|in_proj_z|in_proj_a|in_proj_b|conv1d|out_proj)\.weight" + r"|A_log|dt_bias)$", + key, + ): + return 1 / tp + return 1.0 + + +# ----------------------------------------------------------------- coverage -- + + +@pytest.mark.parametrize("mtp", [False, True], ids=["no-mtp", "mtp"]) +@pytest.mark.parametrize("member", list(MEMBERS)) +def test_every_checkpoint_tensor_is_claimed(stored, configs, member, mtp): + everything = set(stored) + read_by = collections.Counter() + for rank in RANKS: + d = _dims(configs, member, rank, mtp) + claimed = _claimed(_rank_manifest(stored, d, member)) + unread = WT.MODEL_WEIGHTS.expected_unconsumed(d, stored.keys()) + assert claimed <= everything, f"rank {rank} reads keys the checkpoint lacks: " + str( + sorted(claimed - everything)[:8] + ) + assert not claimed & unread, f"rank {rank}: read and predicted unread: " + str( + sorted(claimed & unread)[:8] + ) + assert claimed | unread == everything, ( + f"rank {rank}: neither read nor predicted {sorted(everything - claimed - unread)[:8]}" + ) + read_by.update(claimed) + + mtp_keys = set() if mtp else {k for k in everything if k.startswith("mtp.")} + assert set(read_by) == everything - mtp_keys, "tensors no rank reads: " + str( + sorted(everything - mtp_keys - set(read_by))[:8] + ) + + def readers(key): + # tp4: every rank reads every routed expert, at its slice (checked by the shape test). + return len(RANKS) if member == "tp4" and _TRUNK_EXPERT_TENSOR.match(key) else 1 + + wrong = sorted(k for k in read_by if _EXPERT_TENSOR.search(k) and read_by[k] != readers(k)) + assert not wrong, f"expert tensors not read by their expected rank count: {wrong[:8]}" + + +# -------------------------------------------------------------------- shapes -- + + +@pytest.mark.parametrize("prefill_copy", [False, True], ids=["trtllm-gen", "with-cutedsl-copy"]) +@pytest.mark.parametrize("rank", SHAPE_RANKS) +@pytest.mark.parametrize("member", list(MEMBERS)) +def test_rank_operands_have_their_parameter_shapes(stored, configs, member, rank, prefill_copy): + """Every source, read through its cut and run through its transform on meta tensors, + fills its parameter exactly; the rank reads each tensor at its share, and once per + layout it is kept in (the routed experts twice with the CuteDSL prefill copy).""" + d = _dims(configs, member, rank, mtp=True, prefill_copy=prefill_copy) + with torch.device("meta"): + params = dict(WT.MODEL_WEIGHTS.declare(d)) + # What `load` adds beside the table: the n-gram table and the shell's lm_head. + params["ple_table"] = torch.empty( + d.ple_shards * _rows_per_shard(stored, d), d.ple_head_dim, dtype=d.ple_table_dtype + ) + params["lm_head.weight"] = torch.empty( + _lm_head_rows(member, d.vocab), d.hidden, dtype=torch.bfloat16 + ) + entries = {key: _HeaderEntry(*spec) for key, spec in stored.items()} + manifest = _rank_manifest(stored, d, member) + assert manifest.keys() == params.keys() + + read = collections.Counter() + results: dict = {} + for param_key, sources in manifest.items(): + param = params[param_key] + for ckpt_key, index, cut, transform in sources: + key_cuts = _keys_and_cuts(ckpt_key, cut) + for key, key_cut in key_cuts: + read[key] += _cut_numel(stored[key][0], key_cut) + dst = param if index is None else param[index] + # Sources that agree on the stored tensors, their cuts and the transform agree + # on the result: run each distinct one once. + signature = (transform, tuple((stored[k], c) for k, c in key_cuts)) + if signature not in results: + src = WT.MODEL_WEIGHTS.read(entries, ckpt_key, cut, transform, d) + results[signature] = (tuple(src.shape), src.dtype) + assert results[signature] == (tuple(dst.shape), dst.dtype), ( + param_key, + ckpt_key, + results[signature], + (tuple(dst.shape), dst.dtype), + ) + + def copies(key): + return 2 if prefill_copy and _TRUNK_EXPERT_TENSOR.match(key) else 1 + + wrong = { + key: (n, _share(member, key), copies(key)) + for key, n in read.items() + if n != round(torch.Size(stored[key][0]).numel() * _share(member, key)) * copies(key) + } + assert not wrong, f"tensors not read at the rank's share, once: {sorted(wrong.items())[:8]}" + + +#: Per-rank shapes of the attention and Gated-DeltaNet weights the forward is written +#: against: (under attention TP, under attention DP). +_SPLIT_SHAPES = { + # [q 6x256 | k 256 | v 256 | gate 6x256 | index 4x128 + 128 | pad 128] + "l3_attn_in_w": ((4352, 2560), (14080, 2560)), + "l3_o_w": ((2560, 1536), (2560, 6144)), + # [q 4x128 | k 4x128 | v 12x128 | z 12x128 | a 12 | b 12] + "l0_in_proj_w": ((4120, 2560), (16480, 2560)), + "l0_conv_w": ((2560, 1, 4), (10240, 1, 4)), + "l0_A_log": ((12,), (48,)), + "l0_dt_bias": ((12,), (48,)), + "l0_out_w": ((2560, 1536), (2560, 6144)), + "l48_mtp_q_w": ((3072, 2560), (12288, 2560)), + "l48_mtp_k_w": ((256, 2560), (512, 2560)), + "l48_mtp_v_w": ((256, 2560), (512, 2560)), + "l48_mtp_o_w": ((2560, 1536), (2560, 6144)), +} + +#: The same at every member and rank: replicated weights and the draft layer's 128-expert +#: FP8 window. +_COMMON_SHAPES = { + "l3_idx_qn": (128,), + "l0_gnorm": (128,), + "l0_router": (512, 2560), + "l0_sh_gu": (2 * 640 + 8, 2560), + "l0_sh_dn": (2560, 640), + "l0_e_isc1": (512,), + "l0_ahc_down": (320, 10240), + "fhc_up": (10240, 320), + "embed": (248320, 2560), + "ple_key_w": (10240, 2560), + "l48_mtp_idx_w": (640, 2560), + "l48_mtp_fc_e": (2560, 2560), + "l48_mtp_fc1_w": (128, 1280, 2560), + "l48_mtp_fc1_s": (128, 10, 20), + "l48_mtp_fc2_w": (128, 2560, 640), + "l48_mtp_fc2_s": (128, 20, 5), +} + +#: The routed experts a rank holds: a 128-expert window of whole experts (tep4, dep4), or +#: all 512 at a 160-row intermediate slice stored 256 wide (tp4). +_EXPERT_SHAPES = { + "tep4": { + "l0_fc1_w": (128, 1280, 1280), + "l0_fc1_s": (128, 1280, 160), + "l0_fc2_w": (128, 2560, 320), + "l0_fc2_s": (128, 2560, 40), + }, + "tp4": { + "l0_fc1_w": (512, 512, 1280), + "l0_fc1_s": (512, 512, 160), + "l0_fc2_w": (512, 2560, 128), + "l0_fc2_s": (512, 2560, 16), + }, +} +_EXPERT_SHAPES["dep4"] = _EXPERT_SHAPES["tep4"] + + +@pytest.mark.parametrize("rank", SHAPE_RANKS) +@pytest.mark.parametrize("member", list(MEMBERS)) +def test_rank_shapes_are_what_the_modules_expect(stored, configs, member, rank): + d = _dims(configs, member, rank) + with torch.device("meta"): + params = WT.MODEL_WEIGHTS.declare(d) + dp = MEMBERS[member]["enable_attention_dp"] + for key, shapes in _SPLIT_SHAPES.items(): + assert tuple(params[key].shape) == shapes[dp], (key, tuple(params[key].shape)) + for key, shape in _COMMON_SHAPES.items(): + assert tuple(params[key].shape) == shape, (key, tuple(params[key].shape)) + for key, shape in _EXPERT_SHAPES[member].items(): + assert tuple(params[key].shape) == shape, (key, tuple(params[key].shape)) + window = range(512) if member == "tp4" else range(128 * rank, 128 * (rank + 1)) + experts = [s[0] for s in WT.MODEL_WEIGHTS.manifest(d)["l0_fc2_w"]] + assert experts == [ + f"model.language_model.layers.0.mlp.experts.{e}.down_proj.weight" for e in window + ] + drafts = [s[0] for s in WT.MODEL_WEIGHTS.manifest(d)["l48_mtp_fc2_w"]] + assert drafts == [ + f"mtp.layers.0.mlp.experts.{e}.down_proj.weight" + for e in range(128 * rank, 128 * (rank + 1)) + ] + + # The heads this rank's weights hold are the heads its pools hold. + g = P.pool_geometry( + configs[0], + _mapping(member, rank), + tokens_per_block=32, + kv_cache_quant_algo=None, + dtype=torch.bfloat16, + ) + assert (d.heads_kv, d.kv_head_start) == (g.heads_kv, g.kv_head_start) + assert (d.num_k_heads, d.k_head_start) == (g.num_k_heads, g.k_head_start) + assert (d.num_v_heads, d.v_head_start) == (g.num_v_heads, g.v_head_start) + assert params["l0_conv_w"].shape[0] == g.conv_dim + assert params["l0_A_log"].shape[0] == g.num_v_heads + + +def test_lm_head_reads_the_rows_the_shell_built(): + vocab = 248320 + for rank in RANKS: + cut, _ = WT.lm_head_share(vocab // 4, vocab, rank) + assert cut == (slice(rank * 62080, (rank + 1) * 62080),) + cut, _ = WT.lm_head_share(vocab, vocab, rank) + assert cut == (slice(0, vocab),) + + +# ---------------------------------------------------------------------- load -- + + +def _stub_model(configs, member: str, rank: int, mtp: bool): + """The shell and core `load` reads, on the meta device: the core's configs, mapping and + engine choices, its declared parameters, and the shell's lm_head.""" + text, vision = configs + mapping = _mapping(member, rank) + core = SimpleNamespace( + model_config=SimpleNamespace( + pretrained_config=text, + mapping=mapping, + spec_config=object() if mtp else None, + torch_dtype=torch.bfloat16, + ), + vision_config=vision, + trunk_moe_backend="TRTLLM", + moe_prefill_copy=False, + ple_host_offload=False, + ) + d = WT.MODEL_WEIGHTS.dims(core) + with torch.device("meta"): + core.w = WT.MODEL_WEIGHTS.declare(d) + lm_head = torch.nn.Parameter( + torch.empty(_lm_head_rows(member, d.vocab), d.hidden, dtype=torch.bfloat16), + requires_grad=False, + ) + + def named_parameters(): + yield from ((f"model.w.{k}", v) for k, v in core.w.items()) + yield "lm_head.weight", lm_head + + return SimpleNamespace( + model=core, lm_head=SimpleNamespace(weight=lm_head), named_parameters=named_parameters + ) + + +def _meta_checkpoint(stored) -> dict: + """Every checkpoint tensor as its header: 0-dim scalars as meta tensors, the rest as + entries that index into meta tensors.""" + return { + key: torch.empty((), dtype=dtype, device="meta") + if not shape + else _HeaderEntry(shape, dtype) + for key, (shape, dtype) in stored.items() + } + + +@pytest.mark.parametrize( + "member, rank, mtp", + [("tep4", 3, True), ("dep4", 0, True), ("tep4", 1, False), ("tp4", 2, True)], +) +def test_load_fills_every_parameter_and_covers_the_checkpoint( + stored, configs, member, rank, mtp, monkeypatch +): + """`load` end to end at the real shapes, on meta tensors: the n-gram table declared from + its shards, every manifest row filled through its cut, the shell's lm_head rows, and + both coverage asserts.""" + monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda *a, **k: None) + model = _stub_model(configs, member, rank, mtp) + WT.MODEL_WEIGHTS.load(model, _meta_checkpoint(stored)) + assert tuple(model.model.w["ple_table"].shape) == (320_001_536, 160) + + +# --------------------------------------------------- the cuts pick the heads -- + +_HIDDEN = 2 # the projections' input axis, shrunk: no cut touches it +_FAKE_OFFSETS = dict(q=1, k=2, v=3, idx=4, o=5, qkv=10, z=11, a=12, b=13, conv=14, a_log=15, + dtb=16, out=17) # fmt: skip + + +def _ids(tag: str, rows: int, *trailing: int) -> torch.Tensor: + """`rows` rows whose every value is `tag`'s offset plus the row's index.""" + ids = torch.arange(rows, dtype=torch.int64) + _FAKE_OFFSETS[tag] * 10**7 + return ids.view(rows, *([1] * len(trailing))).expand(rows, *trailing).contiguous() + + +def _col_ids(tag: str, cols: int) -> torch.Tensor: + return _ids(tag, cols, _HIDDEN).t().contiguous() + + +def _fake_checkpoint(stored) -> dict[str, torch.Tensor]: + """Attention layer 3, Gated-DeltaNet layer 0 and the MTP attention, at their stored row + counts and a two-wide input axis.""" + fake = {} + + def rows_of(key): + return stored[key][0][0] + + for p in ("model.language_model.layers.3.self_attn", "mtp.layers.0.self_attn"): + for tag, name in (("q", "q_proj"), ("k", "k_proj"), ("v", "v_proj")): + fake[f"{p}.{name}.weight"] = _ids(tag, rows_of(f"{p}.{name}.weight"), _HIDDEN) + key = f"{p}.indexer.index_qk_proj.weight" + fake[key] = _ids("idx", rows_of(key), _HIDDEN) + fake[f"{p}.o_proj.weight"] = _col_ids("o", stored[f"{p}.o_proj.weight"][0][1]) + p = "model.language_model.layers.0.linear_attn" + for tag, name in (("qkv", "in_proj_qkv"), ("z", "in_proj_z"), ("a", "in_proj_a"), + ("b", "in_proj_b")): # fmt: skip + fake[f"{p}.{name}.weight"] = _ids(tag, rows_of(f"{p}.{name}.weight"), _HIDDEN) + fake[f"{p}.conv1d.weight"] = _ids("conv", rows_of(f"{p}.conv1d.weight"), 1, 4) + fake[f"{p}.A_log"] = _ids("a_log", rows_of(f"{p}.A_log")) + fake[f"{p}.dt_bias"] = _ids("dtb", rows_of(f"{p}.dt_bias")) + fake[f"{p}.out_proj.weight"] = _col_ids("out", stored[f"{p}.out_proj.weight"][0][1]) + return fake + + +def _assemble(d, fake, sources) -> torch.Tensor: + """A parameter filled from its sources the way `load` fills it.""" + if len(sources) == 1 and sources[0][1] is None: + ckpt_key, _, cut, transform = sources[0] + return WT.MODEL_WEIGHTS.read(fake, ckpt_key, cut, transform, d) + parts = [ + (index, WT.MODEL_WEIGHTS.read(fake, ckpt_key, cut, transform, d)) + for ckpt_key, index, cut, transform in sources + ] + rows = max(index[0].stop for index, _ in parts) + out = torch.full((rows, *parts[0][1].shape[1:]), -1, dtype=parts[0][1].dtype) + for index, part in parts: + out[index] = part + return out + + +def _expected_rows(tag: str, starts) -> torch.Tensor: + return torch.tensor([_FAKE_OFFSETS[tag] * 10**7 + r for r in starts], dtype=torch.int64) + + +def _block(start: int, width: int) -> list[int]: + return list(range(start, start + width)) + + +@pytest.mark.parametrize("rank", RANKS) +@pytest.mark.parametrize("member", ["tep4", "dep4"]) +def test_cuts_select_the_rank_heads(stored, configs, member, rank): + """tp4 cuts attention and Gated-DeltaNet exactly as tep4 does: the members differ only + in the routed experts.""" + d = _dims(configs, member, rank) + fake = _fake_checkpoint(stored) + manifest = WT.MODEL_WEIGHTS.manifest(d) + tp = _attention_tp(member) + # Under attention DP each rank holds every head: the split is over one rank. + r = rank if tp > 1 else 0 + q_heads = _block(r * 24 // tp, 24 // tp) + kv_heads = [r * 2 // tp] if tp > 1 else [0, 1] + k_heads = _block(r * 16 // tp, 16 // tp) + v_heads = _block(r * 48 // tp, 48 // tp) + + def rows(heads, width, base=0): + return [base + h * width + j for h in heads for j in range(width)] + + q_rows = [h * 512 + j for h in q_heads for j in range(256)] + gate_rows = [h * 512 + 256 + j for h in q_heads for j in range(256)] + attn_in = _assemble(d, fake, manifest["l3_attn_in_w"]) + expected = torch.cat( + ( + _expected_rows("q", q_rows), + _expected_rows("k", rows(kv_heads, 256)), + _expected_rows("v", rows(kv_heads, 256)), + _expected_rows("q", gate_rows), + _expected_rows("idx", range(640)), + torch.zeros(128, dtype=torch.int64), + ) + ) + assert torch.equal(attn_in[:, 0], expected) and torch.equal(attn_in[:, 1], expected) + + o_cols = _expected_rows("o", rows(q_heads, 256)) + for key in ("l3_o_w", "l48_mtp_o_w"): + assert torch.equal(_assemble(d, fake, manifest[key])[0], o_cols), key + assert torch.equal(_assemble(d, fake, manifest["l48_mtp_q_w"])[:, 0], + _expected_rows("q", rows(q_heads, 512))) # fmt: skip + assert torch.equal(_assemble(d, fake, manifest["l48_mtp_k_w"])[:, 0], + _expected_rows("k", rows(kv_heads, 256))) # fmt: skip + + def qkv(tag): + return torch.cat( + ( + _expected_rows(tag, rows(k_heads, 128)), + _expected_rows(tag, rows(k_heads, 128, 2048)), + _expected_rows(tag, rows(v_heads, 128, 4096)), + ) + ) + + in_proj = torch.cat( + ( + qkv("qkv"), + _expected_rows("z", rows(v_heads, 128)), + _expected_rows("a", v_heads), + _expected_rows("b", v_heads), + ) + ) + assert torch.equal(_assemble(d, fake, manifest["l0_in_proj_w"])[:, 0], in_proj) + assert torch.equal(_assemble(d, fake, manifest["l0_conv_w"])[:, 0, 0], qkv("conv")) + assert torch.equal(_assemble(d, fake, manifest["l0_A_log"]), _expected_rows("a_log", v_heads)) + assert torch.equal(_assemble(d, fake, manifest["l0_dt_bias"]), _expected_rows("dtb", v_heads)) + assert torch.equal(_assemble(d, fake, manifest["l0_out_w"])[0], + _expected_rows("out", rows(v_heads, 128))) # fmt: skip + + +# ------------------------------------------ tp4: the padded intermediate slice -- + + +def _blk32_gather(rows: int) -> torch.Tensor: + """The 32-row block shuffle's gather index, written out here: inside each block of 32 + rows, source row `4u + v` lands at destination row `8v + u`.""" + j = torch.arange(32) + idx = torch.empty(32, dtype=torch.long) + idx[(j % 4) * 8 + j // 4] = j + return idx.repeat(rows // 32) + torch.arange(rows // 32).repeat_interleave(32) * 32 + + +def _interleave_gather(rows: int) -> torch.Tensor: + """`[up | gate]` -> up0, gate0, up1, gate1, ...""" + p = torch.empty(rows, dtype=torch.long) + p[0::2] = torch.arange(rows // 2) + p[1::2] = torch.arange(rows // 2, rows) + return p + + +def _swizzle_flat(s: torch.Tensor) -> torch.Tensor: + """The 128x4 block-scale swizzle of one `[M, C]` matrix, flat: scale `(m, c)` lands at + `(m // 128) * 512 * (C // 4) + (c // 4) * 512 + (m % 32) * 16 + ((m % 128) // 32) * 4 + + c % 4`.""" + m, c = s.shape + v = s.reshape(m // 128, 4, 32, c // 4, 4).permute(0, 3, 2, 1, 4) + return v.reshape(-1).contiguous() + + +def _zero_rows(x: torch.Tensor, rows: int) -> torch.Tensor: + return torch.cat((x, x.new_zeros(rows - x.shape[0], *x.shape[1:])), 0) + + +def _zero_cols(x: torch.Tensor, cols: int) -> torch.Tensor: + return torch.cat((x, x.new_zeros(x.shape[0], cols - x.shape[1])), 1) + + +@pytest.mark.parametrize("expert", [0, 511]) +@pytest.mark.parametrize("rank", RANKS) +def test_tp4_experts_hold_the_rank_slice_zero_padded_per_half(configs, rank, expert, monkeypatch): + """At tp4 rank `r` reads rows `[160r, 160r + 160)` of an expert's `up_proj` and + `gate_proj` (codes and block scales) and the same columns of its `down_proj`; each half is + zero-extended to 256 rows (FC2 to 256 columns) *before* the `[up ; gate]` concatenation, + then interleaved, block-shuffled and swizzled as at the whole width. The bytes are + nonzero, so every zero in the operand is padding. The scale swizzle is the formula + `fp4_block_scale_moe_runner`'s test writes out, standing in for the native op on + the CPU; the shrunk axes (FC1's input bytes, FC2's output rows) take no cut.""" + monkeypatch.setattr(WT, "_swizzle", _swizzle_flat) + d = _dims(configs, "tp4", rank, mtp=False) + assert (d.expert_inter_share, d.expert_inter_start, d.expert_inter) == (160, 160 * rank, 256) + gen = torch.Generator().manual_seed(rank * 1000 + expert) + + def nonzero(*shape): + return torch.randint(1, 256, shape, generator=gen, dtype=torch.uint8) + + p = f"model.language_model.layers.0.mlp.experts.{expert}" + fake = { + f"{p}.up_proj.weight": nonzero(640, 4), + f"{p}.gate_proj.weight": nonzero(640, 4), + f"{p}.up_proj.weight_scale": nonzero(640, 4).view(torch.float8_e4m3fn), + f"{p}.gate_proj.weight_scale": nonzero(640, 4).view(torch.float8_e4m3fn), + f"{p}.down_proj.weight": nonzero(128, 320), + f"{p}.down_proj.weight_scale": nonzero(128, 40).view(torch.float8_e4m3fn), + } + manifest = WT.MODEL_WEIGHTS.manifest(d) + + def operand(param_key): + (row,) = [s for s in manifest[param_key] if s[1] == (expert,)] + ckpt_key, _, cut, transform = row + return WT.MODEL_WEIGHTS.read(fake, ckpt_key, cut, transform, d).view(torch.uint8) + + rows = slice(160 * rank, 160 * rank + 160) + + def fc1(kind): + up = fake[f"{p}.up_proj.{kind}"].view(torch.uint8)[rows] + gate = fake[f"{p}.gate_proj.{kind}"].view(torch.uint8)[rows] + both = torch.cat((_zero_rows(up, 256), _zero_rows(gate, 256)), 0) + return both[_interleave_gather(512)[_blk32_gather(512)]] + + assert torch.equal(operand("l0_fc1_w"), fc1("weight")) + assert torch.equal(operand("l0_fc1_s"), _swizzle_flat(fc1("weight_scale")).view(512, 4)) + + down = fake[f"{p}.down_proj.weight"][:, 80 * rank : 80 * rank + 80] + assert torch.equal(operand("l0_fc2_w"), _zero_cols(down, 128)[_blk32_gather(128)]) + down_s = fake[f"{p}.down_proj.weight_scale"].view(torch.uint8)[:, 10 * rank : 10 * rank + 10] + expected = _swizzle_flat(_zero_cols(down_s, 16)[_blk32_gather(128)]).view(128, 16) + assert torch.equal(operand("l0_fc2_s"), expected) + + # The padding is where the runner's `intermediate_size` = 256 puts it: the last 96 rows + # of each half and the last 96 of FC2's 256 input columns, nothing else. + w1 = operand("l0_fc1_w")[torch.argsort(_interleave_gather(512)[_blk32_gather(512)])] + for half in (w1[:256], w1[256:]): + assert bool((half[:160] != 0).all()) and not bool(half[160:].any()) + + +@pytest.mark.parametrize("member", list(MEMBERS)) +def test_the_stored_intermediate_follows_the_engines_storage_rule(configs, member): + """The width a rank stores each expert's intermediate at is the one the engine's own + NVFP4 trtllm-gen MoE stores that shard at, and the width the runner is certified at.""" + from tensorrt_llm._torch.moe.fused_moe.quantization import NVFP4TRTLLMGenFusedMoEMethod + + d = _dims(configs, member, 1) + align, _ = NVFP4TRTLLMGenFusedMoEMethod.resolve_alignments(d.hidden, d.expert_inter_share) + assert d.expert_inter == -(-d.expert_inter_share // align) * align + assert (d.expert_inter_share, d.expert_inter) == ((160, 256) if member == "tp4" else (640, 640)) + + +# ---------------------------------------------------------------------- dep2 -- + + +def test_dep2_ranks_read_the_checkpoint_in_two_256_expert_windows(stored, configs): + """dep2 reads what dep4 reads, except the routed experts: rank `r` reads the window + `[256r, 256r + 256)`, so across the two ranks every tensor is read, each expert's by one + rank, and the `mtp.*` subtree by none (no speculation at dep2).""" + everything = set(stored) + mtp_keys = {k for k in everything if k.startswith("mtp.")} + read_by = collections.Counter() + for rank in range(DEP2["world_size"]): + d = _dims(configs, "dep2", rank, mtp=False) + claimed = _claimed(_rank_manifest(stored, d, "dep2")) + unread = WT.MODEL_WEIGHTS.expected_unconsumed(d, stored.keys()) + assert claimed <= everything and not claimed & unread, rank + assert claimed | unread == everything, ( + f"rank {rank}: neither read nor predicted {sorted(everything - claimed - unread)[:8]}" + ) + read_by.update(claimed) + + dep4 = _dims(configs, "dep4", rank, mtp=False) + non_expert = {k for k in claimed if not _TRUNK_EXPERT_TENSOR.match(k)} + assert non_expert == { + k + for k in _claimed(_rank_manifest(stored, dep4, "dep4")) + if not _TRUNK_EXPERT_TENSOR.match(k) + }, rank + + with torch.device("meta"): + params = WT.MODEL_WEIGHTS.declare(d) + dep4_params = WT.MODEL_WEIGHTS.declare(dep4) + # Every parameter dep4 declares, at dep4's shape, but the expert window's: 256 + # experts where dep4 holds 128. + assert set(params.keys()) == set(dep4_params.keys()) + for key in dep4_params: + ours, theirs = tuple(params[key].shape), tuple(dep4_params[key].shape) + if ours != theirs: + assert (ours[0], theirs[0], ours[1:]) == (256, 128, theirs[1:]), (key, ours) + for key, shape in _EXPERT_SHAPES["dep4"].items(): + assert tuple(params[key].shape) == (256, *shape[1:]), (key, tuple(params[key].shape)) + experts = [s[0] for s in WT.MODEL_WEIGHTS.manifest(d)["l0_fc2_w"]] + assert experts == [ + f"model.language_model.layers.0.mlp.experts.{e}.down_proj.weight" + for e in range(256 * rank, 256 * (rank + 1)) + ] + g = P.pool_geometry( + configs[0], + _mapping("dep2", rank), + tokens_per_block=32, + kv_cache_quant_algo=None, + dtype=torch.bfloat16, + ) + assert (d.heads_kv, d.num_k_heads, d.num_v_heads) == ( + g.heads_kv, + g.num_k_heads, + g.num_v_heads, + ) + + assert set(read_by) == everything - mtp_keys + wrong = sorted(k for k in read_by if _EXPERT_TENSOR.search(k) and read_by[k] != 1) + assert not wrong, f"expert tensors not read by exactly one rank: {wrong[:8]}" + + +def test_dep2_load_fills_every_parameter_and_covers_the_checkpoint(stored, configs, monkeypatch): + """`load` end to end at dep2's last rank, without speculation, on meta tensors.""" + monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda *a, **k: None) + model = _stub_model(configs, "dep2", 1, False) + WT.MODEL_WEIGHTS.load(model, _meta_checkpoint(stored)) + assert model.lm_head.weight.shape[0] == 248320 diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision/test_modeling_v2_triton_pos_embed_interpolate.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision/test_modeling_v2_triton_pos_embed_interpolate.py new file mode 100644 index 000000000000..6a8b63a8f50e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision/test_modeling_v2_triton_pos_embed_interpolate.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the default target's triton_pos_embed_interpolate entry (its vision +tower). + +`op.reference` is the explicit four-corner sum in fp64, returned unrounded as +fp32, so the band measures the op's own bf16 arithmetic. `_resize64` runs +torch's `F.interpolate` in fp64 and shares no arithmetic with either: each cell +holds the reference to it (the structure: which axis is which, corner +alignment, merge row order), and its wrong variants (`align_corners=False`, +raster rows, swapped table axes) must fail the band. +""" + +import zlib + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.vision.triton_pos_embed_interpolate import ( # noqa: E501 + triton_pos_embed_interpolate as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import ULP, assert_within_ulp, certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "triton_pos_embed_interpolate requires a CUDA device" + +NUM_GRID = 48 +HIDDEN = 1152 +MERGE = 2 + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """4 bf16 ulp per element against the row scale, 1 ulp relative RMS. + + The op rounds its four corner weights to bf16 and multiply-adds in bf16 (the + HF reference's arithmetic); `ref` is the exact value. The worst element + reaches about 3.5 ulp and the RMS 0.75, on `_table` and on the checkpoint's + own table alike. The nearest wrong variant, `align_corners=False`, lands an + order of magnitude outside the band, raster order and swapped axes two. + """ + assert out.dtype == torch.bfloat16, out.dtype + assert_within_ulp(out.float(), ref, element_ulp=4.0, rms_ulp=1.0, ulp=ULP[torch.bfloat16]) + + +def _table(seed: int, num_grid: int = NUM_GRID, hidden: int = HIDDEN) -> torch.Tensor: + """A smooth `[num_grid**2, hidden]` bf16 table, the way a learned 2-D + position table is: a 12x12 normal field upsampled bicubically, plus 0.1 + noise. Neighbouring rows have cosine 0.911; the checkpoint table's have 0.907. + + Contiguous on purpose: the upsampled field's permuted view reshapes to a + strided table, which the kernel reads wrongly (`is_valid` refuses it). + """ + g = torch.Generator(device="cuda").manual_seed(seed) + coarse = torch.randn(1, hidden, 12, 12, generator=g, device="cuda") + up = F.interpolate(coarse, size=(num_grid, num_grid), mode="bicubic", align_corners=True) + up = up[0].permute(1, 2, 0) + 0.1 * torch.randn( + num_grid, num_grid, hidden, generator=g, device="cuda" + ) + return up.reshape(num_grid * num_grid, hidden).to(torch.bfloat16).contiguous() + + +def _resize64(table, t, h, w, m_size, *, align=True, raster=False, swap=False): + """torch's bilinear resize in fp64, in merge order unless `raster`.""" + num_grid = int(round(table.shape[0] ** 0.5)) + grid = table.double().view(num_grid, num_grid, -1) + if swap: + grid = grid.transpose(0, 1) + src = grid.permute(2, 0, 1).unsqueeze(0).contiguous() + out = F.interpolate(src, size=(h, w), mode="bilinear", align_corners=align)[0].permute(1, 2, 0) + if raster: + rows = out.reshape(h * w, -1) + else: + rows = ( + out.reshape(h // m_size, m_size, w // m_size, m_size, -1) + .permute(0, 2, 1, 3, 4) + .reshape(h * w, -1) + ) + return rows.repeat(t, 1) if t > 1 else rows + + +def _args(spec, table): + return ( + table, + spec["t"], + spec["h"], + spec["w"], + spec["num_grid_per_side"], + spec["m_size"], + spec["dtype"], + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + spec = cell.spec + t, h, w, m_size = spec["t"], spec["h"], spec["w"], spec["m_size"] + table = _table(zlib.crc32(cell.why.encode()) % 2**31, spec["num_grid_per_side"], spec["hidden"]) + before = table.clone() + + out = op(*_args(spec, table)) + assert out.shape == (t * h * w, spec["hidden"]) and out.dtype == spec["dtype"] + assert torch.equal(table, before) and out.data_ptr() != table.data_ptr() + + ref = op.reference(*_args(spec, table)) + _compare(out, ref) + + # The reference's structure, from an implementation that shares nothing with it. + torch.testing.assert_close(ref, _resize64(table, t, h, w, m_size).float()) + + # The band refuses each wrong variant of the function. + for wrong in (dict(align=False), dict(raster=True), dict(swap=True)): + with pytest.raises(AssertionError): + _compare(out, _resize64(table, t, h, w, m_size, **wrong).float()) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py deleted file mode 100644 index 1cfb3de32ee3..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_fused_sigmoid_mul_inplace.py +++ /dev/null @@ -1,72 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the fused_sigmoid_mul_inplace entry.""" - -import zlib - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.activation.fused_sigmoid_mul_inplace import ( # noqa: E501 - fused_sigmoid_mul_inplace as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "fused_sigmoid_mul_inplace requires a CUDA device" - - -def _seed(text: str) -> int: - return zlib.crc32(text.encode()) - - -def _build(spec, seed): - """`(out, gate_buffer, gate)` for one spec. - - `gate` is a view of `gate_buffer`: the second half of each - `[T, heads, 2 * D]` q_proj head, or the 2-D gate window of the target's - fused projection row. - """ - g = torch.Generator(device="cuda").manual_seed(seed) - t, heads, d, dtype = spec["tokens"], spec["heads"], spec["head_dim"], spec["dtype"] - hidden = heads * d - - def randn(*shape): - return torch.randn(*shape, generator=g, device="cuda").to(dtype) - - out = randn(t, hidden) - if spec["gate"] == "interleaved": - gate_buffer = randn(t, heads, 2 * d) - gate = gate_buffer[:, :, d:] - else: - assert spec["gate"] == "fused_row", spec["gate"] - # The target's fused input projection row [q | k | v | gate | index | pad]: - # 14080 columns, the gate a 2-D [T, hidden] window at column 7168. - gate_buffer = randn(t, 14080) - gate = gate_buffer[:, 7168 : 7168 + hidden] - return out, gate_buffer, gate - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name - certifying(op) - out, gate_buffer, gate = _build(cell.spec, seed=_seed(cell.why)) - gate_before = gate_buffer.clone() - expected = op.reference(out, gate) - - result = op(out, gate) - - assert result.data_ptr() == out.data_ptr() and result.shape == out.shape - torch.testing.assert_close(result, expected) - assert torch.equal(gate_buffer, gate_before) - - -def test_metadata_misuse_raises() -> None: - """A gate with a strided last dim is refused loudly, as is every layout - violation; hence the entry has no `is_valid`.""" - g = torch.Generator(device="cuda").manual_seed(8) - attention_output = torch.randn(4, 128, generator=g, device="cuda").to(torch.bfloat16) - gate = torch.randn(4, 256, generator=g, device="cuda").to(torch.bfloat16)[:, ::2] - with pytest.raises(AssertionError): - op(attention_output, gate) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_canonical_selection.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_canonical_selection.py deleted file mode 100644 index 0561d585eeae..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_canonical_selection.py +++ /dev/null @@ -1,167 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qsa_canonical_selection entry. - -The rewrite fixes every slot of every row, so `_compare` is index equality against a -composite-key top-k. The input rows come from an adversary that keeps a random part of the -tie class at the k-th value in a random order, or from the radix engine the target calls. -""" - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_canonical_selection import ( # noqa: E501 - qsa_canonical_selection as op, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_select import ( # noqa: E501 - qsa_topk_select, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 - -assert torch.cuda.is_available(), "qsa_canonical_selection requires a CUDA device" - -DEV = "cuda" -COMPRESS_RATIO = 4 -BLOCK_TOPK = 512 -COLS = 2080 - - -def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: - """Index for index. The rule fixes every slot of every row, so nothing is tolerated: a - selection does no arithmetic.""" - assert out.dtype == torch.int32, out.dtype - assert out.shape == ref.shape, (out.shape, ref.shape) - wrong = (out != ref).any(1) - if bool(wrong.any()): - row = int(wrong.nonzero()[0, 0]) - raise AssertionError( - f"{int(wrong.sum())} of {out.shape[0]} rows differ from the canonical " - f"selection; first is row {row}" - ) - - -def _logits(kind: str, rows: int, cols: int, gen: torch.Generator) -> torch.Tensor: - if kind == "relu": - return torch.relu(torch.randn(rows, cols, generator=gen, device=DEV)) * 3.0 - if kind == "relu_sparse": - x = torch.relu(torch.randn(rows, cols, generator=gen, device=DEV)) * 3.0 - return x * (torch.rand(rows, cols, generator=gen, device=DEV) < 0.2) - raise ValueError(kind) - - -def _windows(spec: dict, rows: int, cols: int, gen: torch.Generator): - kind = spec["windows"][0] - starts = torch.zeros(rows, dtype=torch.int32, device=DEV) - if kind == "chunk": - positions = spec["windows"][1] + torch.arange(rows, device=DEV) - else: - assert kind == "positions_random", kind - lo, hi = spec["windows"][1:] - positions = torch.randint(lo, hi + 1, (rows,), generator=gen, device=DEV) - ends = ((positions + 1) // COMPRESS_RATIO).clamp_max(cols).to(torch.int32) - return starts, ends - - -def _poison(logits, starts, ends, kind) -> None: - columns = torch.arange(logits.shape[1], device=DEV)[None, :] - outside = (columns < starts[:, None]) | (columns >= ends[:, None]) - if kind == "nan": - logits.masked_fill_(outside, float("nan")) - return - # Values that would win any comparison they were allowed into. - pattern = torch.tensor([float("nan"), float("inf"), 3e38, float("nan")], device=DEV) - logits.copy_(torch.where(outside, pattern[columns % 4].expand_as(logits), logits)) - - -def _order_key(logits: torch.Tensor) -> torch.Tensor: - """int64 that orders as the fp32 values do, -0.0 folded into +0.0.""" - bits = (logits.float() + 0.0).view(torch.int32).long() - return torch.where(bits < 0, bits ^ 0x7FFFFFFF, bits) - - -def _adversary(logits, starts, ends, index_topk, gen) -> torch.Tensor: - """An exact top-k of each window that keeps a random part of the tie class at the - k-th value, in a random order, `-1` last.""" - rows, cols = logits.shape - columns = torch.arange(cols, device=DEV)[None, :] - in_window = (columns >= starts.long()[:, None]) & (columns < ends.long()[:, None]) - tiebreak = torch.rand(rows, cols, generator=gen, device=DEV).argsort(1) - key = torch.where( - in_window, (_order_key(logits) << 32) | tiebreak, torch.iinfo(torch.int64).min - ) - width = min(index_topk, cols) - top = torch.topk(key, width, dim=1).indices - local = torch.where(torch.gather(in_window, 1, top), top - starts.long()[:, None], -1) - shuffle = torch.where(local >= 0, torch.rand(local.shape, generator=gen, device=DEV), 2.0) - local = torch.gather(local, 1, shuffle.argsort(1)) - out = torch.full((rows, index_topk), -1, dtype=torch.int32, device=DEV) - out[:, :width] = local.to(torch.int32) - return out - - -def _radix(logits, starts, ends) -> torch.Tensor: - out = torch.empty(logits.shape[0], BLOCK_TOPK, dtype=torch.int32, device=DEV) - return qsa_topk_select(logits, starts, ends, out, index_topk=BLOCK_TOPK) - - -def _build(spec: dict, rows: int, seed: int): - gen = torch.Generator(device=DEV).manual_seed(seed) - logits = _logits(spec["logits"], rows, spec["cols"], gen) - starts, ends = _windows(spec, rows, spec["cols"], gen) - _poison(logits, starts, ends, spec["tail"]) - if spec["selection"] == "random": - indices = _adversary(logits, starts, ends, spec["index_topk"], gen) - else: - indices = _radix(logits, starts, ends) - return logits, starts, ends, indices - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 - certifying(op) - spec = cell.spec - seed = sum(map(ord, cell.why)) - row_counts = spec["rows"] if isinstance(spec["rows"], tuple) else (spec["rows"],) - for rows in row_counts: - args = _build(spec, rows, seed + rows) - expected = op.reference(*args) - out = op(*args) - assert out is args[3], "the rewrite writes `indices` and returns it" - _compare(out, expected) - - -def test_narrowing_the_table_to_a_bound_on_the_windows_changes_nothing() -> None: - """A caller that knows every window ends by column 1500 passes `logits[:, :1500]` of a - wide table: the rows are the full table's, and a 2^24-element table is tiled.""" - gen = torch.Generator(device=DEV).manual_seed(21) - rows, cols, bound = 4096, 8192, 1500 - logits = _logits("relu_sparse", rows, cols, gen) - starts = torch.zeros(rows, dtype=torch.int32, device=DEV) - ends = torch.randint(0, bound + 1, (rows,), generator=gen, device=DEV).to(torch.int32) - _poison(logits, starts, ends, "poison") - raw = _radix(logits, starts, ends) - wide = op(logits, starts, ends, raw.clone()) - narrow = op(logits[:, :bound], starts, ends, raw.clone()) - _compare(wide, op.reference(logits, starts, ends, raw)) - assert torch.equal(wide, narrow) - - -def test_guard_refuses_a_broadcast_row_start() -> None: - """A one-element `row_starts` broadcasts over every row without an error and frames - every row by row 0's start. Shown unguarded first, then refused.""" - gen = torch.Generator(device=DEV).manual_seed(9) - rows = 4 - logits = torch.relu(torch.randn(rows, COLS, generator=gen, device=DEV)) - starts = torch.tensor([0, 10, 20, 30], dtype=torch.int32, device=DEV) - ends = torch.full((rows,), 2000, dtype=torch.int32, device=DEV) - raw = _adversary(logits, starts, ends, BLOCK_TOPK, gen) - expected = op.reference(logits, starts, ends, raw) - out = op(logits, starts[:1], ends, raw.clone()) - assert out.shape == expected.shape and not torch.equal(out, expected), ( - "the broadcast hazard stopped reproducing" - ) - with pytest.raises(AssertionError, match="one entry per logits row"): - with validating(op): - op(logits, starts[:1], ends, raw.clone()) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py deleted file mode 100644 index 8498ad6cded2..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_scores.py +++ /dev/null @@ -1,284 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qsa_paged_index_scores catalog entry. - -The paged state is real: a `[pages, tokens_per_block, 1, dim]` index cache -(a view of a wider page under the NVFP4 cells) sized from the pages the table maps, and an -int32 block table whose rows are request *slots* -- more of them than the batch has requests, assigned by a -seeded permutation -- with shuffled physical pages and `-1` past each -sequence on the serving layouts, which is what the engine's table carries. - -The call has two kernels behind it. `_watch_launches` swaps the two Triton -entry points for recording proxies that still launch the real kernel, so every -cell asserts the branch it names. -""" - -import math - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_index_scores import ( # noqa: E501 - qsa_paged_index_scores as op, -) -from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 - -assert torch.cuda.is_available(), "qsa_paged_index_scores requires a CUDA device" - - -def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: - """torch's default fp32 band on every specified column; `-inf` exactly. - - bf16 and fp16 products are exact in the kernel's fp32 accumulator, so - only summation order separates it from the reference: measured at most - 2.9e-6 over every cell (2.4e-6 at 8191 rows x 2080 columns, scores to - 11.6), inside the default rtol 1.3e-6 / atol 1e-5. Wrong variants (no - relu, mean over heads, first group member as key, unscaled, - 1/head_dim scale, block table ignored) were measured 3e5x to 4e6x - outside it. Columns the reference marks NaN are unspecified and not read. - """ - assert out.dtype == ref.dtype == torch.float32, (out.dtype, ref.dtype) - assert out.shape == ref.shape, (out.shape, ref.shape) - specified = ~torch.isnan(ref) - torch.testing.assert_close(out[specified], ref[specified]) - - -DEV = torch.device("cuda") -#: The callable's own threshold: at or below it every call is row-wise. -DECODE_ROW_LIMIT = 256 -#: Index heads must divide this into more than one row for the query tile. -DOT_TILE_WIDTH = 128 - - -class _LaunchRecorder: - """Transparent proxy over a Triton entry point: records the grid, then launches.""" - - def __init__(self, inner, name, log): - self._inner, self._name, self._log = inner, name, log - - def __getitem__(self, grid): - self._log.append((self._name, tuple(grid))) - return self._inner[grid] - - def __getattr__(self, item): - return getattr(self._inner, item) - - -class _watch_launches: - """Record which of the entry's two Triton kernels a call launches.""" - - _NAMES = { - "_qsa_paged_index_scores_kernel": "row_wise", - "_qsa_paged_index_scores_qtile_kernel": "query_tile", - } - - def __enter__(self): - self.log = [] - self._saved = {name: getattr(_qsa_kernels, name) for name in self._NAMES} - for name, tag in self._NAMES.items(): - setattr(_qsa_kernels, name, _LaunchRecorder(self._saved[name], tag, self.log)) - return self - - def __exit__(self, *exc): - for name, kernel in self._saved.items(): - setattr(_qsa_kernels, name, kernel) - return False - - @property - def branch(self): - assert len(self.log) == 1, f"expected exactly one kernel launch, saw {self.log}" - return self.log[0][0] - - @property - def grid(self): - assert len(self.log) == 1, f"expected exactly one kernel launch, saw {self.log}" - return self.log[0][1] - - -def _case(rows, requests, **overrides): - """A narrow-table spec for the non-cell tests, same geometry as the entry's own.""" - spec = dict( - layout="random", - rows=rows, - requests=requests, - slots=requests + 4, - pages=8, - heads=4, - head_dim=128, - tokens_per_block=32, - compress_ratio=4, - dtype=torch.bfloat16, - positions_dtype=torch.int32, - request_dtype=torch.int32, - q_layout="packed", - only_visible_blocks=False, - context_rows=False, - ) - spec.update(overrides) - return spec - - -def _rows_and_requests(spec, gen): - """Per-row positions, per-row request ordinal, and each request's mapped pages.""" - layout, tpb, pages = spec["layout"], spec["tokens_per_block"], spec["pages"] - - def pages_for(seq_len): - return math.ceil(seq_len / tpb) - - if layout == "random": - rows, requests = spec["rows"], spec["requests"] - req = torch.arange(rows, device=DEV) * requests // max(rows, 1) - positions = torch.randint(0, pages * tpb, (rows,), generator=gen, device=DEV) - return positions, req, [pages] * requests - - # (first position, row count) per request, in row order - if layout == "capture": - segments = [(0, n) for n in spec["lengths"]] - elif layout == "decode": - segments = [(p, spec["rows_per_request"]) for p in spec["gen_positions"]] - else: - raise AssertionError(f"unknown layout {layout!r}") - positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments]) - req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)]) - mapped = [pages_for(s + n) for s, n in segments] - assert max(mapped) <= pages - return positions, req, mapped - - -def _build(spec, seed): - """Real paged state for one spec, keyed by `__call__`'s parameter names.""" - gen = torch.Generator(device=DEV).manual_seed(seed) - heads, head_dim, dtype = spec["heads"], spec["head_dim"], spec["dtype"] - positions, req, mapped = _rows_and_requests(spec, gen) - rows, slots = positions.numel(), spec["slots"] - assert len(mapped) <= slots - - slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(mapped)] - pool_pages = sum(mapped) + 3 - page_ids = torch.randperm(pool_pages, generator=gen, device=DEV).to(torch.int32) - block_table = torch.full((slots, spec["pages"]), -1, dtype=torch.int32, device=DEV) - cursor = 0 - for slot, count in zip(slot_of.tolist(), mapped): - block_table[slot, :count] = page_ids[cursor : cursor + count] - cursor += count - tpb = spec["tokens_per_block"] - keys = torch.randn(pool_pages, tpb, 1, head_dim, generator=gen, device=DEV).to(dtype) - if "page_pad" in spec: - # The NVFP4 pool's page: the other index buffers, this layer's, then the position - # page, so the plane is a view at page stride buffers * tpb * dim + page_pad. - buffers, buffer = spec["pool_buffers"], tpb * head_dim - page = torch.zeros(pool_pages, buffers * buffer + spec["page_pad"], dtype=dtype, device=DEV) - index_cache = page[:, (buffers - 1) * buffer : buffers * buffer].view( - pool_pages, tpb, 1, head_dim - ) - index_cache.copy_(keys) - assert index_cache.stride() == (page.stride(0), head_dim, head_dim, 1) - else: - index_cache = keys - - if spec["q_layout"] == "scratch_slice": - # [T, index heads | 1 k | 1 v, D], q is the first `heads` slots. - scratch = torch.randn(rows, heads + 2, head_dim, generator=gen, device=DEV).to(dtype) - q = scratch[:, :heads] - else: - q = torch.randn(rows, heads, head_dim, generator=gen, device=DEV).to(dtype) - if spec["q_layout"] != "packed" and rows > 1: # one row is trivially contiguous - assert not q.is_contiguous() - - return dict( - q=q, - index_cache=index_cache, - block_table=block_table, - query_positions=positions.to(spec["positions_dtype"]), - request_indices=slot_of.to(DEV)[req].to(spec["request_dtype"]), - tokens_per_block=spec["tokens_per_block"], - compress_ratio=spec["compress_ratio"], - only_visible_blocks=spec["only_visible_blocks"], - context_rows=spec["context_rows"], - ) - - -def _branch_rule(rows, spec): - heads = spec["heads"] - packs = DOT_TILE_WIDTH % heads == 0 and DOT_TILE_WIDTH // heads > 1 - tiled = spec["context_rows"] and rows > DECODE_ROW_LIMIT and packs - return "query_tile" if tiled else "row_wise" - - -def _columns(args): - return args["block_table"].shape[1] * args["tokens_per_block"] // args["compress_ratio"] - - -def _visible(args): - blocks = torch.arange(_columns(args), device=DEV) - bound = (args["query_positions"].long() + 1) // args["compress_ratio"] - return blocks[None, :] < bound[:, None] - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 - certifying(op) - args = _build(cell.spec, seed=sum(map(ord, cell.why))) - rows = args["q"].shape[0] - assert cell.spec["branch"] == _branch_rule(rows, cell.spec), "the cell names the wrong kernel" - with _watch_launches() as watcher: - out = op(**args) - torch.cuda.synchronize() - assert watcher.branch == cell.spec["branch"] - assert out.shape == (rows, _columns(args)) and out.dtype == torch.float32 - _compare(out, op.reference(**args)) - - -def test_prefetch_reads_the_fresh_anchor_key_and_q_after_the_wait() -> None: - """`prefetch` reads pages and keys before the PDL wait. Under graph replay a producer - rewrites q and, at rows whose position ends a group, the last visible block's anchor - key -- what the decode prologue writes -- right before the op; both must be seen. - The other keys are static, as they are in the decode step.""" - args = _build(_case(4, 4, only_visible_blocks=True), seed=411) - cr, tpb = args["compress_ratio"], args["tokens_per_block"] - pos = args["query_positions"].clone() - assert bool((pos >= cr).all()) - ends = torch.tensor([True, False, True, True], device=DEV) - args["query_positions"] = torch.where(ends, pos - (pos + 1) % cr, pos).to(pos.dtype) - pos = args["query_positions"].long() - assert bool(((pos + 1) % cr == 0).eq(ends).all()) - pages = args["block_table"][args["request_indices"].long(), pos // tpb].long() - anchors = (pages[ends], (pos % tpb)[ends]) - visible = _visible(args) - plain = op(**args) - assert torch.equal(op(**{**args, "use_pdl": True, "prefetch": True})[visible], plain[visible]) - src_q = args["q"].clone() - src_k = args["index_cache"][anchors].clone() - - def step(): - args["q"].copy_(src_q * 1.0) - args["index_cache"][anchors] = src_k * 1.0 - return op(**{**args, "use_pdl": True, "prefetch": True}) - - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - step() - torch.cuda.current_stream().wait_stream(warmup) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_out = step() - for trial in range(3): - fresh = _build(_case(4, 4, only_visible_blocks=True), seed=412 + trial) - src_q.copy_(fresh["q"]) - src_k.copy_(torch.randn(src_k.shape, device=DEV).to(src_k.dtype)) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(graph_out[visible], op(**args)[visible]), f"replay {trial}" - - -def test_rejects_metadata_shorter_than_the_rows() -> None: - """The callable's own refusal, message as observed -- the reason this entry has no - `is_valid`: positions for fewer rows than `q` carries.""" - args = _build(_case(16, 2), seed=1701) - args["query_positions"] = args["query_positions"][:4] - with pytest.raises(ValueError, match="one entry per row"): - op(**args) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py deleted file mode 100644 index caf83f14f1c7..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_kv_store.py +++ /dev/null @@ -1,190 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qsa_paged_kv_store catalog entry. - -The paged state is real: one `[pages, 2, Hkv, tpb, D]` HND pool whose K and V -role views are the caches, and the engine's `[slots, width]` table -- a row -per request slot, more slots than requests, each live row's pages drawn from a -seeded permutation of the pool, `-1` past its allocation, and the rows past -the batch holding a previous step's plausible page ids. Pool sizes are the old -receipts'; the pool is pre-filled with a sentinel so a stray write shows. -""" - -import zlib - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_kv_store import ( # noqa: E501 - qsa_paged_kv_store as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 - -assert torch.cuda.is_available(), "qsa_paged_kv_store requires a CUDA device" - - -def _compare(out, ref) -> None: - """Bitwise over both role planes, compared as raw bits. - - A store is a move with one conversion, and `reference` performs the - same conversion, so there is no rounding to spend a tolerance on. In - the target both planes together are the whole pool allocation, so this - is also the stray-write check. - """ - for name, got, want in zip(("k_cache", "v_cache"), out, ref): - assert got.dtype == want.dtype, (name, got.dtype, want.dtype) - assert got.shape == want.shape, (name, got.shape, want.shape) - torch.testing.assert_close( - _bits(got), _bits(want), rtol=0, atol=0, msg=lambda m, n=name: f"{n}: {m}" - ) - - -DEV = "cuda" -FILL = -2.0 - - -def _bits(tensor: torch.Tensor) -> torch.Tensor: - as_int = {1: torch.uint8, 2: torch.int16, 4: torch.int32}[tensor.element_size()] - return tensor.contiguous().view(as_int) - - -def _hole(spec, table, allocated_cols): - """One row that resolves to no slot: its (request, position).""" - assert spec["hole"] == "unallocated_page", spec["hole"] - assert int(table[0, allocated_cols[0]]) == -1 - return 0, allocated_cols[0] * spec["tokens_per_block"] - - -def _build(spec, seed: int): - """Real pool, table and rows for one cell; returns the op's kwargs and the pool.""" - gen = torch.Generator(device="cpu").manual_seed(seed) - tpb, hkv, head_dim = spec["tokens_per_block"], spec["hkv"], spec["head_dim"] - slots, width, pages = spec["slots"], spec["width"], spec["pages"] - segments = spec["segments"] - requests = [r for r, (_, n) in enumerate(segments) for _ in range(n)] - positions = [p for start, n in segments for p in range(start, start + n)] - - table = torch.randint(0, pages, (slots, width), generator=gen, dtype=torch.int32) - order = torch.randperm(pages, generator=gen).to(torch.int32) - used, allocated_cols = 0, [] - for r, (start, n) in enumerate(segments): - cols = (start + n - 1) // tpb + 1 - table[r] = -1 - table[r, :cols] = order[used : used + cols] - used += cols - allocated_cols.append(cols) - assert used <= pages and len(segments) < slots, "pool or table too small for the batch" - if spec["hole"] is not None: - request, position = _hole(spec, table, allocated_cols) - requests.append(request) - positions.append(position) - - rows = len(positions) - qkv = torch.randn( - rows, - spec["heads_q"] + 2 * hkv + spec.get("row_pad_heads", 0), - head_dim, - generator=gen, - dtype=torch.float32, - ).to(device=DEV, dtype=spec["row_dtype"]) - k = qkv[:, spec["heads_q"] : spec["heads_q"] + hkv] - v = qkv[:, spec["heads_q"] + hkv : spec["heads_q"] + 2 * hkv] - table = table.to(DEV) - - pool = torch.full((pages, 2, hkv, tpb, head_dim), FILL, dtype=torch.bfloat16, device=DEV).to( - spec["cache_dtype"] - ) - args = dict( - k=k, - v=v, - k_cache=pool[:, 0], - v_cache=pool[:, 1], - request_indices=torch.tensor(requests, dtype=spec["request_dtype"], device=DEV), - logical_positions=torch.tensor(positions, dtype=spec["position_dtype"], device=DEV), - block_table=table, - tokens_per_block=tpb, - ) - return args, pool - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying): # noqa: F811 - certifying(op) - args, pool = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) - before = pool.clone() - ref = op.reference(**args) - out = op(**args) - assert out[0] is args["k_cache"] and out[1] is args["v_cache"] - _compare(out, ref) - wrote = not torch.equal(_bits(pool), _bits(before)) - assert wrote == bool(cell.spec["segments"]), ( - "a cell with rows must write, an empty one must not" - ) - - -def _decode_args(seed: int, **overrides): - """The decode cell's layout over a bf16 pool, without its skipped row.""" - spec = {**op.CELLS[0].spec, "hole": None} - return _build({**spec, **overrides}, seed=seed) - - -#: The cell of 32 MTP7 verify blocks (256 rows): the segments, pool and table it runs over. -_MTP7_32 = next(c.spec for c in op.CELLS if c.why.startswith("fp8 decode MTP7: 32 requests")) -_MTP7_BATCH = {key: _MTP7_32[key] for key in ("segments", "pages", "slots")} - - -def test_cuda_graph_replay_follows_updated_positions() -> None: - """Captured once at 32 MTP7 verify blocks, replayed against new positions written - into the same buffer.""" - args, pool = _decode_args(seed=9, **_MTP7_BATCH) - positions = args["logical_positions"] - saved = positions.clone() - positions.zero_() - - def call(): - op(**args) - - for _ in range(3): # compile the Triton kernel before capture - call() - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - call() - - pool.fill_(FILL) - positions.copy_(saved) - ref = op.reference(**args) - graph.replay() - torch.cuda.synchronize() - _compare((args["k_cache"], args["v_cache"]), ref) - - -def test_guard_refuses_strided_row_metadata() -> None: - """`is_valid`'s one guard, after the hazard it exists for. - - The kernel receives strides for k, v, both caches and the table, but reads - the two metadata vectors as flat unit-stride buffers. A column of a 2-wide - buffer is read through its base pointer: accepted, and rows land elsewhere. - """ - args, pool = _decode_args(seed=10) - requests = args["request_indices"] - rows = requests.numel() - wide_requests = torch.full((rows, 2), rows - 1, dtype=requests.dtype, device=DEV) - wide_requests[:, 0] = requests - ref = op.reference(**args) - strided = {**args, "request_indices": wide_requests[:, 0]} - - pool.fill_(FILL) - out = op(**strided) # unguarded: accepted - assert not all(torch.equal(_bits(got), _bits(want)) for got, want in zip(out, ref)), ( - "strided metadata was expected to place rows elsewhere" - ) - - pool.fill_(FILL) - before = pool.clone() - with pytest.raises(AssertionError, match="must be a contiguous 1-D tensor"): - with validating(op): - op(**strided) - assert torch.equal(_bits(pool), _bits(before)), "a refused call still wrote" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py deleted file mode 100644 index 28ed2deb6c8a..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa.py +++ /dev/null @@ -1,250 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qsa_paged_sparse_gqa catalog entry. - -The paged state is real: one `[pages, 2, kv, tpb, D]` HND pool sized from the -pages its requests map, and the engine's slot table -- a row per request slot, -more slots than requests, assigned by a seeded permutation, `-1` past each -sequence. Selection rows are built the way QSA's top-k + expansion lay them -out: complete groups expanded to tokens, then the trailing partial group, then -`-1`, all packed at the front. - -The call has two kernels behind it (and a merge). The `launches` fixture -records which ones a call indexed with a grid, so every cell asserts its -branch. -""" - -import math -import re - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_sparse_gqa import ( # noqa: E501 - qsa_paged_sparse_gqa as op, -) -from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels - -__extra_import_path__ = ["../../../.."] -from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 - -assert torch.cuda.is_available(), "qsa_paged_sparse_gqa requires a CUDA device" - - -def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: - """4 ulp of each (row, head)'s magnitude, 2 ulp relative RMS, in q's dtype. - - The kernel rounds its online-softmax probabilities to the cache dtype - before P.V, so the error tracks the magnitude of the values averaged, - not of the output: against an fp32 reference it measured 1.2e-3 on rows - averaging a thousand tokens (outputs 0.14 to 0.4) and 1.2e-2 on rows - averaging a handful (outputs to 5). Scaled by each head row's own - magnitude that is one band. Against this reference (fp32, rounded once) - every cell measured at most 2.0 element ulp -- a single final-rounding - flip, since bf16's step is 2 of these units at the bottom of a binade -- - and at most 0.62 ulp RMS: 2x and 3.2x margin. Six wrong references (no - scale, interleaved heads, padding attended, identity pages, K/V - swapped, page/slot transposed) were each measured more than 0.1 off. - """ - assert out.dtype == ref.dtype, (out.dtype, ref.dtype) - assert out.shape == ref.shape, (out.shape, ref.shape) - if out.numel() == 0: - return - assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) - - -DEV = "cuda" - -_KERNELS = { - "_qsa_paged_sparse_gqa_kernel": "prefill", - "_qsa_paged_sparse_gqa_splitk_kernel": "splitk", - "_qsa_merge_splitk_kernel": "merge", -} - - -class _LaunchRecorder: - """Records a Triton kernel's grid indexing, then delegates to the kernel.""" - - def __init__(self, inner, name, log): - self._inner, self._name, self._log = inner, name, log - - def __getitem__(self, grid): - self._log.append((self._name, tuple(grid))) - return self._inner[grid] - - -@pytest.fixture -def launches(monkeypatch): - """Which branch kernels the next call launches, in order, with grids.""" - log: list = [] - for name, tag in _KERNELS.items(): - inner = getattr(_qsa_kernels, name) - monkeypatch.setattr(_qsa_kernels, name, _LaunchRecorder(inner, tag, log)) - return log - - -# -------------------------------------------------------------------------- -# building real state from a cell spec -# -------------------------------------------------------------------------- - - -def _selected_tokens(positions, seed, compress_ratio, token_topk): - """QSA's expansion layout, built in one pass. - - A random `min(available, block_topk)`-subset of the complete visible - groups, each expanded to its `compress_ratio` tokens, then the incomplete - trailing group, then `-1` -- packed at the front, as the producer emits. - """ - block_topk = token_topk // compress_ratio - final_topk = token_topk + compress_ratio - 1 - gen = torch.Generator(device=DEV).manual_seed(seed) - rows = positions.numel() - visible = positions.long() + 1 - available = visible // compress_ratio - taken = available.clamp(max=block_topk) - width = max(block_topk, int(available.max().item()) if rows else 0) - - scores = torch.rand(rows, width, generator=gen, device=DEV) - group_ids = torch.arange(width, device=DEV)[None, :] - scores = torch.where(group_ids < available[:, None], scores, torch.inf) - chosen = scores.argsort(dim=1)[:, :block_topk] - - column = torch.arange(final_topk, device=DEV)[None, :] - group_of_column = (column // compress_ratio).clamp(max=block_topk - 1).expand(rows, -1) - body = torch.gather(chosen, 1, group_of_column) * compress_ratio + column % compress_ratio - in_body = column < (taken * compress_ratio)[:, None] - tail_column = column - (taken * compress_ratio)[:, None] - tail = (available * compress_ratio)[:, None] + tail_column - in_tail = ~in_body & (tail_column >= 0) & (tail < visible[:, None]) - selected = torch.where(in_body, body, torch.where(in_tail, tail, torch.full_like(body, -1))) - return selected.to(torch.int32) - - -def _build(spec, seed): - """Real paged state for one spec, keyed by `__call__`'s parameter names.""" - gen = torch.Generator(device=DEV).manual_seed(seed) - tpb, kv, d, q_heads = ( - spec["tokens_per_block"], - spec["kv_heads"], - spec["head_dim"], - spec["q_heads"], - ) - segments = list(spec["segments"]) - positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments]) - req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)]) - rows, slots = positions.numel(), spec["slots"] - assert len(segments) <= slots - - slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(segments)] - mapped = [math.ceil((s + n) / tpb) for s, n in segments] - assert max(mapped) <= spec["pages"] - pool_pages = sum(mapped) + 7 - page_ids = torch.randperm(pool_pages, generator=gen, device=DEV).to(torch.int32) - block_table = torch.full((slots, spec["pages"]), -1, dtype=torch.int32, device=DEV) - cursor = 0 - for slot, count in zip(slot_of.tolist(), mapped): - block_table[slot, :count] = page_ids[cursor : cursor + count] - cursor += count - - pool = torch.randn(pool_pages, 2, kv, tpb, d, generator=gen, device=DEV).to(spec["cache_dtype"]) - k_cache, v_cache = pool[:, 0], pool[:, 1] - - q_dtype = spec["q_dtype"] - row_heads = q_heads + 2 * kv if spec["q_layout"] == "qkv_slice" else 55 - qkv = torch.randn(rows, row_heads, d, generator=gen, device=DEV).to(q_dtype) - q = qkv[:, :q_heads] - - selected = _selected_tokens(positions, seed + 3, spec["compress_ratio"], spec["token_topk"]) - request_indices = slot_of.to(DEV)[req].to(spec["request_dtype"]) - - gate = None - if spec.get("gate") == "fused_row": - # The gate columns of the target's fused projection row (55 heads of - # D): heads 28..51, after q, k and v. - fused = torch.randn(rows, 55 * d, generator=gen, device=DEV).to(q_dtype) - gate = fused[:, 28 * d : (28 + q_heads) * d] - - args = dict( - q=q, - k_cache=k_cache, - v_cache=v_cache, - block_table=block_table, - selected_tokens=selected, - request_indices=request_indices, - tokens_per_block=tpb, - softmax_scale=d**-0.5, - ) - if spec["prefill_pair"]: - args.update(query_positions=positions, compress_ratio=spec["compress_ratio"]) - if gate is not None: - args.update(output_gate=gate) - return args - - -def _bound(positions, spec): - """The prefill kernel's per-row scan bound, from the positions alone.""" - cr = spec["compress_ratio"] - visible = positions.long() + 1 - return (visible // cr).clamp(max=spec["token_topk"] // cr) * cr + visible % cr - - -def _is_prefix_packed(table): - valid = table >= 0 - columns = torch.arange(table.shape[1], device=table.device)[None, :] - return (valid == (columns < valid.sum(1)[:, None])).all(1) - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811 - certifying(op) - spec = cell.spec - args = _build(spec, seed=sum(map(ord, cell.why))) - rows = args["q"].shape[0] - if spec["prefill_pair"]: - # The precondition the op cannot check: front-packed, and exactly as - # many valid entries as the prefill kernel will scan. - selected = args["selected_tokens"] - assert bool(_is_prefix_packed(selected).all()) - valid = (selected >= 0).sum(1) - assert torch.equal(valid, _bound(args["query_positions"], spec)) - - out = op(**args) - assert [name for name, _ in launches] == list(spec["launches"]) - assert launches[0][1][:2] == (rows, spec["kv_heads"]) - assert out.shape == (rows, spec["q_heads"], spec["head_dim"]) - assert out.dtype == args["q"].dtype and out.is_contiguous() - assert not bool(out.isnan().any()) - _compare(out, op.reference(**args)) - - -#: The first cell is this checkpoint's own decode geometry; the tests below -#: start from it. -CKPT = op.CELLS[0].spec -assert (CKPT["q_heads"], CKPT["kv_heads"], CKPT["head_dim"]) == (24, 2, 256) - - -def _with(**overrides): - return {**CKPT, **overrides} - - -def test_calls_are_bitwise_deterministic() -> None: - """Including a 32-split launch whose merge is a second kernel.""" - spec = _with(segments=tuple((4000 + 37 * i, 2) for i in range(4)), gate=None) - args = _build(spec, seed=8800) - assert torch.equal(op(**args), op(**args)) - - -# -------------------------------------------------------------------------- -# what the op refuses -# -------------------------------------------------------------------------- - - -def test_rejects_positions_without_a_compress_ratio() -> None: - """The op's own loud refusal, verbatim -- the reason `is_valid` is empty: half of - the prefill pair.""" - args = _build(_with(segments=((600, 1),) * 4, gate=None), seed=9100) - args["selected_tokens"] = args["selected_tokens"][:, :64].contiguous() - args["query_positions"] = torch.zeros(4, dtype=torch.int64, device=DEV) - message = "QSA causal-prefix bounds require positions and compression ratio" - with pytest.raises(ValueError, match=re.escape(message)): - op(**args) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_rollback_planes.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_rollback_planes.py deleted file mode 100644 index 59c29b3fccec..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_rollback_planes.py +++ /dev/null @@ -1,84 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test: `qsa_rollback_gather` / `qsa_rollback_restore` against torch's advanced indexing, -bitwise, over planes that are strided views of one per-page buffer (as an NVFP4 KV cache's -manager packs the side planes) and over contiguous planes.""" - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_rollback_planes import ( # noqa: E501 - ROW, - plane_table, - qsa_rollback_gather, - qsa_rollback_restore, -) - -assert torch.cuda.is_available(), "qsa_rollback_planes requires a CUDA device" - -DEV = "cuda" -TPB = 32 -PAGES = 600 - - -def _planes(count, strided, g): - if strided: - page = torch.zeros(PAGES, count * TPB * ROW * 2 + 384, dtype=torch.uint8, device=DEV) - planes = [ - page[:, i * TPB * ROW * 2 : (i + 1) * TPB * ROW * 2] - .view(torch.bfloat16) - .view(PAGES, TPB, 1, ROW) - for i in range(count) - ] - else: - planes = [ - torch.empty(PAGES, TPB, 1, ROW, dtype=torch.bfloat16, device=DEV) for _ in range(count) - ] - for t in planes: - t.copy_(torch.randn(t.shape, generator=g, device=DEV).to(torch.bfloat16)) - return planes - - -def _slots(rows, g): - pages = torch.randperm(PAGES, generator=torch.Generator().manual_seed(rows))[:rows].to(DEV) - within = torch.randint(0, TPB, (rows,), generator=g, device=DEV) - return pages.to(torch.int64), within.to(torch.int64) - - -@pytest.mark.parametrize("strided", [False, True], ids=["contiguous", "shared_page"]) -@pytest.mark.parametrize( - "cell", qsa_rollback_gather.CELLS, ids=[c.why[:40] for c in qsa_rollback_gather.CELLS] -) -def test_gather_cells(cell, strided): - g = torch.Generator(device=DEV).manual_seed(7) - planes = _planes(cell.spec["planes"], strided, g) - pages, within = _slots(cell.spec["rows"], g) - ptrs, strides = plane_table(planes) - got = qsa_rollback_gather(ptrs, strides, pages, within, len(planes)) - want = torch.stack([t[pages, within].reshape(-1, ROW) for t in planes]) - assert torch.equal(got.view(torch.int16), want.view(torch.int16)) - - -@pytest.mark.parametrize("strided", [False, True], ids=["contiguous", "shared_page"]) -@pytest.mark.parametrize( - "cell", qsa_rollback_restore.CELLS, ids=[c.why[:40] for c in qsa_rollback_restore.CELLS] -) -def test_restore_cells(cell, strided): - g = torch.Generator(device=DEV).manual_seed(9) - planes = _planes(cell.spec["planes"], strided, g) - pages, within = _slots(cell.spec["rows"], g) - ptrs, strides = plane_table(planes) - old = qsa_rollback_gather(ptrs, strides, pages, within, len(planes)) - # The step overwrites every slot; the restore puts back the rejected ones only. - for t in planes: - t[pages, within] = torch.randn(t[pages, within].shape, generator=g, device=DEV).to(t.dtype) - after_step = [t.clone() for t in planes] - # Every other row rejected: each cell restores some slots and keeps the rest. - rejected = (torch.arange(pages.shape[0], device=DEV) % 2).to(torch.int32) - qsa_rollback_restore(ptrs, strides, pages, within, rejected, old) - for i, (t, step) in enumerate(zip(planes, after_step)): - want = step.clone() - want[pages, within] = torch.where( - rejected.bool()[:, None, None], old[i].view(-1, 1, ROW), step[pages, within] - ) - assert torch.equal(t.view(torch.int16), want.view(torch.int16)), f"plane {i}" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py deleted file mode 100644 index 0013c41db1bb..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_expand_decode.py +++ /dev/null @@ -1,108 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qsa_topk_expand_decode entry. - -Cells against the entry's reference (a stable descending sort, ties to the higher block, -then the expansion entry's reference): the block list and the token selection exactly. -Columns past each row's visible count carry garbage (huge values, NaN) that must never -be read. Beyond the cells: a stale-read check under graph replay. -""" - -import zlib - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_expand_decode import ( # noqa: E501 - qsa_topk_expand_decode as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "qsa_topk_expand_decode requires a CUDA device" - -DEV = "cuda" -W, K, CR, TOKEN_TOPK = 2560, 512, 4, 2048 - - -def _build(spec, seed): - g = torch.Generator(device=DEV).manual_seed(seed) - if "chunk" in spec: - # One request's consecutive context rows: every row's KV length is the chunk's end. - start, count = spec["chunk"] - pos = torch.arange(start, start + count, dtype=torch.int64, device=DEV) - seqlen = torch.full((count,), start + count, dtype=torch.int32, device=DEV) - else: - pos = torch.tensor(spec["pos"], dtype=torch.int64, device=DEV) - seqlen = (pos + 1).to(torch.int32) - rows = pos.numel() - scores = torch.randn(rows, W, generator=g, device=DEV).abs() * 3.0 - if spec.get("relu_sparse"): - # Fewer positive scores than the budget, so the zero class straddles it. - scores = torch.relu(torch.randn(rows, W, generator=g, device=DEV) - 2.2) - scores[:, ::7] = -0.0 # negative zero ties with zero - if spec.get("narrow"): - scores = 1.0 + 0.05 * torch.rand(rows, W, generator=g, device=DEV) - if spec.get("lognormal"): - scores = torch.exp(1.74 + 0.2 * torch.randn(rows, W, generator=g, device=DEV)) - # bf16-like resolution, so equal scores tie; a zero in every other row. - scores = scores.to(torch.bfloat16).float() - scores[::2, 5] = 0.0 - ends = ((pos + 1) // CR).to(torch.int32) - for r in range(rows): # garbage past the visible blocks - n = int(ends[r]) - scores[r, n:] = 1.0e30 - if n + 3 < W: - scores[r, n + 3] = float("nan") - topk = torch.full((rows, K), -7, dtype=torch.int32, device=DEV) - return [scores, ends, pos, seqlen, topk, CR, TOKEN_TOPK] - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name - certifying(op) - args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) - want_topk, want_sel = op.reference(*args) - margin = cell.spec.get("out_margin") - if margin is None: - sel = op(*args) - else: - # Rows [margin, margin + rows) of a step-sized selection; the rows around stay put. - rows = args[0].shape[0] - whole = torch.full( - (rows + 2 * margin, TOKEN_TOPK + CR - 1), -9, dtype=torch.int32, device=DEV - ) - sel = op(*args, out=whole[margin : margin + rows]) - assert sel.data_ptr() == whole[margin].data_ptr() - assert (whole[:margin] == -9).all() and (whole[margin + rows :] == -9).all(), ( - "rows outside out" - ) - assert sel.dtype == torch.int32 and sel.shape == want_sel.shape - assert torch.equal(args[4], want_topk), "block list" - assert torch.equal(sel, want_sel.to(torch.int32)), "token selection" - - -def test_reads_after_the_wait_under_graph_replay() -> None: - args = _build(dict(pos=(8700,)), seed=21) - src = args[0].clone() - - def step(): - args[0].copy_(src * 1.0) - return op(*args) - - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - step() - torch.cuda.current_stream().wait_stream(warmup) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_sel = step() - for trial in range(3): - src.copy_(_build(dict(pos=(8700,)), seed=100 + trial)[0]) - graph.replay() - torch.cuda.synchronize() - got_topk = args[4].clone() - sel = op(src.clone(), *args[1:]) - assert torch.equal(graph_sel, sel) and torch.equal(got_topk, args[4]), f"replay {trial}" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_spec_rollback_commit.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_spec_rollback_commit.py deleted file mode 100644 index 9d8456e59f62..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_spec_rollback_commit.py +++ /dev/null @@ -1,101 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU tests for the spec_rollback_commit entries, driven by their CELLS. - -`spec_rollback_stage` is compared bitwise with its reference. `spec_rollback_commit` has no -reference of its own (test_modeling_v2_spec_rollback drives it through the handler); here each -cell is checked bitwise against the docstring's per-request rewrite, written out in torch. -""" - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.spec_rollback_commit import ( # noqa: E501 - ROW, - spec_rollback_commit, - spec_rollback_stage, -) - -assert torch.cuda.is_available(), "spec_rollback_commit requires a CUDA device" - -DEV = "cuda" -CTX = 2 -CHANNELS = 640 -STATE = 3 -TPB = 32 - - -def _slots(n: int, g: torch.Generator) -> torch.Tensor: - return (torch.randperm(2 * n + 3, generator=g)[:n] + 1).to(DEV, torch.int32) - - -@pytest.mark.parametrize( - "cell", spec_rollback_stage.CELLS, ids=[c.why for c in spec_rollback_stage.CELLS] -) -def test_stage_cells(cell) -> None: - n, p = cell.spec["n"], cell.spec["p"] - g = torch.Generator().manual_seed(n * 31 + p) - slots = _slots(n, g) - pool_slots = 2 * n + 4 - ngram_pool = torch.randint(0, 1 << 40, (pool_slots, CTX), generator=g).to(DEV) - conv_pool = torch.randn(pool_slots, CHANNELS, STATE, generator=g).to(DEV, torch.bfloat16) - ids = torch.randint(0, 1 << 20, (n * p,), generator=g).to(DEV, torch.int32) - spec_rollback_stage.is_valid(ngram_pool, conv_pool, slots, ids, p) - stream, conv_prev = spec_rollback_stage(ngram_pool, conv_pool, slots, ids, p) - want_stream, want_prev = spec_rollback_stage.reference(ngram_pool, conv_pool, slots, ids, p) - assert torch.equal(stream, want_stream) - assert torch.equal(conv_prev, want_prev) - - -@pytest.mark.parametrize( - "cell", spec_rollback_commit.CELLS, ids=[c.why for c in spec_rollback_commit.CELLS] -) -def test_commit_cells(cell) -> None: - n, p, planes = cell.spec["n"], cell.spec["p"], cell.spec["planes"] - g = torch.Generator().manual_seed(n * 17 + planes) - rows = n * p - slots = _slots(n, g) - pool_slots = 2 * n + 4 - accepted = torch.randint(1, p + 1, (n,), generator=g).to(DEV, torch.int32) - ngram_pool = torch.randint(0, 1 << 40, (pool_slots, CTX), generator=g).to(DEV) - stream = torch.randint(0, 1 << 40, (n, CTX + p), generator=g).to(DEV) - conv_pool = torch.randn(pool_slots, CHANNELS, STATE, generator=g).to(DEV, torch.bfloat16) - conv_prev = torch.randn(n, CHANNELS, STATE, generator=g).to(DEV, torch.bfloat16) - conv_x = torch.randn(n, p, CHANNELS, generator=g).to(DEV, torch.bfloat16) - # Each row its own (page, slot within the page). - flat = torch.randperm(rows * 2, generator=g)[:rows] - pages, within = (flat // TPB).to(DEV), (flat % TPB).to(DEV) - num_pages = int(pages.max()) + 1 - position = torch.randint(0, 1 << 20, (num_pages, TPB, 3), generator=g).to(DEV, torch.int32) - position_old = torch.randint(0, 1 << 20, (rows, 3), generator=g).to(DEV, torch.int32) - plane_list = [ - torch.randn(num_pages, TPB, ROW, generator=g).to(DEV, torch.bfloat16) for _ in range(planes) - ] - plane_ptrs = torch.tensor([t.data_ptr() for t in plane_list], dtype=torch.int64, device=DEV) - plane_strides = torch.tensor([t.stride(0) for t in plane_list], dtype=torch.int64, device=DEV) - index_old = torch.randn(planes, rows, ROW, generator=g).to(DEV, torch.bfloat16) - - want_ngram, want_conv = ngram_pool.clone(), conv_pool.clone() - want_position, want_planes = position.clone(), [t.clone() for t in plane_list] - for i in range(n): - a, slot = int(accepted[i]), int(slots[i]) - want_ngram[slot] = stream[i, a : a + CTX] - cat = torch.cat([conv_prev[i], conv_x[i].T], 1) - want_conv[slot] = cat[:, a : a + STATE] - for t in range(a, p): - r = i * p + t - pg, tok = int(pages[r]), int(within[r]) - want_position[pg, tok] = position_old[r] - for plane, want in zip(range(planes), want_planes): - want[pg, tok] = index_old[plane, r] - - args = (accepted, slots, ngram_pool, stream, conv_pool, conv_prev, conv_x, position, - position_old, pages, within, plane_ptrs, plane_strides, index_old) # fmt: skip - spec_rollback_commit.is_valid(*args) - spec_rollback_commit(*args) - torch.cuda.synchronize() - assert torch.equal(ngram_pool, want_ngram) - assert torch.equal(conv_pool, want_conv) - assert torch.equal(position, want_position) - for got, want in zip(plane_list, want_planes): - assert torch.equal(got, want) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_vision_rope_2d.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_vision_rope_2d.py deleted file mode 100644 index 2a83d03b3fc4..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_vision_rope_2d.py +++ /dev/null @@ -1,120 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qwen4_exp vision_rope_2d entry. - -`op.reference` shares no code with the entry: positions from nested loops, -angles in fp64 from `theta ** (-2j / rope_dim)`, rotation as a complex -multiply. -""" - -import zlib -from typing import Tuple - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.vision_rope_2d import ( # noqa: E501 - vision_rope_2d as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 - -assert torch.cuda.is_available(), "vision_rope_2d requires a CUDA device" - - -def _compare( - out: Tuple[torch.Tensor, torch.Tensor], - ref: Tuple[torch.Tensor, torch.Tensor], -) -> None: - """Default band on each of `(q_out, k_out)`; nothing is loosened.""" - for got, want in zip(out, ref, strict=True): - torch.testing.assert_close(got, want) - - -HEADS = 16 -HEAD_DIM = 72 -MERGE = 2 -THETA = 10000.0 - - -def _num_tokens(grid) -> int: - return sum(t * h * w for t, h, w in grid) - - -def _operands(tokens, heads, heads_k, head_dim, dtype, seed, layout="separate"): - g = torch.Generator(device="cuda").manual_seed(seed) - if layout == "packed_qkv": - assert heads == heads_k - qkv = torch.randn(tokens, 3 * heads, head_dim, generator=g, device="cuda").to(dtype) - q, k, _ = torch.split(qkv, [heads] * 3, 1) - assert not q.is_contiguous() and q.stride(-1) == 1 - return q, k - q = torch.randn(tokens, heads, head_dim, generator=g, device="cuda").to(dtype) - k = torch.randn(tokens, heads_k, head_dim, generator=g, device="cuda").to(dtype) - return q, k - - -def _build(spec, seed): - grid = spec["grid"] - heads = spec.get("heads", HEADS) - q, k = _operands( - _num_tokens(grid), - heads, - spec.get("heads_k", heads), - spec.get("head_dim", HEAD_DIM), - spec["dtype"], - seed, - spec.get("layout", "separate"), - ) - grid_thw = torch.tensor( - [list(row) for row in grid], - dtype=spec.get("grid_dtype", torch.int32), - device=spec.get("grid_device", "cpu"), - ) - return q, k, grid_thw, MERGE, spec.get("theta", THETA) - - -def _max_deviation(lhs, rhs) -> float: - return (lhs.float() - rhs.float()).abs().max().item() - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying): # noqa: F811 - certifying(op) - q, k, grid_thw, merge, theta = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) - q_before, k_before = q.clone(), k.clone() - q_out, k_out = op(q, k, grid_thw, merge, theta) - assert q_out.shape == q.shape and q_out.dtype == q.dtype - assert k_out.shape == k.shape and k_out.dtype == k.dtype - assert torch.equal(q, q_before) and torch.equal(k, k_before), "inputs must not be mutated" - assert q_out.data_ptr() != q.data_ptr() and k_out.data_ptr() != k.data_ptr() - q_ref, k_ref = op.reference(q, k, grid_thw, merge, theta) - print( - f"[{cell.why}] q_dev={_max_deviation(q_out, q_ref):.3e} " - f"k_dev={_max_deviation(k_out, k_ref):.3e}" - ) - _compare((q_out, k_out), (q_ref, k_ref)) - - -def test_grid_placement_and_integer_dtype_do_not_change_the_result() -> None: - """The target hands over an int64 grid on the device; the cells pass host grids.""" - grid = [(1, 16, 24), (1, 8, 12)] - q, k = _operands(_num_tokens(grid), HEADS, HEADS, HEAD_DIM, torch.bfloat16, seed=29) - host = torch.tensor(grid, dtype=torch.int32) - q_host, k_host = op(q, k, host, MERGE, THETA) - for other in (host.to("cuda"), host.to(torch.int64), host.to(torch.int64).to("cuda")): - q_other, k_other = op(q, k, other, MERGE, THETA) - assert torch.equal(q_host, q_other) and torch.equal(k_host, k_other) - - -def test_guard_refuses_a_one_token_k() -> None: - """The hazard, unguarded: a [1, H, D] k broadcasts to [T, H, D].""" - grid_thw = torch.tensor([(1, 4, 4)], dtype=torch.int32) - q, k = _operands(16, HEADS, HEADS, HEAD_DIM, torch.bfloat16, seed=5) - _, k_out = op(q, k[:1], grid_thw, MERGE, THETA) - assert k_out.shape == (16, HEADS, HEAD_DIM) - - with pytest.raises(AssertionError, match="must agree on tokens and head_dim"): - with validating(op): - op(q, k[:1], grid_thw, MERGE, THETA) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py deleted file mode 100644 index 6d7963ed32db..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_weight_path.py +++ /dev/null @@ -1,371 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The routed experts of the Qwen3.8-Flash-Next target, from checkpoint-format tensors to output. - -The CuteDSL op tests build their operands with layout code of their own. This file builds -the checkpoint's expert tensors (512 experts: packed NVFP4 `up_proj` / `gate_proj` / -`down_proj` codes, linear e4m3 block scales, one `weight_scale_2` per expert, per-expert -`input_scale`s) and runs them through the target's own load-time transforms -(`weights._expert_fc1_w` / `_expert_fc1_s` / `_expert_fc2_w` / `_expert_fc2_s`) and scalar -derivation (`modeling._moe_layer_scalars`), once per routed backend. The operands then go -through the call sequence of the target's `_moe_block` -- CuteDSL: sort, FC1 with SwiGLU, -FC2, unpermute; trtllm-gen: the one-call block-scale runner -- with explicit routing. - -Two routings hit all 512 experts. In the first every token's only non-zero routing weight is -on its own expert, so each expert's output is checked alone: a layout error confined to one -expert (a permutation, a scale swizzle, an alpha) cannot hide in a sum over ten. On it both -backends are held to a float64 reference over the dequantized checkpoint tensors, with -`down_proj.input_scale` spread across the experts as a calibrated checkpoint's is. The second -is a weighted top-10: with a uniform `down_proj.input_scale` the two backends requantize the -SwiGLU output identically, so on it they are held to each other -- the area's one -cross-backend check. - -The reference does not model the kernels' accumulation. It models the one rounding step -that is not exact: the SwiGLU output is requantized to NVFP4 before FC2, at the backend's -own global scale (the layer's single `global_sf` on CuteDSL, each expert's -`1 / down_proj.input_scale` on trtllm-gen). -""" - -import functools -from types import SimpleNamespace - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( - fp4_block_scale_moe_runner, - quantize_nvfp4, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( - fp4_quantize, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - modeling as M, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - weights as W, -) - -assert torch.cuda.is_available(), "the CuteDSL experts require a CUDA device" - -DEV = "cuda" -E, K, H, INTER, VEC = 512, 10, 2560, 640, 16 -#: The activation's input scale, one number for all experts: |x| = 5 maps to the e4m3 maximum. -ISC1 = 5.0 / (448 * 6) -#: The SwiGLU output maximum `down_proj.input_scale` is calibrated to. With `spread` each -#: expert's own value ranges over 2^-1/2 .. 2^1/2 of it, so the layer's `max` is not every -#: expert's. -ISC2 = 64.0 / (448 * 6) -E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) -#: The routed runners read activation scales as a linear buffer, not the 128x4-swizzled -#: one of a dense NVFP4 GEMM; byte counts match at multiples of 128 rows, so the wrong -#: one passes silently. -SF_LINEAR = False - -#: Tokens of the two routings: one per expert, and a weighted top-10. -ONE_HOT_TOKENS = E -WEIGHTED_TOKENS = 1024 - -BACKENDS = ("CUTEDSL", "TRTLLM") -#: The float64 reference's gate, relative per token. Both backends sit at 2.5e-3 or below: -#: the bf16 rounding of the output is 1e-3 of it, the rest is the reference's fp32 / tie -#: differences. A wrong layout is ~1; a wrong per-expert scalar is its own relative error. -REFERENCE_TOLERANCE = 5e-3 -#: The same, between the backends under a uniform `down_proj.input_scale`. Each backend is -#: within `REFERENCE_TOLERANCE` of the reference, so the two are within twice that of each -#: other. Both requantize the SwiGLU output to NVFP4, and an fp32 difference in the last -#: place moves an element at an e2m1 rounding boundary by a whole code: the token holding -#: it differs by ~5e-3 while the mean stays near 5e-6. A wrong layout is still ~1. -BACKEND_TOLERANCE = 2 * REFERENCE_TOLERANCE - - -# ── the checkpoint ──────────────────────────────────────────────────────── - - -def _quantize(dense: torch.Tensor): - """A bf16 `[R, C]` matrix as the checkpoint stores it: packed e2m1 codes `[R, C/2]` - uint8, e4m3 block scales `[R, C/16]`, and the `weight_scale_2` that undoes the global - scale the block scales were derived under.""" - rows, cols = dense.shape - gs = (448 * 6 / dense.abs().amax().float()).reshape(1) - codes, sf = torch.ops.trtllm.fp4_quantize(dense, gs, VEC, False, False) - return ( - codes.view(torch.uint8).view(rows, cols // 2), - sf.view(torch.uint8).view(rows, cols // VEC).view(torch.float8_e4m3fn), - (1.0 / gs).squeeze(0), - ) - - -@functools.lru_cache(maxsize=None) -def _checkpoint(spread: bool) -> SimpleNamespace: - """All 512 experts of one layer, with the checkpoint's tensor names and dtypes. - - `up_proj` and `gate_proj` are quantized together, so they share one `weight_scale_2`, as - the fused FC1 requires. Expert magnitudes spread over 2^-1 .. 2^1/2 (FC1) and 2^-1 .. 2^1 - (FC2), so a wrong per-expert scalar shows.""" - g = torch.Generator(device=DEV).manual_seed(20261014) - spread1 = torch.exp2(torch.rand(E, generator=g, device=DEV) * 1.5 - 1) - spread2 = torch.exp2(torch.rand(E, generator=g, device=DEV) * 2 - 1) - fields = {k: [] for k in ("up_w", "up_s", "gate_w", "gate_s", "down_w", "down_s")} - ws2_gu, ws2_dn = [], [] - for e in range(E): - gate_up = torch.randn(2 * INTER, H, generator=g, device=DEV) * (spread1[e] / H**0.5) - down = torch.randn(H, INTER, generator=g, device=DEV) * (spread2[e] / INTER**0.5) - codes, sf, ws2 = _quantize(gate_up.to(torch.bfloat16)) - fields["up_w"].append(codes[:INTER]) - fields["gate_w"].append(codes[INTER:]) - fields["up_s"].append(sf[:INTER]) - fields["gate_s"].append(sf[INTER:]) - ws2_gu.append(ws2) - codes, sf, ws2 = _quantize(down.to(torch.bfloat16)) - fields["down_w"].append(codes) - fields["down_s"].append(sf) - ws2_dn.append(ws2) - ws2_gu, ws2_dn = torch.stack(ws2_gu), torch.stack(ws2_dn) - isc1 = torch.full((E,), ISC1, dtype=torch.float32, device=DEV) - isc2 = torch.full((E,), ISC2, dtype=torch.float32, device=DEV) - if spread: - isc2 = isc2 * torch.exp2(torch.rand(E, generator=g, device=DEV) - 0.5) - return SimpleNamespace( - **fields, - ws2_gu=ws2_gu, - ws2_dn=ws2_dn, - isc2=isc2, - # The rows `modeling._moe_layer_scalars` reads, under layer 0. - scalars={ - "l0_e_isc1": isc1, - "l0_e_isc1_up": isc1.clone(), - "l0_e_ws2_1": ws2_gu, - "l0_e_ws2_1_up": ws2_gu.clone(), - "l0_e_isc2": isc2, - "l0_e_ws2_2": ws2_dn, - }, - ) - - -# ── what the target builds from it ──────────────────────────────────────── - - -@functools.lru_cache(maxsize=None) -def _layer(backend: str, spread: bool) -> SimpleNamespace: - """The layer's expert operands for `backend`, built as the weight table builds them: - one transform call per expert and tensor, stacked, then viewed as `_layer_views` - views them.""" - ck = _checkpoint(spread) - core = SimpleNamespace(trunk_moe_backend=backend) - - def stack(fn, *tensors): - return torch.stack([fn(tuple(t[e] for t in tensors), core) for e in range(E)]) - - def stack1(fn, tensor): - return torch.stack([fn(tensor[e], core) for e in range(E)]) - - fc1_w = stack(W._expert_fc1_w, ck.up_w, ck.gate_w) - fc1_s = stack(W._expert_fc1_s, ck.up_s, ck.gate_s) - fc2_w = stack1(W._expert_fc2_w, ck.down_w) - fc2_s = stack1(W._expert_fc2_s, ck.down_s) - if backend == "CUTEDSL": - w_dtype, s_dtype = torch.float4_e2m1fn_x2, torch.uint8 - else: - w_dtype, s_dtype = torch.uint8, torch.float8_e4m3fn - return SimpleNamespace( - fc1_w=fc1_w.view(w_dtype), - fc1_s=fc1_s.view(s_dtype), - fc2_w=fc2_w.view(w_dtype), - fc2_s=fc2_s.view(s_dtype), - scalars=M._moe_layer_scalars(ck.scalars, 0, backend), - ) - - -def _cutedsl_routed(layer, x, ids, weights) -> torch.Tensor: - """The CuteDSL branch of `_moe_block`, from the quantization to the unpermute.""" - sc = layer.scalars - tokens = x.shape[0] - xq, xsf = fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, SF_LINEAR) - tile_expert, tile_limit, expanded_to_permuted, permuted_to_expanded, _, live_tiles = M.moe_sort( - token_selected_experts=ids, - token_final_scales=weights, - num_experts=E, - top_k=K, - local_expert_offset=0, - local_num_experts=E, - tile_tokens_dim=M._CUTEDSL_TILE_SIZE, - ) - h, h_sf = M.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion( - input=xq.view(torch.float4_e2m1fn_x2), - weight=layer.fc1_w, - input_scale=xsf.view(tokens, H // M._SF_VEC), - weight_scale=layer.fc1_s, - alpha=sc["alpha1"], - tile_idx_to_group_idx=tile_expert, - tile_idx_to_mn_limit=tile_limit, - permuted_idx_to_expanded_idx=permuted_to_expanded, - num_non_exiting_tiles=live_tiles, - global_sf=sc["global_sf"], - num_experts=E, - top_k=K, - num_local_experts=E, - local_expert_offset=0, - tile_size=M._CUTEDSL_TILE_SIZE, - scaling_vector_size=M._SF_VEC, - activation_type=int(M.ActivationType.Swiglu), - ) - y = M.cute_dsl_nvfp4_grouped_gemm( - input=h, - weight=layer.fc2_w, - input_scale=h_sf, - weight_scale=layer.fc2_s, - alpha=sc["alpha2"], - tile_idx_to_group_idx=tile_expert, - num_non_exiting_tiles=live_tiles, - num_experts=E, - top_k=K, - num_local_experts=E, - local_expert_offset=0, - tile_size=M._CUTEDSL_TILE_SIZE, - output_dtype=torch.bfloat16, - scaling_vector_size=M._SF_VEC, - ) - return M.moe_unpermute_inplace( - permuted_input=y, - output=torch.empty(tokens, H, dtype=torch.bfloat16, device=x.device), - expanded_idx_to_permuted_idx=expanded_to_permuted, - topk_scales=weights, - ) - - -def _trtllm_gen_routed(layer, x, ids, weights) -> torch.Tensor: - """The trtllm-gen branch of `_moe_block`.""" - sc = layer.scalars - xq, xsf = fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, SF_LINEAR) - return fp4_block_scale_moe_runner( - routing_logits=None, - routing_bias=None, - hidden_states=xq, - hidden_states_scale=xsf.view(torch.float8_e4m3fn), - gemm1_weights=layer.fc1_w, - gemm1_weights_scale=layer.fc1_s, - gemm1_bias=None, - gemm1_alpha=None, - gemm1_beta=None, - gemm1_clamp_limit=None, - gemm2_weights=layer.fc2_w, - gemm2_weights_scale=layer.fc2_s, - gemm2_bias=None, - output1_scale_scalar=sc["output1_scale_scalar"], - output1_scale_gate_scalar=sc["output1_scale_gate_scalar"], - output2_scale_scalar=sc["output2_scale_scalar"], - num_experts=E, - top_k=K, - n_group=None, - topk_group=None, - intermediate_size=INTER, - local_expert_offset=0, - local_num_experts=E, - routed_scaling_factor=None, - routing_method_type=M._ROUTING_METHOD_RENORM, - do_finalize=True, - act_type=M._ACT_TYPE_SWIGLU, - topk_weights=weights.to(torch.bfloat16), - topk_ids=ids, - )[0] - - -# ── the inputs and the reference ────────────────────────────────────────── - - -@functools.lru_cache(maxsize=None) -def _problem(case: str): - """`(x, ids, weights)`: bf16 activations, expert ids `[T, K]` int32 and routing weights - `[T, K]` fp32 (exact in bf16, which is the dtype trtllm-gen takes them in).""" - tokens = ONE_HOT_TOKENS if case == "one_hot" else WEIGHTED_TOKENS - g = torch.Generator(device=DEV).manual_seed(7 if case == "one_hot" else 8) - x = torch.randn(tokens, H, generator=g, device=DEV).to(torch.bfloat16) - scores = torch.rand(tokens, E, generator=g, device=DEV) - if case == "one_hot": - # The forced expert has the top score, so it sits in slot 0. - scores[torch.arange(tokens), torch.arange(tokens) % E] = 2.0 - weights = torch.zeros(tokens, K, device=DEV) - weights[:, 0] = 1.0 - else: - weights = torch.softmax(torch.randn(tokens, K, generator=g, device=DEV), dim=-1) - weights = weights.to(torch.bfloat16).float() - ids = scores.topk(K, dim=-1).indices.to(torch.int32) - return x, ids, weights - - -def _dequant(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: - """`[R, C/2]` packed e2m1 (element 2j in the low nibble) and `[R, C/16]` e4m3 scales - -> float64 `[R, C]`.""" - magnitude = torch.tensor(E2M1, dtype=torch.float64, device=codes.device) - table = torch.cat([magnitude, -magnitude]) - lo, hi = table[(codes & 0xF).long()], table[(codes >> 4).long()] - values = torch.stack([lo, hi], dim=-1).reshape(codes.shape[0], -1) - return values * scales.view(torch.float8_e4m3fn).double().repeat_interleave(VEC, dim=1) - - -def _reference(backend: str, case: str, spread: bool) -> torch.Tensor: - """The routed output in float64, `[T, H]`. - - The activation is what the kernels see, the dequantized NVFP4 of the target's own - quantization. Each expert's weights are dequantized from the checkpoint tensors, times - its `weight_scale_2`. The SwiGLU output is requantized at `backend`'s global scale.""" - ck, layer = _checkpoint(spread), _layer(backend, spread) - x, ids, weights = _problem(case) - sc = layer.scalars - xq, xsf = fp4_quantize(x, sc["global_scale"], M._SF_VEC, False, SF_LINEAR) - x_dq = _dequant(xq.view(torch.uint8), xsf.view(x.shape[0], H // VEC)) * ISC1 - if backend == "CUTEDSL": - g2 = [float(sc["global_sf"])] * E - else: - g2 = (1.0 / ck.isc2).tolist() - out = torch.zeros(x.shape[0], H, dtype=torch.float64, device=DEV) - for e in range(E): - token, slot = (ids == e).nonzero(as_tuple=True) - w = weights[token, slot].double() - token, w = token[w != 0], w[w != 0] - if token.numel() == 0: - continue - ws2 = ck.ws2_gu[e].double() - xe = x_dq[token] - up = xe @ (_dequant(ck.up_w[e], ck.up_s[e]) * ws2).T - gate = xe @ (_dequant(ck.gate_w[e], ck.gate_s[e]) * ws2).T - act = up * gate * torch.sigmoid(gate) - act = quantize_nvfp4(act.float(), g2[e]).double() / g2[e] - y = act @ (_dequant(ck.down_w[e], ck.down_s[e]) * ck.ws2_dn[e].double()).T - out.index_add_(0, token, y * w[:, None]) - return out - - -@functools.lru_cache(maxsize=None) -def _routed(backend: str, case: str, spread: bool) -> torch.Tensor: - x, ids, weights = _problem(case) - run = _cutedsl_routed if backend == "CUTEDSL" else _trtllm_gen_routed - return run(_layer(backend, spread), x, ids, weights).double() - - -def _token_error(out: torch.Tensor, ref: torch.Tensor) -> torch.Tensor: - """Each token's `|out - ref| / |ref|` (L2 over the hidden width).""" - return (out - ref).norm(dim=1) / ref.norm(dim=1).clamp_min(1e-12) - - -# ── tests ───────────────────────────────────────────────────────────────── - - -@pytest.mark.parametrize("backend", BACKENDS) -def test_routed_output_matches_the_float64_reference(backend: str) -> None: - case = "one_hot" - _, ids, weights = _problem(case) - assert torch.unique(ids[weights != 0]).numel() == E, "the routing misses an expert" - err = _token_error(_routed(backend, case, True), _reference(backend, case, True)) - print(f"[err] {backend} {case}: mean {err.mean():.4e} max {err.max():.4e}") - assert err.max() <= REFERENCE_TOLERANCE, (err.mean().item(), err.max().item()) - - -def test_cutedsl_matches_trtllm_gen_under_a_uniform_input_scale() -> None: - case = "weighted" - _, ids, weights = _problem(case) - assert torch.unique(ids[weights != 0]).numel() == E, "the routing misses an expert" - err = _token_error(_routed("CUTEDSL", case, False), _routed("TRTLLM", case, False)) - print(f"[err] CUTEDSL vs TRTLLM {case}: mean {err.mean():.4e} max {err.max():.4e}") - assert err.max() <= BACKEND_TOLERANCE, (err.mean().item(), err.max().item()) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_mtp_moe_nvfp4.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_mtp_moe_nvfp4.py deleted file mode 100644 index 6e47da43f2fe..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_mtp_moe_nvfp4.py +++ /dev/null @@ -1,142 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the NVFP4 copy of the draft layer's routed experts -(`TRTLLM_MODELING_V2_MTP_MOE_NVFP4`). - -`weights.mtp_experts_to_nvfp4` requantizes FP8 128x128 block-scaled experts to NVFP4 in the -trunk's trtllm-gen layout, `modeling._mtp_moe_nvfp4` adds the runner's scalars and -`modeling._mtp_routed_nvfp4` runs them. Checked against the FP8 block-scale path the draft -layer runs by default, on the same random experts, activations and router: the NVFP4 output -must stay inside a relative-error band of the FP8 one, keep its magnitude (a wrong runner -scalar moves the norm, which unbiased quantization noise does not), and leave no FP8 copy -behind. A layout error is O(1) relative, far outside the band. -""" - -from types import SimpleNamespace - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - modeling as m, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.fp8_block_scale_moe_runner import ( # noqa: E501 - fp8_block_scale_moe_runner, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.moe.renorm_moe_routing_op import ( # noqa: E501 - renorm_moe_routing_op, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.quantization.fp8_quantize_1x128 import ( # noqa: E501 - fp8_quantize_1x128, -) - -assert torch.cuda.is_available(), "the NVFP4 draft MoE requires a CUDA device" - -DEV = "cuda" -E, H, INTER, K = 512, 2560, 640, 10 -BLK = 128 -#: A trunk calibration to derive the static draft scales from: the checkpoint's last trunk layer -#: (FC1 input, FC2 input amax), the layer that feeds the draft. -TRUNK_AMAX = (11.1875, 57.75) -#: Relative RMS of NVFP4 - FP8 over the FP8 output, the floor on each row's cosine, and the -#: band on the output norm ratio. Measured here: 0.239-0.251 RMS, row cosine >= 0.961, at every -#: row count; on the checkpoint's real draft experts 0.222-0.231 and >= 0.964 (bench). That is -#: the NVFP4 chain's own noise -- four e2m1 roundings (activation, FC1 weight, FC1 output, FC2 -#: weight) at ~0.1 each, compounded by the SwiGLU product -- not a bias: it is flat across -#: activation scales 0.5-4x, and FP8 sits at 0.07 from the fp32 reference of the same weights. -REL_RMS_BAND = 0.3 -COS_FLOOR = 0.95 -NORM_RATIO_BAND = 0.1 - - -def _fp8_blocks(rows: int, cols: int, gen: torch.Generator): - """`[E, rows, cols]` e4m3 codes and `[E, rows/128, cols/128]` fp32 scales of a random - expert stack whose 128x128 blocks span 16x in magnitude, as trained experts' do.""" - w = torch.randn(E, rows, cols, generator=gen, device=DEV) * cols**-0.5 - spread = torch.exp2( - torch.randint(-2, 3, (E, rows // BLK, 1, cols // BLK, 1), generator=gen, device=DEV) - ).float() - blocks = w.view(E, rows // BLK, BLK, cols // BLK, BLK) * spread - scale = blocks.abs().amax(dim=(2, 4), keepdim=True) / 448.0 - codes = (blocks / scale).to(torch.float8_e4m3fn).view(E, rows, cols) - return codes, scale.view(E, rows // BLK, cols // BLK).contiguous() - - -@pytest.fixture(scope="module") -def experts(): - gen = torch.Generator(device=DEV).manual_seed(20261022) - fc1_w, fc1_s = _fp8_blocks(2 * INTER, H, gen) - fc2_w, fc2_s = _fp8_blocks(H, INTER, gen) - router = (torch.randn(E, H, generator=gen, device=DEV) * H**-0.5).to(torch.bfloat16) - fp8 = SimpleNamespace(fc1_w=fc1_w, fc1_s=fc1_s, fc2_w=fc2_w, fc2_s=fc2_s, router=router.t()) - p = "l1_mtp_" - w = {p + k: getattr(fp8, k).clone() for k in ("fc1_w", "fc1_s", "fc2_w", "fc2_s")} - w["l0_e_isc1"] = torch.full((E,), TRUNK_AMAX[0] / 2688.0, device=DEV) - w["l0_e_isc2"] = torch.full((E,), TRUNK_AMAX[1] / 2688.0, device=DEV) - core = SimpleNamespace( - w=w, mtp_layer_idx=1, num_layers=1, num_experts=E, topk=K, moe_inter=INTER - ) - fields = m._mtp_moe_nvfp4(core) - nvfp4 = SimpleNamespace(router=fp8.router, **fields) - return core, fp8, nvfp4 - - -def _fp8_routed(fp8, x): - """The FP8 block-scale draft path, call for call as `_mtp_moe_block` makes it.""" - logits = m.cublas_mm(x, fp8.router) - ids, wts = renorm_moe_routing_op(logits, topk=K, output_dtype=torch.bfloat16) - xq, xsf = fp8_quantize_1x128(x, use_ue8m0=False) - return fp8_block_scale_moe_runner( - routing_logits=None, - routing_bias=None, - hidden_states=xq, - hidden_states_scale=xsf, - gemm1_weights=fp8.fc1_w, - gemm1_weights_scale=fp8.fc1_s, - gemm2_weights=fp8.fc2_w, - gemm2_weights_scale=fp8.fc2_s, - num_experts=E, - top_k=K, - num_fused_shared_experts=None, - n_group=None, - topk_group=None, - intermediate_size=INTER, - local_expert_offset=0, - local_num_experts=E, - routed_scaling_factor=None, - routing_method_type=m._ROUTING_METHOD_RENORM, - topk_weights=wts, - topk_ids=ids, - act_type=m._ACT_TYPE_SWIGLU, - gemm1_clamp_limit=None, - ) - - -def test_fp8_copies_freed_and_layout(experts): - core, _, nvfp4 = experts - for k in ("fc1_w", "fc1_s", "fc2_w", "fc2_s"): - assert core.w[f"l1_mtp_{k}"].numel() == 0, k - assert nvfp4.fc1_w.shape == (E, 2 * INTER, H // 2) and nvfp4.fc1_w.dtype == torch.uint8 - assert nvfp4.fc1_s.shape == (E, 2 * INTER, H // 16) and nvfp4.fc1_s.dtype == torch.float8_e4m3fn - assert nvfp4.fc2_w.shape == (E, H, INTER // 2) and nvfp4.fc2_s.shape == (E, H, INTER // 16) - for s in (nvfp4.out1_scale, nvfp4.out1_gate_scale, nvfp4.out2_scale): - assert s.shape == (E,) and s.dtype == torch.float32 and torch.isfinite(s).all() - assert nvfp4.moe_gscale.shape == (1,) - - -@pytest.mark.parametrize("rows", [1, 4, 32, 128, 1024]) -def test_nvfp4_tracks_fp8(experts, rows): - core, fp8, nvfp4 = experts - gen = torch.Generator(device=DEV).manual_seed(rows) - x = torch.randn(rows, H, generator=gen, device=DEV).to(torch.bfloat16) - ref = _fp8_routed(fp8, x).float() - out = m._mtp_routed_nvfp4(core, nvfp4, x) - assert out.shape == (rows, H) and out.dtype == torch.bfloat16 - out = out.float() - rel = ((out - ref).norm() / ref.norm()).item() - cos = torch.nn.functional.cosine_similarity(out, ref, dim=1).min().item() - ratio = (out.norm() / ref.norm()).item() - print(f"rows={rows}: rel RMS {rel:.4f}, min row cosine {cos:.4f}, norm ratio {ratio:.4f}") - assert rel < REL_RMS_BAND, rel - assert cos > COS_FLOOR, cos - assert abs(ratio - 1.0) < NORM_RATIO_BAND, ratio diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_gated_ngram_injection.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_gated_ngram_injection.py deleted file mode 100644 index f93fdf147cd6..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_gated_ngram_injection.py +++ /dev/null @@ -1,164 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the gated_ngram_injection entry: the cells (the sub-floor and -zero dots are a cell), the leading-dim guard and CUDA-graph replay.""" - -import zlib - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.gated_ngram_injection import ( # noqa: E501 - gate_ratio, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.gated_ngram_injection import ( # noqa: E501 - gated_ngram_injection as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "gated_ngram_injection requires a CUDA device" - - -def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: - """16-bit: `rtol = eps`, `atol` one subnormal spacing (`gate_ratio <= 1`) - -- half torch's bf16 `rtol`. Nothing cancels (each element is a sigmoid - times one `value` element), so the budget is the final rounding; with - the reference also rounded once, op and reference agree or sit one ulp - apart. One ulp is at most `eps * |ref|`, so this band always admits it - and never two. Measured at 2048 tokens: 0.992 (a one-ulp element just - above a binade boundary), eight wrong readings at 105x..1115x. The - `atol` is for Gaussian `value` elements that land subnormal (measured - on fp16). - - fp32: torch's default band. The 2560-term signed dot is accumulated at - the width it is reported in, so its conditioning, not the cast, sets - the error; a custom pair would be fitted to a result. - """ - assert out.dtype == ref.dtype, (out.dtype, ref.dtype) - assert out.shape == ref.shape, (out.shape, ref.shape) - assert out.is_contiguous() - if ref.dtype == torch.float32: - torch.testing.assert_close(out, ref) - return - ratio = gate_ratio(out, ref) - assert ratio <= 1.0, f"{ratio:.4g}x the band" - - -DEV = "cuda" -HC_COUNT = 4 -HIDDEN = 2560 -WIDTH = HC_COUNT * HIDDEN -RMS_NORM_EPS = 1e-6 -CLAMP_MIN = 1e-6 - -# This checkpoint's PLE tensors: norm_key / norm_query gammas (element RMS -# 0.1358, extremes 0.863 and -0.326), and the value's scale derived as -# sqrt(2560) * 0.015295 (value_proj RMS) * 0.00747 (dequantized table RMS). -# The entry is exactly linear in `value`, so its absolute scale decides nothing. -CKPT_NORM_GAMMA_RMS = 0.1358 -CKPT_NORM_GAMMA_MAX = 0.863281 -CKPT_NORM_GAMMA_MIN = -0.326172 -CKPT_VALUE_RMS = 0.0058 - - -def _gemma_grouped_rmsnorm(x, gamma, hc_count): - """Where key/query come from -- here only for their distribution.""" - hidden = x.shape[-1] // hc_count - xf = x.float().reshape(*x.shape[:-1], hc_count, hidden) - xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + RMS_NORM_EPS) - return (xf.flatten(-2) * (1.0 + gamma.float())).to(x.dtype) - - -def _make_operands(seed, dtype, leading, hc_count=HC_COUNT, value_dtype=None): - g = torch.Generator(device=DEV).manual_seed(seed) - width = hc_count * HIDDEN - - def normal(*dims, scale, dt=dtype): - return (torch.randn(*dims, generator=g, device=DEV) * scale).to(dt) - - gammas = [] - for _ in range(2): - gamma = normal(width, scale=CKPT_NORM_GAMMA_RMS) - gamma[0] = CKPT_NORM_GAMMA_MAX - gamma[width // 2] = CKPT_NORM_GAMMA_MIN - gammas.append(gamma) - key = _gemma_grouped_rmsnorm(normal(*leading, width, scale=1.0), gammas[0], hc_count) - query = _gemma_grouped_rmsnorm(normal(*leading, width, scale=1.0), gammas[1], hc_count) - value = normal(*leading, HIDDEN, scale=CKPT_VALUE_RMS, dt=value_dtype or dtype) - return key, query, value - - -def _single_channel_operands(dots, value_dtype): - """One nonzero channel per stream: each dot is exactly `dots[s]`.""" - hc_count = len(dots) - key = torch.zeros(8, hc_count * HIDDEN, dtype=torch.bfloat16, device=DEV) - query = torch.zeros_like(key) - for stream, dot in enumerate(dots): - key[:, stream * HIDDEN] = 1.0 - query[:, stream * HIDDEN] = dot - value = torch.full((8, HIDDEN), 0.5, dtype=value_dtype, device=DEV) - return key, query, value - - -def _build(spec, seed): - if spec.get("kind") == "single_channel": - dots = spec["dots"] - key, query, value = _single_channel_operands(dots, spec["value_dtype"]) - return key, query, value, len(dots), spec["clamp_min"] - key, query, value = _make_operands( - seed, torch.bfloat16, spec["leading"], value_dtype=spec.get("value_dtype") - ) - return key, query, value, HC_COUNT, spec.get("clamp_min", CLAMP_MIN) - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name - certifying(op) - args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) - before = [a.clone() for a in args[:3]] - out = op(*args) - key, _, value, hc_count, _ = args - assert out.shape == (*key.shape[:-1], hc_count * value.shape[-1]) - assert out.dtype == value.dtype - for a, b in zip(args[:3], before): - assert torch.equal(a, b), "the operands are read-only" - _compare(out, op.reference(*args)) - - -def test_guard_refuses_unequal_leading_dims() -> None: - """Unguarded, one query (or value) row is broadcast over every key row.""" - key, query, value = _make_operands(60, torch.bfloat16, (64,)) - unguarded = op(key, query[:1], value, HC_COUNT) - assert unguarded.shape == (64, WIDTH), "the hazard this guard covers is gone" - assert not torch.equal(unguarded, op.reference(key, query, value, HC_COUNT)) - assert op(key, query, value[:1], HC_COUNT).shape == (64, WIDTH) - - with pytest.raises(AssertionError, match="broadcasts silently"): - with validating(op): - op(key, query[:1], value, HC_COUNT) - with pytest.raises(AssertionError, match="broadcasts silently"): - with validating(op): - op(key, query, value[:1], HC_COUNT) - - -def test_cuda_graph_capture_and_replay() -> None: - key, query, value = _make_operands(100, torch.bfloat16, (8,)) - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - for _ in range(3): - op(key, query, value, HC_COUNT, CLAMP_MIN) - torch.cuda.current_stream().wait_stream(warmup) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured = op(key, query, value, HC_COUNT, CLAMP_MIN) - - fresh = _make_operands(101, torch.bfloat16, (8,)) - for static, new in zip((key, query, value), fresh): - static.copy_(new) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(captured, op(*fresh, HC_COUNT, CLAMP_MIN)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_dequantize_e4m3_per_tensor.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_dequantize_e4m3_per_tensor.py deleted file mode 100644 index ef52eab0be66..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_dequantize_e4m3_per_tensor.py +++ /dev/null @@ -1,140 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the dequantize_e4m3_per_tensor catalog entry. - -The entry owns the contract. This file turns a cell into operands and holds -the claims no cell carries: graph capture with the scale read at replay, and -the loud rejection of an empty batch. -""" - -import zlib - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.quantization.dequantize_e4m3_per_tensor import ( # noqa: E501 - PLE_SCALE, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.quantization.dequantize_e4m3_per_tensor import ( # noqa: E501 - dequantize_e4m3_per_tensor as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "dequantize_e4m3_per_tensor requires a CUDA device" - - -_BIT_VIEW = { - torch.bfloat16: torch.int16, - torch.float16: torch.int16, - torch.float32: torch.int32, -} - - -def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: - """Bit-exact, NaN for NaN. - - Both sides form the same exact fp32 product and round it once, so any - difference is a defect. Measured bit-exact on every case, including all - 256 codes x three output dtypes x five scales, with -0.0 and the - subnormals compared on bit patterns. NaN payloads are not compared. - """ - assert out.dtype == ref.dtype, (out.dtype, ref.dtype) - assert out.shape == ref.shape, (out.shape, ref.shape) - nan = torch.isnan(ref) - assert torch.equal(torch.isnan(out), nan), "NaN lanes differ" - view = _BIT_VIEW[ref.dtype] - assert torch.equal(out.contiguous().view(view)[~nan], ref.contiguous().view(view)[~nan]), ( - "dequantization is exact; a bit difference is a defect" - ) - - -DEV = "cuda" - -_TABLE = {} - - -def _fp8(x: torch.Tensor) -> torch.Tensor: - return x.to(torch.float8_e4m3fn) - - -def _scale(v: float, rank: int, dtype=torch.bfloat16) -> torch.Tensor: - """The only shape the op accepts: `input`'s rank, every dimension 1.""" - return torch.full((1,) * rank, v, dtype=dtype, device=DEV) - - -def _input(spec, seed: int) -> torch.Tensor: - g = torch.Generator(device=DEV).manual_seed(seed) - source, shape = spec["source"], spec["shape"] - if source == "gather": - if "table" not in _TABLE: - tg = torch.Generator(device=DEV).manual_seed(11) - _TABLE["table"] = _fp8(torch.randn(1 << 20, 160, device=DEV, generator=tg) * 40.0) - idx = torch.randint(0, 1 << 20, shape[:2], device=DEV, generator=g) - x = _TABLE["table"][idx] - else: - assert source == "all_codes", source - x = torch.arange(256, dtype=torch.uint8, device=DEV).view(shape).view(torch.float8_e4m3fn) - assert x.is_contiguous() and x.shape == shape - return x - - -def _bits(t: torch.Tensor) -> torch.Tensor: - view = {torch.bfloat16: torch.int16, torch.float16: torch.int16, torch.float32: torch.int32} - return t.view(view[t.dtype]) - - -# ── the certified cells ─────────────────────────────────────────────────── - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name - certifying(op) - x = _input(cell.spec, zlib.crc32(cell.why.encode())) - for value in cell.spec["scales"]: - scales = _scale(value, x.dim(), cell.spec["dtype"]) - out = op(x, scales) - assert out.shape == x.shape and out.dtype == scales.dtype and out.is_contiguous() - _compare(out, op.reference(x, scales)) - if cell.spec["source"] == "gather": - flat = out.reshape(x.shape[0], -1) - assert flat.shape[1] == 2560 and flat.data_ptr() == out.data_ptr(), "flatten is a view" - - -# ── what a cell cannot carry ────────────────────────────────────────────── - - -def test_cuda_graph_capture_and_replay(): - """Both operands stay live inputs: the scale is read at replay.""" - g = torch.Generator(device=DEV).manual_seed(16) - static_in = _fp8(torch.randn(256, 16, 160, device=DEV, generator=g) * 8.0) - static_scales = _scale(PLE_SCALE, 3) - eager = op(static_in, static_scales).clone() - warm = torch.cuda.Stream() - warm.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warm): - for _ in range(3): - op(static_in, static_scales) - torch.cuda.current_stream().wait_stream(warm) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - static_out = op(static_in, static_scales) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(_bits(static_out), _bits(eager)) - fresh = _fp8(torch.randn(256, 16, 160, device=DEV, generator=g) * 8.0) - static_in.copy_(fresh) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(_bits(static_out), _bits(op(fresh, static_scales))) - static_scales.copy_(_scale(0.5, 3)) - graph.replay() - torch.cuda.synchronize() - _compare(static_out, op.reference(fresh, _scale(0.5, 3))) - - -def test_rejects_an_empty_batch(): - """The op raises on zero rows rather than launching; the caller branches first.""" - with pytest.raises((RuntimeError, NotImplementedError)): - op(_fp8(torch.randn(0, 128, device=DEV)), _scale(PLE_SCALE, 2)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_chunk_gated_delta_rule.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_chunk_gated_delta_rule.py deleted file mode 100644 index 6a8cf2fa2286..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_chunk_gated_delta_rule.py +++ /dev/null @@ -1,242 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the chunk_gated_delta_rule entry. - -The first indexed in-place call per autotune key in a process must be as correct as a -warm one: the state kernel's autotuner restores the indexed state its benchmark launches -write, and the pool's dtype is part of the key. "First call in the process" is a property -of the process, so the cold-call test re-runs this file as a child: -`python `. -""" - -import os -import subprocess -import sys -import zlib -from pathlib import Path - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.chunk_gated_delta_rule import ( # noqa: E501 - ENGINE_SLOTS, - HEAD_K_DIM, - HEAD_V_DIM, - NUM_K_HEADS, - NUM_V_HEADS, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.chunk_gated_delta_rule import ( # noqa: E501 - chunk_gated_delta_rule as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "chunk_gated_delta_rule requires a CUDA device" - - -#: `_compare`'s band; see its docstring. -RTOL, ATOL_OUTPUT, ATOL_STATE = 1.6e-2, 2e-3, 2e-2 - - -def _compare(out: tuple, ref: tuple) -> None: - """`(o, final_state, pool)`: the call's two returns plus the - `initial_state` it was handed, against `reference`'s triple; `None` - components must match. - - Wider than bf16's default because the op is a chunked reformulation - whose WY operands are rounded to bf16 -- the same recurrence, not the - same sum. rtol stays at the bf16 default; atol is 2e-3 on `o` and 2e-2 - on states, about 4x and 6x the largest deviations measured warm over - every cell against this reference (`o` 4.88e-4, one bf16 ulp, on - magnitudes up to 7.9e-2; states 3.57e-3 on up to 8.3e-1). - Discriminated when measured: seven wrong references (no gating, beta = 1, q/k swapped, - initial state dropped, scale = 1, key heads reversed, `[K, V]` - transposed state) all fail the `o` gate, the closest 11x outside it. - """ - for got, want, atol in zip(out, ref, (ATOL_OUTPUT, ATOL_STATE, ATOL_STATE)): - assert (got is None) == (want is None), (got is None, want is None) - if got is None: - continue - assert got.dtype == want.dtype and got.shape == want.shape, ( - (got.dtype, got.shape), - (want.dtype, want.shape), - ) - torch.testing.assert_close(got.float(), want.float(), rtol=RTOL, atol=atol) - - -DEV = "cuda" - - -def _l2norm(x: torch.Tensor) -> torch.Tensor: - xf = x.float() - return xf / torch.sqrt(xf.pow(2).sum(-1, keepdim=True) + 1e-6) - - -def _operands(seqlens: tuple, seed: int): - """Operands in the domain the GDN prefill path produces: unit-norm q/k, - `g = -exp(A_log) * softplus(a + dt_bias) < 0`, `beta` a sigmoid. The delta - rule is only contractive there.""" - gen = torch.Generator(device=DEV).manual_seed(seed) - total = sum(seqlens) - cu_seqlens = torch.tensor( - [0] + torch.tensor(seqlens).cumsum(0).tolist(), dtype=torch.long, device=DEV - ) - shape_qk = (1, total, NUM_K_HEADS, HEAD_K_DIM) - q = torch.randn(shape_qk, generator=gen, device=DEV) - k = torch.randn(shape_qk, generator=gen, device=DEV) - q, k = _l2norm(q).to(torch.bfloat16), _l2norm(k).to(torch.bfloat16) - v = torch.randn( - 1, total, NUM_V_HEADS, HEAD_V_DIM, generator=gen, device=DEV, dtype=torch.bfloat16 - ) - a_log = torch.log(torch.rand(NUM_V_HEADS, generator=gen, device=DEV) * 15 + 1) - dt_bias = torch.randn(NUM_V_HEADS, generator=gen, device=DEV) - a = torch.randn(1, total, NUM_V_HEADS, generator=gen, device=DEV) - g = (-torch.exp(a_log) * torch.nn.functional.softplus(a + dt_bias)).float() - beta = torch.rand(1, total, NUM_V_HEADS, generator=gen, device=DEV).sigmoid().float() - return q, k, v, g, beta, cu_seqlens - - -def _pool(num_slots: int, seed: int, dtype: torch.dtype = torch.float32) -> torch.Tensor: - """A recurrent-state pool, `[slots, H, V, K]` (K innermost), at `dtype`.""" - gen = torch.Generator(device=DEV).manual_seed(seed) - shape = (num_slots, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM) - return (torch.randn(shape, generator=gen, device=DEV) * 0.1).to(dtype) - - -def _gate_output(o: torch.Tensor, ref: tuple) -> None: - """`_compare`, on `o` alone.""" - _compare((o, None, None), (ref[0], None, None)) - - -def _gate_pool(pool: torch.Tensor, ref: tuple) -> None: - """`_compare`, on the pool alone.""" - _compare((ref[0], None, pool), (ref[0], None, ref[2])) - - -def _build(spec, seed): - """`__call__`'s 13 arguments for one spec.""" - q, k, v, g, beta, cu_seqlens = _operands(spec["seqlens"], seed) - pool = _pool(ENGINE_SLOTS, seed=seed + 1, dtype=spec.get("pool_dtype", torch.float32)) - indices = torch.tensor(spec["indices"], dtype=torch.int32, device=DEV) - output = None - if spec.get("output") == "buffer": - shape_o = (q.shape[0], q.shape[1], NUM_V_HEADS, HEAD_V_DIM) - output = torch.empty(shape_o, device=DEV, dtype=torch.bfloat16) - return [q, k, v, g, beta, None, pool, indices, True, False, cu_seqlens, False, output] - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name - certifying(op) - args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) - pool = args[6] - pool_before = pool.clone() - ref = op.reference(*args) - - o, final_state = op(*args) - - _compare((o, final_state, pool), ref) - if args[12] is not None: - assert o.data_ptr() == args[12].data_ptr(), "output= must be written in place" - named = args[7].tolist() - unnamed = [s for s in range(pool.shape[0]) if s not in set(named)] - assert torch.equal(pool[unnamed], pool_before[unnamed]), "unnamed slots written" - assert not torch.equal(pool[named], pool_before[named]), "named slots did not move" - - -def test_guard_refuses_a_noncontiguous_output() -> None: - """Unguarded, a head-slice view of a wider buffer is written through - shape-derived strides and the result is wrong.""" - q, k, v, g, beta, cu = _operands((96,), seed=61) - ref = op.reference(q, k, v, g, beta, cu_seqlens=cu) - head_slice = torch.zeros(1, 96, 2 * NUM_V_HEADS, HEAD_V_DIM, device=DEV, dtype=torch.bfloat16)[ - :, :, :NUM_V_HEADS - ] - assert not head_slice.is_contiguous() - o, _ = op(q, k, v, g, beta, cu_seqlens=cu, output=head_slice) - with pytest.raises(AssertionError): - _gate_output(o, ref) - with pytest.raises(AssertionError, match="output must be contiguous"): - with validating(op): - op(q, k, v, g, beta, cu_seqlens=cu, output=head_slice) - - -# --------------------------------------------------------------------------- -# The cold call. The cells above have already benchmarked this autotune key by -# the time pytest reaches it, so it runs in a fresh interpreter. - -CHILD_TIMEOUT_S = 420 -#: The GDN prefill call shape, packed varlen, state by slot index. -CHILD_SEQLENS = (256, 96) - - -def _state_kernel_autotuner(): - """The `triton.autotune` wrapper around `chunk_gated_delta_rule_fwd_kernel_h_blockdim64`; - a key in its `cache` is a key whose benchmark launches already happened.""" - from tensorrt_llm._torch.modules.fla.chunk_delta_h import ( - chunk_gated_delta_rule_fwd_kernel_h_blockdim64 as node, - ) - - while node is not None: - if isinstance(getattr(node, "cache", None), dict): - return node - node = getattr(node, "fn", None) - raise AssertionError("no triton autotuner found around the chunked state kernel") - - -def _checkpoint_case(seed: int = 101, pool_dtype: torch.dtype = torch.float32): - q, k, v, g, beta, cu = _operands(CHILD_SEQLENS, seed=seed) - pool = _pool(4, seed=seed + 1, dtype=pool_dtype) - indices = torch.tensor([2, 0], dtype=torch.int32, device=DEV) - buf = torch.empty( - 1, sum(CHILD_SEQLENS), NUM_V_HEADS, HEAD_V_DIM, device=DEV, dtype=torch.bfloat16 - ) - return (q, k, v, g, beta, None, pool, indices, True, False, cu, False, buf) - - -def _child_cold_call(pool_dtype: torch.dtype = torch.float32) -> None: - autotuner = _state_kernel_autotuner() - assert not autotuner.cache, "the child must be the first indexed call in its process" - args = _checkpoint_case(pool_dtype=pool_dtype) - ref = op.reference(*args) - o, _ = op(*args) - assert len(autotuner.cache) == 1, "the call must have benchmarked exactly one key" - _gate_output(o, ref) - _gate_pool(args[6], ref) - - -_CHILD_MODES = { - "cold-call-bf16-pool": lambda: _child_cold_call(torch.bfloat16), -} - - -def _run_child(mode: str) -> None: - # `_validating` lives four directories up; the child has no pytest to - # resolve `__extra_import_path__` for it. - root = str(Path(__file__).resolve().parents[4]) - existing = os.environ.get("PYTHONPATH") - env = dict(os.environ, PYTHONPATH=root if not existing else os.pathsep.join([root, existing])) - proc = subprocess.run( - [sys.executable, __file__, mode], - capture_output=True, - text=True, - timeout=CHILD_TIMEOUT_S, - env=env, - ) - print(proc.stdout) - assert proc.returncode == 0, ( - f"child `{mode}` exited {proc.returncode}\n--- stdout ---\n{proc.stdout}\n" - f"--- stderr ---\n{proc.stderr}" - ) - - -def test_cold_indexed_call_over_a_bf16_pool() -> None: - """The first indexed call of a process at 16 / 48 heads, K = V = 128, over a bf16 - pool, in the mixer's shape (256 + 96 packed, by slot, into `output=`), matches the - reference.""" - _run_child("cold-call-bf16-pool") - - -if __name__ == "__main__": - _CHILD_MODES[sys.argv[1]]() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_dilated_short_conv_paged.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_dilated_short_conv_paged.py deleted file mode 100644 index 8908d558b5ad..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_dilated_short_conv_paged.py +++ /dev/null @@ -1,147 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the dilated_short_conv_paged entry.""" - -import zlib - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.dilated_short_conv_paged import ( # noqa: E501 - ENGINE_SLOTS, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.dilated_short_conv_paged import ( # noqa: E501 - dilated_short_conv_paged as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "dilated_short_conv_paged requires a CUDA device" - - -def _compare( - out: tuple[torch.Tensor, torch.Tensor], ref: tuple[torch.Tensor, torch.Tensor] -) -> None: - """`(out, conv_states)` pairs: the returned activation with the pool as - the call left it, against `reference`'s pair. - - Activation at torch's default band: bf16 and fp16 were measured bitwise - equal to the reference (at 10240 channels and 461 slots too); fp32 - differs only by tap order. The pool is a copy, gated bitwise over the - whole pool -- the lock on every slot the batch did not name. - """ - (act, pool), (ref_act, ref_pool) = out, ref - assert act.dtype == ref_act.dtype and act.shape == ref_act.shape - assert act.is_contiguous() - torch.testing.assert_close(act, ref_act) - assert torch.equal(pool, ref_pool), "the published state must be bitwise exact" - - -KERNEL_SIZE = 4 -DILATION = 3 -STATE_LEN = (KERNEL_SIZE - 1) * DILATION -PAD_SLOT_ID = -1 - - -def _build(spec, seed): - """`__call__`'s positional arguments for one spec, freshly allocated, in the - call site's form: `x` a bare `.t()` of a token-major activation, int64 - `query_start_loc`, int32 slots. - - The checkpoint's taps are small (its four positions average 1e-4 .. 0.0185), - so the weight is scaled down rather than left at unit variance. - """ - g = torch.Generator(device="cuda").manual_seed(seed) - channels, lengths = spec["channels"], spec["lengths"] - num_seqs = len(lengths) - - def randn(*shape, scale=1.0): - return (scale * torch.randn(*shape, generator=g, device="cuda")).to(torch.bfloat16) - - starts = [0] - for length in lengths: - starts.append(starts[-1] + length) - qsl = torch.tensor(starts, dtype=torch.int64, device="cuda") - x = randn(starts[-1], channels).t() - weight = randn(channels, 1, KERNEL_SIZE, scale=0.3) - pool = randn(spec.get("pool_slots", 8), channels, STATE_LEN) - slots = spec.get("slots") or tuple(range(num_seqs)) - cache_indices = torch.tensor(slots, dtype=torch.int32, device="cuda") - initial = spec.get("initial", (True,) * num_seqs) - his = torch.tensor(initial, dtype=torch.bool, device="cuda") - return [x, weight, None, pool, qsl, cache_indices, his, DILATION, True, PAD_SLOT_ID] - - -def _run(args): - """Drive the op on a copy of the pool; return (out, pool after).""" - got = list(args) - got[3] = args[3].clone() - return op(*got), got[3] - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name - certifying(op) - args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) - x, weight = args[0].clone(), args[1].clone() - ref = op.reference(*args) - out, pool = _run(args) - assert out.data_ptr() != args[0].data_ptr() - assert torch.equal(args[0], x) and torch.equal(args[1], weight), "x and weight are read-only" - _compare((out, pool), ref) - if cell.spec.get("pool_slots") == ENGINE_SLOTS: - for slot in cell.spec["slots"]: # non-vacuity of the whole-pool gate - assert not torch.equal(pool[slot], args[3][slot]), slot - - -def _small(lengths, slots, seed, **spec): - return _build(dict(channels=256, lengths=lengths, slots=slots, **spec), seed=seed) - - -def test_cuda_graph_capture_and_replay() -> None: - """Captured once at the MTP-7 verify's 32 requests of 8 tokens, replayed after - the batch tensors' values change in place.""" - batch, p = 32, 8 - args = _small((p,) * batch, tuple(range(batch)), seed=25, pool_slots=batch + 2) - x, weight, _, pool, qsl, cidx, his = args[:7] - live = pool.clone() - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - for _ in range(3): - op(x, weight, None, live, qsl, cidx, his, DILATION, True, PAD_SLOT_ID) - torch.cuda.current_stream().wait_stream(stream) - - graph = torch.cuda.CUDAGraph() - live.copy_(pool) - with torch.cuda.graph(graph): - captured = op(x, weight, None, live, qsl, cidx, his, DILATION, True, PAD_SLOT_ID) - - g = torch.Generator(device="cuda").manual_seed(26) - # Every row but the last shrinks to one token; the slot map reverses and shifts by one. - qsl.copy_(torch.tensor([*range(batch), batch * p], dtype=torch.int32, device="cuda")) - cidx.copy_(torch.tensor(list(range(batch, 0, -1)), dtype=torch.int32, device="cuda")) - his.copy_(torch.tensor([i % 2 == 0 for i in range(batch)], dtype=torch.bool, device="cuda")) - x.copy_(torch.randn(x.shape, generator=g, device="cuda").to(x.dtype)) - live.copy_(pool) - graph.replay() - torch.cuda.synchronize() - - ref_out, ref_pool = op.reference( - x, weight, None, pool, qsl, cidx, his, DILATION, True, PAD_SLOT_ID - ) - torch.testing.assert_close(captured, ref_out) - assert torch.equal(live, ref_pool) - - -def test_mismatched_operands_raise() -> None: - """Refused loudly by the torch calls the entry is composed from -- why it - has no `is_valid`. Probed at a pool of another dtype (the index copy).""" - args = _small((4, 4), None, seed=27) - args[3] = args[3].float() - with pytest.raises(RuntimeError) as excinfo: - with validating(op): - op(*args) - fragment = "index_copy_(): self and source expected to have the same dtype" - assert fragment in str(excinfo.value), str(excinfo.value) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_extract_transpose_prefill_slice.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_extract_transpose_prefill_slice.py deleted file mode 100644 index d0e4fff2e5ad..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_extract_transpose_prefill_slice.py +++ /dev/null @@ -1,65 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the extract_transpose_prefill_slice entry.""" - -import zlib - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.extract_transpose_prefill_slice import ( # noqa: E501 - extract_transpose_prefill_slice as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "extract_transpose_prefill_slice requires a CUDA device" - - -def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: - """Bitwise: the kernel moves elements with no cast and no arithmetic, - so any tolerance would only hide a wrong index. Compared as raw bytes so - dtypes without an equality kernel (`float8_e4m3fn`) are gated the same.""" - assert out.dtype == ref.dtype, (out.dtype, ref.dtype) - assert out.shape == ref.shape, (out.shape, ref.shape) - assert out.is_contiguous() - if out.numel(): - assert torch.equal(out.reshape(-1).view(torch.uint8), ref.reshape(-1).view(torch.uint8)) - - -def _src(num_rows: int, num_cols: int, dtype: torch.dtype, seed: int) -> torch.Tensor: - g = torch.Generator(device="cuda").manual_seed(seed) - return torch.randn(num_rows, num_cols, generator=g, device="cuda").to(dtype) - - -def _build(spec, seed): - src = _src(spec["num_rows"], spec["num_cols"], spec["dtype"], seed) - if "row_stop" in spec: - src = src[: spec["row_stop"]] - if "col_slice" in spec: - lo, hi = spec["col_slice"] - src = src[:, lo:hi] - assert not src.is_contiguous() and src.stride(1) == 1 - return src, spec["n"], spec["start_col"], spec["width"] - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name - certifying(op) - src, n, start_col, width = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) - before = src.contiguous().view(torch.uint8).clone() - out = op(src, n, start_col, width) - assert out.data_ptr() != src.data_ptr() or out.numel() == 0 - assert torch.equal(src.contiguous().view(torch.uint8), before), "src is read-only" - _compare(out, op.reference(src, n, start_col, width)) - - -def test_guard_refuses_a_range_outside_src() -> None: - """The kernel masks only against the requested extents and reads past - `src` rather than failing; the guard bounds both axes (probed here at 128 - rows of a 64-row `src`).""" - src = torch.randn(64, 256, dtype=torch.bfloat16, device="cuda") - with pytest.raises(AssertionError, match="outside src"): - with validating(op): - op(src, 128, 0, 256) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py deleted file mode 100644 index c02361fd0d2e..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py +++ /dev/null @@ -1,143 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the flashinfer_chunk_gated_delta_rule entry (FlashInfer GDN prefill).""" - -import zlib - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.flashinfer_chunk_gated_delta_rule import ( # noqa: E501 - flashinfer_chunk_gated_delta_rule as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "flashinfer_chunk_gated_delta_rule requires a CUDA device" - -DEV = "cuda" -#: The band the Triton chunk op is certified at (rtol bf16 default; atol 2e-3 -#: on `o`, 2e-2 on states), derived for a chunked WY reformulation with bf16 -#: operands. This op is the same class of sum: measured against fp64 at -#: 1 x 8192 tokens, o max abs 1.26e-3 and state max abs 8.6e-3 (Triton op: -#: 1.13e-3 / 7.7e-3), inside it. -RTOL, ATOL_OUTPUT, ATOL_STATE = 1.6e-2, 2e-3, 2e-2 - - -def _compare(out: tuple, ref: tuple) -> None: - """`(o, final_state, pool)` against `reference`'s triple; `None` must match.""" - for got, want, atol in zip(out, ref, (ATOL_OUTPUT, ATOL_STATE, ATOL_STATE)): - assert (got is None) == (want is None), (got is None, want is None) - if got is None: - continue - assert got.dtype == want.dtype and got.shape == want.shape, ( - (got.dtype, got.shape), - (want.dtype, want.shape), - ) - assert torch.isfinite(got).all() - torch.testing.assert_close(got.float(), want.float(), rtol=RTOL, atol=atol) - - -def _build(spec, seed): - """Keyword arguments in the domain a GDN post-conv step produces: - unit-norm q/k, g < 0 in log space, beta a sigmoid.""" - gen = torch.Generator(device=DEV).manual_seed(seed) - num_k_heads, num_v_heads, head_k_dim, head_v_dim = spec["geometry"] - seqlens = spec["seqlens"] - total = sum(seqlens) - q = torch.randn(1, total, num_k_heads, head_k_dim, generator=gen, device=DEV) - k = torch.randn(1, total, num_k_heads, head_k_dim, generator=gen, device=DEV) - q = q / torch.sqrt(q.pow(2).sum(-1, keepdim=True) + 1e-6) - k = k / torch.sqrt(k.pow(2).sum(-1, keepdim=True) + 1e-6) - q, k = q.bfloat16(), k.bfloat16() - v = torch.randn( - 1, total, num_v_heads, head_v_dim, generator=gen, device=DEV, dtype=torch.bfloat16 - ) - a_log = torch.log(torch.rand(num_v_heads, generator=gen, device=DEV) * 15 + 1) - x = torch.randn(1, total, num_v_heads, generator=gen, device=DEV) + torch.randn( - num_v_heads, generator=gen, device=DEV - ) - g = (-torch.exp(a_log) * torch.nn.functional.softplus(x)).float() - beta = torch.rand(1, total, num_v_heads, generator=gen, device=DEV).sigmoid().float() - pool = ( - torch.randn(spec["slots"], num_v_heads, head_v_dim, head_k_dim, generator=gen, device=DEV) - * 0.1 - ).to(spec.get("pool_dtype", torch.float32)) - cu = torch.tensor([0] + torch.tensor(seqlens).cumsum(0).tolist(), dtype=torch.long, device=DEV) - indices = torch.tensor(spec["indices"], dtype=torch.int32, device=DEV) - output = None - if spec.get("output") == "buffer": - output = torch.full( - (1, total, num_v_heads, head_v_dim), -7.0, device=DEV, dtype=torch.bfloat16 - ) - extra = {} - if "initial" in spec: - extra["has_initial_state"] = torch.tensor(spec["initial"], dtype=torch.bool, device=DEV) - if spec.get("linear_gate"): - extra["linear_gate"] = True - g = torch.exp(g) - return dict( - **extra, - q=q, - k=k, - v=v, - g=g, - beta=beta, - scale=None, - initial_state=pool, - initial_state_indices=indices, - inplace_indexed_state_update=True, - output_final_state=False, - cu_seqlens=cu, - head_first=False, - use_qk_l2norm_in_kernel=False, - output=output, - state_workspace=None, - ) - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name - certifying(op) - args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) - pool = args["initial_state"] - pool_before = pool.clone() - ref = op.reference(**args) - - o, final_state = op(**args) - - _compare((o, final_state, pool), ref) - if args["output"] is not None: - assert o.data_ptr() == args["output"].data_ptr(), "output= must be written in place" - named = set(args["initial_state_indices"].tolist()) - unnamed = [s for s in range(pool.shape[0]) if s not in named] - assert torch.equal(pool[unnamed], pool_before[unnamed]), "unnamed slots written" - - -_CARRIED = next(c for c in op.CELLS if "initial" in c.spec) - - -def test_has_initial_state_is_the_reset_then_the_plain_call() -> None: - """The flags replace `reset_recurrent_state_rows` before the call: zeroing the fresh - slots and running the plain in-place call gives the same output and pool, bitwise; a - linear gate that is `torch.exp` of the plain call's log gate changes nothing.""" - args = _build(_CARRIED.spec, seed=5) - flags = args.pop("has_initial_state") - linear_gate = args.pop("linear_gate", False) - pool = args["initial_state"] - fresh = args["initial_state_indices"][~flags].long() - pool[fresh] = torch.randn_like(pool[fresh]) # stale rows the flags must not read - plain = dict(args, initial_state=pool.clone(), output=args["output"].clone()) - plain["initial_state"][fresh] = 0.0 - if linear_gate: - plain["g"] = torch.log(args["g"]) - # The linear gate of the target call is torch.exp of the log gate, bitwise. - args["g"] = torch.exp(plain["g"]) - - o_plain, _ = op(**plain) - o, final_state = op(**args, has_initial_state=flags, linear_gate=linear_gate) - - assert final_state is None - assert torch.equal(o, o_plain) - assert torch.equal(pool, plain["initial_state"]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_gating.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_gating.py deleted file mode 100644 index 4022900963ef..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_gating.py +++ /dev/null @@ -1,115 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the fused_gdn_gating entry. - -Geometry is a Qwen3.8-Flash-Next linear-attention layer: 48 value heads, `a` -bf16, `A_log` / `dt_bias` fp32 drawn across the checkpoint's own ranges. -""" - -import math -import zlib - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_gdn_gating import ( # noqa: E501 - A_LOG_RANGE, - APPROX_ULP, - BA_WIDTH, - DT_BIAS_RANGE, - NUM_V_HEADS, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_gdn_gating import ( # noqa: E501 - fused_gdn_gating as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "fused_gdn_gating requires a CUDA device" - - -#: The softplus `beta` the absolute floor is sized for; the tail error scales as -#: 1 / beta. The target calls 1.0, so 0.5 leaves a factor 2 of headroom. -_MIN_CERTIFIED_BETA = 0.5 - - -def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: - """Both fp32. `g` multiplies two approximate transcendentals, so the - relative term is `4 * 2^-22`. The absolute floor is the unstabilised - softplus's tail error -- about one approximate ulp of the log's - argument -- scaled by `1 / beta` and by the `exp(A_log)` it multiplies; - `_compare` cannot see either, so it takes exp(A_log) <= 139.42 (the - checkpoint's extreme) and beta >= 0.5 (`_MIN_CERTIFIED_BETA`): atol = - 4 * 2^-22 * 139.42 / 0.5 = 2.66e-04. - - Measured against this gate on 256 checkpoint-range rows: worst error / - allowance 0.215. Five wrong formulas (dt_bias dropped, sign dropped, - A_log without exp, exp(-A_log), sigmoid for softplus) land 8.7e+05 to - 1.5e+07 allowances away. - """ - assert out.dtype == ref.dtype == torch.float32, (out.dtype, ref.dtype) - assert out.shape == ref.shape, (out.shape, ref.shape) - atol = 4.0 * APPROX_ULP * math.exp(A_LOG_RANGE[1]) / _MIN_CERTIFIED_BETA - torch.testing.assert_close(out, ref, rtol=4.0 * APPROX_ULP, atol=atol) - - -DEV = "cuda" -SOFTPLUS_THRESHOLD = 20.0 - - -def _params(g, num_heads=NUM_V_HEADS, dtype=torch.float32): - lo, hi = A_LOG_RANGE - A_log = torch.rand(num_heads, generator=g, device=DEV) * (hi - lo) + lo - lo, hi = DT_BIAS_RANGE - dt_bias = torch.rand(num_heads, generator=g, device=DEV) * (hi - lo) + lo - return A_log.to(dtype), dt_bias.to(dtype) - - -def _build(spec, seed): - """`__call__`'s positional arguments for one spec.""" - g = torch.Generator(device=DEV).manual_seed(seed) - rows, heads = spec["rows"], NUM_V_HEADS - - def randn(*shape): - return torch.randn(*shape, generator=g, device=DEV).to(torch.bfloat16) - - if "fill" in spec: - a = torch.full((rows, heads), spec["fill"], dtype=torch.bfloat16, device=DEV) - elif spec.get("layout", "contiguous") == "target_rows": - # `_gdn_block` splits the [T, 48] projection into prefill and decode rows. - a = randn(rows + 23, heads)[23:] - assert a.is_contiguous() and a.storage_offset() == 23 * heads - else: - a = randn(rows, heads) - A_log, dt_bias = _params(g, heads) - return A_log, a, dt_bias, 1.0, SOFTPLUS_THRESHOLD - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name - certifying(op) - args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) - g = op(*args) - _compare(g, op.reference(*args)) - assert g.is_contiguous() and torch.isfinite(g).all() - - -def test_guard_refuses_what_the_kernel_reads_at_the_wrong_addresses() -> None: - """Unguarded, an `a` with head stride 2 returns a finite `g` -- exactly the - one for the memory the kernel actually addressed (`row * stride(0) + - head`: the first 48 columns of each row), not the tensor it was handed; - the guard refuses it.""" - g = torch.Generator(device=DEV).manual_seed(111) - A_log, dt_bias = _params(g) - wide = torch.randn(32, BA_WIDTH, generator=g, device=DEV).to(torch.bfloat16) - bad = (A_log, wide[:, ::2], dt_bias) - addressed = (A_log, wide[:, :NUM_V_HEADS], dt_bias) - unguarded = op(*bad) - assert torch.isfinite(unguarded).all() - assert torch.equal(unguarded, op(*addressed)), "the hazard this guard covers is gone" - with pytest.raises(AssertionError): - _compare(unguarded, op.reference(*bad)) - with pytest.raises(AssertionError, match="head axis must be unit-stride"): - with validating(op): - op(*bad) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_post_conv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_post_conv.py deleted file mode 100644 index 01b35ad0d378..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_gdn_post_conv.py +++ /dev/null @@ -1,143 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the fused_gdn_post_conv entry. - -Geometry is a Qwen3.8-Flash-Next linear-attention layer: 16 x 128 key heads, -48 x 128 value heads, qkv_dim 10240, bf16 activations. -""" - -import math -import zlib - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_gdn_post_conv import ( # noqa: E501 - A_LOG_RANGE, - DT_BIAS_RANGE, - HEAD_K_DIM, - HEAD_V_DIM, - NUM_K_HEADS, - NUM_V_HEADS, - QKV_DIM, - SOFTPLUS_THRESHOLD, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_gdn_post_conv import ( # noqa: E501 - fused_gdn_post_conv as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "fused_gdn_post_conv requires a CUDA device" - - -#: Relative error of the SASS approximate transcendentals (`ex2.approx.f32`, -#: `lg2.approx.f32`) Triton lowers `tl.exp` / `tl.log` to. -_APPROX_ULP = 2.0**-22 - - -def _compare(out: tuple, ref: tuple) -> None: - """Per output, each gate derived rather than tuned. - - q, k -- torch's bf16 band: both sides round one fp32 value, so they - differ by at most one bf16 ulp. v -- bitwise: the kernel moves bits. - g -- a product of two SASS approximate transcendentals (relative error - ~2^-22 each) with the softplus magnitude bounded by the threshold and - every term scaled by `exp(A_log)`: - `atol = 4 * 2^-22 * softplus_threshold * exp(max A_log)`, `rtol = 8 * 2^-22`. - At this checkpoint's A_log ceiling that is 6.7e-04 against a measured - worst case of 7.1e-05; five wrong gate formulas land at 1.2e+02 and up. - beta -- fp32: torch's fp32 band (one well-conditioned transcendental); - narrower: bitwise, the fp32 sigmoid rounded exactly once. - """ - assert len(out) == len(ref) == 5 - for got, want in zip(out, ref): - assert got.dtype == want.dtype and got.shape == want.shape, ( - (got.dtype, got.shape), - (want.dtype, want.shape), - ) - assert got.is_contiguous() - q, k, v, g, beta = out - q_ref, k_ref, v_ref, g_ref, beta_ref = ref - torch.testing.assert_close(q, q_ref) - torch.testing.assert_close(k, k_ref) - assert torch.equal(v, v_ref), "v must be a bit-exact copy" - g_atol = 4 * _APPROX_ULP * SOFTPLUS_THRESHOLD * math.exp(A_LOG_RANGE[1]) - torch.testing.assert_close(g, g_ref, rtol=8 * _APPROX_ULP, atol=g_atol) - if beta.dtype == torch.float32: - torch.testing.assert_close(beta, beta_ref) - else: - assert torch.equal(beta, beta_ref), "beta must be the fp32 sigmoid rounded once" - - -DEV = "cuda" -DTYPE = torch.bfloat16 -HEADS = (NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM) -L2_NORM_EPS = 1e-6 -#: The `z` columns of the input projection. -Z_WIDTH = NUM_V_HEADS * HEAD_V_DIM -#: The target's one GDN input-projection output: [q|k|v | z | a | b]. -IN_PROJ_WIDTH = QKV_DIM + Z_WIDTH + 2 * NUM_V_HEADS # 16480 - - -def _build(spec, seed): - """`__call__`'s positional arguments (its keyword `beta_dtype` left at fp32).""" - g = torch.Generator(device=DEV).manual_seed(seed) - n_p, n_d = spec["num_prefill"], spec["num_decode"] - n = n_p + n_d - - def randn(*shape): - return torch.randn(*shape, generator=g, device=DEV).to(DTYPE) - - # `_gdn_block`: one [T, 16480] GEMM output; q|k|v, a and b are its column - # views, the decode half the row split of q|k|v, the prefill half the - # transposing extraction of the rest. - proj = randn(n, IN_PROJ_WIDTH) - mixed = proj[:, :QKV_DIM] - prefill = mixed[:n_p].t().contiguous() - decode = mixed[n_p:] if n_d else None - a = proj[:, QKV_DIM + Z_WIDTH : QKV_DIM + Z_WIDTH + NUM_V_HEADS] - b = proj[:, QKV_DIM + Z_WIDTH + NUM_V_HEADS :] - assert a.stride(0) == b.stride(0) == IN_PROJ_WIDTH - assert decode is None or decode.stride() == (IN_PROJ_WIDTH, 1) - - lo, hi = A_LOG_RANGE - A_log = torch.rand(NUM_V_HEADS, generator=g, device=DEV) * (hi - lo) + lo - lo, hi = DT_BIAS_RANGE - dt_bias = torch.rand(NUM_V_HEADS, generator=g, device=DEV) * (hi - lo) + lo - if "dt_bias_fill" in spec: - dt_bias = torch.full((NUM_V_HEADS,), spec["dt_bias_fill"], device=DEV) - return [prefill, decode, a, b, A_log, dt_bias, *HEADS, L2_NORM_EPS, SOFTPLUS_THRESHOLD] - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name - certifying(op) - args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) - if "dt_bias_fill" in cell.spec: - # Every head strictly on the side of the threshold the cell names. - x = args[2].float() + args[5].float() - above = cell.spec["dt_bias_fill"] > SOFTPLUS_THRESHOLD - assert bool((x > SOFTPLUS_THRESHOLD).all() if above else (x <= SOFTPLUS_THRESHOLD).all()) - before = [t.clone() if isinstance(t, torch.Tensor) else t for t in args] - out = op(*args) - for arg, old in zip(args, before): - if isinstance(arg, torch.Tensor): - assert torch.equal(arg, old), "inputs are read-only" - _compare(out, op.reference(*args)) - - -def test_guard_refuses_strided_gate_parameters() -> None: - """Unguarded, a stride-2 `A_log` view is read as `ptr + head` -- the - interleaved zeros -- and returns a finite wrong g.""" - args = _build(dict(num_prefill=32, num_decode=0), seed=108) - strided = torch.stack([args[4], torch.zeros_like(args[4])], dim=1)[:, 0] - assert strided.stride() == (2,) - bad = list(args) - bad[4] = strided - with pytest.raises(AssertionError): - _compare(op(*bad), op.reference(*bad)) - with pytest.raises(AssertionError, match="must be contiguous 1-D"): - with validating(op): - op(*bad) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_update.py deleted file mode 100644 index 3b2041b63919..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_update.py +++ /dev/null @@ -1,250 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the fused_recurrent_gated_delta_rule_update entry. - -The axis this entry exists for is more than one token per sequence per step -- -the speculative verify call -- so the cells are that call as the target makes it -(per-step states into the rollback buffer, the pool untouched) at both pool -dtypes. The write-back sink is driven only by the cold-call test below. - -The neighbouring `chunk_gated_delta_rule` was once wrong on the first indexed -in-place call of a process (an autotuner trial launch consumes the state), so -`test_cold_call_in_a_fresh_process_is_correct` measures the first call of a -process that has made none. -""" - -from __future__ import annotations - -import json -import subprocess -import sys -import zlib -from pathlib import Path - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_recurrent_gated_delta_rule_update import ( # noqa: E501 - HEAD_K_DIM, - HEAD_V_DIM, - NUM_K_HEADS, - NUM_V_HEADS, - SLOTS, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_recurrent_gated_delta_rule_update import ( # noqa: E501 - fused_recurrent_gated_delta_rule_update as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "fused_recurrent_gated_delta_rule_update requires a CUDA device" - - -def _compare(out: tuple, ref: tuple) -> None: - """`(returned, initial_state_source, intermediate_states_buffer)` as the - call left them, against `reference`'s triple, each at torch's default - dtype-aware band -- nothing is loosened. - - The kernel and the reference both accumulate in fp32 and differ only in - reduction order. Measured over every cell shape: output <= 2.27e-4 at a - bf16 absmax of 3.9e-2 .. 8.2e-2 (under one bf16 ulp there), written-back - state <= 2.98e-7 at absmax ~1, per-step states <= 1.8e-7; a bf16 - per-step buffer costs one bf16 ulp (1.94e-3 at absmax 0.89). Eight wrong - recurrences (q/k heads tiled not interleaved, q/k swapped, beta 1, no - gating, reversed tokens, every token from the initial state, a - transposed or absent initial state) land 1.56e-2 .. 8.17e-2 away on the - output -- 128x the correct result's own deviation at the closest. - """ - (o, pool, cache), (o_ref, pool_ref, cache_ref) = out, ref - assert o.dtype == o_ref.dtype and o.shape == o_ref.shape, ( - (o.dtype, o.shape), - (o_ref.dtype, o_ref.shape), - ) - torch.testing.assert_close(o, o_ref) - assert (pool is None) == (pool_ref is None) - if pool is not None: - torch.testing.assert_close(pool, pool_ref) - assert (cache is None) == (cache_ref is None) - if cache is not None: - torch.testing.assert_close(cache, cache_ref) - - -DEV = "cuda" -MIXED_QKV_WIDTH = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240 -CACHE_SENTINEL = -3.0 - - -def _build(spec, seed): - """`__call__`'s arguments by name for one spec, plus the live slots. - - q/k/v are views over one packed `[T, 10240]` row block (or contiguous - copies, as the target passes them); `g` is built the way the gating step - builds it -- negative, so `exp(g) < 1` -- and `beta` is a sigmoid. - """ - gen = torch.Generator(device=DEV).manual_seed(seed) - mode, n, tokens = spec["mode"], spec["n"], spec["p"] - total = n * tokens - - def randn(*shape, dtype=torch.float32): - return torch.randn(*shape, generator=gen, device=DEV, dtype=dtype) - - mixed = randn(total, MIXED_QKV_WIDTH, dtype=torch.bfloat16) - kd = NUM_K_HEADS * HEAD_K_DIM - q = mixed[:, :kd].view(n, tokens, NUM_K_HEADS, HEAD_K_DIM) - k = mixed[:, kd : 2 * kd].view(n, tokens, NUM_K_HEADS, HEAD_K_DIM) - v = mixed[:, 2 * kd :].view(n, tokens, NUM_V_HEADS, HEAD_V_DIM) - if spec.get("qkv") == "contiguous": - q, k, v = q.contiguous(), k.contiguous(), v.contiguous() - - a, a_log, dt_bias = randn(total, NUM_V_HEADS), randn(NUM_V_HEADS), randn(NUM_V_HEADS) - g = -torch.exp(a_log) * torch.nn.functional.softplus(a + dt_bias) - g = g.view(n, tokens, NUM_V_HEADS) - b = randn(total, NUM_V_HEADS, dtype=torch.bfloat16) - beta = b.sigmoid() if spec.get("beta", "fp32") == "bf16" else b.float().sigmoid() - beta = beta.view(n, tokens, NUM_V_HEADS) - - pool_slots = spec.get("pool_slots", SLOTS) - pool = randn(pool_slots, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM) * 0.05 - pool = pool.to(spec.get("pool_dtype", torch.float32)) - - slots = spec.get("slots") or tuple(range(n)) - indices = torch.tensor(slots, dtype=torch.int32, device=DEV) - - buffer = cache_steps = None - if mode == "verify": - rows = spec.get("cache_rows", pool_slots) - cache_steps = tokens - buffer = torch.full( - (rows, cache_steps, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM), - CACHE_SENTINEL, - dtype=spec.get("cache_dtype", torch.float32), - device=DEV, - ) - - kwargs = dict( - q=q, - k=k, - v=v, - g=g, - beta=beta, - initial_state_source=pool, - initial_state_indices=indices, - scale=None, - use_qk_l2norm_in_kernel=True, - cu_seqlens=None, - disable_state_update=mode == "verify", - disable_output_calculation=False, - intermediate_states_buffer=buffer, - cache_steps=cache_steps, - output=None, - ) - return kwargs, list(slots) - - -def _seed(why): - return zlib.crc32(why.encode()) - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name - certifying(op) - kw, live = _build(cell.spec, seed=_seed(cell.why)) - pool, buffer = kw["initial_state_source"], kw["intermediate_states_buffer"] - pool_before, buffer_before = pool.clone(), buffer.clone() - - ref = op.reference(**kw) - got = op(**kw) - _compare((got, pool, buffer), ref) - - # The in-place locks, bitwise, independent of the reference. - assert torch.equal(pool, pool_before), "the pool must not move on a verify call" - for s in live: - assert not (buffer[s] == CACHE_SENTINEL).any(), f"row {s} not fully cached" - buffer_before[live] = buffer[live] - assert torch.equal(buffer, buffer_before), "a buffer row outside the indices was written" - - -# ── the mechanism the neighbouring entry failed on ──────────────────────── - - -_COLD_SPEC = dict(mode="write_back", n=2, p=3, slots=(5, 2)) - - -def cold_call_measurement() -> dict: - """The first indexed in-place call of the process, then the second. - - Public because the child process below calls it by name. - """ - kw, _ = _build(_COLD_SPEC, seed=7) - pool = kw["initial_state_source"] - pool_before = pool.clone() - ref = op.reference(**kw) - cold_out = op(**kw).clone() - cold_pool = pool.clone() - try: - _compare((cold_out, cold_pool, None), ref) - cold_ok = True - except AssertionError: - cold_ok = False - pool.copy_(pool_before) - warm_out = op(**kw) - return { - "cold_compare_ok": cold_ok, - "cold_out_err": (cold_out.float() - ref[0].float()).abs().max().item(), - "cold_pool_err": (cold_pool - ref[1]).abs().max().item(), - "out_bitwise_equal": bool(torch.equal(cold_out, warm_out)), - "pool_bitwise_equal": bool(torch.equal(cold_pool, pool)), - } - - -_COLD_CHILD = """ -import importlib.util, json, sys -sys.path.insert(0, {root!r}) -spec = importlib.util.spec_from_file_location("_cold_entry_test", {path!r}) -mod = importlib.util.module_from_spec(spec) -sys.modules["_cold_entry_test"] = mod -spec.loader.exec_module(mod) -print("COLD_RESULT " + json.dumps(mod.cold_call_measurement())) -""" - - -def test_cold_call_in_a_fresh_process_is_correct() -> None: - """Measured in a child that has launched nothing: a warm-up here would - hide exactly the defect the sibling prefill core has. Cold must pass the - entry's own gate *and* be bitwise the warm call.""" - root = str(Path(__file__).resolve().parents[4]) # where `_validating` lives - proc = subprocess.run( - [sys.executable, "-c", _COLD_CHILD.format(root=root, path=__file__)], - capture_output=True, - text=True, - timeout=480, - check=False, - ) - assert proc.returncode == 0, ( - f"cold-call child exited {proc.returncode}\n{proc.stdout[-4000:]}\n{proc.stderr[-4000:]}" - ) - line = [ln for ln in proc.stdout.splitlines() if ln.startswith("COLD_RESULT ")] - assert line, f"child produced no result\n{proc.stdout[-4000:]}" - result = json.loads(line[0][len("COLD_RESULT ") :]) - print(f"\ncold call: {result}") - assert result["cold_compare_ok"], result - assert result["out_bitwise_equal"] and result["pool_bitwise_equal"], result - - -# ── what the guard refuses ───────────────────────────────────────────── - - -def test_cache_steps_off_the_buffer_step_axis_misplaces_rows_and_is_refused() -> None: - """Unguarded, a 4-step buffer driven at cache_steps 2: row 1's states land in - row 0's steps 2..3 and row 1 is never written -- the rows the promotion - reads are not the ones the kernel filled.""" - kw = _build(dict(mode="verify", n=2, p=2), seed=75)[0] - deep = torch.full((SLOTS, 4, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM), CACHE_SENTINEL, device=DEV) - op(**dict(kw, intermediate_states_buffer=deep, cache_steps=2)) - assert (deep[1] == CACHE_SENTINEL).all() and not (deep[0, 2:] == CACHE_SENTINEL).any(), ( - "the hazard this guard covers is gone" - ) - with pytest.raises(AssertionError, match="must equal intermediate_states_buffer's step axis"): - with validating(op): - op(**dict(kw, intermediate_states_buffer=deep.clone(), cache_steps=2)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py deleted file mode 100644 index b1039ed71fd6..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_sigmoid_gating_delta_rule_update.py +++ /dev/null @@ -1,231 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the fused_sigmoid_gating_delta_rule_update entry. - -Everything here certifies the vendored Triton branch: the target's decode call -- the -rows form (no `cu_seqlens`) over a bf16 pool -- and the packed call over an fp32 pool that -the GDN block test uses as its reference. `test_rows_form_is_batch_invariant` shows the -property the bf16 call takes the rows form for. -""" - -import zlib -from typing import Optional - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_sigmoid_gating_delta_rule_update import ( # noqa: E501 - ENGINE_SLOTS, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_sigmoid_gating_delta_rule_update import ( # noqa: E501 - fused_sigmoid_gating_delta_rule_update as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "fused_sigmoid_gating_delta_rule_update requires a CUDA device" - - -def _compare( - out: tuple[torch.Tensor, Optional[torch.Tensor]], - ref: tuple[torch.Tensor, Optional[torch.Tensor]], -) -> None: - """`(o, pool)`: the returned output with the pool as the call left it. - - Both at torch's default dtype-aware band: the kernel and the reference - accumulate the same fp32 sequence and round once (bf16 `o`: measured - 3.05e-05 max abs at N=32, 461 slots; fp32 pool 2.4e-07). Over a bf16 pool the - state is rounded once on store, as the reference rounds it, so a stored value - can sit one bf16 step (2^-8 relative) from the reference's, inside the band. Five wrong - variants (a/b swapped, norm dropped, pool `(V, K)` transposed, initial - state ignored, q-head order reversed) land 1.7e-2 to 4.5 away, three - orders of magnitude outside. The pool is compared whole, so the - unnamed slots are gated too (the reference leaves them bit-identical). - """ - (o, pool), (ref_o, ref_pool) = out, ref - assert o.dtype == ref_o.dtype and o.shape == ref_o.shape, ( - (o.dtype, o.shape), - (ref_o.dtype, ref_o.shape), - ) - assert torch.isfinite(o).all() - torch.testing.assert_close(o, ref_o) - assert (pool is None) == (ref_pool is None) - if pool is not None: - assert pool.dtype == ref_pool.dtype and pool.shape == ref_pool.shape - torch.testing.assert_close(pool, ref_pool) - - -DEV = "cuda" -NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM = 16, 48, 128, 128 -MIXED_QKV_WIDTH = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240 -#: The target's one GDN input-projection output: [q|k|v | z | a | b], 16480 wide. -IN_PROJ_A = MIXED_QKV_WIDTH + NUM_V_HEADS * HEAD_V_DIM # 16384 -IN_PROJ_WIDTH = IN_PROJ_A + 2 * NUM_V_HEADS -SOFTPLUS_BETA, SOFTPLUS_THRESHOLD = 1.0, 20.0 - - -def _rows_form(mixed, a, b): - """The target's bf16 decode operands: q/k/v `[N, 1, heads, 128]` views of the packed - `[N, 10240]` block, reshaped whole so the size-1 token axis carries the row stride, and - a/b `[N, 1, 48]`.""" - n = mixed.shape[0] - qkv = mixed.view(n, 1, MIXED_QKV_WIDTH // HEAD_K_DIM, HEAD_K_DIM) - q, k, v = torch.split(qkv, [NUM_K_HEADS, NUM_K_HEADS, NUM_V_HEADS], 2) - return q, k, v, a.reshape(n, 1, NUM_V_HEADS), b.reshape(n, 1, NUM_V_HEADS) - - -def _build(spec, seed): - """`__call__`'s 15 positional arguments, plus the packed block the rows form - views. q/k/v are views over one packed `[T, 10240]` conv output, exactly as - the decode call reshapes its column splits.""" - g = torch.Generator(device=DEV).manual_seed(seed) - seq_lens = spec["seq_lens"] - total = sum(seq_lens) - if spec["ab"] == "in_proj": - proj = torch.randn(total, IN_PROJ_WIDTH, generator=g, device=DEV, dtype=torch.bfloat16) - mixed = proj[:, :MIXED_QKV_WIDTH] - else: - mixed = torch.randn(total, MIXED_QKV_WIDTH, generator=g, device=DEV, dtype=torch.bfloat16) - kd = NUM_K_HEADS * HEAD_K_DIM - q = mixed[:, :kd].view(1, total, NUM_K_HEADS, HEAD_K_DIM) - k = mixed[:, kd : 2 * kd].view(1, total, NUM_K_HEADS, HEAD_K_DIM) - v = mixed[:, 2 * kd :].view(1, total, NUM_V_HEADS, HEAD_V_DIM) - if spec["ab"] == "in_proj": - a = proj[:, IN_PROJ_A : IN_PROJ_A + NUM_V_HEADS] - b = proj[:, IN_PROJ_A + NUM_V_HEADS :] - assert mixed.stride(0) == a.stride(0) == b.stride(0) == IN_PROJ_WIDTH - else: - a = torch.randn(total, NUM_V_HEADS, generator=g, device=DEV, dtype=torch.bfloat16) - b = torch.randn(total, NUM_V_HEADS, generator=g, device=DEV, dtype=torch.bfloat16) - A_log = torch.randn(NUM_V_HEADS, generator=g, device=DEV) - dt_bias = torch.randn(NUM_V_HEADS, generator=g, device=DEV) - - state_shape = (spec["slots"], NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM) - pool = torch.randn(state_shape, generator=g, device=DEV) * 0.05 - pool = pool.to(spec.get("pool_dtype", torch.float32)) - indices = spec.get("indices") or tuple(range(len(seq_lens))) - indices = torch.tensor(indices, dtype=torch.int32, device=DEV) - cu_seqlens = torch.tensor( - [0] + torch.tensor(seq_lens).cumsum(0).tolist(), dtype=torch.int64, device=DEV - ) - if spec.get("layout", "varlen") == "rows": - q, k, v, a, b = _rows_form(mixed, a, b) - cu_seqlens = None - args = [ - A_log, - a, - dt_bias, - SOFTPLUS_BETA, - SOFTPLUS_THRESHOLD, - q, - k, - v, - b, - pool, - indices, - None, - True, - cu_seqlens, - None, - ] - return args, {"mixed": mixed} - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name - certifying(op) - spec = cell.spec - args, _ = _build(spec, seed=zlib.crc32(cell.why.encode())) - pool = args[9] - pool_before = pool.clone() - ref = op.reference(*args) - - out = op(*args) - - _compare((out, pool), ref) - # The whole-pool compare is the lock on unnamed slots; make it bitwise. - named = {s for s in args[10].tolist() if s >= 0} - unnamed = [s for s in range(pool.shape[0]) if s not in named] - assert torch.equal(pool[unnamed], pool_before[unnamed]), "unnamed slots written" - if spec["slots"] == ENGINE_SLOTS: - for slot in named: - assert not torch.equal(pool[slot], pool_before[slot]), slot - - -#: Batch sizes either side of every FlashInfer dispatch boundary at 48 value heads (its -#: kernel and tile change at 2|3, 6|7, 10|11 and 21|22) and the engine's CUDA-graph batch -#: sizes up to 32. -_BATCHES = (1, 2, 3, 10, 11, 16, 21, 22, 32) -_FIXED_SLOT = 301 -_FIXED_ROW_SLOTS = 512 - - -def test_rows_form_is_batch_invariant() -> None: - """The property the target's bf16 decode call relies on: a row's output and state are - bitwise independent of the batch it rides in, at batch sizes either side of every - FlashInfer dispatch boundary and at both ends of the batch. One fixed row (activations - and slot 301 of a bf16 pool) sits at position 0 and N-1 of random N-row batches.""" - gen = torch.Generator(device=DEV).manual_seed(3) - cpu = torch.Generator().manual_seed(3) - pool0 = ( - torch.randn( - _FIXED_ROW_SLOTS, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM, generator=gen, device=DEV - ) - * 0.05 - ).to(torch.bfloat16) - row = torch.randn(MIXED_QKV_WIDTH, generator=gen, device=DEV, dtype=torch.bfloat16) - row_a, row_b = ( - torch.randn(NUM_V_HEADS, generator=gen, device=DEV, dtype=torch.bfloat16) for _ in range(2) - ) - A_log = torch.randn(NUM_V_HEADS, generator=gen, device=DEV) - dt_bias = torch.randn(NUM_V_HEADS, generator=gen, device=DEV) - first = None - for n in _BATCHES: - for pos in sorted({0, n - 1}): - g = torch.Generator(device=DEV).manual_seed(1000 * n + pos) - mixed = torch.randn(n, MIXED_QKV_WIDTH, generator=g, device=DEV, dtype=torch.bfloat16) - a = torch.randn(n, NUM_V_HEADS, generator=g, device=DEV, dtype=torch.bfloat16) - b = torch.randn(n, NUM_V_HEADS, generator=g, device=DEV, dtype=torch.bfloat16) - mixed[pos], a[pos], b[pos] = row, row_a, row_b - others = [s for s in torch.randperm(_FIXED_ROW_SLOTS, generator=cpu).tolist()] - others = [s for s in others if s != _FIXED_SLOT][:n] - others[pos] = _FIXED_SLOT - slots = torch.tensor(others, dtype=torch.int32, device=DEV) - pool = pool0.clone() - q, k, v, a_in, b_in = _rows_form(mixed, a, b) - out = op(A_log, a_in, dt_bias, 1.0, 20.0, q, k, v, b_in, pool, slots, None, True, None) - got = (out.reshape(n, NUM_V_HEADS, HEAD_V_DIM)[pos].clone(), pool[_FIXED_SLOT].clone()) - first = first or got - assert torch.equal(got[0], first[0]) and torch.equal(got[1], first[1]), ( - f"the fixed row differs at batch {n}, position {pos}" - ) - - -def test_guard_refuses_rows_views_whose_token_stride_is_not_the_row_stride() -> None: - """Unguarded, q/k/v reshaped to `[N, 1, heads, 128]` from column slices of the packed - block carry `stride(1) = heads * 128`: without `cu_seqlens` the kernel reads row n at - `n * stride(1)`, from other rows' columns.""" - n, kd = 8, NUM_K_HEADS * HEAD_K_DIM - spec = dict( - seq_lens=(1,) * n, slots=64, ab="separate", layout="rows", pool_dtype=torch.bfloat16 - ) - args, extras = _build(spec, seed=600) - mixed = extras["mixed"] - pool_before = args[9].clone() - ref = op.reference(*args) - bad = list(args) - bad[5] = mixed[:, :kd].reshape(n, 1, NUM_K_HEADS, HEAD_K_DIM) - bad[6] = mixed[:, kd : 2 * kd].reshape(n, 1, NUM_K_HEADS, HEAD_K_DIM) - bad[7] = mixed[:, 2 * kd :].reshape(n, 1, NUM_V_HEADS, HEAD_V_DIM) - assert all(bad[i].stride(1) != bad[i].stride(0) for i in (5, 6, 7)) - for i in (5, 6, 7): - assert torch.equal(bad[i], args[i]), "the same values, other strides" - got = op(*bad) - with pytest.raises(AssertionError): - _compare((got, bad[9]), ref) - guarded = list(bad) - guarded[9] = pool_before - with pytest.raises(AssertionError, match=r"stride\(0\) must equal"): - with validating(op): - op(*guarded) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py deleted file mode 100644 index 6447b4cf57fd..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py +++ /dev/null @@ -1,416 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test: the target's whole Gated-DeltaNet block against the all-Triton path. - -Drives `DecodeTarget._gdn_block` and `PrefillTarget._gdn_block` at the -checkpoint's geometry (hidden 2560, 16 k-heads, 48 v-heads, head dim 128, -fp32 or bf16 recurrent pool of 461 slots) on a decode and a mixed (context + one-token -generation rows) batch, and compares the block output -and both state pools with a reference that composes the Triton recurrences -(`fused_sigmoid_gating_delta_rule_update`'s rows form for decode, bitwise its -packed call at either pool dtype; `chunk_gated_delta_rule` over every row of a -prefill step) around the same projections, convolutions and epilogue. - -The reference's context rows run the unfused `extract_transpose_prefill_slice` --> `causal_conv1d_fwd` -> `fused_gdn_post_conv` chain the target's -`gdn_prefill_conv_post` replaces. - -Bands: the conv pools are bitwise (both sides publish the raw projection). The -recurrent pool is held to the looser of the two recurrence entries' state -bands (`flashinfer_chunk_gated_delta_rule`: rtol 1.6e-2, atol 2e-2), and the -block output, after the gated norm and the output projection, to a relative -L2 distance of 1e-2 and a max-abs distance of 3e-2 of its own peak. -""" - -import types - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - modeling as M, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_fwd import ( # noqa: E501 - causal_conv1d_fwd, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.chunk_gated_delta_rule import ( # noqa: E501 - chunk_gated_delta_rule as triton_chunk, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.extract_transpose_prefill_slice import ( # noqa: E501 - extract_transpose_prefill_slice, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_gdn_post_conv import ( # noqa: E501 - fused_gdn_post_conv, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_sigmoid_gating_delta_rule_update import ( # noqa: E501 - fused_sigmoid_gating_delta_rule_update as triton_decode, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.reset_recurrent_state_rows import ( # noqa: E501 - reset_recurrent_state_rows, -) - -assert torch.cuda.is_available(), "the Gated-DeltaNet block requires a CUDA device" - -DEV = "cuda" -BF16 = torch.bfloat16 -HIDDEN, H, HV, K, V, WIDTH = 2560, 16, 48, 128, 128, 4 -KD, VD = H * K, HV * V -CONV = 2 * KD + VD -SLOTS = 461 -MAX_BATCH = 256 -RTOL_STATE, ATOL_STATE = 1.6e-2, 2e-2 -OUT_REL_L2, OUT_MAX_FRAC = 1e-2, 3e-2 - - -def _slots(n: int) -> list[int]: - """`n` distinct slots of the pool, both ends first.""" - rest = [(i * 37 + 4) % (SLOTS - 2) + 1 for i in range(max(n - 2, 0))] - return ([SLOTS - 1, 0] + rest)[:n] - - -def _core(seed: int, state: torch.dtype = torch.float32) -> types.SimpleNamespace: - """The attributes `_gdn_block` reads off the core, with layer 0's pools (the recurrent - one in `state`).""" - g = torch.Generator(device=DEV).manual_seed(seed) - conv = torch.randn(SLOTS, CONV, WIDTH - 1, generator=g, device=DEV, dtype=BF16) - ssm = (torch.randn(SLOTS, HV, V, K, generator=g, device=DEV) * 0.05).to(state) - return types.SimpleNamespace( - conv_dim=CONV, - key_dim=KD, - value_dim=VD, - num_k_heads=H, - head_k_dim=K, - num_v_heads=HV, - head_v_dim=V, - eps=1e-6, - conv_kernel=WIDTH, - in_proj_split=(CONV, VD, HV, HV), - _bound=types.SimpleNamespace(conv={0: conv}, ssm={0: ssm}), - ) - - -def _clone_core(core: types.SimpleNamespace) -> types.SimpleNamespace: - out = types.SimpleNamespace(**vars(core)) - out._bound = types.SimpleNamespace( - conv={0: core._bound.conv[0].clone()}, ssm={0: core._bound.ssm[0].clone()} - ) - return out - - -def _layer(seed: int) -> types.SimpleNamespace: - """Layer weights as the target views them: HF `[out, in]` storage read - through `.t()`, fp32 gating vectors over the checkpoint's ranges.""" - g = torch.Generator(device=DEV).manual_seed(seed) - - def lin(n_out, n_in): - return (torch.randn(n_out, n_in, generator=g, device=DEV) * n_in**-0.5).to(BF16).t() - - # The target's one [q|k|v | z | a | b] input projection, read through `.t()`. - # The reference blocks run the same GEMM, so the conv pools stay bitwise. - return types.SimpleNamespace( - in_proj=lin(CONV + VD + 2 * HV, HIDDEN), - out=lin(HIDDEN, VD), - conv_w=(torch.randn(CONV, WIDTH, generator=g, device=DEV) * 0.5).to(BF16), - a_log=torch.empty(HV, device=DEV).uniform_(-5.5, 3.5, generator=g), - dt_bias=torch.empty(HV, device=DEV).uniform_(-8.9, 19.2, generator=g), - gnorm=(1.0 + 0.1 * torch.randn(VD // HV, generator=g, device=DEV)).to(BF16), - # No FP8 copies: the projections stay on the bf16 entries (dense FP8 mode off). - in_proj8=None, - out8=None, - in_proj_g=None, - out_g=None, - in_proj_m=None, - out_m=None, - in_proj_nv=None, - ) - - -def _x(tokens: int, seed: int) -> torch.Tensor: - g = torch.Generator(device=DEV).manual_seed(seed) - return torch.randn(tokens, HIDDEN, generator=g, device=DEV, dtype=BF16) - - -def _target(cls, core, L=None, step=None): - """A target over `core` alone; a prefill target's `_gdn_core` break reads layer 0's - weights and the step's bookkeeping off it.""" - target = object.__new__(cls) - target.core = core - target._layers = [L] - target._step = step - return target - - -def _prefill_block(core, x, L, st): - """`PrefillTarget._gdn_block` over the step `st`, the break's output buffer allocated - as the forward allocates it.""" - tokens = x.shape[0] - step = types.SimpleNamespace(gdn=st, num_tokens=tokens) - out = torch.empty([1, tokens, HV, V], dtype=BF16, device=DEV) - return _target(M.PrefillTarget, core, L, step)._gdn_block(x, 0, L, out) - - -def _decode_state(m: int): - """`DecodeTarget.forward`'s `gdn_state` for `m` generation sequences: the - slot vector is the head of an engine-sized buffer.""" - buf = torch.zeros(MAX_BATCH, dtype=torch.int32, device=DEV) - buf[:m] = torch.tensor(_slots(m), dtype=torch.int32, device=DEV) - return (buf[:m], False, m, 1, None, None) - - -def _prefill_state(ctx_lens: tuple, n_dec: int, has_init: tuple): - """`PrefillTarget.forward`'s `gdn_state` for a context prefix plus `n_dec` - one-token generation rows, built the way the forward builds it from the - engine's `Mamba2Metadata` buffers.""" - n_ctx = len(ctx_lens) - batch = n_ctx + n_dec - n_p = sum(ctx_lens) - si_buf = torch.zeros(MAX_BATCH, dtype=torch.int32, device=DEV) - si_buf[:batch] = torch.tensor(_slots(batch), dtype=torch.int32, device=DEV) - hi_buf = torch.zeros(MAX_BATCH, dtype=torch.bool, device=DEV) - hi_buf[:n_ctx] = torch.tensor(has_init, dtype=torch.bool, device=DEV) - si_all, hi_all = si_buf[:batch], hi_buf[:batch] - si_p, si_d = torch.split(si_all, [n_ctx, n_dec], 0) - si_d = torch.clone(si_d) - hi_p = hi_all[:n_ctx] - bounds = [0] - for n in ctx_lens: - bounds.append(bounds[-1] + n) - bounds += [n_p + j for j in range(1, n_dec + 1)] - cu_buf = torch.zeros(MAX_BATCH + 1, dtype=torch.int32, device=DEV) - cu_buf[: batch + 1] = torch.tensor(bounds, dtype=torch.int32, device=DEV) - cu_long_buf = cu_buf.to(torch.int64) - qsl_all = cu_buf[: batch + 1] - qsl_p = qsl_all[: n_ctx + 1] - qsl_long = cu_long_buf[: batch + 1] - qsl_long_p = qsl_long[: n_ctx + 1] - return ( - n_p, - n_dec, - si_p, - si_all, - si_d, - hi_p, - hi_all, - qsl_p, - qsl_all, - qsl_long, - False, - n_dec, - 1, - None, - qsl_long_p, - None, - ) - - -def _post_conv(prefill, decode, a, b, L): - return fused_gdn_post_conv( - prefill=prefill, - decode=decode, - a=a, - b=b, - A_log=L.a_log, - dt_bias=L.dt_bias, - num_k_heads=H, - head_k_dim=K, - num_v_heads=HV, - head_v_dim=V, - l2_norm_eps=M._L2_NORM_EPS, - softplus_threshold=M._SOFTPLUS_THRESHOLD, - ) - - -def _triton_decode_block(core, x, L, si_all): - """The decode block with the Triton single-token recurrence in its rows form (q/k/v - `[m, 1, heads, 128]` from a contiguous copy, `cu_seqlens` None), which stays on the - Triton kernel at either pool dtype.""" - m = x.shape[0] - mixed, z, a, b = torch.split(M.cublas_mm(x, L.in_proj), [CONV, VD, HV, HV], 1) - M.causal_conv1d_update( - x=mixed, - conv_state=core._bound.conv[0], - weight=L.conv_w, - bias=None, - silu_activation=True, - cache_seqlens=None, - conv_state_indices=si_all, - pad_slot_id=M._PAD_SLOT_ID, - ) - qv, kv, vv = torch.split(torch.reshape(mixed.contiguous(), [m, 1, -1, K]), [H, H, HV], 2) - out = triton_decode( - A_log=L.a_log, - a=torch.reshape(a.contiguous(), [m, 1, HV]), - dt_bias=L.dt_bias, - softplus_beta=M._SOFTPLUS_BETA, - softplus_threshold=M._SOFTPLUS_THRESHOLD, - q=qv, - k=kv, - v=vv, - b=torch.reshape(b.contiguous(), [m, 1, HV]), - initial_state_source=core._bound.ssm[0], - initial_state_indices=si_all, - scale=None, - use_qk_l2norm_in_kernel=True, - cu_seqlens=None, - output=None, - ) - return M.DecodeTarget._gdn_epilogue(_target(M.DecodeTarget, core), out, z, L, m) - - -def _triton_prefill_block(core, x, L, st): - """The prefill block with every row -- context and generation -- in one - Triton chunk call.""" - n_p, n_d, si_p, si_all, si_d, hi_p, hi_all, qsl_p, qsl_all, qsl_long = st[:10] - T = x.shape[0] - conv, ssm = core._bound.conv[0], core._bound.ssm[0] - mixed, z, a, b = torch.split(M.cublas_mm(x, L.in_proj), [CONV, VD, HV, HV], 1) - reset_recurrent_state_rows(ssm, si_p, hi_p, conv) - conv_args = dict( - weight=L.conv_w, - bias=None, - conv_states=conv, - silu_activation=True, - pad_slot_id=M._PAD_SLOT_ID, - out=None, - ) - if n_d > 0: - mq_p, mq_d = torch.split(mixed, [n_p, n_d], 0) - mq_p_t = extract_transpose_prefill_slice(mq_p, n_p, 0, CONV) - causal_conv1d_fwd( - x=mq_p_t, query_start_loc=qsl_p, cache_indices=si_p, has_initial_state=hi_p, **conv_args - ) - M.causal_conv1d_update( - x=mq_d, - conv_state=conv, - weight=L.conv_w, - bias=None, - silu_activation=True, - cache_seqlens=None, - conv_state_indices=si_d, - pad_slot_id=M._PAD_SLOT_ID, - ) - q, k, v, g, beta = _post_conv(mq_p_t, mq_d, a, b, L) - else: - mq_t = extract_transpose_prefill_slice(mixed, T, 0, CONV) - causal_conv1d_fwd( - x=mq_t, - query_start_loc=qsl_all, - cache_indices=si_all, - has_initial_state=hi_all, - **conv_args, - ) - q, k, v, g, beta = _post_conv(mq_t, None, a, b, L) - out, _ = triton_chunk( - q, - k, - v, - g, - beta, - scale=None, - initial_state=ssm, - initial_state_indices=si_all, - inplace_indexed_state_update=True, - output_final_state=False, - cu_seqlens=qsl_long, - use_qk_l2norm_in_kernel=False, - output=None, - ) - return M.PrefillTarget._gdn_epilogue(_target(M.PrefillTarget, core), out, z, L, T) - - -def _check(y, y_ref, core, core_ref, what): - """Output and both pools against the reference block's, at the module bands.""" - assert y.shape == y_ref.shape and y.dtype == y_ref.dtype - assert torch.isfinite(y).all() - d = y.float() - y_ref.float() - rel = (d.norm() / y_ref.float().norm()).item() - peak = (d.abs().max() / y_ref.float().abs().max()).item() - s_d = core._bound.ssm[0] - core_ref._bound.ssm[0] - s_rel = (s_d.norm() / core_ref._bound.ssm[0].norm()).item() - print( - f"{what}: out rel-L2 {rel:.3e}, out max-abs/peak {peak:.3e}, " - f"state rel-L2 {s_rel:.3e}, state max-abs {s_d.abs().max().item():.3e}" - ) - assert rel <= OUT_REL_L2 and peak <= OUT_MAX_FRAC, (rel, peak) - assert torch.equal(core._bound.conv[0], core_ref._bound.conv[0]), "conv pools differ" - torch.testing.assert_close( - core._bound.ssm[0], core_ref._bound.ssm[0], rtol=RTOL_STATE, atol=ATOL_STATE - ) - - -@pytest.fixture(scope="module") -def warm_triton_chunk(): - """The Triton chunk op's first indexed in-place call per autotune key - returns a wrong answer (its entry's note); spend it on scratch.""" - core = _core(0) - L = _layer(0) - st = _prefill_state((64,), 0, (False,)) - _triton_prefill_block(core, _x(64, 0), L, st) - torch.cuda.synchronize() - - -#: The recurrent pool dtypes the target serves (`mamba_ssm_cache_dtype`). -STATES = {"fp32": torch.float32, "bf16": torch.bfloat16} - - -@pytest.mark.parametrize("state", STATES) -def test_decode_block(state) -> None: - """A con-128 decode step; the graph-replay test below runs 32 rows.""" - m = 128 - core, L, x = _core(100 + m, STATES[state]), _layer(200 + m), _x(m, 300 + m) - core_ref = _clone_core(core) - st = _decode_state(m) - y = _target(M.DecodeTarget, core)._gdn_block(x, 0, L, st) - y_ref = _triton_decode_block(core_ref, x, L, st[0]) - _check(y, y_ref, core, core_ref, f"decode m={m}, {state} pool") - named = set(_slots(m)) - unnamed = [s for s in range(SLOTS) if s not in named] - original = _core(100 + m, STATES[state]) - assert torch.equal(core._bound.ssm[0][unnamed], original._bound.ssm[0][unnamed]), ( - "unnamed slots written" - ) - - -@pytest.mark.parametrize("state", STATES) -def test_prefill_block(state, warm_triton_chunk) -> None: # noqa: F811 — fixture by name - """A mixed step: three variable-length contexts (two carrying state) plus five - one-token generation rows.""" - ctx_lens, n_dec, has_init = (200, 37, 91), 5, (True, False, True) - tokens = sum(ctx_lens) + n_dec - core, L, x = _core(400 + tokens, STATES[state]), _layer(500 + tokens), _x(tokens, 600 + tokens) - core_ref = _clone_core(core) - st = _prefill_state(ctx_lens, n_dec, has_init) - y = _prefill_block(core, x, L, st) - y_ref = _triton_prefill_block(core_ref, x, L, st) - _check(y, y_ref, core, core_ref, f"prefill {ctx_lens}+{n_dec}x1, {state} pool") - - -@pytest.mark.parametrize("state", STATES) -def test_decode_block_replays_under_cuda_graph(state) -> None: - """The decode block captured into a CUDA graph after one eager call (the - eager call compiles the decode kernel and sizes its per-batch scratch, as - the engine's warmup does) replays to exactly the eager result.""" - m = 32 - core0, L, x = _core(700, STATES[state]), _layer(701), _x(m, 702) - st = _decode_state(m) - eager_core = _clone_core(core0) - y_eager = _target(M.DecodeTarget, eager_core)._gdn_block(x, 0, L, st).clone() - graph_core = _clone_core(core0) - target = _target(M.DecodeTarget, graph_core) - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - target._gdn_block(x, 0, L, st) - torch.cuda.current_stream().wait_stream(stream) - torch.cuda.synchronize() - graph_core._bound.conv[0].copy_(core0._bound.conv[0]) - graph_core._bound.ssm[0].copy_(core0._bound.ssm[0]) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - y_graph = target._gdn_block(x, 0, L, st) - assert torch.equal(graph_core._bound.ssm[0], core0._bound.ssm[0]), "capture must not execute" - graph.replay() - torch.cuda.synchronize() - assert torch.equal(y_graph, y_eager) - assert torch.equal(graph_core._bound.ssm[0], eager_core._bound.ssm[0]) - assert torch.equal(graph_core._bound.conv[0], eager_core._bound.conv[0]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_cached_replay_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_cached_replay_cute.py deleted file mode 100644 index 78902d1b7580..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_cached_replay_cute.py +++ /dev/null @@ -1,364 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test: the CuTe DSL cached-replay kernel (`gdn_cached_replay_cute`) over consecutive -verify steps against the Triton entry's fp32 `reference` and the Triton kernel itself, driven by -that entry's test harness (`_Bench`: work items, all-layer commit, `_advance_replay_state` with -random partial acceptance), in both commit modes: (a) no in-kernel commit + the manager's -all-layer commit when partitioned, (b) every wrapping row committed in-kernel. Plus rows with a -negative slot and CUDA-graph replay (PDL on and off) bitwise against eager.""" - -import importlib.util -import os - -import pytest -import torch - -import tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_cached_replay_cute as cr # noqa: E501 - -_spec = importlib.util.spec_from_file_location( - "trep", - os.path.join( - os.path.dirname(os.path.abspath(__file__)), - "test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py", - ), -) -trep = importlib.util.module_from_spec(_spec) -_spec.loader.exec_module(trep) - -assert torch.cuda.is_available(), "gdn_cached_replay_cute requires a CUDA device" - -BF = torch.bfloat16 -ULP = trep.BF16_ULP - - -def blk_err(got, want): - """max |got - want| / absmax of its block (last two dims), in bf16 ulps; relative RMS.""" - block = got.shape[-2] * got.shape[-1] - g, w = got.float().reshape(-1, block), want.float().reshape(-1, block) - keep = ~torch.isnan(w).any(1) - g, w = g[keep], w[keep] - if g.numel() == 0: - return 0.0, 0.0 - scale = w.abs().amax(1, keepdim=True).clamp_min(1e-30) - worst = ((g - w).abs() / scale).nan_to_num(nan=float("inf")).max().item() / ULP - rms = ((g - w).pow(2).mean().sqrt() / w.pow(2).mean().sqrt().clamp_min(1e-30)).item() / ULP - return worst, rms - - -class Stat: - def __init__(self): - self.d = {} - - def add(self, name, val): - a, b = val if isinstance(val, tuple) else (val, 0.0) - o = self.d.get(name, (0.0, 0.0)) - self.d[name] = (max(o[0], a), max(o[1], b)) - - def fmt(self): - return " ".join(f"{k}={v[0]:.2f}/{v[1]:.2f}" for k, v in self.d.items()) - - -def run_cell(name, spec, mode, config=None, seed=1234): - """mode 'a': mine without commit (+ the manager's all-layer commit when partitioned); - mode 'b': mine commits wrapping rows in-kernel (no all-layer commit for mine).""" - spec = dict(spec, state_dtype=BF) - bt = trep._Bench(spec, seed, exact=False) - bm = trep._Bench(spec, seed, exact=False) - dup = tuple(spec.get("dup", ())) - n = bt.n - unique = n - len(dup) - others = torch.ones(bm.slots, dtype=torch.bool, device="cuda") - if dup: - others[int(bm.state_indices[unique])] = False - vs_ref, vs_tri, tri_ref = Stat(), Stat(), Stat() - fails = [] - wraps = 0 - for step in range(bt.steps): - _, wrapped, _ = bm.writes() - wraps += wrapped.numel() - bt.prepare() - bm.prepare() - wi_all = nw_all = None - if bm.path == "all_layer": - # the engine's work items cover every decode row (padding included) - wi_all = torch.zeros(bm.n, 4, dtype=torch.int32, device="cuda") - nw_all = torch.zeros(1, dtype=torch.int32, device="cuda") - trep._build_replay_work_items_triton( - bm.state_indices, bm.pnat, bm.buf, wi_all, nw_all, bm.t, trep.HIST - ) - for layer in range(bt.layers): - at, am = bt.inputs(), bm.inputs() - kt, km = bt.kwargs(layer, at), bm.kwargs(layer, am) - rkw = dict(km) - if mode == "b": - rkw.update(replay_work_items=None, n_writes=None, use_all_layer_commit=False) - if dup: - rkw.update( - q=km["q"][:unique], - k=km["k"][:unique], - v=km["v"][:unique], - g=km["g"][:unique], - beta=km["beta"][:unique], - state_indices=km["state_indices"][:unique], - packed_qkv=am["packed"][: unique * bm.t], - ) - ref = trep.op.reference(**rkw) - tkw = dict(kt) - if mode == "b": - tkw.update(replay_work_items=None, n_writes=None, use_all_layer_commit=False) - rt = None if dup else trep.op.reference(**tkw) - got_t = trep.op(**tkw) - if rt is not None: - tri_ref.add("out", blk_err(got_t, rt[0])) - tri_ref.add("u", blk_err(bt.state(layer)[1], rt[2])[0]) - tri_ref.add("state", blk_err(bt.state(layer)[0], rt[1])[0]) - kmm = dict(km) - if wi_all is not None: - kmm.update(replay_work_items=wi_all, n_writes=nw_all) - got_m = cr.cached_replay_update(**kmm, in_kernel_commit=(mode == "b"), config=config) - pool, u, k, G = bm.state(layer) - gm = (got_m[:unique], pool[others], u[others], k[others], G[others]) - rf = (ref[0], *(x[others] for x in ref[1:])) - try: - trep._compare(gm, rf) - except AssertionError as e: # record, keep going - fails.append(f"step {step} layer {layer}: {str(e)[:200]}") - vs_ref.add("out", blk_err(gm[0], rf[0])) - vs_ref.add("state", blk_err(gm[1], rf[1])[0]) - vs_ref.add("u", blk_err(gm[2], rf[2])[0]) - vs_ref.add("k", blk_err(gm[3], rf[3])[0]) - vs_ref.add("G", (gm[4] - rf[4]).abs().max().item()) - tp, tu, tk, tG = bt.state(layer) - vs_tri.add("out", blk_err(got_m[:unique], got_t[:unique])) - vs_tri.add("u", blk_err(u[others], tu[others])[0]) - vs_tri.add("k", blk_err(k[others], tk[others])[0]) - vs_tri.add("G", (G[others] - tG[others]).abs().max().item()) - if mode == "a": - bt.commit() - bm.commit() - for layer in range(bt.layers): - vs_tri.add("state", blk_err(bm.state(layer)[0][others], bt.state(layer)[0][others])[0]) - acc = bt.draw_acceptance() - bm.draw_acceptance() - bt.accept(acc) - bm.accept(acc) - torch.cuda.synchronize() - ok = not fails - print( - f"[{'PASS' if ok else 'FAIL'}] {name} mode={mode} cfg={config} wraps={wraps}\n" - f" vs reference (max ulp of block absmax / rel-RMS ulp; G abs): {vs_ref.fmt()}\n" - f" vs Triton: {vs_tri.fmt()}\n" - f" (Triton vs reference: {tri_ref.fmt()})", - flush=True, - ) - for f in fails[:4]: - print(" ", f) - return ok - - -def run_negative(config=None, n=8, t=4, seed=99): - """Rows with a negative slot: nothing read or written for them, output rows untouched.""" - spec = dict(path="small", n=n, t=t, steps=1, layers=1, state_dtype=BF) - bm = trep._Bench(spec, seed, exact=False) - sidx = bm.state_indices.clone() - dead = [1, n - 2] - sidx[dead] = -1 - ok = True - for mode in ("b", "a"): - act = bm.inputs() - kw = bm.kwargs(0, act) - sentinel = torch.full((n, t, trep.HV, trep.V), 3.0, dtype=BF, device="cuda") - kw.update(state_indices=sidx, output=sentinel.clone()) - if mode == "a": - # every row listed (row, slot, pending count, half), dead rows with their negative slot - sl = sidx.long() - safe = sl.clamp(min=0) - wi = ( - torch.stack( - [torch.arange(n, device="cuda"), sl, bm.pnat[safe].long(), bm.buf[safe].long()], - 1, - ) - .to(torch.int32) - .flip(0) - .contiguous() - ) - nw = torch.zeros(1, dtype=torch.int32, device="cuda") - kw.update(replay_work_items=wi, n_writes=nw, use_all_layer_commit=True) - ref = trep.op.reference(**kw) - got = cr.cached_replay_update(**kw, config=config) - try: - trep._compare((got, *bm.state(0)), ref) - assert torch.equal(got[dead], sentinel[dead]), "dead rows written" - except AssertionError as e: - ok = False - print(" negative-slot FAIL", mode, str(e)[:200]) - print(f"[{'PASS' if ok else 'FAIL'}] negative slots cfg={config}", flush=True) - return ok - - -def run_graph(config=None, n=32, t=4, mode="b", pdl=False): - """Captured once, replayed on new activations with acceptance in between: bitwise the eager - calls on a twin state (the op test's CUDA-graph check, bf16 state, both commit modes).""" - spec = dict(path="all_layer", n=n, t=t, layers=2, slots=40, state_dtype=BF) - seed = 7000 + 10 * n + t - eager, graph, warm = (trep._Bench(spec, seed, exact=False) for _ in range(3)) - acts = [warm.inputs() for _ in range(2)] - - def step(b): - b.prepare() - outs = [] - for layer in range(2): - kw = b.kwargs(layer, acts[layer]) - kw.update(launch_with_pdl=pdl) - outs.append( - cr.cached_replay_update(**kw, in_kernel_commit=(mode == "b"), config=config) - ) - if mode == "a": - b.commit() - return outs - - step(warm) - torch.cuda.synchronize() - captured = torch.cuda.CUDAGraph() - s_ = torch.cuda.Stream() - s_.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(s_): - step(warm) - torch.cuda.current_stream().wait_stream(s_) - torch.cuda.synchronize() - with torch.cuda.graph(captured): - outs_graph = step(graph) - refill = torch.Generator(device="cuda").manual_seed(seed + 1) - ok = True - for _ in range(3): - for act in acts: - for key in ("packed", "a", "b"): - act[key].copy_( - torch.randn( - act[key].shape, generator=refill, device="cuda", dtype=act[key].dtype - ) - ) - outs_eager = step(eager) - captured.replay() - torch.cuda.synchronize() - for e, g in zip(outs_eager, outs_graph): - ok &= torch.equal(e, g) - for name in ("backing", "old_x", "old_B", "old_dt", "pnat", "buf"): - ok &= torch.equal(getattr(eager, name), getattr(graph, name)) - accepted = eager.draw_acceptance() - eager.accept(accepted) - graph.accept(accepted) - print( - f"[{'PASS' if ok else 'FAIL'}] CUDA graph replay bitwise eager mode={mode} pdl={pdl} cfg={config}", - flush=True, - ) - return ok - - -CELLS = { - "small_n4_t8_full": dict(path="small", n=4, t=8, start="fresh", accept="full", steps=8), - "small_n8_t3_dup": dict(path="small", n=8, t=3, dup=(5, 6, 7), steps=6), - "all_n32_t8_full": dict(path="all_layer", n=32, t=8, start="fresh", accept="full", steps=8), - "all_n16_t3_L8_S461": dict( - path="all_layer", n=16, t=3, layers=8, slots=461, edge_slots=True, steps=3 - ), - "t4_n1": dict(path="small", n=1, t=4, steps=6), - "t4_n4": dict(path="small", n=4, t=4, steps=6), - "t4_n16": dict(path="all_layer", n=16, t=4, steps=5), - "t4_n64": dict(path="all_layer", n=64, t=4, steps=4, slots=96), - "t4_n256": dict(path="all_layer", n=256, t=4, steps=3, slots=300, layers=1), - "t2_n8": dict(path="small", n=8, t=2, steps=6), - "t5_n8": dict(path="small", n=8, t=5, steps=6), - "t4_n32_dup": dict(path="all_layer", n=32, t=4, dup=(28, 29, 30, 31), steps=4), -} - - -_RUNS = [ - (name, mode) - for name, spec in CELLS.items() - for mode in ("a", "b") - if not (mode == "a" and spec["path"] == "small") -] - - -@pytest.mark.parametrize("name,mode", _RUNS, ids=[f"{n}-{m}" for n, m in _RUNS]) -def test_steps_match_the_reference(name, mode) -> None: - assert run_cell(name, CELLS[name], mode) - - -@pytest.mark.parametrize( - "cell", - cr.gdn_cached_replay_cute.CELLS, - ids=[c.why[:40] for c in cr.gdn_cached_replay_cute.CELLS], -) -def test_certified_cells(cell) -> None: - # The mode the target uses at the cell's size: in-kernel commit on the small path, the - # manager's all-layer commit on the partitioned one. - assert run_cell(cell.why, cell.spec, "b" if cell.spec["path"] == "small" else "a") - - -def test_negative_slots() -> None: - assert run_negative() - - -@pytest.mark.parametrize("pdl", [False, True]) -@pytest.mark.parametrize("mode", ["a", "b"]) -def test_graph_replay_matches_eager(mode, pdl) -> None: - assert run_graph(mode=mode, pdl=pdl) - - -#: Launch mappings a batch can take, and the mapping each is checked against: the state tiles -#: all by bulk copy with the state dots before the first barrier (the small-batch default), -#: against the first tile in registers and the dots after it. -_MAPPINGS = ( - (dict(tpc=2, warps=2, pf=0, minb=4, early_acc=True), dict(tpc=2, warps=2, pf=1, minb=4)), - (dict(tpc=4, warps=4, pf=0, minb=4, early_acc=True), dict(tpc=4, warps=4, pf=1, minb=4)), - (dict(tpc=8, warps=4, pf=1, minb=4, early_acc=True), dict(tpc=8, warps=4, pf=1, minb=4)), - (dict(tpc=8, warps=4, pf=0, minb=4), dict(tpc=24, warps=8, pf=1, minb=2)), -) - - -@pytest.mark.parametrize("n,t,pdl", [(1, 4, True), (5, 3, False), (32, 4, True), (8, 8, True)]) -def test_launch_mappings_agree(n, t, pdl) -> None: - """Every launch mapping computes the same instructions per state row. - - One step bitwise equal across mappings (output, checkpoint with in-kernel commits, u / k / - G history), with wrapping and appending rows and a negative slot, with and without the - engine's work items. - """ - spec = dict(path="small", n=n, t=t, steps=1, layers=1, state_dtype=BF) - bm = trep._Bench(spec, 4242 + n, exact=False) - sidx = bm.state_indices.clone() - if n > 2: - sidx[n // 2] = -1 - act = bm.inputs() - snap = [x.clone() for x in bm.state(0)] - # every row listed (row, slot, pending count, half), reversed, the dead row with its -1 - sl = sidx.long() - safe = sl.clamp(min=0) - wi = ( - torch.stack( - [torch.arange(n, device="cuda"), sl, bm.pnat[safe].long(), bm.buf[safe].long()], 1 - ) - .to(torch.int32) - .flip(0) - .contiguous() - ) - nw = torch.zeros(1, dtype=torch.int32, device="cuda") - for new_cfg, old_cfg in _MAPPINGS: - for with_wi in (False, True): - got = [] - for cfg in (new_cfg, old_cfg): - for x, y in zip(bm.state(0), snap): - x.copy_(y) - kw = bm.kwargs(0, act) - # a dead row's output is never written: a sentinel, the same for every call - sentinel = torch.full((n, t, trep.HV, trep.V), 3.0, dtype=BF, device="cuda") - kw.update(state_indices=sidx, launch_with_pdl=pdl, output=sentinel) - if with_wi: - kw.update(replay_work_items=wi, n_writes=nw, use_all_layer_commit=False) - out = cr.cached_replay_update(**kw, in_kernel_commit=True, config=cfg) - torch.cuda.synchronize() - got.append([out.clone()] + [x.clone() for x in bm.state(0)]) - for name, x, y in zip(("out", "state", "u", "k", "G"), *got): - assert torch.equal(x, y), f"{name}: {new_cfg} vs {old_cfg} (work items {with_wi})" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py deleted file mode 100644 index 24ebe8c685a8..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py +++ /dev/null @@ -1,122 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the gdn_decode_gated_norm_pdl entry. - -Cells against the composition's fp64 reference (the gated delta rule's recurrence, -then the gated norm entry's) in torch's default band for the bf16 output and the fp32 -pool, the pool compared whole so slots the call did not name are -gated too. Beyond the cells: a stale-read check under graph replay (a producer -rewrites the projection inside the graph right before the op), and the row guard. -""" - -import zlib - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_decode_gated_norm_pdl import ( # noqa: E501 - gdn_decode_gated_norm_pdl as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "gdn_decode_gated_norm_pdl requires a CUDA device" - -DEV = "cuda" -H, HV, K, V = 16, 48, 128, 128 -SLOTS, EPS = 461, 1e-6 -WIDTH = 2 * H * K + 2 * HV * V + 2 * HV - - -def _scattered(n: int, slots: int) -> tuple[int, ...]: - """`n` distinct slots of a `slots`-row pool, both ends first, then a stride-37 walk - over the interior (37 is coprime to every pool size here).""" - rest = tuple((i * 37 + 4) % (slots - 2) + 1 for i in range(max(n - 2, 0))) - return ((slots - 1, 0) + rest)[:n] - - -def _build(spec, seed): - g = torch.Generator(device=DEV).manual_seed(seed) - rows = spec["rows"] - proj = torch.randn(rows, WIDTH, generator=g, device=DEV).to(torch.bfloat16) - if spec.get("pool") == "coalesced": - # The KV-cache-v2 manager's view: every GDN layer's state of a slot in one page, - # this layer's state one sub-page of it, so the slot stride is 36 layers' states - # and slot 100 sits ~2.8G elements past the base (past int32). - layers, slots = 36, 102 - page = torch.empty(slots * layers * HV * V * K, device=DEV) - page[: HV * V * K * 4].normal_(0.0, 0.05, generator=g) - pool = page.as_strided((slots, HV, V, K), (layers * HV * V * K, V * K, K, 1)) - for s in (100, 101): - pool[s].copy_(torch.randn(HV, V, K, generator=g, device=DEV) * 0.05) - else: - pool = torch.randn(SLOTS, HV, V, K, generator=g, device=DEV) * 0.05 - if spec.get("pool") == "coalesced": - slots = torch.tensor([100, 101][:rows], dtype=torch.int32, device=DEV) - else: - slots = torch.tensor(_scattered(rows, SLOTS), dtype=torch.int32, device=DEV) - if spec.get("pad"): - slots[1] = -1 - A_log = torch.empty(HV, device=DEV).uniform_(-5.5, 1.0, generator=g) - dt_bias = torch.randn(HV, generator=g, device=DEV) * 4.0 - w = (1.0 + 0.2 * torch.randn(V, generator=g, device=DEV)).to(torch.bfloat16) - return [proj, pool, slots, A_log, dt_bias, w, EPS, H, K] - - -def _compare(out, pool, ref_out, ref_pool) -> None: - assert out.dtype == torch.bfloat16 and out.shape == ref_out.shape and out.is_contiguous() - assert torch.isfinite(out).all() - torch.testing.assert_close(out, ref_out) - torch.testing.assert_close(pool, ref_pool) - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name - certifying(op) - args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) - ref_out, ref_pool = op.reference(*args) - out = op(*args) - _compare(out, args[1], ref_out, ref_pool) - if cell.spec.get("pad"): - assert (out[1] == 0).all(), "a pad row's output is zero" - - -def test_reads_after_the_wait_under_graph_replay() -> None: - """A producer rewrites the projection inside the graph right before the op; with PDL - a read of it before `griddepcontrol.wait` would see the previous replay's row. The - pool advances one step per replay, as in decode.""" - args = _build(dict(rows=1), seed=41) - src = args[0].clone() - - def step(): - args[0].copy_(src * 1.0) - return op(*args) - - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - pool0 = args[1].clone() - with torch.cuda.stream(warmup): - step() - torch.cuda.current_stream().wait_stream(warmup) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_out = step() - for trial in range(3): - src.copy_(_build(dict(rows=1), seed=100 + trial)[0]) - before = args[1].clone() - graph.replay() - torch.cuda.synchronize() - after = args[1].clone() - args[1].copy_(before) - want = op(src.clone(), *args[1:]) - assert torch.equal(graph_out, want), f"replay {trial} output" - assert torch.equal(after, args[1]), f"replay {trial} pool" - del pool0 - - -def test_guards() -> None: - """More rows than `MAX_ROWS` are refused.""" - args = _build(dict(rows=1), seed=5) - with validating(op), pytest.raises(AssertionError): - op(args[0].expand(5, -1).contiguous(), args[1], args[2].expand(5).contiguous(), *args[3:]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py deleted file mode 100644 index 206ca93d1d3a..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py +++ /dev/null @@ -1,252 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the gdn_decode_rows_gated_norm entry. - -Cells against the composition's fp64 reference (the gated delta rule's recurrence, -then the gated norm entry's) in torch's default band for the bf16 output and the pool's -dtype (fp32 or bf16), the pool compared whole so slots the call did not name are -gated too. Beyond the cells: a fixed row bitwise across batch sizes over a bf16 pool, a -stale-read check under graph replay (a producer rewrites the projection inside the graph -right before the op), the two launch variants (state rows by bulk copy into shared memory, -or loaded into registers) bitwise each other, and the guards. The `gdn_decode_rows` entry -(the recurrence alone) has its own cells below. -""" - -import zlib - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm import ( # noqa: E501 - gdn_decode_rows_gated_norm as rows_mod, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_decode_rows_gated_norm import ( # noqa: E501 - gdn_decode_rows as rows_op, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_decode_rows_gated_norm import ( # noqa: E501 - gdn_decode_rows_gated_norm as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "gdn_decode_rows_gated_norm requires a CUDA device" - -DEV = "cuda" -H, HV, K, V = 16, 48, 128, 128 -SLOTS, EPS = 461, 1e-6 -WIDTH = 2 * H * K + 2 * HV * V + 2 * HV - - -def _scattered(n: int, slots: int) -> tuple[int, ...]: - """`n` distinct slots of a `slots`-row pool, both ends first, then a stride-37 walk - over the interior (37 is coprime to every pool size here).""" - rest = tuple((i * 37 + 4) % (slots - 2) + 1 for i in range(max(n - 2, 0))) - return ((slots - 1, 0) + rest)[:n] - - -def _build(spec, seed): - g = torch.Generator(device=DEV).manual_seed(seed) - rows = spec["rows"] - state = torch.bfloat16 if spec.get("state") == "bf16" else torch.float32 - proj = torch.randn(rows, WIDTH, generator=g, device=DEV).to(torch.bfloat16) - if spec.get("pool") == "coalesced": - # The KV-cache-v2 manager's view: every GDN layer's state of a slot in one page, - # this layer's state one sub-page of it, so the slot stride is 36 layers' states - # and slot 100 sits ~2.8G elements past the base (past int32). - layers, slots = 36, 102 - page = torch.empty(slots * layers * HV * V * K, device=DEV, dtype=state) - page[: HV * V * K * 4].normal_(0.0, 0.05, generator=g) - pool = page.as_strided((slots, HV, V, K), (layers * HV * V * K, V * K, K, 1)) - for s in (100, 101): - pool[s].copy_(torch.randn(HV, V, K, generator=g, device=DEV) * 0.05) - else: - pool = (torch.randn(SLOTS, HV, V, K, generator=g, device=DEV) * 0.05).to(state) - if spec.get("pool") == "coalesced": - slots = torch.tensor([100, 101][:rows], dtype=torch.int32, device=DEV) - else: - slots = torch.tensor(_scattered(rows, SLOTS), dtype=torch.int32, device=DEV) - if spec.get("pad"): - slots[1] = -1 - A_log = torch.empty(HV, device=DEV).uniform_(-5.5, 1.0, generator=g) - dt_bias = torch.randn(HV, generator=g, device=DEV) * 4.0 - w = (1.0 + 0.2 * torch.randn(V, generator=g, device=DEV)).to(torch.bfloat16) - return [proj, pool, slots, A_log, dt_bias, w, EPS, H, K] - - -def _compare(out, pool, ref_out, ref_pool) -> None: - assert out.dtype == torch.bfloat16 and out.shape == ref_out.shape and out.is_contiguous() - assert torch.isfinite(out).all() - torch.testing.assert_close(out, ref_out) - torch.testing.assert_close(pool, ref_pool) - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name - certifying(op) - args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) - ref_out, ref_pool = op.reference(*args) - out = op(*args) - _compare(out, args[1], ref_out, ref_pool) - if cell.spec.get("pad"): - assert (out[1] == 0).all(), "a pad row's output is zero" - - -#: Decode batch sizes up to the engine's largest CUDA-graph batch. -_BATCHES = (1, 2, 4, 5, 16, 33, 128, 256) -_FIXED_SLOT = 301 - - -def test_bf16_pool_is_batch_invariant() -> None: - """The property the bf16 decode step relies on: a row's output and state are bitwise - independent of the batch it rides in. One fixed row (its projection and slot 301 of a - bf16 pool) sits first and last in random batches of every size.""" - gen = torch.Generator(device=DEV).manual_seed(3) - cpu = torch.Generator().manual_seed(3) - pool0 = (torch.randn(SLOTS, HV, V, K, generator=gen, device=DEV) * 0.05).to(torch.bfloat16) - row = torch.randn(WIDTH, generator=gen, device=DEV).to(torch.bfloat16) - rest = _build(dict(rows=1), seed=3)[3:] - first = None - for n in _BATCHES: - for pos in sorted({0, n - 1}): - g = torch.Generator(device=DEV).manual_seed(1000 * n + pos) - proj = torch.randn(n, WIDTH, generator=g, device=DEV).to(torch.bfloat16) - proj[pos] = row - others = [s for s in torch.randperm(SLOTS, generator=cpu).tolist() if s != _FIXED_SLOT] - others = others[:n] - others[pos] = _FIXED_SLOT - slots = torch.tensor(others, dtype=torch.int32, device=DEV) - pool = pool0.clone() - out = op(proj, pool, slots, *rest) - got = (out[pos].clone(), pool[_FIXED_SLOT].clone()) - first = first or got - assert torch.equal(got[0], first[0]) and torch.equal(got[1], first[1]), ( - f"the fixed row differs at batch {n}, position {pos}" - ) - - -def test_reads_after_the_wait_under_graph_replay() -> None: - """A producer rewrites the projection inside the graph right before the op; with PDL - a read of it before `griddepcontrol.wait` would see the previous replay's row. The - pool advances one step per replay, as in decode.""" - args = _build(dict(rows=16), seed=41) - src = args[0].clone() - - def step(): - args[0].copy_(src * 1.0) - return op(*args) - - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - step() - torch.cuda.current_stream().wait_stream(warmup) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_out = step() - for trial in range(3): - src.copy_(_build(dict(rows=16), seed=100 + trial)[0]) - before = args[1].clone() - graph.replay() - torch.cuda.synchronize() - after = args[1].clone() - args[1].copy_(before) - want = op(src.clone(), *args[1:]) - assert torch.equal(graph_out, want), f"replay {trial} output" - assert torch.equal(after, args[1]), f"replay {trial} pool" - - -_VARIANT_SPECS = ( - dict(rows=1, state="bf16"), - dict(rows=33, state="bf16"), - dict(rows=33, pad=True), - dict(rows=2, pool="coalesced", state="bf16", pad=True), -) - - -@pytest.mark.parametrize("norm", [True, False], ids=["norm", "rows"]) -@pytest.mark.parametrize("spec", _VARIANT_SPECS, ids=lambda s: "-".join(map(str, s.values()))) -def test_bulk_and_register_variants_agree(spec, norm) -> None: - """The bulk-copy and the register variants: outputs and state bitwise equal. - - The bulk-copy variant stages the state rows in shared memory; the register variant holds - them in registers; both compute the same instructions. Over contiguous bf16 and fp32 pools - and the cache manager's coalesced pool (slot 100, ~2.8G elements past the base, by the - bulk copy), with pad rows. - """ - args = _build(spec, seed=zlib.crc32(repr(spec).encode()) % (2**31)) - proj, pool, slots = args[0], args[1], args[2] - assert rows_mod._bulk_eligible(pool, K) - live = slots[slots >= 0].long() - before = pool[live].clone() - got = [] - for bulk in (False, True): - pool[live] = before - w = args[5] if norm else None - out = rows_mod.launch(proj, pool, slots, args[3], args[4], w, EPS, H, K, bulk=bulk) - got.append((out.clone(), pool[live].clone())) - assert torch.equal(got[0][0], got[1][0]), "outputs differ between the variants" - assert torch.equal(got[0][1], got[1][1]), "states differ between the variants" - - -def test_variant_choice_and_unaligned_pool() -> None: - """The variant choice, and the fallback for a pool the bulk copy cannot take. - - The bulk variant from the batch the register variant needs a second wave for; a pool 8- - but not 16-byte aligned stays on the register variant whatever is asked, and matches the - aligned pool's result. - """ - sms = torch.cuda.get_device_properties(DEV).multi_processor_count - edge = (2 * sms) // HV - assert not rows_mod._use_bulk(edge, HV, DEV) and rows_mod._use_bulk(edge + 1, HV, DEV) - args = _build(dict(rows=33, state="bf16"), seed=19) - pool = args[1] - flat = torch.empty(pool.numel() + 4, dtype=pool.dtype, device=DEV) - shifted = flat[4:].view(pool.shape) - shifted.copy_(pool) - assert not rows_mod._bulk_eligible(shifted, K) - want = rows_mod.launch(args[0], pool, *args[2:], bulk=True) - got = rows_mod.launch(args[0], shifted, *args[2:], bulk=True) - assert torch.equal(got, want) and torch.equal(shifted, pool) - - -def test_guards() -> None: - """An fp16 pool, int64 slot indices and a projection of the wrong width are refused.""" - args = _build(dict(rows=4), seed=5) - with validating(op), pytest.raises(AssertionError): - op(args[0], args[1].to(torch.float16), *args[2:]) - with validating(op), pytest.raises(AssertionError): - op(args[0], args[1], args[2].long(), *args[3:]) - with validating(op), pytest.raises(AssertionError): - op(args[0][:, :-2], *args[1:]) - - -# ---- gdn_decode_rows: the recurrence alone, into a row range of a step's output ---------- - - -def _rows_args(spec, seed): - """`(args, buf)`: the entry's arguments, `out` rows 5.. of a `[1, rows + 9, HV, V]` buffer - (the step's output, the call's rows a range of it), and that buffer.""" - proj, pool, slots, A_log, dt_bias, _, _, h, k = _build(spec, seed) - rows = proj.shape[0] - g = torch.Generator(device=DEV).manual_seed(seed + 1) - buf = torch.randn(1, rows + 9, HV, V, generator=g, device=DEV).to(torch.bfloat16) - out = buf[0, 5 : 5 + rows].view(rows, HV * V) - return [proj, pool, slots, A_log, dt_bias, h, k, out], buf - - -@pytest.mark.parametrize("cell", rows_op.CELLS, ids=[c.why[:44] for c in rows_op.CELLS]) -def test_rows_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name - certifying(rows_op) - args, buf = _rows_args(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) - outside = torch.cat([buf[0, :5], buf[0, 5 + args[0].shape[0] :]]).clone() - ref_o, ref_pool = rows_op.reference(*args) - got = rows_op(*args) - assert got.data_ptr() == args[-1].data_ptr(), "written into out" - torch.testing.assert_close(got, ref_o) - torch.testing.assert_close(args[1], ref_pool) - assert torch.equal(torch.cat([buf[0, :5], buf[0, 5 + args[0].shape[0] :]]), outside), ( - "rows outside the range untouched" - ) - if cell.spec.get("pad"): - assert (got[1] == 0).all(), "a pad row's output is zero" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py deleted file mode 100644 index e6a5f55bd730..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_in_proj_conv1d_pdl.py +++ /dev/null @@ -1,103 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the gdn_in_proj_conv1d_pdl entry. - -Cells against the entry's fp32 reference: the output within the band the -decode projections use (4 bf16 ulp element-wise against the row scale, 2 ulp -relative RMS -- the GEMV's summation order is the only difference), the pool -the same, a pad row's slot bitwise untouched. Beyond the cells: two decode -steps under CUDA-graph replay, the second reading the window the first one -wrote; and the row guard. -""" - -import zlib - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_in_proj_conv1d_pdl import ( # noqa: E501 - MAX_ROWS, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_in_proj_conv1d_pdl import ( # noqa: E501 - gdn_in_proj_conv1d_pdl as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import ( # noqa: E402,F401 — path above; fixture by name - ULP, - assert_within_ulp, - certifying, - validating, -) - -assert torch.cuda.is_available(), "gdn_in_proj_conv1d_pdl requires a CUDA device" - -N, K, CONV_DIM, PAD = 16480, 2560, 10240, -1 - - -def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: - assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0, ulp=ULP[torch.bfloat16]) - - -def _build(rows, pool, seed): - g = torch.Generator(device="cuda").manual_seed(seed) - x = torch.randn(rows, K, generator=g, device="cuda").to(torch.bfloat16) - weight = (torch.randn(N, K, generator=g, device="cuda") * K**-0.5).to(torch.bfloat16) - state = torch.randn(pool, CONV_DIM, 3, generator=g, device="cuda").to(torch.bfloat16) - taps = (torch.randn(CONV_DIM, 4, generator=g, device="cuda") * 0.3).to(torch.bfloat16) - return x, weight.t(), state, taps - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 - certifying(op) - s = cell.spec - x, mat_b, state, taps = _build(s["rows"], s["pool"], seed=zlib.crc32(cell.why.encode()) % 2**31) - slots = torch.tensor(s["slots"], dtype=torch.int32, device="cuda") - pool = state.clone() - out = op(x, mat_b, pool, taps, slots, CONV_DIM, PAD, s["pdl"]) - ref_out, ref_state = op.reference(x, mat_b, state, taps, slots, CONV_DIM, PAD, s["pdl"]) - assert out.shape == (s["rows"], N) and out.dtype == torch.bfloat16 and out.is_contiguous() - _compare(out, ref_out) - _compare(pool, ref_state) - touched = {v for v in s["slots"] if v != PAD} - for slot in range(s["pool"]): - if slot not in touched: - assert torch.equal(pool[slot], state[slot]), f"slot {slot} was written" - - -def test_two_steps_under_graph_replay() -> None: - """Decode runs this under CUDA graphs with PDL, layer after layer: capture one - step, replay it twice with new activations. The second replay's conv reads the - window the first replay advanced, so a stale or misplaced window write fails - the second comparison.""" - pdl = True - x, mat_b, state, taps = _build(1, 64, seed=3) - slots = torch.tensor([9], dtype=torch.int32, device="cuda") - pool = state.clone() - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - for _ in range(2): - op(x, mat_b, pool.clone(), taps, slots, CONV_DIM, PAD, pdl) - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - out = op(x, mat_b, pool, taps, slots, CONV_DIM, PAD, pdl) - want_state = state.clone() - for trial in range(2): - g = torch.Generator(device="cuda").manual_seed(50 + trial) - x.copy_(torch.randn(x.shape, generator=g, device="cuda").to(x.dtype)) - graph.replay() - torch.cuda.synchronize() - ref_out, want_state = op.reference(x, mat_b, want_state, taps, slots, CONV_DIM, PAD) - _compare(out, ref_out) - _compare(pool, want_state) - - -def test_guards() -> None: - """More rows than `MAX_ROWS` are refused.""" - x, mat_b, state, taps = _build(MAX_ROWS + 1, 16, seed=21) - slots = torch.zeros(MAX_ROWS + 1, dtype=torch.int32, device="cuda") - with validating(op), pytest.raises(AssertionError): - op(x, mat_b, state, taps, slots, CONV_DIM, PAD) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_replay_verify_cute.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_replay_verify_cute.py deleted file mode 100644 index 322d994709bd..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_replay_verify_cute.py +++ /dev/null @@ -1,128 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the target's cached-replay verify dispatch over a bf16 pool (the CuTe kernel). - -`_replay_verify` as the decode target calls it -- `a`/`b` as column views of the input -projection, PDL after the fused in_proj GEMV -- against the call it replaces (contiguous -`a`/`b`, no PDL), each on its own copy of the pool and the history, over consecutive verify -steps with random partial acceptance and `_advance_replay_state` in between: outputs, the -checkpoint (committed in-kernel on every wrap, as without the manager's all-layer commit) and -the history bitwise equal at every step. The partitioned replay (work items and the manager's -all-layer commit) is the CuTe entry's own test's. -""" - -from types import SimpleNamespace - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - modeling as M, -) -from tensorrt_llm._torch.modules.fla.cache_manager import ( - ReplayStateUpdateMetadata, - _advance_replay_state, -) - -assert torch.cuda.is_available(), "the GDN verify dispatch test requires a CUDA device" - -DEV = "cuda" -BF = torch.bfloat16 -H, HV, K, V = 16, 48, 128, 128 -KEY_DIM, VALUE_DIM = H * K, HV * V -CONV_DIM = 2 * KEY_DIM + VALUE_DIM -HIST = 16 -SLOTS = 40 - - -class _Side: - """One copy of a layer's replay state and the `core` stand-in `_replay_verify` reads. - - The state: the pool, the history and the bookkeeping (pending count, half). - """ - - def __init__(self, pool, old, pnat, buf, t, gating): - self.pool = pool.clone() - self.old = [x.clone() for x in old] - self.pnat, self.buf = pnat.clone(), buf.clone() - self.meta = ReplayStateUpdateMetadata(self.pnat, self.buf, t, HIST) - replay = M._Replay( - old={0: tuple(self.old)}, - cache_buf_idx=self.buf, - prev_num_accepted=self.pnat, - history=HIST, - width=t, - all_layer_commit=False, - cute=True, - ) - self.core = SimpleNamespace( - conv_dim=CONV_DIM, - key_dim=KEY_DIM, - value_dim=VALUE_DIM, - num_k_heads=H, - head_k_dim=K, - num_v_heads=HV, - head_v_dim=V, - _bound=SimpleNamespace(replay=replay), - ) - self.layer = SimpleNamespace(a_log=gating[0], dt_bias=gating[1]) - - def state(self): - return [self.pool, *self.old[:3]] - - -@pytest.mark.parametrize("n,t", [(1, 4), (4, 4), (32, 4), (5, 3)]) -def test_cute_dispatch_matches_the_plain_call(n, t): - gen = torch.Generator(device=DEV).manual_seed(77 + n) - cpu = torch.Generator().manual_seed(77 + n) - slots = torch.randperm(SLOTS, generator=cpu)[:n].to(DEV, torch.int32) - pool = (0.3 * torch.randn(SLOTS, HV, V, K, generator=gen, device=DEV)).to(BF) - old = ( - (0.5 * torch.randn(SLOTS, 2, HIST, HV, V, generator=gen, device=DEV)).to(BF), - torch.nn.functional.normalize( - torch.randn(SLOTS, 2, HIST, H, K, generator=gen, device=DEV), dim=-1 - ).to(BF), - -(0.2 * torch.rand(SLOTS, 2, HV, HIST, generator=gen, device=DEV)).cumsum(-1), - torch.zeros(SLOTS, 2, HV, HIST, device=DEV), - ) - pnat = torch.randint(0, HIST + 1, (SLOTS,), generator=cpu, dtype=torch.int32).to(DEV) - buf = torch.randint(0, 2, (SLOTS,), generator=cpu, dtype=torch.int32).to(DEV) - gating = ( - (torch.randn(HV, generator=gen, device=DEV) * 1.5 + 1.0).clamp(-7.0, 5.0), - (torch.randn(HV, generator=gen, device=DEV) * 1.7 - 4.4).clamp(-11.0, 2.5), - ) - plain = _Side(pool, old, pnat, buf, t, gating) - disp = _Side(pool, old, pnat, buf, t, gating) - wraps = 0 - for step in range(HIST // t + 3): - wraps += int((disp.pnat[slots.long()] + t > HIST).sum()) - replay = M._replay_work(disp.core, SimpleNamespace(), n) - assert replay == (None, None, False), "no all-layer commit: the kernel commits itself" - mq_out = (0.5 * torch.randn(n * t, CONV_DIM, generator=gen, device=DEV)).to(BF) - # z | a | b: column views of one projection output, as the decode target holds them - zab = torch.randn(n * t, VALUE_DIM + 2 * HV, generator=gen, device=DEV).to(BF) - _, a, b = torch.split(zab, [VALUE_DIM, HV, HV], 1) - want = M._replay_verify( - plain.core, - mq_out, - a.contiguous(), - b.contiguous(), - plain.pool, - 0, - plain.layer, - slots, - n, - t, - (None, None, False), - ) - got = M._replay_verify( - disp.core, mq_out, a, b, disp.pool, 0, disp.layer, slots, n, t, replay, pdl=True - ) - torch.cuda.synchronize() - assert torch.equal(got, want), f"step {step}: outputs differ" - for name, x, y in zip(("pool", "u", "k", "G"), disp.state(), plain.state()): - assert torch.equal(x, y), f"step {step}: {name} differs" - accepted = torch.randint(1, t + 1, (n,), generator=cpu, dtype=torch.int32).to(DEV) - for side in (plain, disp): - _advance_replay_state(side.meta, slots, accepted) - assert wraps > 0, "no history wrapped: the in-kernel commit was never exercised" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_verify_fused_conv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_verify_fused_conv.py deleted file mode 100644 index 8476e49feada..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_verify_fused_conv.py +++ /dev/null @@ -1,100 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test: `DecodeTarget._gdn_block`'s verify dispatch under the dense FP8 mode, and over the -in_proj's NVFP4 copy (`TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4`, up to `_INPROJ_NVFP4_MAX_ROWS`). - -A captured verify step (`fused_conv=True`) runs the in_proj with the multi-token conv step as its -epilogue (`fp8_gemv_conv1d_verify` / `nvfp4_split_gemv_conv1d_verify`); an eager one runs the -projection and then `causal_conv1d_triton_update`. Both must hand the recurrence the same thing: the conv output -(contiguous, as the cached replay reads it), `a`, `b` and `z` bitwise, and leave the conv pool -and the per-step snapshot buffer bitwise the same. The recurrence is stubbed to record its -inputs (it is the same call on both sides), the epilogue to record `z`. Over a bf16 pool (the -CuTe replay) `a` and `b` reach the recurrence as views of the projection (no copies), and only -the fused step asks for PDL (its GEMV signals its dependents; the Triton conv does not). -""" - -import types - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - modeling as M, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.fp8_gemv import ( # noqa: E501 - pack_from_bf16, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.gemm.nvfp4_gemv import ( # noqa: E501 - pack_split_from_bf16, -) - -assert torch.cuda.is_available(), "the Gated-DeltaNet block requires a CUDA device" - -DEV, BF16 = "cuda", torch.bfloat16 -HIDDEN, HV, V, CONV, VD, SLOTS, ROWS = 2560, 48, 128, 10240, 6144, 461, 32 - - -@pytest.fixture(scope="module", params=["fp8", "nvfp4"]) -def layer(request): - g = torch.Generator(device=DEV).manual_seed(0) - w = (torch.randn(CONV + VD + 2 * HV, HIDDEN, generator=g, device=DEV) * HIDDEN**-0.5).to(BF16) - nv = request.param == "nvfp4" - return types.SimpleNamespace( - in_proj=w.t(), - in_proj8=None, - in_proj_m=None, - in_proj_g=None if nv else pack_from_bf16(w), - in_proj_nv=pack_split_from_bf16(w) if nv else None, - conv_w=(torch.randn(CONV, 4, generator=g, device=DEV) * 0.5).to(BF16), - ) - - -def _block(layer, x, si, int_idx, conv0, win0, p, fused, monkeypatch): - """One verify `_gdn_block` on copies of the pools; returns what the recurrence and the - epilogue were handed, and the pools.""" - seen = {} - - def replay(core, mq_out, a_d, b_d, ssm_states, i, L, si_d, n, p_, work, pdl=False): - seen.update(mq_out=mq_out.clone(), a=a_d.clone(), b=b_d.clone(), pdl=pdl) - assert mq_out.is_contiguous() and mq_out.shape == (n * p_, CONV) - # the CuTe replay reads a / b through their row stride: views, not copies - assert a_d.stride(0) > HV and b_d.stride(0) > HV - return torch.zeros(1, n * p_, HV, V, dtype=BF16, device=DEV) - - monkeypatch.setattr(M, "_replay_verify", replay) - core = types.SimpleNamespace( - conv_dim=CONV, - in_proj_split=(CONV, VD, HV, HV), - _bound=types.SimpleNamespace( - conv={0: conv0.clone()}, - ssm={0: torch.zeros(1, device=DEV)}, - int_conv={0: win0.clone()}, - replay=types.SimpleNamespace(cute=True), - ), - ) - target = object.__new__(M.DecodeTarget) - target.core = core - target._gdn_epilogue = lambda out, z, L, num_tokens, gemm=None: z.clone() - n = si.shape[0] - gdn_state = (si, True, n, p, int_idx, (None, None, False)) - z = target._gdn_block(x, 0, layer, gdn_state, cublas_mm, fused) - return dict(seen, z=z, conv=core._bound.conv[0], win=core._bound.int_conv[0]) - - -@pytest.mark.parametrize("n,p", [(1, 4), (2, 4), (8, 2), (2, 8)]) -def test_fused_verify_matches_two_launch(layer, n, p, monkeypatch) -> None: - if layer.in_proj_nv is not None and n * p > M._INPROJ_NVFP4_MAX_ROWS: - pytest.skip("past the NVFP4 copy's rows") - g = torch.Generator(device=DEV).manual_seed(10 * n + p) - conv0 = torch.randn(SLOTS, CONV, 3, generator=g, device=DEV).to(BF16) - win0 = torch.randn(ROWS, p, CONV, 3, generator=g, device=DEV).to(BF16) - x = torch.randn(n * p, HIDDEN, generator=g, device=DEV).to(BF16) - si = torch.randperm(SLOTS, generator=g, device=DEV)[:n].to(torch.int32) - int_idx = torch.arange(ROWS, dtype=torch.int32, device=DEV)[:n] - eager = _block(layer, x, si, int_idx, conv0, win0, p, False, monkeypatch) - fused = _block(layer, x, si, int_idx, conv0, win0, p, True, monkeypatch) - for key in ("mq_out", "a", "b", "z", "conv", "win"): - assert torch.equal(eager[key].view(torch.int16), fused[key].view(torch.int16)), key - assert fused["pdl"] and not eager["pdl"], "PDL exactly after the fused GEMV" - assert not torch.equal(fused["conv"], conv0) and not torch.equal(fused["win"], win0) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_pack_gdn_decode_qkv.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_pack_gdn_decode_qkv.py deleted file mode 100644 index 4adc8c137f6f..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_pack_gdn_decode_qkv.py +++ /dev/null @@ -1,117 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the pack_gdn_decode_qkv entry. - -The op is a pure copy, so every comparison is bitwise. -""" - -import zlib - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.pack_gdn_decode_qkv import ( # noqa: E501 - CKPT_HEADS, - HEAD_K_DIM, - HEAD_V_DIM, - NUM_K_HEADS, - NUM_V_HEADS, - QKV_DIM, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.pack_gdn_decode_qkv import ( # noqa: E501 - pack_gdn_decode_qkv as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "pack_gdn_decode_qkv requires a CUDA device" - - -def _compare(out: tuple, ref: tuple) -> None: - """Bitwise, all three: the op moves bits and changes nothing.""" - assert len(out) == len(ref) == 3 - for name, got, want in zip("qkv", out, ref): - assert got.dtype == want.dtype and got.shape == want.shape, ( - name, - (got.dtype, got.shape), - (want.dtype, want.shape), - ) - assert torch.equal(got, want), f"{name} must be a bit-exact copy" - - -DEV = "cuda" -Q_WIDTH = NUM_K_HEADS * HEAD_K_DIM # 2048 -V_WIDTH = NUM_V_HEADS * HEAD_V_DIM # 6144 - - -def _build(spec, seed): - """`(mixed_qkv, conv_out or None)`: the bf16 input, and for a verify batch the - `[B, qkv_dim, P]` conv output it was flattened from.""" - g = torch.Generator(device=DEV).manual_seed(seed) - - def randn(*shape): - return torch.randn(*shape, generator=g, device=DEV).to(torch.bfloat16) - - if "draft" in spec: - b, p = spec["num_decodes"], spec["draft"] - # The target: the conv result carries its `[n, C, p]` input view's - # strides, so transposing back and flattening is a contiguous view. - conv_out = randn(b, p, QKV_DIM).transpose(1, 2) - mixed = conv_out.transpose(1, 2).reshape(b * p, QKV_DIM) - assert mixed.is_contiguous() and mixed.data_ptr() == conv_out.data_ptr() - return mixed, conv_out - return randn(spec["num_tokens"], QKV_DIM), None - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name - certifying(op) - spec = cell.spec - mixed, conv_out = _build(spec, seed=zlib.crc32(cell.why.encode())) - snapshot = mixed.clone() - out = op(mixed, *CKPT_HEADS) - _compare(out, op.reference(mixed, *CKPT_HEADS)) - - assert all(t.is_contiguous() for t in out) - assert torch.equal(mixed, snapshot), "mixed_qkv is read-only" - storages = {t.untyped_storage().data_ptr() for t in out} - assert len(storages) == 3 and mixed.untyped_storage().data_ptr() not in storages - - if "draft" in spec: - # The verify caller regroups each return as [B, P, heads, dim]: that is - # a view of the same storage, and it lands on (sequence, draft) exactly. - b, p = spec["num_decodes"], spec["draft"] - grouped = ( - out[0].reshape(b, p, NUM_K_HEADS, HEAD_K_DIM), - out[1].reshape(b, p, NUM_K_HEADS, HEAD_K_DIM), - out[2].reshape(b, p, NUM_V_HEADS, HEAD_V_DIM), - ) - assert all(gt.data_ptr() == t.data_ptr() for gt, t in zip(grouped, out)) - per_token = conv_out.permute(0, 2, 1) # [B, P, qkv_dim], straight off the conv - want = ( - per_token.narrow(2, 0, Q_WIDTH).reshape(b, p, NUM_K_HEADS, HEAD_K_DIM), - per_token.narrow(2, Q_WIDTH, Q_WIDTH).reshape(b, p, NUM_K_HEADS, HEAD_K_DIM), - per_token.narrow(2, 2 * Q_WIDTH, V_WIDTH).reshape(b, p, NUM_V_HEADS, HEAD_V_DIM), - ) - assert all(torch.equal(gt, w) for gt, w in zip(grouped, want)) - - -def test_a_short_row_reads_the_next_row_and_is_refused() -> None: - """Unguarded, a `[32, 9216]` input at qkv_dim 10240 runs and silently reads - the next row's first 1024 columns into the tail of `v`; the guard refuses it. - - The input is carved from a longer flat buffer so the over-read lands in - known memory rather than in allocator garbage. - """ - rows, width = 32, QKV_DIM - 1024 - flat = torch.randn(rows * width + 1024, device=DEV).to(torch.bfloat16) - mixed = flat[: rows * width].view(rows, width) - _, _, v = op(mixed, *CKPT_HEADS) - torch.cuda.synchronize() - tail = v[0].reshape(rows, V_WIDTH)[:, V_WIDTH - 1024 :] - assert torch.equal(tail[:-1], mixed[1:, :1024]), "the hazard this guard covers is gone" - - with pytest.raises(AssertionError, match="must be at least qkv_dim=10240 wide"): - with validating(op): - op(mixed, *CKPT_HEADS) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_qwen4_exp_causal_conv1d_fwd.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_qwen4_exp_causal_conv1d_fwd.py deleted file mode 100644 index 5ca8d32dfed5..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_qwen4_exp_causal_conv1d_fwd.py +++ /dev/null @@ -1,131 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the causal_conv1d_fwd entry.""" - -import zlib -from typing import Optional - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_fwd import ( # noqa: E501 - ENGINE_SLOTS, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_fwd import ( # noqa: E501 - causal_conv1d_fwd as op, -) - -# Cells sized at the engine's pool get a non-vacuity check below. -assert ENGINE_SLOTS == 461 - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "causal_conv1d_fwd requires a CUDA device" - - -def _compare( - out: tuple[torch.Tensor, Optional[torch.Tensor]], - ref: tuple[torch.Tensor, Optional[torch.Tensor]], -) -> None: - """`(activation, conv_states)` pairs: the returned activation with the - pool as the call left it, against `reference`'s pair. - - Activation at torch's default band: the kernel accumulates the taps in - fp32 and rounds once, so it sits within one ulp of the fp32 reference - (checkpoint geometry: max abs 0.015625, bitwise fraction 0.999996; fp16 - one ulp; fp32 1.4e-6). The state is a copy of the input tail, so it is - gated bitwise over the whole pool -- which is also the lock on every - slot the batch did not name. - """ - (act, states), (ref_act, ref_states) = out, ref - assert act.dtype == ref_act.dtype and act.shape == ref_act.shape - torch.testing.assert_close(act, ref_act) - assert (states is None) == (ref_states is None) - if states is not None: - assert torch.equal(states, ref_states), "the published state must be bitwise exact" - - -PAD_SLOT_ID = -1 - - -def _seed(text: str) -> int: - return zlib.crc32(text.encode()) - - -def _build(spec, seed): - """`__call__`'s positional arguments for one spec, freshly allocated.""" - g = torch.Generator(device="cuda").manual_seed(seed) - dim, width, dtype = spec["dim"], 4, torch.bfloat16 - - def randn(*shape): - return torch.randn(*shape, generator=g, device="cuda").to(dtype) - - lengths = spec["lengths"] - starts = [0] - for length in lengths: - starts.append(starts[-1] + length) - qsl = torch.tensor(starts, dtype=torch.int32, device="cuda") - x = randn(sum(lengths), dim).t().contiguous() - weight = randn(dim, width) - pool = randn(spec["pool_slots"], dim, width - 1) - cache_indices = torch.tensor(spec["slots"], dtype=torch.int32, device="cuda") - his = torch.tensor(spec["initial"], dtype=torch.bool, device="cuda") - return [x, weight, None, pool, qsl, cache_indices, his, True, PAD_SLOT_ID, None] - - -def _clone(args): - """The mutable operands cloned (x, conv_states).""" - got = list(args) - for i in (0, 3): - got[i] = got[i].clone() - return got - - -def _run(args): - """Drive the op on clones; return (returned activation, pool after, clones).""" - got = _clone(args) - returned = op(*got) - return returned, got[3], got - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name - certifying(op) - args = _build(cell.spec, seed=_seed(cell.why)) - ref = op.reference(*args) - returned, pool, got = _run(args) - assert returned is got[0], "returns the tensor it wrote" - _compare((returned, pool), ref) - if cell.spec["pool_slots"] == ENGINE_SLOTS: - # The whole-pool bitwise gate above already locks the unnamed slots; - # this keeps it from being vacuous on the large pool. - for slot in cell.spec["slots"]: - assert not torch.equal(pool[slot], args[3][slot]), slot - - -def _token_strided_args(): - """A fresh argument list whose `x` has token stride 2: the channel-major - kernel ignores that stride and answers silently wrong.""" - dim, tokens = 512, 64 - spec = dict(dim=dim, lengths=(tokens,), pool_slots=1, slots=(0,), initial=(True,)) - args = _build(spec, seed=15) - g = torch.Generator(device="cuda").manual_seed(16) - wide = torch.randn(2 * tokens, dim, generator=g, device="cuda").to(torch.bfloat16) - args[0] = wide.t().contiguous()[:, ::2] - assert args[0].stride(1) == 2 - return args - - -def test_guard_refuses_what_the_op_answers_silently_wrong() -> None: - """First the hazard, unguarded: the op returns without raising and its - answer misses the reference. Then the guard refuses the same call.""" - ref = op.reference(*_token_strided_args()) - unguarded = _token_strided_args() # strided views must not be cloned dense - returned = op(*unguarded) - with pytest.raises(AssertionError): - _compare((returned, unguarded[3]), ref) - - with pytest.raises(AssertionError): - with validating(op): - op(*_token_strided_args()) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_reset_recurrent_state_rows.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_reset_recurrent_state_rows.py deleted file mode 100644 index 13a81044a508..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_reset_recurrent_state_rows.py +++ /dev/null @@ -1,117 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the reset_recurrent_state_rows entry.""" - -import zlib -from typing import Optional - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.reset_recurrent_state_rows import ( # noqa: E501 - CONV_SHAPE, - REC_SHAPE, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.reset_recurrent_state_rows import ( # noqa: E501 - reset_recurrent_state_rows as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name - -assert torch.cuda.is_available(), "reset_recurrent_state_rows requires a CUDA device" - - -def _compare( - out: tuple[torch.Tensor, Optional[torch.Tensor]], - ref: tuple[torch.Tensor, Optional[torch.Tensor]], -) -> None: - """Bitwise, both pools whole: every element is its old value or +0.0, - so a tolerance could only hide a wrong row -- and comparing the whole - pool is the lock on every slot the batch did not name. A zeroed row - must be +0.0 by bit pattern; `torch.equal` alone would accept -0.0.""" - for got, want in zip(out, ref): - assert (got is None) == (want is None) - if got is None: - continue - assert got.dtype == want.dtype and got.shape == want.shape - assert torch.equal(got, want), "pool differs from the reference" - zero = want.reshape(want.shape[0], -1).eq(0).all(dim=1) - if zero.any(): - bits = {1: torch.uint8, 2: torch.int16, 4: torch.int32}[got.element_size()] - assert (got[zero].contiguous().view(bits) == 0).all(), "a reset row is not +0.0" - - -DEV = "cuda" -_LAYERS = 3 - - -def _build(spec, seed): - """`__call__`'s arguments, plus the backing allocations a layer view views.""" - g = torch.Generator(device=DEV).manual_seed(seed) - slots = spec["slots"] - rec_dtype = spec.get("rec_dtype", torch.float32) - layer_view = spec.get("layout", "plain") == "layer_view" - - def randn(*shape, dtype): - return torch.randn(*shape, generator=g, device=DEV).to(dtype) - - backing = {} - if layer_view: - backing["rec"] = randn(_LAYERS, slots, *REC_SHAPE, dtype=rec_dtype) - backing["conv"] = randn(_LAYERS, slots, *CONV_SHAPE, dtype=torch.bfloat16) - rec, conv = backing["rec"][1], backing["conv"][1] - else: - rec = randn(slots, *REC_SHAPE, dtype=rec_dtype) - conv = randn(slots, *CONV_SHAPE, dtype=torch.bfloat16) - - indices = torch.tensor(spec["indices"], dtype=torch.int32, device=DEV) - flags = torch.tensor(spec["flags"], dtype=torch.bool, device=DEV) - return [rec, indices, flags, conv], backing - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name - certifying(op) - spec = cell.spec - args, backing = _build(spec, seed=zlib.crc32(cell.why.encode())) - rec, indices, flags, conv = args - before = {name: t.clone() for name, t in backing.items()} - indices_before, flags_before = indices.clone(), flags.clone() - ref = op.reference(*args) - - out = op(*args) - - assert out[0] is rec and out[1] is conv, "returns the pools it was handed" - _compare(out, ref) - assert torch.equal(indices, indices_before) and torch.equal(flags, flags_before) - - named = {s for s in indices.tolist() if 0 <= s < spec["slots"]} - reset = {s for s, f in zip(indices.tolist(), flags.tolist()) if not f and s in named} - if spec["slots"] > 64: - # Non-vacuity: the engine-width cells reset at least one row, so the - # whole-pool compare is also the lock on the ~430 unnamed slots. - assert reset - if spec.get("layout") == "layer_view": - for name in backing: - for other in (0, 2): - assert torch.equal(backing[name][other], before[name][other]), (name, other) - - -def _pools(): - torch.manual_seed(7) - rec = torch.randn(6, *REC_SHAPE[1:], dtype=torch.float32, device=DEV) - conv = torch.randn(6, *CONV_SHAPE, dtype=torch.bfloat16, device=DEV) - return rec, conv - - -def test_guard_refuses_index_tensors_the_op_misreads() -> None: - """Unguarded, a stride-2 `state_indices` view is read through its base pointer - as a flat buffer: indices [0, 1, 2] arrive as [0, 9, 1], so rows 0 and 1 are reset - and row 2 is not (measured). The guard refuses it.""" - rec, conv = _pools() - indices = torch.tensor([0, 9, 1, 9, 2, 9], dtype=torch.int32, device=DEV)[::2] - flags = torch.zeros(3, dtype=torch.bool, device=DEV) - with pytest.raises(AssertionError, match="state_indices must be a contiguous 1-D tensor"): - with validating(op): - op(rec, indices, flags, conv) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py deleted file mode 100644 index b67a01545790..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py +++ /dev/null @@ -1,357 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the target's Gated-DeltaNet verify step in both recurrent modes. - -Drives the two targets' own `_gdn_verify` -- the prefill target's (generation rows -behind a context prefix) and the decode target's -- over consecutive verify steps of -this checkpoint's GDN geometry, with the cache manager's bookkeeping run in between as -`update_mamba_states` runs it: the conv window promotion in both modes, then either the -legacy per-step state promotion, or the cached replay's all-layer commit (partitioned -batches only) and `_advance_replay_state`. - -Three things are checked at every step, with a random partial acceptance (or, in the -full-acceptance cases, every draft accepted): - -* the replay mode's output against the legacy mode's on identical inputs and identical - accepted histories: the two modes compute one recurrence with different kernels and - state disciplines, so a mis-bound history buffer, a wrong partition condition or a - skipped commit drifts off within a step; -* the prefill and the decode target's replay calls bitwise equal to each other (each on - its own copy of the state); -* the conv pool identical across the modes (it is promoted the same way). -""" - -from types import SimpleNamespace - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - modeling as M, -) -from tensorrt_llm._torch.modules.fla.cache_manager import ( - ReplayStateUpdateMetadata, - _advance_replay_state, -) -from tensorrt_llm._torch.modules.fla.cached_replay import ( - CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, - commit_gdn_cached_replay_history_layers, -) -from tensorrt_llm._torch.modules.mamba.mamba2_metadata import _build_replay_work_items_torch - -assert torch.cuda.is_available(), "the GDN verify test requires a CUDA device" - -DEV = "cuda" -H, HV, K, V = 16, 48, 128, 128 -KEY_DIM, VALUE_DIM = H * K, HV * V -CONV_DIM = 2 * KEY_DIM + VALUE_DIM -CONV_KERNEL = 4 -HIST = 16 -LAYERS = 2 -SLOTS = 40 -STEPS = 7 - - -def _steps(t): - """Enough steps that every slot fills its history at least once: each step accepts - at least one token, and a slot wraps once `pnat + t > HIST`.""" - return HIST - t + 3 - - -BF16_ULP = 2.0**-8 -#: Replay against legacy, per token block: the replay kernel rounds the history's update -#: vectors and keys to bf16 and rebuilds the state from them, where the legacy kernel -#: carries an fp32 state, so the two agree to a few bf16 ulps of the block. -ELEMENT_GATE = 6 * BF16_ULP -RMS_GATE = 1.5 * BF16_ULP - - -class _Mode: - """One cache manager's worth of state for `LAYERS` GDN layers, in one mode, plus the - `core` / `L` stand-ins `_gdn_verify` reads.""" - - def __init__(self, replay, t, conv0, ssm0, weights, gating, max_batch): - self.conv = [c.clone() for c in conv0] - self.int_conv = [ - torch.zeros(max_batch, t, CONV_DIM, CONV_KERNEL - 1, dtype=torch.bfloat16, device=DEV) - for _ in range(LAYERS) - ] - # V2 lays the layers out inside each slot; the all-layer commit is handed that - # affine layout. - self.backing = torch.stack(ssm0, 1).contiguous() - self.ssm = [self.backing[:, i] for i in range(LAYERS)] - bound = dict( - conv=dict(enumerate(self.conv)), - int_conv=dict(enumerate(self.int_conv)), - int_ssm=None, - replay=None, - ) - if replay: - self.old_x = torch.zeros( - LAYERS, SLOTS, 2, HIST, HV, V, dtype=torch.bfloat16, device=DEV - ) - self.old_B = torch.zeros(LAYERS, SLOTS, 2, HIST, H, K, dtype=torch.bfloat16, device=DEV) - self.old_dt = torch.zeros(LAYERS, SLOTS, 2, HV, HIST, device=DEV) - self.old_dA = torch.zeros(LAYERS, SLOTS, 2, HV, HIST, device=DEV) - self.pnat = torch.zeros(SLOTS, dtype=torch.int32, device=DEV) - self.buf = torch.zeros(SLOTS, dtype=torch.int32, device=DEV) - bound["replay"] = M._Replay( - old={ - i: (self.old_x[i], self.old_B[i], self.old_dt[i], self.old_dA[i]) - for i in range(LAYERS) - }, - cache_buf_idx=self.buf, - prev_num_accepted=self.pnat, - history=HIST, - width=t, - all_layer_commit=True, - # The Triton replay and the manager's all-layer commit, as an fp32 pool runs. - cute=False, - ) - self.meta = ReplayStateUpdateMetadata(self.pnat, self.buf, t, HIST) - self.mm = SimpleNamespace( - replay_work_items=torch.zeros(max_batch, 4, dtype=torch.int32, device=DEV), - replay_n_writes=torch.zeros(1, dtype=torch.int32, device=DEV), - ) - else: - self.int_ssm = [torch.zeros(max_batch, t, HV, V, K, device=DEV) for _ in range(LAYERS)] - bound["int_ssm"] = dict(enumerate(self.int_ssm)) - self.core = SimpleNamespace( - conv_dim=CONV_DIM, - key_dim=KEY_DIM, - value_dim=VALUE_DIM, - num_k_heads=H, - head_k_dim=K, - num_v_heads=HV, - head_v_dim=V, - _bound=SimpleNamespace(**bound), - ) - self.layers = [ - SimpleNamespace(conv_w=weights[i], a_log=gating[i][0], dt_bias=gating[i][1]) - for i in range(LAYERS) - ] - - def work(self, slots, n, t): - """What the forward hands `_gdn_verify` as `replay`: the metadata's prepare builds - the work items, then `_replay_work` applies the manager's partition condition.""" - partitioned = n >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE - if partitioned: - _build_replay_work_items_torch( - slots, - self.pnat, - self.buf, - self.mm.replay_work_items, - self.mm.replay_n_writes, - t, - HIST, - ) - replay = M._replay_work(self.core, self.mm, n) - assert replay[2] == partitioned - return replay - - def verify(self, target_cls, layer, mq, a, b, slots, int_idx, n, t, replay): - return target_cls._gdn_verify( - SimpleNamespace(core=self.core), - mq, - a, - b, - self.ssm[layer], - layer, - self.layers[layer], - slots, - int_idx, - n, - t, - replay, - ) - - def commit(self, slots, n): - """The manager's all-layer commit, then the replay bookkeeping advance is the - caller's (it needs the accepted counts).""" - commit_gdn_cached_replay_history_layers( - ssm_states=self.ssm[0], - ssm_state_num_layers=LAYERS, - ssm_state_layer_stride=self.backing.stride(1), - ssm_state_slot_stride=self.backing.stride(0), - old_u=self.old_x, - old_k=self.old_B, - old_G=self.old_dt, - replay_work_items=self.mm.replay_work_items[:n], - n_writes=self.mm.replay_n_writes, - history_size=HIST, - ) - - -def _rows(gen, rows): - """One layer's verify-row activations: the pre-conv `[q|k|v]` projection and the raw - `a` and `b` projections, each its own contiguous GEMM output.""" - return tuple( - torch.randn(rows, width, generator=gen, device=DEV).to(torch.bfloat16) - for width in (CONV_DIM, HV, HV) - ) - - -def _behind_prefix(x, prefix=5): - """`x` as the trailing rows of a buffer that starts with `prefix` context rows.""" - buf = torch.zeros(prefix + x.shape[0], x.shape[1], dtype=x.dtype, device=DEV) - buf[prefix:] = x - return torch.split(buf, [prefix, x.shape[0]], 0)[1] - - -def _promote(pool, inter, slots, accepted): - """`update_mamba_states`' copy: step `accepted - 1` of each batch row's per-step - buffer into the row's slot.""" - rows = torch.arange(slots.shape[0], device=DEV) - pool[slots.long()] = inter[rows, accepted.long() - 1] - - -def _gate(got, want, what): - block = got.shape[-2] * got.shape[-1] - g, w = got.float().reshape(-1, block), want.float().reshape(-1, block) - scale = w.abs().amax(1, keepdim=True).clamp_min(1e-30) - worst = ((g - w).abs() / scale).max().item() - rms = ((g - w).pow(2).mean().sqrt() / w.pow(2).mean().sqrt()).item() - assert worst <= ELEMENT_GATE and rms <= RMS_GATE, ( - f"{what}: replay vs legacy {worst / BF16_ULP:.2f} ulps of the block, " - f"RMS {rms / BF16_ULP:.2f} ulps" - ) - - -def _state(gen, cpu, n, t, gating=None): - """Random weights, gating parameters and pre-step pools shared by the modes, and a - random slot assignment of `n` requests.""" - slots = torch.randperm(SLOTS, generator=cpu)[:n].to(DEV, torch.int32) - weights = [ - (0.3 * torch.randn(CONV_DIM, CONV_KERNEL, generator=gen, device=DEV)).to(torch.bfloat16) - for _ in range(LAYERS) - ] - if gating is None: - gating = [ - ( - (torch.randn(HV, generator=gen, device=DEV) * 1.5 + 1.0).clamp(-7.0, 5.0), - (torch.randn(HV, generator=gen, device=DEV) * 1.7 - 4.4).clamp(-11.0, 2.5), - ) - for _ in range(LAYERS) - ] - conv0 = [ - torch.randn(SLOTS, CONV_DIM, CONV_KERNEL - 1, generator=gen, device=DEV).to(torch.bfloat16) - for _ in range(LAYERS) - ] - ssm0 = [0.3 * torch.randn(SLOTS, HV, V, K, generator=gen, device=DEV) for _ in range(LAYERS)] - return slots, weights, gating, conv0, ssm0 - - -@pytest.mark.parametrize( - "n,t,accept", - [ - # At the partition threshold, at max_draft_len 7 (T = 8, the replay kernel's - # ceiling): the history wraps through the all-layer commit. - (CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, 8, "random"), - # Below the threshold, every draft accepted: from a fresh slot at T = 8 the second - # step appends exactly up to the history's end (`pnat + T == 16`) and every second - # step after it wraps in-kernel. - (4, 8, "full"), - ], - ids=lambda v: str(v), -) -def test_replay_verify_tracks_legacy_in_both_targets(n, t, accept): - gen = torch.Generator(device=DEV).manual_seed(1000 * n + t) - cpu = torch.Generator().manual_seed(1000 * n + t) - slots, weights, gating, conv0, ssm0 = _state(gen, cpu, n, t) - int_idx = torch.arange(n, dtype=torch.int32, device=DEV) - legacy = _Mode(False, t, conv0, ssm0, weights, gating, n) - prefill = _Mode(True, t, conv0, ssm0, weights, gating, n) - decode = _Mode(True, t, conv0, ssm0, weights, gating, n) - partitioned = n >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE - - wraps = to_the_end = 0 - for step in range(_steps(t)): - pending = prefill.pnat[slots.long()] - wraps += int((pending + t > HIST).sum()) - to_the_end += int((pending + t == HIST).sum()) - replays = {id(mode): mode.work(slots, n, t) for mode in (prefill, decode)} - for layer in range(LAYERS): - mq, a, b = _rows(gen, n * t) - # The prefill target's generation rows sit behind a context prefix: row views - # at an offset into the batch's projections. - mq_p, a_p, b_p = (_behind_prefix(x) for x in (mq, a, b)) - want = legacy.verify(M.DecodeTarget, layer, mq, a, b, slots, int_idx, n, t, None) - got_p = prefill.verify( - M.PrefillTarget, layer, mq_p, a_p, b_p, slots, int_idx, n, t, replays[id(prefill)] - ) - got_d = decode.verify( - M.DecodeTarget, layer, mq, a, b, slots, int_idx, n, t, replays[id(decode)] - ) - assert got_p.shape == got_d.shape == want.shape == (1, n * t, HV, V) - assert torch.equal(got_p, got_d), ( - f"step {step} layer {layer}: the prefill and decode targets' replay calls disagree" - ) - _gate(got_d, want, f"step {step} layer {layer}") - if accept == "full": - accepted = torch.full((n,), t, dtype=torch.int32, device=DEV) - else: - accepted = torch.randint(1, t + 1, (n,), generator=cpu, dtype=torch.int32).to(DEV) - for layer in range(LAYERS): - for mode in (legacy, prefill, decode): - _promote(mode.conv[layer], mode.int_conv[layer], slots, accepted) - _promote(legacy.ssm[layer], legacy.int_ssm[layer], slots, accepted) - for mode in (prefill, decode): - if partitioned: - mode.commit(slots, n) - _advance_replay_state(mode.meta, slots, accepted) - for layer in range(LAYERS): - assert torch.equal(prefill.conv[layer], legacy.conv[layer]) - assert torch.equal(decode.conv[layer], legacy.conv[layer]) - assert wraps > 0, "the steps never filled a history, so the checkpoint never advanced" - if accept == "full": - assert wraps >= 3 * n and to_the_end >= 3 * n, ( - f"full acceptance at T={t} should wrap every row at least 3 times and append up to " - f"the history's end as often ({wraps}, {to_the_end})" - ) - - -def test_a_wrong_partition_drifts_off_legacy(): - """The control: the work items withheld from a partitioned step while the manager - still runs its all-layer commit -- the checkpoint committed twice -- must fail the - same gate.""" - n, t = CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, 4 - gen = torch.Generator(device=DEV).manual_seed(7) - cpu = torch.Generator().manual_seed(7) - gating = [ - (torch.full((HV,), 0.5, device=DEV), torch.full((HV,), -4.0, device=DEV)) - for _ in range(LAYERS) - ] - slots, weights, gating, conv0, ssm0 = _state(gen, cpu, n, t, gating) - int_idx = torch.arange(n, dtype=torch.int32, device=DEV) - legacy = _Mode(False, t, conv0, ssm0, weights, gating, n) - wrong = _Mode(True, t, conv0, ssm0, weights, gating, n) - failed = False - for _ in range(STEPS): - # The manager's prepare builds the work items; the step does not receive them. - _build_replay_work_items_torch( - slots, - wrong.pnat, - wrong.buf, - wrong.mm.replay_work_items, - wrong.mm.replay_n_writes, - t, - HIST, - ) - for layer in range(LAYERS): - mq, a, b = _rows(gen, n * t) - want = legacy.verify(M.DecodeTarget, layer, mq, a, b, slots, int_idx, n, t, None) - got = wrong.verify( - M.DecodeTarget, layer, mq, a, b, slots, int_idx, n, t, (None, None, False) - ) - try: - _gate(got, want, "control") - except AssertionError: - failed = True - accepted = torch.randint(1, t + 1, (n,), generator=cpu, dtype=torch.int32).to(DEV) - for layer in range(LAYERS): - _promote(legacy.conv[layer], legacy.int_conv[layer], slots, accepted) - _promote(wrong.conv[layer], wrong.int_conv[layer], slots, accepted) - _promote(legacy.ssm[layer], legacy.int_ssm[layer], slots, accepted) - wrong.commit(slots, n) - _advance_replay_state(wrong.meta, slots, accepted) - assert failed, "a doubled checkpoint commit passed the replay-vs-legacy gate" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py deleted file mode 100644 index 548008dacdf4..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py +++ /dev/null @@ -1,168 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU tests for the target's MTP IndexShare plumbing. - -The target's canonical selection, the capture of each request's draft-extend row -into the shared buffers (decode and mixed-batch forms), and the reuse entry's -order preservation over canonical captures. The targets are built without weights: -the capture methods read only their arguments, the metadata and the core's -draft binding. -""" - -from types import SimpleNamespace - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - modeling as m, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_canonical_selection import ( # noqa: E501 - qsa_canonical_selection as canonical, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_mtp_reuse_topk import ( # noqa: E501 - qsa_mtp_reuse_topk as reuse, -) - -assert torch.cuda.is_available(), "the IndexShare plumbing runs on a CUDA device" - -DEV = "cuda" -K = 512 -SLOT = 1 - - -def _target(cls): - target = object.__new__(cls) - target.core = SimpleNamespace(_draft_bound=SimpleNamespace(share_slot=SLOT)) - return target - - -def _metadata(batch: int, accepted, seq_lens=None, num_contexts: int = 0): - return SimpleNamespace( - num_seqs=batch, - num_contexts=num_contexts, - seq_lens_cuda=None - if seq_lens is None - else torch.tensor(seq_lens, dtype=torch.int32, device=DEV), - qsa_mtp_num_accepted=torch.tensor(accepted, dtype=torch.int32, device=DEV), - qsa_shared_topk_indices=torch.full((3, 8, K), -7, dtype=torch.int32, device=DEV), - qsa_shared_topk_visible_blocks=torch.full((3, 8), -7, dtype=torch.int32, device=DEV), - qsa_shared_topk_captured_slots=set(), - ) - - -COLS = 2080 - - -def _scores_and_windows(rows: int, gen: torch.Generator): - """Relu-shaped scores (half zeros: ties at the budget's last score whenever fewer than - K are positive) over windows of 0 to 3K groups; the columns past a window are NaN.""" - ends = torch.randint(0, COLS + 1, (rows,), generator=gen).to(torch.int32).to(DEV) - scores = torch.relu(torch.randn(rows, COLS, generator=gen)).to(DEV) - scores = scores * (torch.rand(rows, COLS, generator=gen).to(DEV) < 0.3) - outside = torch.arange(COLS, device=DEV)[None, :] >= ends[:, None] - scores.masked_fill_(outside, float("nan")) - return scores, torch.zeros(rows, dtype=torch.int32, device=DEV), ends - - -def _select(scores, starts, ends) -> torch.Tensor: - out = torch.empty((scores.shape[0], K), dtype=torch.int32, device=DEV) - m._select(SimpleNamespace(block_topk=K), scores, starts, ends, out) - return out - - -def test_select_is_the_canonical_row_ascending_with_padding_last(): - gen = torch.Generator().manual_seed(3) - scores, starts, ends = _scores_and_windows(64, gen) - - out = _select(scores, starts, ends) - - assert out.dtype == torch.int32 - assert torch.equal(out, canonical.reference(scores, starts, ends, out)) - for row, end in zip(out.cpu().tolist(), ends.tolist()): - ids = [i for i in row if i >= 0] - assert ids == sorted(ids) and len(set(ids)) == len(ids) == min(end, K) - assert row == ids + [-1] * (K - len(ids)), "padding must stay last" - - -def test_decode_capture_at_draft_seven_takes_row_s_times_8_plus_a_minus_1(): - """max_draft_len 7: eight rows per request, accepted 1..8, so the captured row is - `s * 8 + a - 1` -- the first and the last row of a block included.""" - p, batch = 8, 4 - accepted = [1, 8, 5, 3] - topk = torch.arange(batch * p * K, dtype=torch.int32, device=DEV).reshape(batch * p, K) - visible = torch.arange(300, 300 + batch * p, dtype=torch.int32, device=DEV) - md = _metadata(batch, accepted) - - _target(m.DecodeTarget)._capture_selection(md, topk, visible, batch, p) - - rows = [s * p + a - 1 for s, a in enumerate(accepted)] - assert rows == [0, 15, 20, 26] - torch.testing.assert_close(md.qsa_shared_topk_indices[SLOT, :batch], topk[rows]) - torch.testing.assert_close(md.qsa_shared_topk_visible_blocks[SLOT, :batch], visible[rows]) - # Other slots and rows past the batch are untouched. - assert torch.all(md.qsa_shared_topk_indices[SLOT, batch:] == -7) - assert torch.all(md.qsa_shared_topk_indices[0] == -7) - assert md.qsa_shared_topk_captured_slots == {SLOT} - - -def test_prefill_capture_at_draft_seven(): - """A mixed batch at max_draft_len 7: two context chunks of 5 and 3 rows, then three - generation requests of 8 rows with 8, 1 and 6 accepted. A context request captures its - chunk's last row whatever accepted count the worker left there.""" - seq_lens = [5, 3, 8, 8, 8] - accepted = [2, 3, 8, 1, 6] - rows_total = sum(seq_lens) - topk = torch.arange(rows_total * K, dtype=torch.int32, device=DEV).reshape(rows_total, K) - visible = torch.arange(400, 400 + rows_total, dtype=torch.int32, device=DEV) - md = _metadata(5, accepted, seq_lens=seq_lens, num_contexts=2) - - _target(m.PrefillTarget)._capture_selection(md, topk, visible) - - rows = [4, 7, 8 + 7, 16 + 0, 24 + 5] - torch.testing.assert_close(md.qsa_shared_topk_indices[SLOT, :5], topk[rows]) - torch.testing.assert_close(md.qsa_shared_topk_visible_blocks[SLOT, :5], visible[rows]) - assert md.qsa_shared_topk_captured_slots == {SLOT} - - -def _reuse_case(fresh: int): - """A canonical capture of 16 requests and the reuse call's arguments for it.""" - gen = torch.Generator().manual_seed(11 + fresh) - batch = 16 - scores, starts, ends = _scores_and_windows(batch, gen) - captured_visible = ends.cpu() - captured = _select(scores, starts, ends) - shared = torch.full((8, 32, K), -1, dtype=torch.int32, device=DEV) - shared[SLOT, :batch] = captured - shared_visible = torch.zeros((8, 32), dtype=torch.int32, device=DEV) - shared_visible[SLOT, :batch] = captured_visible.to(DEV) - args = dict( - shared_indices=shared[SLOT], - shared_visible_blocks=shared_visible[SLOT], - request_indices=torch.arange(batch, dtype=torch.int32, device=DEV), - visible_blocks=(captured_visible + fresh).to(DEV), - output=torch.empty((batch, K), dtype=torch.int32, device=DEV), - ) - return args, captured_visible - - -def test_reuse_of_a_canonical_capture_stays_ascending(): - """A canonical capture plus the groups completed since is still ascending: appended ids - are newer than every captured one, and a full row gives up its oldest ids to them, so - the newest captured groups -- the ones next to the new ones -- are kept.""" - for fresh in (0, 2): - args, captured_visible = _reuse_case(fresh) - - out = reuse(**args) - - assert torch.equal(out, reuse.reference(**args)) - for row, shared_row, vis in zip( - out.cpu().tolist(), args["shared_indices"].cpu().tolist(), captured_visible.tolist() - ): - ids = [i for i in row if i >= 0] - assert row == ids + [-1] * (K - len(ids)), "padding must stay last" - assert ids == sorted(ids) and len(set(ids)) == len(ids) - assert len(ids) == min(vis + fresh, K) - assert all(i < vis + fresh for i in ids) - captured = [i for i in shared_row if i >= 0] - kept = captured[max(0, len(captured) + fresh - K) :] - assert ids == kept + list(range(vis, vis + fresh)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_kv_range_probe.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_kv_range_probe.py deleted file mode 100644 index 2603aaf63e00..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_kv_range_probe.py +++ /dev/null @@ -1,82 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The FP8 KV range probe (`TRTLLM_MODELING_V2_VALIDATE` over an fp8 pool). - -The pool stores K/V at the unit scale, so a magnitude past the e4m3 maximum saturates. The probe -reports the largest magnitude of the first eager forward's attention layers, once, and never reads -a tensor under CUDA graph capture. -""" - -from unittest.mock import Mock - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - modeling as m, -) - -DEVICE = "cuda" if torch.cuda.is_available() else "cpu" -LAST = 11 - - -def _kv(peak: float = 1.0): - """K and V of one layer, 5 rows x 2 heads x 8 dims, with `peak` as V's largest value.""" - k = torch.linspace(-1.0, 1.0, 80, device=DEVICE).reshape(5, 2, 8).to(torch.bfloat16) - v = torch.linspace(-0.5, 0.5, 80, device=DEVICE).reshape(5, 2, 8).to(torch.bfloat16) - v[3, 1, 4] = peak - return k, v - - -def _run_forward(probe, peaks): - """One forward's attention layers 3, 7 and LAST; returns each `observe` result.""" - return [probe.observe(layer, *_kv(peak)) for layer, peak in zip((3, 7, LAST), peaks)] - - -def test_a_forward_within_the_e4m3_range_is_silent(monkeypatch): - logger = Mock() - monkeypatch.setattr(m, "logger", logger) - - assert _run_forward(m._KvRangeProbe(LAST), (1.0, 300.0, 448.0)) == [True, True, False] - - logger.warning.assert_not_called() - - -@pytest.mark.parametrize("peak", [-600.0, float("nan")]) -def test_a_magnitude_past_the_e4m3_maximum_is_reported_once(monkeypatch, peak): - """A negative value counts by its magnitude; a NaN counts as out of range.""" - logger = Mock() - monkeypatch.setattr(m, "logger", logger) - - # The oversized value is in a middle layer: the report waits for the forward's last one. - assert _run_forward(m._KvRangeProbe(LAST), (1.0, peak, 1.0)) == [True, True, False] - - logger.warning.assert_called_once() - assert "e4m3 maximum 448" in logger.warning.call_args.args[0] - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_observing_under_capture_reads_nothing(monkeypatch): - logger = Mock() - monkeypatch.setattr(m, "logger", logger) - probe = m._KvRangeProbe(LAST) - k, v = _kv(1000.0) - scratch = torch.zeros(1, device="cuda") - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - scratch.add_(1) # warm the allocator off the capture stream - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - kept = [] - with torch.cuda.graph(graph): - scratch.add_(1) - # The last layer, with an out-of-range value: a read here would fail the capture - # (a host sync) and a report would fire. - kept.append(probe.observe(LAST, k, v)) - - assert kept == [True], "a captured forward does not end the probe" - logger.warning.assert_not_called() - # The first eager forward after capture is what reports. - assert _run_forward(probe, (1.0, 1.0, 1000.0)) == [True, True, False] - logger.warning.assert_called_once() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py deleted file mode 100644 index 62dc201db2db..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py +++ /dev/null @@ -1,100 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The target's logits processor returns the shell's logits bit for bit: the same row gather, -the same fp32 values, its head's bf16 logits widened by one vectorized kernel.""" - -from types import SimpleNamespace - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - modeling as T, -) -from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor - -assert torch.cuda.is_available(), "the logits processor test requires a CUDA device" - -VOCAB, HIDDEN = 248320, 2560 - - -@pytest.fixture(scope="module") -def head(): - g = torch.Generator(device="cuda").manual_seed(0) - lin = torch.nn.Linear(HIDDEN, VOCAB, bias=False, device="cuda", dtype=torch.bfloat16) - with torch.no_grad(): - lin.weight.copy_(torch.randn(VOCAB, HIDDEN, generator=g, device="cuda") * 0.02) - return lin - - -@pytest.mark.parametrize("seq_lens", [[1] * 128, [5000, 3000, 1] + [1] * 121]) -def test_matches_the_shell_processor_bitwise(head, seq_lens) -> None: - g = torch.Generator(device="cuda").manual_seed(len(seq_lens)) - lens = torch.tensor(seq_lens, dtype=torch.int32, device="cuda") - hidden = torch.randn(int(lens.sum()), HIDDEN, generator=g, device="cuda").to(torch.bfloat16) - md = SimpleNamespace(seq_lens_cuda=lens) - with torch.no_grad(): - want = LogitsProcessor().forward(hidden, head, md) - got = T._WideningLogitsProcessor().forward(hidden, head, md) - assert got.dtype == torch.float32 and got.shape == (len(seq_lens), VOCAB) - assert torch.equal(got.view(torch.int32), want.view(torch.int32)) - - -def test_context_logits_and_no_metadata(head) -> None: - g = torch.Generator(device="cuda").manual_seed(9) - hidden = torch.randn(33, HIDDEN, generator=g, device="cuda").to(torch.bfloat16) - with torch.no_grad(): - for md, ctx in ((None, False), (SimpleNamespace(seq_lens_cuda=None), True)): - want = LogitsProcessor().forward(hidden, head, md, ctx) - got = T._WideningLogitsProcessor().forward(hidden, head, md, ctx) - assert torch.equal(got.view(torch.int32), want.view(torch.int32)) - - -@pytest.mark.parametrize("rows", [1, 4, 256]) -def test_unwidened_logits_give_the_same_greedy_tokens(head, rows) -> None: - """`widen=False` returns the head's own bf16 logits, whose greedy samplers (the verify - sampler's `greedy_argmax`, the draft sampler's int32 one) pick the widened logits' tokens.""" - from tensorrt_llm._torch.speculative.greedy_argmax import greedy_argmax - - g = torch.Generator(device="cuda").manual_seed(rows) - hidden = torch.randn(rows, HIDDEN, generator=g, device="cuda").to(torch.bfloat16) - proc = T._WideningLogitsProcessor() - with torch.no_grad(): - wide = proc.forward(hidden, head, None, True) - narrow = proc.forward(hidden, head, None, True, widen=False) - # Without return_context_logits the shell's gather and fp32 contract stay. - gathered = proc.forward(hidden, head, None, False, widen=False) - assert narrow.dtype == torch.bfloat16 and wide.dtype == torch.float32 - assert gathered.dtype == torch.float32 - assert torch.equal(narrow.float().view(torch.int32), wide.view(torch.int32)) - narrow[0, 17] = narrow[0].max() # an exact tie, broken to the lower index either way - wide = narrow.float() - assert torch.equal(greedy_argmax(narrow), greedy_argmax(wide)) - assert torch.equal(greedy_argmax(narrow, torch.int32), greedy_argmax(wide, torch.int32)) - assert torch.equal(greedy_argmax(narrow), torch.argmax(wide, dim=-1)) - - -def _gate(guided=None, greedy=True, rejection=False, penalty=False, relaxed=False): - worker = SimpleNamespace( - guided_decoder=guided, - spec_config=SimpleNamespace(use_relaxed_acceptance_for_thinking=relaxed), - ) - md = SimpleNamespace( - is_all_greedy_sample=greedy, - wants_advanced_draft_sampling=rejection and not greedy, - enable_penalty=penalty, - ) - return T._spec_argmax_only(worker, md) - - -def test_argmax_only_gate() -> None: - assert _gate() == (True, True) - # A non-greedy batch: the verify sampler samples; the draft stays argmax unless - # rejection sampling needs its distribution. - assert _gate(greedy=False) == (False, True) - assert _gate(greedy=False, rejection=True) == (False, False) - # The occurrence penalties and relaxed acceptance read the verify logits only. - assert _gate(penalty=True) == (False, True) - assert _gate(relaxed=True) == (False, True) - # A guided decoder edits both. - assert _gate(guided=object()) == (False, False) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_nvfp4_context_groups.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_nvfp4_context_groups.py deleted file mode 100644 index 847d5d8abc5e..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_nvfp4_context_groups.py +++ /dev/null @@ -1,136 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test: the target's context pass over an NVFP4 pool, a group of requests at a time. - -The context rows read a bf16 scratch of their requests' pages (`_nvfp4_context_plan`), filled -and attended one group of consecutive requests at a time, so it holds the largest group's -pages rather than the step's. Under a budget of a few pages, with one request longer than -the budget, the rows equal, bitwise, one group over the whole step -- for a prefill step's -rows and for a mixed step's one kept row a request. The prefill requests' rows come in -multiples of the row-tile pass's 10-row groups, so a launch per group splits none of them and -every row takes the same pass as in the one launch (a split group's rows go to the per-row -pass, which rounds differently). -""" - -import math -from types import SimpleNamespace - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - modeling as M, -) - -assert torch.cuda.is_available(), "the NVFP4 context pass requires a CUDA device" - -DEV = "cuda" -TPB, HKV, D, QH, RATIO, TOPK = 32, 2, 256, 24, 4, 2048 -# (cached, new) tokens a context request: 10, 30, 1, 47 and 4 pages. -SEGMENTS = ((0, 300), (900, 40), (20, 10), (1400, 100), (64, 60)) -BUDGET_PAGES = 32 # request 3's 47 pages form a group of their own - - -def _step(one_row_per_request: bool, seed: int = 5): - gen = torch.Generator(device=DEV).manual_seed(seed) - kv_lens = [s + n for s, n in SEGMENTS] - pages = [math.ceil(k / TPB) for k in kv_lens] - pool_pages = sum(pages) + 3 - shape = (pool_pages, 2, HKV, TPB) - data = torch.randint(0, 256, (*shape, D // 2), generator=gen, device=DEV).to(torch.uint8) - scales = torch.randint(0x08, 0x31, (*shape, D // 16), generator=gen, device=DEV) - scales = scales.to(torch.uint8) - slots = 2 * len(SEGMENTS) - slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed)) - slot_of = slot_of[: len(SEGMENTS)].to(DEV, torch.int32) - page_ids = torch.randperm(pool_pages, generator=gen, device=DEV).to(torch.int32) - block_table = torch.full((slots, max(pages)), -1, dtype=torch.int32, device=DEV) - cursor = 0 - for slot, count in zip(slot_of.tolist(), pages): - block_table[slot, :count] = page_ids[cursor : cursor + count] - cursor += count - new = torch.tensor([n for _, n in SEGMENTS], dtype=torch.int32) - md = SimpleNamespace( - num_contexts=len(SEGMENTS), - kv_lens_runtime=torch.tensor(kv_lens, dtype=torch.int32), - kv_lens_cuda_runtime=torch.tensor(kv_lens, dtype=torch.int32, device=DEV), - seq_lens=new, - seq_lens_cuda=new.to(DEV), - ) - # The step's rows: a request's new tokens, or the last of them alone. - kept = [1 if one_row_per_request else n for _, n in SEGMENTS] - logical = torch.cat( - [torch.arange(s + n - k, s + n, device=DEV) for (s, n), k in zip(SEGMENTS, kept)] - ) - rows = logical.numel() - # Every visible token, as the expansion lays out a selection under its top-k. - column = torch.arange(TOPK + RATIO - 1, device=DEV)[None, :] - selected = torch.where(column <= logical[:, None], column, -1).to(torch.int32) - core = SimpleNamespace( - heads_kv=HKV, - head_dim=D, - q_width=QH * D, - kv_scale_quant_orig=torch.ones(3, dtype=torch.float32, device=DEV), - softmax_scale=D**-0.5, - compress_ratio=RATIO, - ) - kv = SimpleNamespace(k=data[:, 0], v=data[:, 1], k_scale=scales[:, 0], v_scale=scales[:, 1]) - return dict( - core=core, - kv=kv, - md=md, - block_table=block_table, - req_all=slot_of.repeat_interleave(new.to(DEV).long()), - req_rows=slot_of.repeat_interleave(torch.tensor(kept, device=DEV)), - q=torch.randn(rows, QH, D, generator=gen, device=DEV).to(torch.bfloat16), - selected=selected, - logical=logical.to(torch.int32), - rows=rows, - ) - - -def _context_pass(step, one_row_per_request): - plan = M._nvfp4_context_plan( - step["core"], - step["md"], - step["req_all"], - step["block_table"], - TPB, - one_row_per_request=one_row_per_request, - ) - out = torch.empty(step["rows"], QH * D, dtype=torch.bfloat16, device=DEV) - M._sparse_gqa_nvfp4_context( - step["core"], - step["kv"], - plan, - step["q"], - step["block_table"], - step["selected"], - step["req_rows"], - step["logical"], - TPB, - step["rows"], - out, - ) - return plan, out - - -@pytest.mark.parametrize("one_row_per_request", [False, True], ids=["prefill", "mixed_kept"]) -def test_grouped_context_pass_equals_one_group(one_row_per_request, monkeypatch) -> None: - step = _step(one_row_per_request) - whole, want = _context_pass(step, one_row_per_request) - assert len(whole[0]) == 1 and whole[2].shape[0] == 92 - - page_bytes = 2 * HKV * TPB * D * 2 - monkeypatch.setattr(M, "_NVFP4_CONTEXT_SCRATCH_BYTES", BUDGET_PAGES * page_bytes) - (groups, _, k_scratch, _), got = _context_pass(step, one_row_per_request) - ends = [1, 2, 3, 4, 5] if one_row_per_request else [300, 340, 350, 450, 510] - # Requests [0], [1, 2], [3] (over the budget, alone), [4]. - assert [(g[0], g[1]) for g in groups] == [ - (0, ends[0]), - (ends[0], ends[2]), - (ends[2], ends[3]), - (ends[3], ends[4]), - ] - assert k_scratch.shape[0] == 47 - assert torch.equal(got.view(torch.int16), want.view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_patch_embed_conv3d.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_patch_embed_conv3d.py deleted file mode 100644 index 21bdba2b0d0c..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_patch_embed_conv3d.py +++ /dev/null @@ -1,118 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qwen4_exp patch_embed_conv3d entry. - -The reference is a true-fp32 matmul of the flattened patch against the -flattened filter: a different op, accumulation order and memory layout from -the convolution under test, so it checks the contraction pairing rather than -restating it. Beyond the cells: the one-patch-per-row guard and graph capture. -""" - -import zlib - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.vision.patch_embed_conv3d import ( # noqa: E501 - patch_embed_conv3d as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import ULP, certifying, validating # noqa: E402,F401 - -assert torch.cuda.is_available(), "patch_embed_conv3d requires a CUDA device" - - -def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: - """Default rtol; atol raised to 4 ulps of the dtype at the result's RMS. - - Summing 1536 signed products leaves a few outputs per thousand near - zero, where the relative gate measures nothing and torch's 1e-5 atol is - three orders below one bf16 ulp. Measured at the checkpoint's geometry: - max |conv - fp32 reference| was 4.31 / 4.38 / 7.81 ulps at the RMS for - 64 / 256 / 4096 bf16 patches (1-2 ulps of the element's own magnitude, - which the untouched rtol covers); fp16 4.36 / 4.39 / 8.15. Wrong - variants land 325-392x outside this gate (bias dropped: 10.9x, bounded - by `max|bias|` ~ 0.02). For fp32 the 1e-5 floor binds, so nothing is - loosened there. - """ - assert out.dtype == ref.dtype, (out.dtype, ref.dtype) - assert out.shape == ref.shape, (out.shape, ref.shape) - rms = ref.float().pow(2).mean().sqrt().item() - rtol = {torch.bfloat16: 1.6e-2, torch.float16: 1e-3, torch.float32: 1.3e-6}[ref.dtype] - atol = max(1e-5, 4.0 * rms * ULP[ref.dtype]) - torch.testing.assert_close(out, ref, rtol=rtol, atol=atol) - - -IN_CHANNELS = 3 -KT, KH, KW = 2, 16, 16 -EMBED_DIM = 1152 -PATCH_NUMEL = IN_CHANNELS * KT * KH * KW # 1536 - - -def _operands( - num_patches, - dtype=torch.bfloat16, - seed=0, - in_channels=IN_CHANNELS, - kernel=(KT, KH, KW), - embed_dim=EMBED_DIM, - bias=True, -): - """Pixels in [-1, 1] (processor mean/std 0.5), weights at the tower's - initializer_range 0.02.""" - g = torch.Generator(device="cuda").manual_seed(seed) - kt, kh, kw = kernel - pixel_values = torch.empty(num_patches, in_channels * kt * kh * kw, device="cuda") - pixel_values = pixel_values.uniform_(-1.0, 1.0, generator=g).to(dtype) - weight = ( - 0.02 * torch.randn(embed_dim, in_channels, kt, kh, kw, generator=g, device="cuda") - ).to(dtype) - bias_t = (0.02 * torch.randn(embed_dim, generator=g, device="cuda")).to(dtype) if bias else None - return pixel_values, weight, bias_t - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying): # noqa: F811 - certifying(op) - spec = cell.spec - pixel_values, weight, bias = _operands( - spec["num_patches"], dtype=spec["dtype"], seed=zlib.crc32(cell.why.encode()) - ) - before = pixel_values.clone() - out = op(pixel_values, weight, bias) - assert out.shape == (pixel_values.shape[0], weight.shape[0]) - assert out.dtype == pixel_values.dtype and out.is_contiguous() - assert out.data_ptr() != pixel_values.data_ptr() - assert torch.equal(pixel_values, before), "the pixels are read-only" - _compare(out, op.reference(pixel_values, weight, bias)) - - -def test_guard_refuses_rows_that_are_not_one_patch() -> None: - """Two patches per row are split into twice the rows, without an error.""" - pixel_values, weight, bias = _operands(8, seed=8) - doubled = pixel_values.reshape(4, 2 * PATCH_NUMEL) - assert op(doubled, weight, bias).shape == (8, EMBED_DIM) - with pytest.raises(AssertionError, match="each row must be exactly one patch"): - with validating(op): - op(doubled, weight, bias) - - -def test_cuda_graph_capture_and_replay() -> None: - pixel_values, weight, bias = _operands(64, seed=12) - - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - for _ in range(3): - op(pixel_values, weight, bias) - torch.cuda.current_stream().wait_stream(warmup) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured = op(pixel_values, weight, bias) - - pixel_values.copy_(_operands(64, seed=13)[0]) - graph.replay() - torch.cuda.synchronize() - _compare(captured, op.reference(pixel_values, weight, bias)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_pos_embed_interpolate.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_pos_embed_interpolate.py deleted file mode 100644 index 3517e2f3d18e..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_pos_embed_interpolate.py +++ /dev/null @@ -1,110 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qwen4_exp pos_embed_interpolate entry. - -`op.reference` is the explicit four-corner sum in fp64 -- exact, but sharing -the entry's sample positions, so it gates numerics. `_ref_resize` here runs -torch's own `F.interpolate` in fp64 and shares no arithmetic with the entry, so -it gates the structure: which axis is which, how the table is read, and the -spatial-merge row order. It sits up to 8.6e-6 from the exact value (its sample -coordinates are float): far inside a bf16 cell's band, though not an fp32 one's. -Beyond the cells: the square-table guard and graph capture. -""" - -import zlib - -import pytest -import torch -import torch.nn.functional as F - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.vision.pos_embed_interpolate import ( # noqa: E501 - pos_embed_interpolate as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import certifying, validating # noqa: E402,F401 - -assert torch.cuda.is_available(), "pos_embed_interpolate requires a CUDA device" - -NUM_GRID = 48 -NUM_POS = NUM_GRID * NUM_GRID -HIDDEN = 1152 -MERGE = 2 - - -def _table(rows=NUM_POS, hidden=HIDDEN, dtype=torch.bfloat16, seed=0): - g = torch.Generator(device="cuda").manual_seed(seed) - return torch.randn(rows, hidden, generator=g, device="cuda").to(dtype) - - -def _merge_order(x, height, width, merge): - hidden = x.shape[-1] - return ( - x.reshape(height // merge, merge, width // merge, merge, hidden) - .permute(0, 2, 1, 3, 4) - .reshape(height * width, hidden) - ) - - -def _ref_resize(table, grid_thw, merge): - """Structural reference via torch's resize kernel, in fp64.""" - num_grid = int(round(table.shape[0] ** 0.5)) - hidden = table.shape[1] - src = table.double().view(num_grid, num_grid, hidden).permute(2, 0, 1).unsqueeze(0).contiguous() - outs = [] - for frames, height, width in grid_thw: - resized = F.interpolate(src, size=(height, width), mode="bilinear", align_corners=True) - rows = _merge_order(resized[0].permute(1, 2, 0), height, width, merge) - outs.append(rows.repeat(frames, 1) if frames > 1 else rows) - return torch.cat(outs) if len(outs) > 1 else outs[0] - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell, certifying): # noqa: F811 - certifying(op) - spec = cell.spec - table = _table(spec["rows"], spec["hidden"], spec["dtype"], seed=zlib.crc32(cell.why.encode())) - grid = [list(g) for g in spec["grid"]] - before = table.clone() - out = op(table, grid, spec["merge"]) - assert out.shape == (sum(t * h * w for t, h, w in grid), spec["hidden"]) - assert out.dtype == table.dtype - assert torch.equal(table, before) and out.data_ptr() != table.data_ptr() - torch.testing.assert_close(out, op.reference(table, grid, spec["merge"])) - # The structural gate, from a reference that shares no arithmetic. - torch.testing.assert_close(out, _ref_resize(table, grid, spec["merge"]).to(out.dtype)) - - -def test_guard_refuses_a_non_square_table() -> None: - """The hazard, unguarded: 2300 rows -> isqrt 47, every index in range, a - correctly shaped result off the wrong grid.""" - table = _table(rows=2300, hidden=64, seed=10) - out = op(table, [(1, 4, 4)], MERGE) - assert out.shape == (16, 64) - assert torch.equal(out, op(table[: 47 * 47], [(1, 4, 4)], MERGE)), "read as a 47x47 grid" - - with pytest.raises(AssertionError, match="must hold a square grid"): - with validating(op): - op(table, [(1, 4, 4)], MERGE) - - -def test_cuda_graph_capture_and_replay() -> None: - """The grid is frozen at capture; the table's contents are not.""" - table = _table(seed=12) - grid = [(1, 8, 12)] - - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - for _ in range(3): - op(table, grid, MERGE) - torch.cuda.current_stream().wait_stream(warmup) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured = op(table, grid, MERGE) - - table.copy_(_table(seed=13)) - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(captured, op.reference(table, grid, MERGE)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_vision_tower_calls.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_vision_tower_calls.py deleted file mode 100644 index d8a1f9396844..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_vision_tower_calls.py +++ /dev/null @@ -1,165 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the qwen4_exp vision tower's call grouping. - -A request's frames run in calls of whole frames under `_VISION_CALL_MAX_PATCHES` rows. -The tower here has the checkpoint's geometry (1152 wide, 16 heads of 72, 2x2 merge) with -two random blocks, so every op the real tower runs runs here; the cap is a module -constant, lowered so a few hundred rows span several calls. The reference is the same -tower in one call. - -Cutting between frames leaves attention, the position embedding, the rope and every -row-wise op unchanged; what can differ is which kernel a call picks (a GEMM's split, the -fused attention kernel for a call whose longest frame is at most 64 rows), within the ops' -reference bands. -""" - -from types import SimpleNamespace - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - modeling as M, -) - -assert torch.cuda.is_available(), "the vision tower requires a CUDA device" - -HIDDEN = 1152 -HEADS = 16 -HEAD_DIM = HIDDEN // HEADS -INTER = 4304 -MERGE = 2 -MERGE_DIM = HIDDEN * MERGE * MERGE -OUT = 2560 -POS_ROWS = 48 * 48 -PIXEL_WIDTH = 3 * 2 * 16 * 16 -DEPTH = 2 -DT = torch.bfloat16 - - -def _core(): - g = torch.Generator(device="cuda").manual_seed(0) - - def w(*shape, scale=0.02): - return (scale * torch.randn(*shape, generator=g, device="cuda")).to(DT) - - def norm_w(): - return (1.0 + 0.1 * torch.randn(HIDDEN, generator=g, device="cuda")).to(DT) - - weights = { - "embed": torch.zeros(1, dtype=DT, device="cuda"), - "vis_patch_w": w(HIDDEN, 3, 2, 16, 16), - "vis_patch_b": w(HIDDEN), - "vis_pos": w(POS_ROWS, HIDDEN, scale=0.5), - "vis_m_n_w": norm_w(), - "vis_m_n_b": w(HIDDEN), - "vis_m_fc1_w": w(MERGE_DIM, MERGE_DIM), - "vis_m_fc1_b": w(MERGE_DIM), - "vis_m_fc2_w": w(OUT, MERGE_DIM), - "vis_m_fc2_b": w(OUT), - } - for b in range(DEPTH): - weights.update( - { - f"l{b}_vis_n1_w": norm_w(), - f"l{b}_vis_n1_b": w(HIDDEN), - f"l{b}_vis_qkv_w": w(3 * HIDDEN, HIDDEN), - f"l{b}_vis_qkv_b": w(3 * HIDDEN), - f"l{b}_vis_proj_w": w(HIDDEN, HIDDEN), - f"l{b}_vis_proj_b": w(HIDDEN), - f"l{b}_vis_n2_w": norm_w(), - f"l{b}_vis_n2_b": w(HIDDEN), - f"l{b}_vis_fc1_w": w(INTER, HIDDEN), - f"l{b}_vis_fc1_b": w(INTER), - f"l{b}_vis_fc2_w": w(HIDDEN, INTER), - f"l{b}_vis_fc2_b": w(HIDDEN), - } - ) - return SimpleNamespace( - w=weights, - v_depth=DEPTH, - v_hidden=HIDDEN, - v_heads=HEADS, - v_head_dim=HEAD_DIM, - v_merge=MERGE, - v_merge_dim=MERGE_DIM, - v_pixel_width=PIXEL_WIDTH, - ) - - -def _pixels(grid, seed=1): - patches = sum(t * h * w for t, h, w in grid) - g = torch.Generator(device="cuda").manual_seed(seed) - return torch.empty(patches, PIXEL_WIDTH, device="cuda").uniform_(-1, 1, generator=g).to(DT) - - -def _count_calls(tower, monkeypatch): - """Wrap `tower._run`; the returned list gets one `(rows, grid)` entry per call.""" - calls = [] - run = tower._run - - def counted(pixel_values, grid, grid_thw): - calls.append((pixel_values.shape[0], list(grid))) - return run(pixel_values, grid, grid_thw) - - monkeypatch.setattr(tower, "_run", counted) - return calls - - -def _assert_close(out, ref): - """Two runs of the tower over different row counts differ by bf16 rounding: another - call picks another GEMM kernel, and the difference compounds through the blocks and the - merger. Measured here: 0.6% RMS relative error and at most 8 ulps (of the RMS) in one - element, against 117% RMS between unrelated rows. The gates sit between the two.""" - assert out.dtype == ref.dtype and out.shape == ref.shape, (out.shape, ref.shape) - error = out.float() - ref.float() - rms = ref.float().pow(2).mean().sqrt() - assert error.pow(2).mean().sqrt() <= 0.02 * rms, (error.pow(2).mean().sqrt(), rms) - assert error.abs().max() <= 16 * 2**-8 * rms, (error.abs().max(), rms) - - -@pytest.fixture(scope="module") -def tower(): - return M.VisionTower(_core()) - - -# Frames of 256, 144 and 32 rows under a 512-row cap: three calls -- the first item, the -# second with two of the third's six frames, the other four -- so a cut between items and one -# inside an item; the last call's longest frame is 32 rows, where the fused attention kernel -# is the short-frame one. -GRID, CAP = [(2, 16, 16), (3, 12, 12), (6, 4, 8)], 512 - - -def test_calls_of_whole_frames_equal_one_call(tower, monkeypatch): - grid, cap = GRID, CAP - pixels = _pixels(grid) - grid_thw = torch.tensor(grid) - monkeypatch.setattr(M, "_VISION_CALL_MAX_PATCHES", 1 << 20) - whole = tower(pixels, grid_thw) - - monkeypatch.setattr(M, "_VISION_CALL_MAX_PATCHES", cap) - calls = _count_calls(tower, monkeypatch) - chunked = tower(pixels, grid_thw) - - expected = M._frame_calls([tuple(row) for row in grid], cap) - assert len(calls) == len(expected) > 1 - assert all(rows <= cap for rows, _ in calls) - _assert_close(chunked, whole) - - -def test_a_request_past_the_old_whole_request_ceiling_runs(tower): - # 66 frames of 512x512 pixels: 33 temporal groups of 32x32 patches, 33792 rows in two - # calls (32 groups, then 1). The last group's rows, alone in its call, and the first - # group's, in a call with 31 others, equal each group run on its own. - grid = [(33, 32, 32)] - pixels = _pixels(grid) - out = tower(pixels, torch.tensor(grid)) - assert out.shape == (33 * 32 * 32 // 4, OUT) - assert torch.isfinite(out).all() - group_rows, group_out = 32 * 32, 32 * 32 // 4 - for group in (0, 32): - alone = tower( - pixels[group * group_rows : (group + 1) * group_rows], torch.tensor([(1, 32, 32)]) - ) - _assert_close(out[group * group_out : (group + 1) * group_out], alone) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_triton_update.py similarity index 51% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_triton_update.py index 7f8a993b6785..a289411fbafb 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_triton_update.py @@ -7,14 +7,14 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_triton_update import ( # noqa: E501 - ENGINE_SLOTS, +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.causal_conv1d_triton_update import ( # noqa: E501 + CKPT_X_STRIDE, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_triton_update import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.causal_conv1d_triton_update import ( # noqa: E501 causal_conv1d_triton_update as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "causal_conv1d_triton_update requires a CUDA device" @@ -29,8 +29,10 @@ def _compare(out: tuple, ref: tuple) -> None: the reference marks the rest): against the operand-dtype-product reference bf16/fp16 were bit-exact at dim 512 and at dim 10240 off by one ulp in at most 9 of 2.6M elements (the fp32 `exp`); fp32 is off by - 1.4e-6. The pool and the snapshots are copies, gated bitwise over the - whole buffer -- the lock on every unnamed slot and snapshot row. + 1.4e-6. At a tep4 rank's 2560 channels the same one-ulp differences + touch a handful of elements. The pool and the snapshots are copies, gated + bitwise over the whole buffer -- the lock on every unnamed slot and + snapshot row. """ (act, state, inter), (ref_act, ref_state, ref_inter) = out, ref assert act.dtype == ref_act.dtype and act.shape == ref_act.shape @@ -46,27 +48,44 @@ def _compare(out: tuple, ref: tuple) -> None: def _build(spec, seed): - """`__call__`'s positional arguments for one spec, freshly allocated.""" + """`__call__`'s positional arguments for one spec, freshly allocated, plus the + buffers its views are cut from. + + x is the q|k|v columns of an `[n * p, x_row_stride]` projection output + (z | a | b past them), reshaped `[n, p, C]` and seen as `[n, C, p]`; the + pool is the last layer's view of a `[slots, pool_layers, C, 3]` page, as the + V2 manager coalesces the GDN layers' conv states inside each slot; the + snapshot buffer is the first layer's view of a `[snapshot_layers, rows, + steps, C, 3]` buffer, as the GDN speculative state holds every layer's. + """ g = torch.Generator(device="cuda").manual_seed(seed) dim, batch, seqlen, width = spec["dim"], spec["batch"], spec["seqlen"], 4 + wide = spec.get("x_row_stride", CKPT_X_STRIDE) + layers = spec.get("pool_layers", 1) def randn(*shape): return torch.randn(*shape, generator=g, device="cuda").to(torch.bfloat16) - # The call site under the one input GEMM: the q|k|v columns of an - # [n * p, C + 6240] projection, reshaped [n, p, C] and seen as [n, C, p]. - wide = dim + 6144 + 2 * 48 - x = randn(batch * seqlen, wide)[:, :dim].reshape(batch, seqlen, dim).transpose(1, 2) + buf = randn(batch * seqlen, wide) + x = buf[:, :dim].reshape(batch, seqlen, dim).transpose(1, 2) assert x.stride() == (seqlen * wide, 1, wide) - conv_state = randn(spec["pool_slots"], dim, width - 1) + pool_buf = randn(spec["pool_slots"], layers, dim, width - 1) + conv_state = pool_buf[:, layers - 1] + assert conv_state.stride() == (layers * dim * (width - 1), width - 1, 1) weight = randn(dim, width) indices = torch.tensor(spec["slots"], dtype=torch.int32, device="cuda") steps = spec.get("snapshot_steps", seqlen) - window = torch.zeros( - spec["snapshots"], steps, dim, width - 1, dtype=torch.bfloat16, device="cuda" + window_buf = torch.zeros( + spec.get("snapshot_layers", 1), + spec["snapshots"], + steps, + dim, + width - 1, + dtype=torch.bfloat16, + device="cuda", ) rows_idx = torch.arange(batch, dtype=torch.int32, device="cuda") - return [ + args = [ x, conv_state, weight, @@ -75,42 +94,37 @@ def randn(*shape): indices, PAD_SLOT_ID, None, - window, + window_buf[0], rows_idx, - False, + spec.get("launch_dependent_kernels", False), ] + return args, dict(buf=buf, pool_buf=pool_buf, window_buf=window_buf) -def _run(args): - """Drive the op on copies of the two buffers; return the triple compare takes.""" +def _run(args, extras): + """Drive the op on copies of the pool page and the snapshot buffer, through + the same views; return the triple compare takes.""" got = list(args) - got[1], got[8] = args[1].clone(), args[8].clone() + pool_buf, window_buf = extras["pool_buf"].clone(), extras["window_buf"].clone() + got[1], got[8] = pool_buf[:, pool_buf.shape[1] - 1], window_buf[0] + buf = extras["buf"].clone() returned = op(*got) + assert torch.equal(extras["buf"], buf), "x (and the columns past it) are read-only" + assert torch.equal(pool_buf[:, :-1], extras["pool_buf"][:, :-1]), "another layer written" + assert torch.equal(window_buf[1:], extras["window_buf"][1:]), "another layer's snapshots" return returned, got[1], got[8] @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name certifying(op) - args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) - x, weight = args[0].clone(), args[2].clone() + args, extras = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + weight = args[2].clone() ref = op.reference(*args) - got = _run(args) + got = _run(args, extras) assert got[0] is not args[0] and got[0].shape == args[0].shape - assert torch.equal(args[0], x) and torch.equal(args[2], weight), "x and weight are read-only" + assert torch.equal(args[2], weight), "weight is read-only" _compare(got, ref) - if cell.spec["pool_slots"] == ENGINE_SLOTS: - for slot in cell.spec["slots"]: # non-vacuity of the whole-pool gate - assert not torch.equal(got[1][slot], args[1][slot]), slot - - -def test_guard_refuses_undersized_buffers() -> None: - """A snapshot buffer shallower than seqlen is accepted by the kernel and - written past the row it was given (measured); the guard refuses it - before the launch.""" - spec = dict(dim=512, batch=2, seqlen=2, pool_slots=8, slots=(0, 2), snapshots=2) - args = _build(spec, seed=15) - args[8] = torch.zeros(2, 1, 512, 3, dtype=torch.bfloat16, device="cuda") - with pytest.raises(AssertionError, match="intermediate_conv_window must be at least"): - with validating(op): - op(*args) + # Non-vacuity of the whole-pool gate: every slot the batch named moved. + for slot in cell.spec["slots"]: + assert not torch.equal(got[1][slot], args[1][slot]), slot diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_update.py similarity index 68% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_update.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_update.py index bbd3b86e9575..d958f21cb0e3 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_update.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_causal_conv1d_update.py @@ -7,14 +7,18 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_update import ( # noqa: E501 - ENGINE_SLOTS, +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.causal_conv1d_update import ( # noqa: E501 + CKPT_DIM, + CKPT_X_STRIDE, + RANK_DIM, + RANK_X_STRIDE, + _decode_conv_ok, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_update import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.causal_conv1d_update import ( # noqa: E501 causal_conv1d_update as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "causal_conv1d_update requires a CUDA device" @@ -41,23 +45,29 @@ def _compare( PAD_SLOT_ID = -1 #: Rows of prefill ahead of the decode rows in the in_proj output. _PREFILL_ROWS = 37 -#: Columns past q|k|v in the target's one GDN input-projection output: z | a | b. -_IN_PROJ_TAIL = 6144 + 2 * 48 def _build(spec, seed): - """`__call__`'s arguments for one spec, plus the buffers x and the pool view.""" + """`__call__`'s arguments for one spec, plus the buffers x and the pool view. + + x is the first `dim` columns of the decode rows of a `[rows, x_row_stride]` + projection output (its z | a | b columns past them); the pool is the last + layer's view of a `[slots, pool_layers, dim, 3]` page, as the V2 manager + coalesces the GDN layers' conv states inside each slot. + """ g = torch.Generator(device="cuda").manual_seed(seed) dim, batch, layers = spec["dim"], spec["batch"], spec.get("pool_layers", 1) + row_stride = spec.get("x_row_stride", CKPT_X_STRIDE) def randn(*shape): return torch.randn(*shape, generator=g, device="cuda").to(torch.bfloat16) - buf = randn(_PREFILL_ROWS + batch, dim + _IN_PROJ_TAIL) + buf = randn(_PREFILL_ROWS + batch, row_stride) x = buf[_PREFILL_ROWS:, :dim] - assert x.stride() == (dim + _IN_PROJ_TAIL, 1) + assert x.stride() == (row_stride, 1) pool_buf = randn(spec["pool_slots"], layers, dim, 3) conv_state = pool_buf[:, layers - 1] + assert conv_state.stride() == (layers * dim * 3, 3, 1) weight = randn(dim, 4) slots = spec["slots"] if slots == "random": @@ -93,19 +103,21 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res returned, state = _run(args, extras) assert torch.equal(args[2], weight), "weight is read-only" _compare((returned, state), ref) - if cell.spec["pool_slots"] == ENGINE_SLOTS: - # Non-vacuity for the large pool: every named slot moved. - for slot in cell.spec["slots"]: - if slot != PAD_SLOT_ID: - assert not torch.equal(state[slot], args[1][slot]), slot - - -def test_op_rejects_metadata_violations_loudly() -> None: - """Why the entry has no `is_valid`: the op raises on metadata violations - itself (probed here at int64 slot indices).""" - spec = dict(dim=512, batch=4, pool_slots=16, slots=(0, 1, 2, 3)) - args = _build(spec, seed=14)[0] - args[6] = args[6].long() - with pytest.raises(RuntimeError, match="conv_state_indices.scalar_type"): - with validating(op): - op(*args) + # Non-vacuity: every slot the batch named moved (the whole-pool gate above + # holds every other slot, and every pad row, unchanged). + for slot in args[6].tolist(): + if slot != PAD_SLOT_ID: + assert not torch.equal(state[slot], args[1][slot]), slot + + +def test_which_kernel_each_member_runs() -> None: + """The path the note states each member's cells certify: a tp4/tep4 rank's 2560 + channels never reach the CuTe kernel (2560 is not a multiple of 1024), so its 128-row + cells certify the TRT-LLM kernel; a dep4 rank's 10240 do from 64 rows.""" + for dim, row_stride, expect in ( + (RANK_DIM, RANK_X_STRIDE, False), + (CKPT_DIM, CKPT_X_STRIDE, True), + ): + spec = dict(dim=dim, x_row_stride=row_stride, batch=128, pool_slots=256, slots="random") + args = _build(spec, seed=dim)[0] + assert _decode_conv_ok(*args[:7]) is expect, (dim, expect) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py similarity index 63% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py index 64926845e39c..cc65c056494b 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py @@ -12,11 +12,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.cu_seqlens_from_seq_lens import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.cu_seqlens_from_seq_lens import ( # noqa: E501 cu_seqlens_from_seq_lens as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "cu_seqlens_from_seq_lens requires a CUDA device" @@ -70,36 +70,3 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res assert bounds[-1] == packed.shape[0] for s, piece in enumerate(pieces): assert torch.equal(packed[bounds[s] : bounds[s + 1]], piece), f"sequence {s} mis-sliced" - - -def test_cuda_graph_capture_and_replay() -> None: - """A captured graph follows lengths rewritten in place: only the shape - reaches Python, the values are read on device and never synchronized on. - 32 requests of 8 rows (the MTP-7 verify's 256 rows).""" - batch, p = 32, 8 - static_lens = _seq_lens([p] * batch) - side = torch.cuda.Stream() - side.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(side): - for _ in range(3): - op(static_lens, torch.int32) - torch.cuda.current_stream().wait_stream(side) - torch.cuda.synchronize() - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - static_out = op(static_lens, torch.int32) - for lengths in ([1] * batch, [p] * batch, [p + 1] + [1] * (batch - 1)): - static_lens.copy_(_seq_lens(lengths)) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(static_out, op.reference(static_lens, torch.int32)) - - -def test_guard_refuses_non_integer_seq_lens() -> None: - """Unguarded, a float32 seq_lens is truncated toward zero before the scan.""" - floats = torch.tensor([7.9, 1.2, 13.6], dtype=torch.float32, device=DEV) - assert op(floats, torch.int32).tolist() == [0, 7, 8, 21], "hazard gone" - with pytest.raises(AssertionError, match="seq_lens must have an integer dtype"): - with validating(op): - op(floats, torch.int32) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py new file mode 100644 index 000000000000..da97829e7f5d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_flashinfer_chunk_gated_delta_rule.py @@ -0,0 +1,377 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the flashinfer_chunk_gated_delta_rule entry (FlashInfer GDN prefill). + +A cell's pool is one of two: + +- a contiguous pool (tp1's cells), compared whole so slots the call did not name are + gated too; +- with `page_layers`, the stock cache manager's bf16 state page at a member's measured + slot count: `page_layers` GDN layers' states a slot, the view at the last layer's + sub-page. A tep4 page spans 26.5 GB and a dep4 page 48 GB, so it is a reserved virtual + range with physical memory mapped under the named slots only (`_SparsePage`). The + reference runs on a dense copy of those slots, and every other byte of their pages is + checked unchanged; a write to any other slot would fault. + +Beyond the cells: the has_initial_state equivalence. +""" + +import warnings +import zlib + +import cuda.bindings.driver as cuda_drv +import pytest +import torch +from flashinfer.gdn_kernels.delta_rule_dsl.varlen_helper import should_use_cp_host +from flashinfer.utils import get_compute_capability, get_device_name, get_device_sm_count + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.flashinfer_chunk_gated_delta_rule import ( # noqa: E501 + INPLACE_MAX_TOKENS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.flashinfer_chunk_gated_delta_rule import ( # noqa: E501 + flashinfer_chunk_gated_delta_rule as op, +) + +__extra_import_path__ = ["../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "flashinfer_chunk_gated_delta_rule requires a CUDA device" + +DEV = "cuda" +#: The band the Triton chunk op is certified at (rtol bf16 default; atol 2e-3 +#: on `o`, 2e-2 on states), derived for a chunked WY reformulation with bf16 +#: operands. This op is the same class of sum and sits inside it against fp64, CP +#: pipeline and chunked kernel alike; a value head reading the wrong key head at 4/12 +#: lands far outside it (measured). +RTOL, ATOL_OUTPUT, ATOL_STATE = 1.6e-2, 2e-3, 2e-2 +#: Rows past the call's own in the output buffer of a `slice` cell: the target passes +#: `out[:, :num_tokens]` of a wider buffer, and the rows past it must stay as they were. +_OUTPUT_TAIL = 16 + + +def _compare(out: tuple, ref: tuple) -> None: + """`(o, final_state, pool)` against `reference`'s triple; `None` must match.""" + for got, want, atol in zip(out, ref, (ATOL_OUTPUT, ATOL_STATE, ATOL_STATE)): + assert (got is None) == (want is None), (got is None, want is None) + if got is None: + continue + assert got.dtype == want.dtype and got.shape == want.shape, ( + (got.dtype, got.shape), + (want.dtype, want.shape), + ) + assert torch.isfinite(got).all() + torch.testing.assert_close(got.float(), want.float(), rtol=RTOL, atol=atol) + + +def _stats(name: str, got: torch.Tensor, want: torch.Tensor) -> str: + """Relative L2, cosine, p99.9 and max absolute error of `got` against `want`.""" + g, w = got.double().flatten(), want.double().flatten() + err = (g - w).abs() + rel = (err.norm() / w.norm().clamp_min(1e-30)).item() + cos = torch.nn.functional.cosine_similarity(g, w, dim=0).item() + p999 = torch.quantile(err[:: max(1, err.numel() // 2_000_000)].float(), 0.999).item() + return f"{name} rel-L2 {rel:.3e} cos {cos:.6f} p99.9 {p999:.3e} max {err.max().item():.3e}" + + +def _flashinfer_path(spec) -> str: + """Which FlashInfer kernel the call reaches: `inplace` (the chunked kernel over the pool + itself) for a `has_initial_state` call of at most `INPLACE_MAX_TOKENS` tokens, else `cp` + (the context-parallel pipeline) when FlashInfer's own heuristic picks it for + `num_seqs x num_v_heads` on this device.""" + if "initial" in spec and sum(spec["seqlens"]) <= INPLACE_MAX_TOKENS: + return "inplace" + device = torch.device(DEV) + cp = should_use_cp_host( + len(spec["seqlens"]) * spec["geometry"][1], + get_device_sm_count(device), + get_device_name(device), + device_capability=get_compute_capability(device), + ) + return "cp" if cp else "chunked" + + +def _check(result): + err, *value = result + assert err == cuda_drv.CUresult.CUDA_SUCCESS, err + return value[0] if len(value) == 1 else value + + +class _ArrayInterface: + def __init__(self, ptr: int, shape, strides_bytes): + self.__cuda_array_interface__ = dict( + shape=tuple(shape), + strides=tuple(strides_bytes), + typestr=" torch.Tensor: + """bf16 tensor at element `offset`, element `strides`.""" + raw = torch.as_tensor( + _ArrayInterface(self.va + self.ELEM * offset, shape, [self.ELEM * s for s in strides]), + device=DEV, + ) + return raw.view(torch.bfloat16) + + def close(self) -> None: + torch.cuda.synchronize() + for ptr, size, handle in self.maps: + _check(cuda_drv.cuMemUnmap(ptr, size)) + _check(cuda_drv.cuMemRelease(handle)) + _check(cuda_drv.cuMemAddressFree(cuda_drv.CUdeviceptr(self.va), self.size)) + + +class _EnginePool: + """The stock manager's coalesced bf16 state page at `slots` slots (slot stride: + `layers` `[HV, V, K]` states), this layer's view the last layer's sub-page, mapped + under the named slots: each one's whole page and the next slot's first layer. The + view's own base (slot 0's sub-page) is mapped too: torch reads the device off a + tensor's base pointer.""" + + def __init__(self, slots: int, layers: int, used, v_heads: int, head_v: int, head_k: int, gen): + layer = v_heads * head_v * head_k + stride = layers * layer + total = slots * stride + base = (layers - 1) * layer + self.used = sorted(set(used)) + self.spans = [(s * stride, min((s + 1) * stride + layer, total)) for s in self.used] + self.page = _SparsePage(total, [*self.spans, (base, base + layer)]) + self.own = base + self.layer = layer + self.pool = self.page.view( + base, (slots, v_heads, head_v, head_k), (stride, head_v * head_k, head_k, 1) + ) + self.flat = [self.page.view(lo, (hi - lo,), (1,)) for lo, hi in self.spans] + for f in self.flat: + f.copy_(torch.randn(f.numel(), generator=gen, device=DEV) * 0.1) + torch.cuda.synchronize() + self.before = [f.clone() for f in self.flat] + + def dense(self, slots: torch.Tensor): + """A dense copy of the named slots and `slots` renumbered into it.""" + index = {s: i for i, s in enumerate(self.used)} + compact = torch.stack([self.pool[s] for s in self.used]).contiguous() + renamed = [index[s] for s in slots.tolist()] + return compact, torch.tensor(renamed, dtype=slots.dtype, device=DEV) + + def touched(self) -> torch.Tensor: + return torch.stack([self.pool[s] for s in self.used]) + + def assert_rest_unchanged(self) -> None: + """Every byte of a named slot's page but this layer's state, and the next slot's + first layer, is what it was.""" + torch.cuda.synchronize() + for f, b, s in zip(self.flat, self.before, self.used): + keep = torch.ones(f.numel(), dtype=torch.bool, device=DEV) + keep[self.own : self.own + self.layer] = False + assert torch.equal(f[keep], b[keep]), f"slot {s}: the call wrote outside its state" + + def close(self) -> None: + del self.pool, self.flat + self.page.close() + + +def _build(spec, seed): + """`(args, engine)`: keyword arguments in the domain a GDN post-conv step produces + (unit-norm q/k, g < 0 in log space, beta a sigmoid), and the `_EnginePool` behind the + pool (or None).""" + gen = torch.Generator(device=DEV).manual_seed(seed) + num_k_heads, num_v_heads, head_k_dim, head_v_dim = spec["geometry"] + seqlens = spec["seqlens"] + total = sum(seqlens) + q = torch.randn(1, total, num_k_heads, head_k_dim, generator=gen, device=DEV) + k = torch.randn(1, total, num_k_heads, head_k_dim, generator=gen, device=DEV) + q = q / torch.sqrt(q.pow(2).sum(-1, keepdim=True) + 1e-6) + k = k / torch.sqrt(k.pow(2).sum(-1, keepdim=True) + 1e-6) + q, k = q.bfloat16(), k.bfloat16() + v = torch.randn( + 1, total, num_v_heads, head_v_dim, generator=gen, device=DEV, dtype=torch.bfloat16 + ) + a_log = torch.log(torch.rand(num_v_heads, generator=gen, device=DEV) * 15 + 1) + x = torch.randn(1, total, num_v_heads, generator=gen, device=DEV) + torch.randn( + num_v_heads, generator=gen, device=DEV + ) + g = (-torch.exp(a_log) * torch.nn.functional.softplus(x)).float() + beta = torch.rand(1, total, num_v_heads, generator=gen, device=DEV).sigmoid().float() + engine = None + if spec.get("page_layers"): + engine = _EnginePool( + spec["slots"], + spec["page_layers"], + spec["indices"], + num_v_heads, + head_v_dim, + head_k_dim, + gen, + ) + pool = engine.pool + else: + pool = ( + torch.randn( + spec["slots"], num_v_heads, head_v_dim, head_k_dim, generator=gen, device=DEV + ) + * 0.1 + ).to(spec.get("pool_dtype", torch.float32)) + cu = torch.tensor([0] + torch.tensor(seqlens).cumsum(0).tolist(), dtype=torch.long, device=DEV) + indices = torch.tensor(spec["indices"], dtype=torch.int32, device=DEV) + output = None + if spec.get("output") == "buffer": + output = torch.full( + (1, total, num_v_heads, head_v_dim), -7.0, device=DEV, dtype=torch.bfloat16 + ) + elif spec.get("output") == "slice": + output = torch.full( + (1, total + _OUTPUT_TAIL, num_v_heads, head_v_dim), + -7.0, + device=DEV, + dtype=torch.bfloat16, + )[:, :total] + extra = {} + if "initial" in spec: + extra["has_initial_state"] = torch.tensor(spec["initial"], dtype=torch.bool, device=DEV) + if spec.get("linear_gate"): + extra["linear_gate"] = True + g = torch.exp(g) + args = dict( + **extra, + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=None, + initial_state=pool, + initial_state_indices=indices, + inplace_indexed_state_update=True, + output_final_state=False, + cu_seqlens=cu, + head_first=False, + use_qk_l2norm_in_kernel=False, + output=output, + state_workspace=None, + ) + return args, engine + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + args, engine = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + try: + if engine is None: + pool = args["initial_state"] + pool_before = pool.clone() + ref = op.reference(**args) + else: + compact, renamed = engine.dense(args["initial_state_indices"]) + ref = op.reference(**dict(args, initial_state=compact, initial_state_indices=renamed)) + output = args["output"] + # A `slice` cell's output is a prefix view of the wider buffer the target owns. + full = None if output is None else output._base + tail_before = None if full is None else full[:, output.shape[1] :].clone() + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + o, final_state = op(**args) + # FlashInfer falls back from a CP heuristic match with a RuntimeWarning, so a CP + # cell certifies the CP pipeline only if none was raised. + fallback = [str(w.message) for w in caught if "falling back to non-CP" in str(w.message)] + path = "chunked" if fallback else _flashinfer_path(cell.spec) + + got_pool = pool if engine is None else engine.touched() + _compare((o, final_state, got_pool), ref) + print( + f"\n[{cell.why}] path={path} sms={get_device_sm_count(torch.device(DEV))}; " + f"{_stats('o', o, ref[0])}; {_stats('state', got_pool, ref[2])}" + ) + if "fi_path" in cell.spec: + assert path == cell.spec["fi_path"], (path, fallback) + if output is not None: + assert o.data_ptr() == output.data_ptr(), "output= must be written in place" + if tail_before is not None: + assert torch.equal(full[:, output.shape[1] :], tail_before), "rows past output written" + if engine is None: + named = set(args["initial_state_indices"].tolist()) + unnamed = [s for s in range(pool.shape[0]) if s not in named] + assert torch.equal(pool[unnamed], pool_before[unnamed]), "unnamed slots written" + else: + engine.assert_rest_unchanged() + finally: + if engine is not None: + engine.close() + + +_CARRIED = next(c for c in op.CELLS if "initial" in c.spec) + + +def test_has_initial_state_is_the_reset_then_the_plain_call() -> None: + """The flags replace `reset_recurrent_state_rows` before the call: zeroing the fresh + slots and running the plain in-place call gives the same output and pool, bitwise; a + linear gate that is `torch.exp` of the plain call's log gate changes nothing.""" + args, _ = _build(_CARRIED.spec, seed=5) + flags = args.pop("has_initial_state") + linear_gate = args.pop("linear_gate", False) + pool = args["initial_state"] + fresh = args["initial_state_indices"][~flags].long() + pool[fresh] = torch.randn_like(pool[fresh]) # stale rows the flags must not read + plain = dict(args, initial_state=pool.clone(), output=args["output"].clone()) + plain["initial_state"][fresh] = 0.0 + if linear_gate: + plain["g"] = torch.log(args["g"]) + # The linear gate of the target call is torch.exp of the log gate, bitwise. + args["g"] = torch.exp(plain["g"]) + + o_plain, _ = op(**plain) + o, final_state = op(**args, has_initial_state=flags, linear_gate=linear_gate) + + assert final_state is None + assert torch.equal(o, o_plain) + assert torch.equal(pool, plain["initial_state"]) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py similarity index 61% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py index 5ebd22d61d5e..d811aa537d70 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_fused_recurrent_gated_delta_rule_cached_replay_update.py @@ -19,17 +19,24 @@ token-by-token recurrence over every accepted token, independent of the wrap, the double buffer and the commit, so a wrong commit or wrong bookkeeping surfaces as an output that drifts off it. + +A cell with `page_layers` runs over the stock manager's coalesced state page at an +engine's slot count (36 layers' states a slot): 26.5 GB at tep4's 1874 bf16 slots, so a +reserved virtual range with memory mapped under the touched slots only (`_SparsePage`). +The reference and the twin run on a dense copy of those slots, and every byte of their +pages outside the cell's sub-pages is checked unchanged. """ from __future__ import annotations import zlib +import cuda.bindings.driver as cuda_drv import pytest import torch import torch.nn.functional as F -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_recurrent_gated_delta_rule_cached_replay_update import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.fused_recurrent_gated_delta_rule_cached_replay_update import ( # noqa: E501 ALL_LAYER_COMMIT_MIN_BATCH, HEAD_K_DIM, HEAD_V_DIM, @@ -38,7 +45,7 @@ NUM_V_HEADS, PACKED_QKV_WIDTH, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_recurrent_gated_delta_rule_cached_replay_update import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.fused_recurrent_gated_delta_rule_cached_replay_update import ( # noqa: E501 fused_recurrent_gated_delta_rule_cached_replay_update as op, ) from tensorrt_llm._torch.modules.fla.cache_manager import ( @@ -51,7 +58,7 @@ ) from tensorrt_llm._torch.modules.mamba.mamba2_metadata import _build_replay_work_items_triton -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import ULP, certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), ( @@ -80,6 +87,14 @@ #: Gate of the stored cumulative log-decays: |dG| <= atol + rtol * |G|. _G_ATOL, _G_RTOL = 2.0**-11, 2.0**-18 +#: The largest deviation each gated member reached in the running cell, in units of +#: its gate's measure (printed per cell, so the margins are on record). +_MEASURED: dict[str, float] = {} + + +def _record(key: str, value: float) -> None: + _MEASURED[key] = max(_MEASURED.get(key, 0.0), value) + def _gate(name, got, want, element, rms=None): """`got` against `want` per block of the last two dims (a token's @@ -92,12 +107,14 @@ def _gate(name, got, want, element, rms=None): scale = w.abs().amax(1, keepdim=True).clamp_min(torch.finfo(torch.float32).tiny) ratio = ((g - w).abs() / scale).nan_to_num(nan=float("inf")) worst = ratio.max().item() if ratio.numel() else 0.0 + _record(f"{name} element/gate", worst / element) assert worst <= element, ( f"{name}: an element is {worst:.3e} of its block's absmax from the reference " f"(gate {element:.3e})" ) if rms is not None and w.numel(): rel = ((g - w).pow(2).mean().sqrt() / w.pow(2).mean().sqrt().clamp_min(1e-30)).item() + _record(f"{name} rms/gate", rel / rms) assert rel <= rms, f"{name}: relative RMS {rel:.3e} exceeds {rms:.3e}" @@ -137,6 +154,8 @@ def _compare(out: tuple, ref: tuple) -> None: (want.dtype, tuple(want.shape)), ) if name == "old_G": + bound = _G_ATOL + _G_RTOL * want.abs() + _record("old_G |dG|/gate", ((got - want).abs() / bound).max().item()) torch.testing.assert_close(got, want, rtol=_G_RTOL, atol=_G_ATOL, msg=name) elif name == "ssm_states": bf16 = got.dtype == torch.bfloat16 @@ -150,13 +169,13 @@ def _seed(why): return zlib.crc32(why.encode()) -def _gating(gen, layers): +def _gating(gen, layers, hv=HV): """Per-layer A_log / dt_bias drawn around the checkpoint's own statistics; head 0 pinned to the corner that forgets everything every token (g ~ -400), head 1 to the one that hardly decays (g ~ -1e-8).""" (a_mean, a_std, a_range), (d_mean, d_std, d_range) = A_LOG_STATS, DT_BIAS_STATS - a_log = torch.randn(layers, HV, generator=gen, device=DEV) * a_std + a_mean - dt_bias = torch.randn(layers, HV, generator=gen, device=DEV) * d_std + d_mean + a_log = torch.randn(layers, hv, generator=gen, device=DEV) * a_std + a_mean + dt_bias = torch.randn(layers, hv, generator=gen, device=DEV) * d_std + d_mean a_log, dt_bias = a_log.clamp(*a_range), dt_bias.clamp(*d_range) a_log[:, 0], dt_bias[:, 0] = a_range[1], d_range[1] a_log[:, 1], dt_bias[:, 1] = a_range[0], d_range[0] @@ -167,6 +186,141 @@ def _unit(x): return (x / x.norm(dim=-1, keepdim=True)).to(torch.bfloat16) +def _scattered(n: int, slots: int) -> list[int]: + """`n` distinct slots of a `slots`-slot pool: both ends first, then a stride-37 walk + over the interior (37 is coprime to `slots - 2` for every pool here).""" + rest = [(i * 37 + 4) % (slots - 2) + 1 for i in range(max(n - 2, 0))] + picked = ([slots - 1, 0] + rest)[:n] + assert len(set(picked)) == n, (n, slots) + return picked + + +def _check(result): + err, *value = result + assert err == cuda_drv.CUresult.CUDA_SUCCESS, err + return value[0] if len(value) == 1 else value + + +#: Array-interface type of each pool dtype; bf16 has none, so it is read as int16 and viewed. +_TYPESTR = {torch.float32: (" torch.Tensor: + """Tensor of the page's dtype at element `offset`, element `strides`.""" + typestr, carrier = _TYPESTR[self.dtype] + es = self.es + t = torch.as_tensor( + _ArrayInterface(self.va + es * offset, shape, [es * s for s in strides], typestr), + device=DEV, + ) + assert t.dtype == carrier + return t.view(self.dtype) + + def close(self) -> None: + torch.cuda.synchronize() + for ptr, size, handle in self.maps: + _check(cuda_drv.cuMemUnmap(ptr, size)) + _check(cuda_drv.cuMemRelease(handle)) + _check(cuda_drv.cuMemAddressFree(cuda_drv.CUdeviceptr(self.va), self.size)) + + +class _EnginePage: + """The stock manager's coalesced state page at `slots` slots: a slot's page holds + `page_layers` layers' `[HV, V, K]` states, and the bench's `layers` are its top + sub-pages, the page's highest addresses. Memory is mapped under each touched slot's + page plus the next slot's first sub-page (where a write past the page would land), + and under slot 0's page, which holds the views' base.""" + + def __init__(self, slots, page_layers, layers, used, hv, dtype, gen): + state = hv * V * K + stride = page_layers * state + total = slots * stride + self.state, self.stride, self.layers = state, stride, layers + self.top = (page_layers - layers) * state + self.used = sorted(set(used)) + self.mapped = sorted(set(self.used) | {0}) + self.spans = [(s * stride, min((s + 1) * stride + state, total)) for s in self.mapped] + self.page = _SparsePage(total, self.spans, dtype) + self.pools = [ + self.page.view(self.top + j * state, (slots, hv, V, K), (stride, V * K, K, 1)) + for j in range(layers) + ] + self.flat = [self.page.view(lo, (hi - lo,), (1,)) for lo, hi in self.spans] + for f in self.flat: + f.normal_(0.0, 0.5, generator=gen) + torch.cuda.synchronize() + self.before = [f.clone() for f in self.flat] + + def assert_rest_unchanged(self) -> None: + """Every mapped byte outside the touched slots' own sub-pages is what it was.""" + torch.cuda.synchronize() + lo, hi = self.top, self.top + self.layers * self.state + for s, f, b in zip(self.mapped, self.flat, self.before): + if s in self.used: + same = torch.equal(f[:lo], b[:lo]) and torch.equal(f[hi:], b[hi:]) + else: + same = torch.equal(f, b) + assert same, f"slot {s}: the page was written outside the cell's sub-pages" + + def close(self) -> None: + del self.pools, self.flat, self.before + self.page.close() + + class _Bench: """The state the cache manager keeps for `layers` GDN layers, laid out as `MambaHybridCacheManagerV2` lays it out, driven the way the engine drives @@ -176,19 +330,28 @@ class _Bench: `_advance_replay_state`. With `exact=True` it also carries the exact twin (see the module docstring). + Over an engine page (`page_layers`) the reference and the twin see a dense copy + of the touched slots: `idx` the slots in order, `position` a slot's row in it. """ def __init__(self, spec, seed, exact=True): self.spec = spec self.path, self.n, self.t = spec["path"], spec["n"], spec["t"] + self.h, self.hv = spec.get("H", H), spec.get("HV", HV) + assert spec.get("history_size", HIST) == HIST, spec + self.width = 2 * self.h * K + self.hv * V self.layers = spec.get("layers", 2) self.slots = spec.get("slots", 64) self.steps = spec.get("steps", 8) self.gen = torch.Generator(device=DEV).manual_seed(seed) self.cpu = torch.Generator().manual_seed(seed) n, t, L, S = self.n, self.t, self.layers, self.slots + h, hv = self.h, self.hv + page_layers = spec.get("page_layers") - if spec.get("edge_slots"): + if page_layers: + chosen = torch.tensor(_scattered(n, S)) + elif spec.get("edge_slots"): inner = torch.randperm(S - 2, generator=self.cpu)[: n - 2] + 1 chosen = torch.cat([torch.tensor([S - 1]), inner, torch.tensor([0])]) else: @@ -203,23 +366,32 @@ def __init__(self, spec, seed, exact=True): self.state_indices = chosen.to(DEV) self.live = self.state_indices[live_rows.to(DEV)].clone() - # [slots, layers, HV, V, K]: each layer's state is a slot-strided view, - # the affine layout whose strides the all-layer commit is handed. - self.backing = (self._randn(S, L, HV, V, K) * 0.5).to( - spec.get("state_dtype", torch.float32) - ) - self.pool = [self.backing[:, i] for i in range(L)] - self.old_x = torch.zeros(L, S, 2, HIST, HV, V, device=DEV, dtype=torch.bfloat16) - self.old_B = torch.zeros(L, S, 2, HIST, H, K, device=DEV, dtype=torch.bfloat16) - self.old_dt = torch.zeros(L, S, 2, HV, HIST, device=DEV) - self.old_dA = self._randn(L, S, 2, HV, HIST) # handed over as old_beta; never read + state_dtype = spec.get("state_dtype", torch.float32) + self.page = self.idx = self.position = None + if page_layers: + self.page = _EnginePage(S, page_layers, L, chosen.tolist(), hv, state_dtype, self.gen) + self.pool = self.page.pools + self.layer_stride, self.slot_stride = self.page.state, self.page.stride + self.idx = torch.tensor(self.page.used, dtype=torch.long, device=DEV) + self.position = torch.full((S,), -1, dtype=torch.long, device=DEV) + self.position[self.idx] = torch.arange(self.idx.numel(), device=DEV) + else: + # [slots, layers, HV, V, K]: each layer's state is a slot-strided view, + # the affine layout whose strides the all-layer commit is handed. + self.backing = (self._randn(S, L, hv, V, K) * 0.5).to(state_dtype) + self.pool = [self.backing[:, i] for i in range(L)] + self.layer_stride, self.slot_stride = self.backing.stride(1), self.backing.stride(0) + self.old_x = torch.zeros(L, S, 2, HIST, hv, V, device=DEV, dtype=torch.bfloat16) + self.old_B = torch.zeros(L, S, 2, HIST, h, K, device=DEV, dtype=torch.bfloat16) + self.old_dt = torch.zeros(L, S, 2, hv, HIST, device=DEV) + self.old_dA = self._randn(L, S, 2, hv, HIST) # handed over as old_beta; never read self.pnat = torch.zeros(S, dtype=torch.int32, device=DEV) self.buf = torch.zeros(S, dtype=torch.int32, device=DEV) - self.a_log, self.dt_bias = _gating(self.gen, L) + self.a_log, self.dt_bias = _gating(self.gen, L, hv) if spec.get("start", "warm") == "warm": - self.old_x.copy_(self._randn(L, S, 2, HIST, HV, V) * 0.5) - self.old_B.copy_(_unit(self._randn(L, S, 2, HIST, H, K))) - raw_a = self._randn(L, S, 2, HV, HIST) + self.old_x.copy_(self._randn(L, S, 2, HIST, hv, V) * 0.5) + self.old_B.copy_(_unit(self._randn(L, S, 2, HIST, h, K))) + raw_a = self._randn(L, S, 2, hv, HIST) g = -torch.exp(self.a_log)[:, None, None, :, None] * F.softplus( raw_a + self.dt_bias[:, None, None, :, None] ) @@ -240,32 +412,51 @@ def __init__(self, spec, seed, exact=True): self.exact = exact if exact: self.depth = HIST + (self.steps + 1) * t - self.x_pool = [p.clone() for p in self.pool] - self.x_u = torch.zeros(L, S, 1, self.depth, HV, V, device=DEV) - self.x_k = torch.zeros(L, S, 1, self.depth, H, K, device=DEV) - self.x_G = torch.zeros(L, S, 1, HV, self.depth, device=DEV) - self.x_pnat = self.pnat.clone() - self.x_buf = torch.zeros_like(self.buf) + self.x_pool = [self.compact(p).clone() for p in self.pool] + rows = self.x_pool[0].shape[0] + self.x_u = torch.zeros(L, rows, 1, self.depth, hv, V, device=DEV) + self.x_k = torch.zeros(L, rows, 1, self.depth, h, K, device=DEV) + self.x_G = torch.zeros(L, rows, 1, hv, self.depth, device=DEV) + self.x_pnat = self.compact(self.pnat).clone() + self.x_buf = torch.zeros_like(self.x_pnat) for s in self.live.tolist(): p, b = int(self.pnat[s]), int(self.buf[s]) - self.x_u[:, s, 0, :p] = self.old_x[:, s, b, :p].float() - self.x_k[:, s, 0, :p] = self.old_B[:, s, b, :p].float() - self.x_G[:, s, 0, :, :p] = self.old_dt[:, s, b, :, :p] + c = s if self.position is None else int(self.position[s]) + self.x_u[:, c, 0, :p] = self.old_x[:, s, b, :p].float() + self.x_k[:, c, 0, :p] = self.old_B[:, s, b, :p].float() + self.x_G[:, c, 0, :, :p] = self.old_dt[:, s, b, :, :p] def _randn(self, *shape, dtype=torch.float32): return torch.randn(*shape, generator=self.gen, device=DEV, dtype=dtype) + def compact(self, per_slot: torch.Tensor) -> torch.Tensor: + """A per-slot tensor as the reference sees it: whole over a dense pool, the touched + slots' rows (a copy, in `idx` order) over an engine page.""" + return per_slot if self.idx is None else per_slot[self.idx] + + def compact_indices(self, slots: torch.Tensor) -> torch.Tensor: + """Slot ids as rows of `compact`'s output (negative ids stay negative).""" + if self.position is None: + return slots + renamed = torch.where(slots >= 0, self.position[slots.long().clamp(min=0)], -1) + return renamed.to(slots.dtype) + + def close(self) -> None: + if self.page is not None: + self.pool = [] + self.page.close() + def inputs(self): """One layer's fresh activations: post-conv `[q|k|v]` rows and the raw `a`/`b` projections, as the verify step hands them over.""" - n, t = self.n, self.t - packed = self._randn(n * t, W, dtype=torch.bfloat16) - rows = packed.view(n, t, W) - q = rows[..., : H * K].view(n, t, H, K) - k = rows[..., H * K : 2 * H * K].view(n, t, H, K) - v = rows[..., 2 * H * K :].view(n, t, HV, V) - a = self._randn(n, t, HV, dtype=torch.bfloat16) - b = self._randn(n, t, HV, dtype=torch.bfloat16) + n, t, h, hv = self.n, self.t, self.h, self.hv + packed = self._randn(n * t, self.width, dtype=torch.bfloat16) + rows = packed.view(n, t, self.width) + q = rows[..., : h * K].view(n, t, h, K) + k = rows[..., h * K : 2 * h * K].view(n, t, h, K) + v = rows[..., 2 * h * K :].view(n, t, hv, V) + a = self._randn(n, t, hv, dtype=torch.bfloat16) + b = self._randn(n, t, hv, dtype=torch.bfloat16) return dict(q=q, k=k, v=v, a=a, b=b, packed=packed) def kwargs(self, layer, act): @@ -296,10 +487,30 @@ def kwargs(self, layer, act): output=None, ) + def reference(self, kw): + """The entry's `reference` on a call's arguments, every per-slot buffer through + `compact`: the touched slots' rows, the rows' slots renumbered into them.""" + if self.idx is None: + return op.reference(**kw) + per_slot = ( + "ssm_states", + "old_u", + "old_k", + "old_G", + "old_beta", + "cache_buf_idx", + "prev_num_accepted_tokens", + ) + compact = {name: self.compact(kw[name]) for name in per_slot} + return op.reference( + **dict(kw, **compact, state_indices=self.compact_indices(kw["state_indices"])) + ) + def exact_kwargs(self, layer, kw): return dict( kw, ssm_states=self.x_pool[layer], + state_indices=self.compact_indices(kw["state_indices"]), old_u=self.x_u[layer], old_k=self.x_k[layer], old_G=self.x_G[layer], @@ -321,6 +532,19 @@ def adopt_exact(self, layer, ref): def state(self, layer): return self.pool[layer], self.old_x[layer], self.old_B[layer], self.old_dt[layer] + def observed(self, layer): + """`state(layer)` through `compact`, as `reference`'s results are laid out.""" + return tuple(self.compact(x) for x in self.state(layer)) + + def unwritten(self, wrapped): + """Rows of `observed`'s pool the call must leave as they were: all of them on the + partitioned path, all but the wrapping slots' when the call commits in-kernel.""" + rows = self.slots if self.idx is None else self.idx.numel() + keep = torch.ones(rows, dtype=torch.bool, device=DEV) + if self.path == "small": + keep[self.compact_indices(wrapped).long()] = False + return keep + def prepare(self): if self.path == "all_layer": self.work_items.zero_() @@ -333,8 +557,8 @@ def commit(self): commit_gdn_cached_replay_history_layers( ssm_states=self.pool[0], ssm_state_num_layers=self.layers, - ssm_state_layer_stride=self.backing.stride(1), - ssm_state_slot_stride=self.backing.stride(0), + ssm_state_layer_stride=self.layer_stride, + ssm_state_slot_stride=self.slot_stride, old_u=self.old_x, old_k=self.old_B, old_G=self.old_dt, @@ -355,7 +579,7 @@ def draw_acceptance(self): def accept(self, accepted): _advance_replay_state(self.meta, self.live, accepted) if self.exact: - self.x_pnat[self.live.long()] += accepted + self.x_pnat[self.compact_indices(self.live).long()] += accepted def writes(self): """Where this step's calls append, by the documented rule, which slots wrap, and how @@ -412,9 +636,18 @@ def _check_duplicate_padding_slot(cell) -> None: @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name certifying(op) + _MEASURED.clear() if "dup" in cell.spec: return _check_duplicate_padding_slot(cell) bench = _Bench(cell.spec, seed=_seed(cell.why)) + try: + _drive_cell(cell, bench) + finally: + bench.close() + print(f"\n[{cell.why}] " + ", ".join(f"{k} {v:.3f}" for k, v in sorted(_MEASURED.items()))) + + +def _drive_cell(cell, bench) -> None: wraps = appends = to_the_end = 0 for _ in range(bench.steps): mask, wrapped, at_end = bench.writes() @@ -426,34 +659,32 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res for layer in range(bench.layers): act = bench.inputs() kw = bench.kwargs(layer, act) - pool, u, k, G = bench.state(layer) + _, u, k, G = bench.state(layer) before = dict( - pool=pool.clone(), + pool=bench.observed(layer)[0].clone(), u=u.clone(), k=k.clone(), G=G.clone(), old_beta=bench.old_dA[layer].clone(), ) - ref = op.reference(**kw) + ref = bench.reference(kw) exact = op.reference(**bench.exact_kwargs(layer, kw)) committed = None if bench.path == "all_layer": - committed = op.reference( - **dict(kw, replay_work_items=None, n_writes=None, use_all_layer_commit=False) + committed = bench.reference( + dict(kw, replay_work_items=None, n_writes=None, use_all_layer_commit=False) ) got = op(**kw) - _compare((got, *bench.state(layer)), ref) + _compare((got, *bench.observed(layer)), ref) # The locks: in place, and nothing outside the documented writes. - pool, u, k, G = bench.state(layer) + _, u, k, G = bench.state(layer) assert torch.equal(bench.old_dA[layer], before["old_beta"]), "old_beta is never written" assert torch.equal(u[~mask], before["u"][~mask]), "old_u written outside the append" assert torch.equal(k[~mask], before["k"][~mask]), "old_k written outside the append" Gt, Gt0 = G.transpose(-1, -2), before["G"].transpose(-1, -2) assert torch.equal(Gt[~mask], Gt0[~mask]), "old_G written outside the append" - untouched = torch.ones(bench.slots, dtype=torch.bool, device=DEV) - if bench.path == "small": - untouched[wrapped] = False - assert torch.equal(pool[untouched], before["pool"][untouched]), ( + kept = bench.unwritten(wrapped) + assert torch.equal(bench.observed(layer)[0][kept], before["pool"][kept]), ( "the pool moved where the call commits nothing (a partitioned call commits none)" ) bench.adopt_exact(layer, exact) @@ -463,8 +694,10 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res if committed is not None: # The partitioned call plus the manager's commit is the call # that commits in-kernel. - _compare((got, *bench.state(layer)), committed) + _compare((got, *bench.observed(layer)), committed) _compare((got, None, None, None, None), (exact[0], None, None, None, None)) + if bench.page is not None: + bench.page.assert_rest_unchanged() bench.accept(bench.draw_acceptance()) assert wraps > 0 and appends > 0, ( f"the cell never wrapped or never appended ({wraps}, {appends})" @@ -524,24 +757,3 @@ def test_cuda_graph_replay_on_new_data_is_bitwise_eager() -> None: accepted = eager.draw_acceptance() eager.accept(accepted) graph.accept(accepted) - - -# ── guards: accepted and silently wrong ────────────────────────────────── - - -def test_a_short_state_indices_reads_past_its_end_and_is_refused() -> None: - """Unguarded, three indices for four rows: the fourth row reads whatever - follows the view -- here a slot the call never named -- and writes its - history.""" - bench = _Bench(dict(path="small", n=4, t=2, layers=1), seed=101, exact=False) - stranger = next(s for s in range(bench.slots) if s not in bench.live.tolist()) - parent = torch.cat( - [bench.state_indices[:3], torch.tensor([stranger], dtype=torch.int32, device=DEV)] - ) - kw = dict(bench.kwargs(0, bench.inputs()), state_indices=parent[:3]) - before = bench.old_x[0, stranger].clone() - op(**kw) - assert not torch.equal(bench.old_x[0, stranger], before), "the hazard this guard covers is gone" - with pytest.raises(AssertionError, match="state_indices has shape"): - with validating(op): - op(**kw) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py new file mode 100644 index 000000000000..fafe33a39b78 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_gated_norm_pdl.py @@ -0,0 +1,311 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gdn_decode_gated_norm_pdl entry. + +Cells against the composition's fp64 reference (the gated delta rule's recurrence, +then the gated norm entry's) in torch's default band for the bf16 output and the fp32 +pool. A cell's pool is one of three: + +- a contiguous 461-slot pool, compared whole so slots the call did not name are gated too; +- the tp1 coalesced page (36 layers' states a slot, 102 slots), the view at layer 0; +- with `pool_slots`, the stock cache manager's page at an engine's measured slot count: + 36 GDN layers' states a slot, the view at the last layer's sub-page. Such a page spans + 53 GB (tep4) or 96 GB (dep4) in fp32, so it is a reserved virtual range with physical + memory mapped under the touched slots only (`_SparsePage`). The reference runs on a dense + copy of those slots, and every other byte of their pages is checked unchanged. + +Beyond the cells: a stale-read check under graph replay at both head layouts (a producer +rewrites the projection inside the graph right before the op). +""" + +import zlib + +import cuda.bindings.driver as cuda_drv +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_decode_gated_norm_pdl import ( # noqa: E501 + gdn_decode_gated_norm_pdl as op, +) + +__extra_import_path__ = ["../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "gdn_decode_gated_norm_pdl requires a CUDA device" + +DEV = "cuda" +K, V = 128, 128 +SLOTS, EPS = 461, 1e-6 +#: GDN layers whose states share one slot page in the stock manager (48 layers, every 4th +#: full attention); the view is at the last one's sub-page, the page's highest addresses. +PAGE_LAYERS = 36 + + +def _heads(spec) -> tuple[int, int]: + """(k heads, v heads): tp1's 16 / 48 where the cell names none.""" + return spec.get("k_heads", 16), spec.get("v_heads", 48) + + +def _width(k_heads: int, v_heads: int) -> int: + """`[q | k | v | z | a | b]`: 16480 at 16 / 48, 4120 at a tp4/tep4 rank's 4 / 12.""" + return 2 * k_heads * K + 2 * v_heads * V + 2 * v_heads + + +def _scattered(n: int, slots: int) -> tuple[int, ...]: + """`n` distinct slots of a `slots`-row pool, both ends first, then a stride-37 walk + over the interior (37 is coprime to every pool size here).""" + rest = tuple((i * 37 + 4) % (slots - 2) + 1 for i in range(max(n - 2, 0))) + return ((slots - 1, 0) + rest)[:n] + + +def _check(result): + err, *value = result + assert err == cuda_drv.CUresult.CUDA_SUCCESS, err + return value[0] if len(value) == 1 else value + + +class _ArrayInterface: + def __init__(self, ptr: int, shape, strides_bytes): + self.__cuda_array_interface__ = dict( + shape=tuple(shape), + strides=tuple(strides_bytes), + typestr=" torch.Tensor: + """fp32 tensor at element `offset`, element `strides`.""" + return torch.as_tensor( + _ArrayInterface(self.va + 4 * offset, shape, [4 * s for s in strides]), device=DEV + ) + + def close(self) -> None: + torch.cuda.synchronize() + for ptr, size, handle in self.maps: + _check(cuda_drv.cuMemUnmap(ptr, size)) + _check(cuda_drv.cuMemRelease(handle)) + _check(cuda_drv.cuMemAddressFree(cuda_drv.CUdeviceptr(self.va), self.size)) + + +class _EnginePool: + """The stock manager's coalesced state page at `slots` slots (measured slot stride: 36 + layers' `[HV, V, K]` states), this layer's view the last layer's sub-page, mapped under + the touched slots: each one's whole page and the next slot's first layer. The view's own + base (slot 0's sub-page) is mapped too: torch reads the device off a tensor's base + pointer.""" + + def __init__(self, slots: int, used, v_heads: int, gen): + layer = v_heads * V * K + stride = PAGE_LAYERS * layer + total = slots * stride + base = (PAGE_LAYERS - 1) * layer + self.used = sorted(set(used)) + self.spans = [(s * stride, min((s + 1) * stride + layer, total)) for s in self.used] + self.page = _SparsePage(total, [*self.spans, (base, base + layer)]) + # This layer's state within each mapped span: every span starts at a slot's page. + self.own = base + self.layer = layer + self.pool = self.page.view(base, (slots, v_heads, V, K), (stride, V * K, K, 1)) + self.flat = [self.page.view(lo, (hi - lo,), (1,)) for lo, hi in self.spans] + for f in self.flat: + f.normal_(0.0, 0.05, generator=gen) + torch.cuda.synchronize() + self.before = [f.clone() for f in self.flat] + + def dense(self, slots: torch.Tensor): + """A dense copy of the touched slots and `slots` renumbered into it (-1 stays).""" + index = {s: i for i, s in enumerate(self.used)} + compact = torch.stack([self.pool[s] for s in self.used]).contiguous() + renamed = [index[s] if s >= 0 else -1 for s in slots.tolist()] + return compact, torch.tensor(renamed, dtype=torch.int32, device=DEV) + + def touched(self) -> torch.Tensor: + return torch.stack([self.pool[s] for s in self.used]) + + def assert_rest_unchanged(self) -> None: + """Every byte of a touched slot's page but this layer's state, and the next slot's + first layer, is what it was.""" + torch.cuda.synchronize() + for f, b, s in zip(self.flat, self.before, self.used): + keep = torch.ones(f.numel(), dtype=torch.bool, device=DEV) + keep[self.own : self.own + self.layer] = False + assert torch.equal(f[keep], b[keep]), f"slot {s}: the call wrote outside its state" + + def close(self) -> None: + del self.pool, self.flat + self.page.close() + + +def _build(spec, seed): + """`(args, engine)`: the call's arguments, and the `_EnginePool` behind them (or None).""" + g = torch.Generator(device=DEV).manual_seed(seed) + rows = spec["rows"] + k_heads, v_heads = _heads(spec) + proj = torch.randn(rows, _width(k_heads, v_heads), generator=g, device=DEV).to(torch.bfloat16) + engine = None + if spec.get("pool_slots"): + n = spec["pool_slots"] + picked = (n - 1, n - 2) if rows == 2 else _scattered(rows, n) + engine = _EnginePool(n, picked, v_heads, g) + pool = engine.pool + slots = torch.tensor(picked, dtype=torch.int32, device=DEV) + elif spec.get("pool") == "coalesced": + # The KV-cache-v2 manager's view: every GDN layer's state of a slot in one page, + # this layer's state one sub-page of it, so the slot stride is 36 layers' states + # and slot 100 sits ~2.8G elements past the base (past int32). + layers, num = PAGE_LAYERS, 102 + page = torch.empty(num * layers * v_heads * V * K, device=DEV) + pool = page.as_strided((num, v_heads, V, K), (layers * v_heads * V * K, V * K, K, 1)) + # The pool is compared whole: no slot of the view is left uninitialized. + pool.zero_() + pool[0].normal_(0.0, 0.05, generator=g) + for s in (100, 101): + pool[s].copy_(torch.randn(v_heads, V, K, generator=g, device=DEV) * 0.05) + slots = torch.tensor([100, 101][:rows], dtype=torch.int32, device=DEV) + else: + pool = torch.randn(SLOTS, v_heads, V, K, generator=g, device=DEV) * 0.05 + slots = torch.tensor(_scattered(rows, SLOTS), dtype=torch.int32, device=DEV) + if spec.get("pad") and rows > 1: + slots[1] = -1 + A_log = torch.empty(v_heads, device=DEV).uniform_(-5.5, 1.0, generator=g) + dt_bias = torch.randn(v_heads, generator=g, device=DEV) * 4.0 + w = (1.0 + 0.2 * torch.randn(V, generator=g, device=DEV)).to(torch.bfloat16) + return [proj, pool, slots, A_log, dt_bias, w, EPS, k_heads, K], engine + + +def _compare(out, pool, ref_out, ref_pool) -> None: + assert out.dtype == torch.bfloat16 and out.shape == ref_out.shape and out.is_contiguous() + assert torch.isfinite(out).all() + torch.testing.assert_close(out, ref_out) + torch.testing.assert_close(pool, ref_pool) + + +def _stats(out, ref_out, pool, ref_pool) -> str: + """What `note` quotes: the output against its bf16-rounded reference, the state against + the fp32-rounded one.""" + o, r = out.float().flatten(), ref_out.float().flatten() + cos = torch.nn.functional.cosine_similarity(o, r, dim=0).item() + rel = ((o - r).norm() / r.norm().clamp_min(1e-30)).item() + differ = (out.view(torch.int16) != ref_out.view(torch.int16)).sum().item() + p, q = pool.flatten(), ref_pool.flatten() + p_rel = ((p - q).norm() / q.norm().clamp_min(1e-30)).item() + return ( + f"out max|d| {(o - r).abs().max().item():.3e}, {differ}/{o.numel()} elements differ, " + f"cos {cos:.8f}, rel L2 {rel:.3e}; state max|d| {(p - q).abs().max().item():.3e}, " + f"rel L2 {p_rel:.3e}" + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + spec = cell.spec + args, engine = _build(spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + try: + if engine is None: + ref_out, ref_pool = op.reference(*args) + out = op(*args) + pool = args[1] + else: + dense, renamed = engine.dense(args[2]) + ref_out, ref_pool = op.reference(args[0], dense, renamed, *args[3:]) + out = op(*args) + torch.cuda.synchronize() + pool = engine.touched() + _compare(out, pool, ref_out, ref_pool) + if engine is not None: + engine.assert_rest_unchanged() + print(f"[{cell.why}] {_stats(out, ref_out, pool, ref_pool)}") + if spec.get("pad"): + if spec["rows"] > 1: + assert (out[1] == 0).all(), "a pad row's output is zero" + else: + # One row: a second call over the pad slot (-1) writes zeros and touches no state. + state = engine.touched().clone() if engine else args[1].clone() + pad = torch.full_like(args[2], -1) + out_pad = op(args[0], args[1], pad, *args[3:]) + torch.cuda.synchronize() + assert (out_pad == 0).all(), "a pad row's output is zero" + after = engine.touched() if engine else args[1] + assert torch.equal(after, state), "a pad row touches no state" + if engine: + engine.assert_rest_unchanged() + finally: + if engine is not None: + engine.close() + + +@pytest.mark.parametrize("k_heads, v_heads", [(16, 48), (4, 12)], ids=["16x48", "4x12"]) +def test_reads_after_the_wait_under_graph_replay(k_heads, v_heads) -> None: + """A producer rewrites the projection inside the graph right before the op; with PDL + a read of it before `griddepcontrol.wait` would see the previous replay's row. The + pool advances one step per replay, as in decode.""" + spec = dict(rows=1, k_heads=k_heads, v_heads=v_heads) + args, _ = _build(spec, seed=41) + src = args[0].clone() + + def step(): + args[0].copy_(src * 1.0) + return op(*args) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + step() + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_out = step() + for trial in range(3): + src.copy_(_build(spec, seed=100 + trial)[0][0]) + before = args[1].clone() + graph.replay() + torch.cuda.synchronize() + after = args[1].clone() + args[1].copy_(before) + want = op(src.clone(), *args[1:]) + assert torch.equal(graph_out, want), f"replay {trial} output" + assert torch.equal(after, args[1]), f"replay {trial} pool" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py new file mode 100644 index 000000000000..d9688b4400f2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_decode_rows_gated_norm.py @@ -0,0 +1,471 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gdn_decode_rows_gated_norm entry. + +Cells against the composition's fp64 reference (the gated delta rule's recurrence, +then the gated norm entry's) in torch's default band for the bf16 output and the pool's +dtype (fp32 or bf16). A cell's pool is one of three: + +- a contiguous 461-slot pool, compared whole so slots the call did not name are gated too; +- the tp1 coalesced page (36 layers' states a slot, 102 slots), the view at layer 0; +- with `pool_slots`, the stock cache manager's page at an engine's slot count: 36 GDN + layers' states a slot, the view at the last layer's sub-page. Such a page spans 26 GB + (tep4) or 48 GB (dep4) in bf16, so it is a reserved virtual range with physical memory + mapped under the touched slots only (`_SparsePage`). The reference runs on a dense copy + of those slots, and every other byte of their pages is checked unchanged. + +Beyond the cells: a fixed row bitwise across batch sizes over a bf16 pool, a stale-read +check under graph replay (a producer rewrites the projection inside the graph right +before the op), the two launch variants (state rows by bulk copy into shared memory, or +loaded into registers) bitwise each other, and the variant choice over an unaligned +pool, each at tp1's 16 / 48 heads and a tp4/tep4 rank's 4 / 12. The `gdn_decode_rows` +entry (the recurrence alone) has its own cells below. +""" + +import zlib + +import cuda.bindings.driver as cuda_drv +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm import ( + gdn_decode_rows_gated_norm as rows_mod, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_decode_rows_gated_norm import ( # noqa: E501 + gdn_decode_rows as rows_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_decode_rows_gated_norm import ( # noqa: E501 + gdn_decode_rows_gated_norm as op, +) + +__extra_import_path__ = ["../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "gdn_decode_rows_gated_norm requires a CUDA device" + +DEV = "cuda" +H, HV, K, V = 16, 48, 128, 128 +SLOTS, EPS = 461, 1e-6 +WIDTH = 2 * H * K + 2 * HV * V + 2 * HV +#: GDN layers whose states share one slot page in the stock manager (48 layers, every 4th +#: full attention); the view is at the last one's sub-page, the page's highest addresses. +PAGE_LAYERS = 36 +#: (k heads, v heads) of a tp4/tep4 rank. +RANK_HEADS = (4, 12) + + +def _heads(spec) -> tuple[int, int]: + """(k heads, v heads): tp1's 16 / 48 where the cell names none.""" + return spec.get("k_heads", H), spec.get("v_heads", HV) + + +def _width(k_heads: int, v_heads: int) -> int: + """`[q | k | v | z | a | b]`: 16480 at 16 / 48, 4120 at a tp4/tep4 rank's 4 / 12.""" + return 2 * k_heads * K + 2 * v_heads * V + 2 * v_heads + + +def _scattered(n: int, slots: int) -> tuple[int, ...]: + """`n` distinct slots of a `slots`-row pool, both ends first, then a stride-37 walk + over the interior (37 is coprime to every pool size here).""" + rest = tuple((i * 37 + 4) % (slots - 2) + 1 for i in range(max(n - 2, 0))) + return ((slots - 1, 0) + rest)[:n] + + +def _check(result): + err, *value = result + assert err == cuda_drv.CUresult.CUDA_SUCCESS, err + return value[0] if len(value) == 1 else value + + +#: Array-interface type of each pool dtype; bf16 has none, so it is read as int16 and viewed. +_TYPESTR = {torch.float32: (" torch.Tensor: + """Tensor of the page's dtype at element `offset`, element `strides`.""" + typestr, carrier = _TYPESTR[self.dtype] + es = self.es + t = torch.as_tensor( + _ArrayInterface(self.va + es * offset, shape, [es * s for s in strides], typestr), + device=DEV, + ) + assert t.dtype == carrier + return t.view(self.dtype) + + def close(self) -> None: + torch.cuda.synchronize() + for ptr, size, handle in self.maps: + _check(cuda_drv.cuMemUnmap(ptr, size)) + _check(cuda_drv.cuMemRelease(handle)) + _check(cuda_drv.cuMemAddressFree(cuda_drv.CUdeviceptr(self.va), self.size)) + + +class _EnginePool: + """The stock manager's coalesced state page at `slots` slots (slot stride: 36 layers' + `[HV, V, K]` states), this layer's view the last layer's sub-page, mapped under the + touched slots: each one's whole page and the next slot's first layer. The view's own + base (slot 0's sub-page) is mapped too: torch reads the device off a tensor's base + pointer.""" + + def __init__(self, slots: int, used, v_heads: int, dtype: torch.dtype, gen): + layer = v_heads * V * K + stride = PAGE_LAYERS * layer + total = slots * stride + base = (PAGE_LAYERS - 1) * layer + self.used = sorted(set(used)) + self.spans = [(s * stride, min((s + 1) * stride + layer, total)) for s in self.used] + self.page = _SparsePage(total, [*self.spans, (base, base + layer)], dtype) + # This layer's state within each mapped span: every span starts at a slot's page. + self.own = base + self.layer = layer + self.pool = self.page.view(base, (slots, v_heads, V, K), (stride, V * K, K, 1)) + self.flat = [self.page.view(lo, (hi - lo,), (1,)) for lo, hi in self.spans] + for f in self.flat: + f.normal_(0.0, 0.05, generator=gen) + torch.cuda.synchronize() + self.before = [f.clone() for f in self.flat] + + def dense(self, slots: torch.Tensor): + """A dense copy of the touched slots and `slots` renumbered into it (-1 stays).""" + index = {s: i for i, s in enumerate(self.used)} + compact = torch.stack([self.pool[s] for s in self.used]).contiguous() + renamed = [index[s] if s >= 0 else -1 for s in slots.tolist()] + return compact, torch.tensor(renamed, dtype=torch.int32, device=DEV) + + def touched(self) -> torch.Tensor: + return torch.stack([self.pool[s] for s in self.used]) + + def assert_rest_unchanged(self) -> None: + """Every byte of a touched slot's page but this layer's state, and the next slot's + first layer, is what it was.""" + torch.cuda.synchronize() + lo, hi = self.own, self.own + self.layer + for f, b, s in zip(self.flat, self.before, self.used): + assert torch.equal(f[:lo], b[:lo]) and torch.equal(f[hi:], b[hi:]), ( + f"slot {s}: the call wrote outside its state" + ) + + def close(self) -> None: + del self.pool, self.flat + self.page.close() + + +def _build(spec, seed): + """`(args, engine)`: the call's arguments, and the `_EnginePool` behind them (or None).""" + g = torch.Generator(device=DEV).manual_seed(seed) + rows = spec["rows"] + k_heads, v_heads = _heads(spec) + width = _width(k_heads, v_heads) + assert spec.get("proj_width", width) == width, (spec, width) + state = torch.bfloat16 if spec.get("state") == "bf16" else torch.float32 + proj = torch.randn(rows, width, generator=g, device=DEV).to(torch.bfloat16) + engine = None + if spec.get("pool_slots"): + n = spec["pool_slots"] + picked = _scattered(rows, n) + engine = _EnginePool(n, picked, v_heads, state, g) + pool = engine.pool + slots = torch.tensor(picked, dtype=torch.int32, device=DEV) + elif spec.get("pool") == "coalesced": + # The KV-cache-v2 manager's view: every GDN layer's state of a slot in one page, + # this layer's state one sub-page of it, so the slot stride is 36 layers' states + # and slot 100 sits ~2.8G elements past the base (past int32). + layers, num = PAGE_LAYERS, 102 + page = torch.empty(num * layers * v_heads * V * K, device=DEV, dtype=state) + page[: v_heads * V * K * 4].normal_(0.0, 0.05, generator=g) + pool = page.as_strided((num, v_heads, V, K), (layers * v_heads * V * K, V * K, K, 1)) + for s in (100, 101): + pool[s].copy_(torch.randn(v_heads, V, K, generator=g, device=DEV) * 0.05) + slots = torch.tensor([100, 101][:rows], dtype=torch.int32, device=DEV) + else: + pool = (torch.randn(SLOTS, v_heads, V, K, generator=g, device=DEV) * 0.05).to(state) + slots = torch.tensor(_scattered(rows, SLOTS), dtype=torch.int32, device=DEV) + if spec.get("pad"): + slots[1] = -1 + A_log = torch.empty(v_heads, device=DEV).uniform_(-5.5, 1.0, generator=g) + dt_bias = torch.randn(v_heads, generator=g, device=DEV) * 4.0 + w = (1.0 + 0.2 * torch.randn(V, generator=g, device=DEV)).to(torch.bfloat16) + return [proj, pool, slots, A_log, dt_bias, w, EPS, k_heads, K], engine + + +def _compare(out, pool, ref_out, ref_pool) -> None: + assert out.dtype == torch.bfloat16 and out.shape == ref_out.shape and out.is_contiguous() + assert torch.isfinite(out).all() + torch.testing.assert_close(out, ref_out) + torch.testing.assert_close(pool, ref_pool) + + +def _stats(out, ref_out, pool, ref_pool) -> str: + """What `note` quotes: the output against its bf16-rounded reference, the state against + the reference rounded to the pool's dtype. Elements are counted by value (a pad row's + -0.0 equals the reference's 0.0).""" + o, r = out.float().flatten(), ref_out.float().flatten() + cos = torch.nn.functional.cosine_similarity(o, r, dim=0).item() + rel = ((o - r).norm() / r.norm().clamp_min(1e-30)).item() + differ = (o != r).sum().item() + p, q = pool.float().flatten(), ref_pool.float().flatten() + p_rel = ((p - q).norm() / q.norm().clamp_min(1e-30)).item() + p_differ = (p != q).sum().item() + return ( + f"out max|d| {(o - r).abs().max().item():.3e}, {differ}/{o.numel()} elements differ, " + f"cos {cos:.8f}, rel L2 {rel:.3e}; state max|d| {(p - q).abs().max().item():.3e}, " + f"{p_differ}/{p.numel()} differ, rel L2 {p_rel:.3e}" + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(op) + args, engine = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + try: + if engine is None: + ref_out, ref_pool = op.reference(*args) + out = op(*args) + pool = args[1] + else: + dense, renamed = engine.dense(args[2]) + ref_out, ref_pool = op.reference(args[0], dense, renamed, *args[3:]) + out = op(*args) + torch.cuda.synchronize() + pool = engine.touched() + _compare(out, pool, ref_out, ref_pool) + if engine is not None: + engine.assert_rest_unchanged() + print(f"[{cell.why}] {_stats(out, ref_out, pool, ref_pool)}") + if cell.spec.get("pad"): + assert (out[1] == 0).all(), "a pad row's output is zero" + finally: + if engine is not None: + engine.close() + + +#: Decode batch sizes up to the engine's largest CUDA-graph batch. +_BATCHES = (1, 2, 4, 5, 16, 33, 128, 256) +_FIXED_SLOT = 301 + + +@pytest.mark.parametrize("k_heads, v_heads", [(H, HV), RANK_HEADS], ids=["16x48", "4x12"]) +def test_bf16_pool_is_batch_invariant(k_heads, v_heads) -> None: + """The property the bf16 decode step relies on: a row's output and state are bitwise + independent of the batch it rides in. One fixed row (its projection and slot 301 of a + bf16 pool) sits first and last in random batches of every size, across the batch where + the launch switches to the bulk-copy variant (7 rows at 48 v heads, 26 at 12).""" + width = _width(k_heads, v_heads) + gen = torch.Generator(device=DEV).manual_seed(3) + cpu = torch.Generator().manual_seed(3) + pool0 = (torch.randn(SLOTS, v_heads, V, K, generator=gen, device=DEV) * 0.05).to(torch.bfloat16) + row = torch.randn(width, generator=gen, device=DEV).to(torch.bfloat16) + rest = _build(dict(rows=1, k_heads=k_heads, v_heads=v_heads), seed=3)[0][3:] + first = None + for n in _BATCHES: + for pos in sorted({0, n - 1}): + g = torch.Generator(device=DEV).manual_seed(1000 * n + pos) + proj = torch.randn(n, width, generator=g, device=DEV).to(torch.bfloat16) + proj[pos] = row + others = [s for s in torch.randperm(SLOTS, generator=cpu).tolist() if s != _FIXED_SLOT] + others = others[:n] + others[pos] = _FIXED_SLOT + slots = torch.tensor(others, dtype=torch.int32, device=DEV) + pool = pool0.clone() + out = op(proj, pool, slots, *rest) + got = (out[pos].clone(), pool[_FIXED_SLOT].clone()) + first = first or got + assert torch.equal(got[0], first[0]) and torch.equal(got[1], first[1]), ( + f"the fixed row differs at batch {n}, position {pos}" + ) + + +@pytest.mark.parametrize( + "spec", + ( + dict(rows=16), + dict(rows=3, k_heads=4, v_heads=12, state="bf16"), + dict(rows=128, k_heads=4, v_heads=12, state="bf16"), + ), + ids=["16x48-16rows", "4x12-3rows-bf16", "4x12-128rows-bf16"], +) +def test_reads_after_the_wait_under_graph_replay(spec) -> None: + """A producer rewrites the projection inside the graph right before the op; with PDL + a read of it before `griddepcontrol.wait` would see the previous replay's row. The + pool advances one step per replay, as in decode. At 4 / 12 heads 3 rows take the + register variant, 128 the bulk copy.""" + args = _build(spec, seed=41)[0] + src = args[0].clone() + + def step(): + args[0].copy_(src * 1.0) + return op(*args) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + step() + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_out = step() + for trial in range(3): + src.copy_(_build(spec, seed=100 + trial)[0][0]) + before = args[1].clone() + graph.replay() + torch.cuda.synchronize() + after = args[1].clone() + args[1].copy_(before) + want = op(src.clone(), *args[1:]) + assert torch.equal(graph_out, want), f"replay {trial} output" + assert torch.equal(after, args[1]), f"replay {trial} pool" + + +_VARIANT_SPECS = ( + dict(rows=1, state="bf16"), + dict(rows=33, state="bf16"), + dict(rows=33, pad=True), + dict(rows=2, pool="coalesced", state="bf16", pad=True), + dict(rows=3, k_heads=4, v_heads=12, state="bf16", pad=True), + dict(rows=128, k_heads=4, v_heads=12, state="bf16"), + dict(rows=3, k_heads=4, v_heads=12, state="bf16", pool_slots=rows_mod.TEP4_SLOTS), +) + + +@pytest.mark.parametrize("norm", [True, False], ids=["norm", "rows"]) +@pytest.mark.parametrize("spec", _VARIANT_SPECS, ids=lambda s: "-".join(map(str, s.values()))) +def test_bulk_and_register_variants_agree(spec, norm) -> None: + """The bulk-copy and the register variants: outputs and state bitwise equal. + + The bulk-copy variant stages the state rows in shared memory; the register variant holds + them in registers; both compute the same instructions. Over contiguous bf16 and fp32 pools, + the tp1 coalesced page (slot 100, ~2.8G elements past the base, by the bulk copy) and + tep4's engine page (slot 1873, ~1.3e10 elements past the base), with pad rows, at 16 / 48 + and 4 / 12 heads. + """ + args, engine = _build(spec, seed=zlib.crc32(repr(spec).encode()) % (2**31)) + try: + proj, pool, slots, k_heads = args[0], args[1], args[2], args[7] + assert rows_mod._bulk_eligible(pool, K) + live = slots[slots >= 0].long() + before = pool[live].clone() + got = [] + for bulk in (False, True): + pool[live] = before + w = args[5] if norm else None + out = rows_mod.launch( + proj, pool, slots, args[3], args[4], w, EPS, k_heads, K, bulk=bulk + ) + got.append((out.clone(), pool[live].clone())) + assert torch.equal(got[0][0], got[1][0]), "outputs differ between the variants" + assert torch.equal(got[0][1], got[1][1]), "states differ between the variants" + if engine is not None: + engine.assert_rest_unchanged() + finally: + if engine is not None: + engine.close() + + +@pytest.mark.parametrize("v_heads", [HV, RANK_HEADS[1]], ids=["48", "12"]) +def test_variant_choice_and_unaligned_pool(v_heads) -> None: + """The variant choice, and the fallback for a pool the bulk copy cannot take. + + The bulk variant from the batch the register variant needs a second wave for; a pool 8- + but not 16-byte aligned stays on the register variant whatever is asked, and matches the + aligned pool's result. + """ + sms = torch.cuda.get_device_properties(DEV).multi_processor_count + edge = (2 * sms) // v_heads + assert not rows_mod._use_bulk(edge, v_heads, DEV) and rows_mod._use_bulk(edge + 1, v_heads, DEV) + print(f"[{v_heads} v heads] {sms} SMs: the bulk-copy variant from {edge + 1} rows") + k_heads = v_heads // 3 + args = _build(dict(rows=33, k_heads=k_heads, v_heads=v_heads, state="bf16"), seed=19)[0] + pool = args[1] + flat = torch.empty(pool.numel() + 4, dtype=pool.dtype, device=DEV) + shifted = flat[4:].view(pool.shape) + shifted.copy_(pool) + assert not rows_mod._bulk_eligible(shifted, K) + want = rows_mod.launch(args[0], pool, *args[2:], bulk=True) + got = rows_mod.launch(args[0], shifted, *args[2:], bulk=True) + assert torch.equal(got, want) and torch.equal(shifted, pool) + + +# ---- gdn_decode_rows: the recurrence alone, into a row range of a step's output ---------- + + +def _rows_args(spec, seed): + """`(args, buf)`: the entry's arguments, `out` rows 5.. of a `[1, rows + 9, HV, V]` buffer + (the step's output, the call's rows a range of it), and that buffer.""" + (proj, pool, slots, A_log, dt_bias, _, _, h, k), _ = _build(spec, seed) + rows, v_heads = proj.shape[0], pool.shape[1] + g = torch.Generator(device=DEV).manual_seed(seed + 1) + buf = torch.randn(1, rows + 9, v_heads, V, generator=g, device=DEV).to(torch.bfloat16) + out = buf[0, 5 : 5 + rows].view(rows, v_heads * V) + return [proj, pool, slots, A_log, dt_bias, h, k, out], buf + + +@pytest.mark.parametrize("cell", rows_op.CELLS, ids=[c.why[:44] for c in rows_op.CELLS]) +def test_rows_certified_cells(cell, certifying) -> None: # noqa: F811 — fixture by name + certifying(rows_op) + args, buf = _rows_args(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + outside = torch.cat([buf[0, :5], buf[0, 5 + args[0].shape[0] :]]).clone() + ref_o, ref_pool = rows_op.reference(*args) + got = rows_op(*args) + assert got.data_ptr() == args[-1].data_ptr(), "written into out" + torch.testing.assert_close(got, ref_o) + torch.testing.assert_close(args[1], ref_pool) + print(f"[{cell.why}] {_stats(got, ref_o, args[1], ref_pool)}") + assert torch.equal(torch.cat([buf[0, :5], buf[0, 5 + args[0].shape[0] :]]), outside), ( + "rows outside the range untouched" + ) + if cell.spec.get("pad"): + assert (got[1] == 0).all(), "a pad row's output is zero" diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_prefill_conv_post.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_prefill_conv_post.py similarity index 54% rename from tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_prefill_conv_post.py rename to tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_prefill_conv_post.py index 9cfca29e1712..a49ba3fde83a 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_prefill_conv_post.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/ssm/test_modeling_v2_gdn_prefill_conv_post.py @@ -3,8 +3,14 @@ """GPU test for the gdn_prefill_conv_post entry. Every cell is checked against the entry's fp32/fp64 reference, the conv pool -bitwise. Geometry is a Qwen3.8-Flash-Next linear-attention layer: 16 x 128 key heads, -48 x 128 value heads, qkv_dim 10240, conv width 4, bf16 activations. +bitwise. Geometry is a Qwen3.8-Flash-Next linear-attention layer, conv width 4, bf16 +activations: the whole layer (16 x 128 key heads, 48 x 128 value heads, qkv_dim 10240, +a 16480-wide projection) at tp1 and on a dep4 rank, a quarter of it (4 / 12 heads, +qkv_dim 2560, a 4120-wide projection) on a tp4 or tep4 rank. The pool is a view of a +bigger buffer, which the test also checks the call leaves alone outside the view: a +two-layer buffer for the tp1 cells, the stock manager's page (36 GDN layers a slot, +the view at the last) for a rank's cells: at an engine's measured slot count, or at +slots whose offset passes 2^31 elements. """ import math @@ -13,7 +19,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_prefill_conv_post import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_prefill_conv_post import ( # noqa: E501 A_LOG_RANGE, CONV_WIDTH, DT_BIAS_RANGE, @@ -22,30 +28,42 @@ IN_PROJ_WIDTH, NUM_K_HEADS, NUM_V_HEADS, - QKV_DIM, SOFTPLUS_THRESHOLD, _engine_slots, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.gdn_prefill_conv_post import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.ssm.gdn_prefill_conv_post import ( # noqa: E501 gdn_prefill_conv_post as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = ["../../.."] from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name assert torch.cuda.is_available(), "gdn_prefill_conv_post requires a CUDA device" DEV = "cuda" DTYPE = torch.bfloat16 -HEADS = (NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM) PAD = -1 L2_NORM_EPS = 1e-6 -#: `g`'s band, derived as in the fused_gdn_post_conv entry test: two SASS +#: `g`'s band: two SASS #: approximate transcendentals (~2^-22 relative each), the softplus bounded by #: the threshold, every term scaled by exp(A_log). _APPROX_ULP = 2.0**-22 G_ATOL = 4 * _APPROX_ULP * SOFTPLUS_THRESHOLD * math.exp(A_LOG_RANGE[1]) -A_COL = QKV_DIM + NUM_V_HEADS * HEAD_V_DIM + + +def _geometry(spec): + """`(num_k_heads, num_v_heads, head_dim, qkv_dim, projection width)` of a cell: its + own, or the tp1 layer's.""" + hk = spec.get("num_k_heads", NUM_K_HEADS) + hv = spec.get("num_v_heads", NUM_V_HEADS) + d = spec.get("head_dim", HEAD_K_DIM) + assert d == HEAD_K_DIM == HEAD_V_DIM + qkv = 2 * hk * d + hv * d + width = spec.get("x_row_stride", IN_PROJ_WIDTH) + assert width == qkv + hv * d + 2 * hv, ( + f"x_row_stride {width} is not the [q|k|v|z|a|b] projection of {hk}/{hv} heads" + ) + return hk, hv, d, qkv, width def _build(spec, seed): @@ -53,22 +71,28 @@ def _build(spec, seed): g = torch.Generator(device=DEV).manual_seed(seed) lengths = spec["lengths"] n_tok = sum(lengths) + hk, hv, d, qkv, width = _geometry(spec) + a_col = qkv + hv * d def randn(*shape): return torch.randn(*shape, generator=g, device=DEV).to(DTYPE) # "in_proj" and "rows" both carve x, a and b out of one projection; "rows" # leaves `extra_rows` trailing rows the call must not read. - proj = randn(n_tok + spec.get("extra_rows", 0), IN_PROJ_WIDTH) - x = proj[:n_tok, :QKV_DIM] - a = proj[:n_tok, A_COL : A_COL + NUM_V_HEADS] - b = proj[:n_tok, A_COL + NUM_V_HEADS :] - assert x.stride(0) == a.stride(0) == IN_PROJ_WIDTH - weight = (torch.randn(QKV_DIM, CONV_WIDTH, generator=g, device=DEV) * 0.5).to(DTYPE) + proj = randn(n_tok + spec.get("extra_rows", 0), width) + x = proj[:n_tok, :qkv] + a = proj[:n_tok, a_col : a_col + hv] + b = proj[:n_tok, a_col + hv : a_col + 2 * hv] + assert x.stride(0) == a.stride(0) == width + weight = (torch.randn(qkv, CONV_WIDTH, generator=g, device=DEV) * 0.5).to(DTYPE) pool_slots = spec["pool_slots"] - # The pool is a per-layer view of a two-layer buffer, as a cache manager - # hands it out: slot stride twice the slot's own size. - conv_states = randn(pool_slots, 2, QKV_DIM, CONV_WIDTH - 1)[:, 1] + # The pool is a per-layer view of a multi-layer buffer, as a cache manager + # hands it out: two layers for the tp1 cells (slot stride twice the slot's own + # size); for a rank's cells the stock page, `page_layers` (36) GDN layers' + # states a slot, viewed at the last -- the strides and offset the manager's + # `get_conv_states` returns, measured. + layers = spec.get("page_layers", 2) + conv_states = randn(pool_slots, layers, qkv, CONV_WIDTH - 1)[:, layers - 1] slots = spec.get("slots") or _engine_slots(len(lengths)) bounds = [0] for n in lengths: @@ -77,9 +101,9 @@ def randn(*shape): si = torch.tensor(slots, dtype=torch.int32, device=DEV) hi = torch.tensor(spec["initial"], dtype=torch.bool, device=DEV) lo, hi_ = A_LOG_RANGE - A_log = torch.rand(NUM_V_HEADS, generator=g, device=DEV) * (hi_ - lo) + lo + A_log = torch.rand(hv, generator=g, device=DEV) * (hi_ - lo) + lo lo, hi_ = DT_BIAS_RANGE - dt_bias = torch.rand(NUM_V_HEADS, generator=g, device=DEV) * (hi_ - lo) + lo + dt_bias = torch.rand(hv, generator=g, device=DEV) * (hi_ - lo) + lo return [ x, weight, @@ -91,7 +115,10 @@ def randn(*shape): b, A_log, dt_bias, - *HEADS, + hk, + d, + hv, + d, PAD, L2_NORM_EPS, SOFTPLUS_THRESHOLD, @@ -123,20 +150,34 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) want = op.reference(*args) inputs = [t.clone() if isinstance(t, torch.Tensor) else t for t in args] + pool = args[2] + page = pool._base + page_before = page.clone() out = op(*args) for i, (arg, before) in enumerate(zip(args, inputs)): if isinstance(arg, torch.Tensor) and i != 2: assert torch.equal(arg, before), f"input {i} must be read-only" - _compare(out, args[2], want, what=f"{cell.why[:30]} vs reference") - - -def test_guard_refuses_channel_strided_x() -> None: - args = _build(dict(lengths=(16,), pool_slots=2, slots=(0,), initial=(True,)), seed=302) - bad = list(args) - bad[0] = torch.empty(QKV_DIM, 16, dtype=DTYPE, device=DEV).t() - with pytest.raises(AssertionError, match="unit channel stride"): - with validating(op): - op(*bad) + _compare(out, pool, want, what=f"{cell.why[:30]} vs reference") + # The rest of the buffer -- the other layers' sub-pages of every slot -- is untouched. + layer = (pool.storage_offset() - page.storage_offset()) // page.stride(1) + page_before[:, layer] = want[5] + assert torch.equal(page, page_before), "the call wrote outside its layer's view of the pool" + + +def test_cells_carry_from_slots_past_the_int32_offset() -> None: + """The cells read a carried window at a slot whose element offset on the stock page + passes 2^31 at both rank geometries, and 2^32 at 16 / 48 heads -- where a 32-bit slot + offset faults or wraps back into the page -- so the 64-bit slot base stays exercised.""" + reach = {} + for cell in op.CELLS: + spec = cell.spec + if "page_layers" not in spec: + continue + hk, _, _, qkv, _ = _geometry(spec) + slot_stride = spec["page_layers"] * qkv * (CONV_WIDTH - 1) + carried = [s for s, init in zip(spec["slots"], spec["initial"]) if init] + reach[hk] = max([reach.get(hk, 0)] + [s * slot_stride for s in carried]) + assert reach[NUM_K_HEADS] >= 2**32 and reach[NUM_K_HEADS // 4] >= 2**31, reach def test_linear_gate_is_torch_exp_of_the_log_gate() -> None: diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_kv_quant.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_kv_quant.py deleted file mode 100644 index 46c8dad1a1ab..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_kv_quant.py +++ /dev/null @@ -1,229 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""qwen4_exp derives its KV storage from the engine's `kv_cache_quant_algo`. - -`None` keeps the checkpoint's bf16 K/V with the attention op's `quant_mode` 0 and no KV -scales; `FP8` stores K/V as e4m3 under `QM_FP8_KV_CACHE` with unit fp32 `[1]` scales; `NVFP4` -stores e2m1 pairs as uint8 half a head wide under `QM_NVFP4_KV_CACHE` with unit fp32 `[3]` -scales; any other algorithm is refused at construction. The core is built on the meta device from the -checkpoint's text and vision config values, so nothing here touches a device; the -`post_load_weights` cases stub out everything that reads weight values. -""" - -from __future__ import annotations - -from types import SimpleNamespace -from unittest import mock - -import pytest -import torch - -__extra_import_path__ = ["../../.."] # noqa: F841 -- repo's file-scoped import hook - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - modeling as M, -) -from tensorrt_llm._torch.configs import Qwen4ExpTextConfig, Qwen4ExpVisionConfig -from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm.llmapi import QSASparseAttentionConfig -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.quantization.mode import QuantAlgo - -pytestmark = pytest.mark.cpu_only - -NUM_LAYERS = 48 - -# Qwen3.8-Flash-Next's text config. -_TEXT = dict( - architectures=["Qwen4ExpForCausalLM"], - attention_bias=False, - attention_dropout=0.0, - bos_token_id=248044, - dtype="bfloat16", - eos_token_id=248044, - full_attention_interval=4, - hc_count=4, - hc_lowrank=320, - head_dim=256, - heads_per_ngram=8, - hidden_act="silu", - hidden_size=2560, - indexer_budget=2048, - indexer_compress_ratio=4, - indexer_head_dim=128, - indexer_kv_heads=1, - indexer_n_heads=4, - initializer_range=0.02, - layer_types=["full_attention" if i % 4 == 3 else "linear_attention" for i in range(NUM_LAYERS)], - linear_conv_kernel_dim=4, - linear_key_head_dim=128, - linear_num_key_heads=16, - linear_num_value_heads=48, - linear_value_head_dim=128, - make_ngram_vocab_size_divisible_by=128, - mamba_ssm_dtype="float32", - max_position_embeddings=262144, - moe_intermediate_size=640, - mtp=dict( - hybrid=True, - layer_types=["full_attention"], - mtp_use_hidden_state_from_layer=None, - num_hidden_layers=1, - rope_theta=10000000, - ), - mtp_num_hidden_layers=1, - mtp_use_dedicated_embeddings=False, - ngram_size=3, - ngram_vocab_size_base=20000000, - norm_topk_prob=True, - num_attention_heads=24, - num_experts=512, - num_experts_per_tok=10, - num_hidden_layers=NUM_LAYERS, - num_key_value_heads=2, - number_of_conv_states=3, - output_gate_type="sigmoid", - output_router_logits=False, - pad_token_id=None, - partial_rotary_factor=0.25, - ple_conv_kernel_size=4, - ple_embed_dim=2560, - ple_layer_ids=[2], - rms_norm_eps=1e-06, - rope_parameters=dict( - mrope_interleaved=True, - mrope_section=[11, 11, 10], - partial_rotary_factor=0.25, - rope_theta=10000000, - rope_type="default", - ), - router_aux_loss_coef=0.001, - seed=1234, - shared_expert_intermediate_size=640, - split_ngram_parts=128, - tie_word_embeddings=False, - use_cache=True, - vocab_size=248320, -) -_VISION = dict( - deepstack_visual_indexes=[], - depth=27, - dtype="bfloat16", - hidden_act="gelu_pytorch_tanh", - hidden_size=1152, - in_channels=3, - intermediate_size=4304, - num_heads=16, - num_position_embeddings=2304, - out_hidden_size=2560, - patch_size=16, - spatial_merge_size=2, - temporal_patch_size=2, -) - - -def _core(kv_cache_quant_algo): - """The core of a Qwen3.8-Flash-Next engine whose resolved `kv_cache_quant_algo` is - `kv_cache_quant_algo`, with every parameter on the meta device.""" - model_config = ModelConfig( - pretrained_config=Qwen4ExpTextConfig(**_TEXT), - mapping=Mapping(world_size=1, tp_size=1, rank=0), - quant_config=QuantConfig( - quant_algo=QuantAlgo.MIXED_PRECISION, kv_cache_quant_algo=kv_cache_quant_algo - ), - moe_backend="CUTLASS", - moe_backend_requested="TRTLLM", - sparse_attention_config=QSASparseAttentionConfig(), - ) - composite = SimpleNamespace( - vision_config=Qwen4ExpVisionConfig(**_VISION), image_token_id=248056, video_token_id=248057 - ) - # The core asserts the device it runs on; the weights it declares stay on meta. - with ( - mock.patch.object(torch.cuda, "get_device_capability", return_value=(10, 3)), - torch.device("meta"), - ): - return M.Qwen4ExpTp1ModelingV2(model_config, composite) - - -def _load(core, monkeypatch) -> None: - """`post_load_weights` on a core with the weights it reads replaced by CPU stubs. - - Everything that reads weight values (the expert scalars, the n-gram table's driver - residency, the per-step targets) is replaced; what is left is the derivation under test - and the shape checks around it. The derived scales land on `embed`'s device, here CPU, - so their values can be read. - """ - table_rows = core.ngram_vocab_divisor - stubs = dict( - embed=torch.empty(1), - ple_table=torch.empty(table_rows, 1), - ple_offs=torch.zeros(1, dtype=torch.int64), - ple_vocab=torch.full((1,), table_rows, dtype=torch.int64), - ) - for key, value in stubs.items(): - core.w[key] = torch.nn.Parameter(value, requires_grad=False) - monkeypatch.setattr(M, "_moe_layer_scalars", lambda w, i, backend: {}) - for target in ("PrefillTarget", "DecodeTarget", "VisionTower"): - monkeypatch.setattr(M, target, lambda core: None) - monkeypatch.setattr(core, "_report_ple_residency", lambda: None) - core.post_load_weights() - - -def test_unquantized_kv_cache_keeps_the_checkpoints_bf16(monkeypatch) -> None: - core = _core(None) - assert core.quant_mode == 0 - assert core.kv_pool_dtype == torch.bfloat16 - geometry = core._pool_geometry(32) - assert geometry.kv_dtype == torch.bfloat16 - assert geometry.kv_width == core.head_dim - assert core.kv_scale_orig_quant is None and core.kv_scale_quant_orig is None - _load(core, monkeypatch) - assert core.kv_scale_orig_quant is None and core.kv_scale_quant_orig is None - - -def test_fp8_kv_cache_stores_e4m3_with_unit_fp32_single_element_scales(monkeypatch) -> None: - """The generation kernels read `kvCacheScale[0]` and the QSA store and sparse GQA take - no scale, so the scale is a real tensor holding exactly 1.0: the dense and the sparse - path then agree on what a cache slot holds.""" - core = _core(QuantAlgo.FP8) - assert core.quant_mode == M.QM_FP8_KV_CACHE - assert core.kv_pool_dtype == torch.float8_e4m3fn - # The pool contract holds every manager's K/V pool to this dtype; the index and position - # planes keep the activation dtype. - geometry = core._pool_geometry(32) - assert geometry.kv_dtype == torch.float8_e4m3fn - assert geometry.kv_width == core.head_dim - assert geometry.dtype == torch.bfloat16 - _load(core, monkeypatch) - for scale in (core.kv_scale_orig_quant, core.kv_scale_quant_orig): - assert scale is not None - assert scale.dtype == torch.float32 and tuple(scale.shape) == (1,) - assert torch.equal(scale, torch.ones(1)) - - -def test_nvfp4_kv_cache_stores_e2m1_pairs_with_unit_fp32_scale_triples(monkeypatch) -> None: - """The attention op reads the q, K and V second-level scales out of one fp32 `[3]` tensor, - and the QSA NVFP4 store and sparse GQA take the same tensor, so every writer and reader - applies the one value, 1.0 (the checkpoint has no KV scales).""" - core = _core(QuantAlgo.NVFP4) - assert core.quant_mode == M.QM_NVFP4_KV_CACHE - # An NVFP4 byte holds two values: the pool is uint8, half a head wide. - assert core.kv_pool_dtype == torch.uint8 - geometry = core._pool_geometry(32) - assert geometry.kv_dtype == torch.uint8 - assert geometry.kv_width == core.head_dim // 2 - assert geometry.dtype == torch.bfloat16 - _load(core, monkeypatch) - for scale in (core.kv_scale_orig_quant, core.kv_scale_quant_orig): - assert scale is not None - assert scale.dtype == torch.float32 and tuple(scale.shape) == (3,) - assert torch.equal(scale, torch.ones(3)) - - -def test_other_kv_cache_algorithms_are_refused() -> None: - """The QSA backend serves its sparse path over half, bf16, fp8 and NVFP4 pools and falls - back to dense for any other, so a pool of another storage must not be accepted.""" - with pytest.raises(AssertionError, match="kv_cache_quant_algo"): - _core(QuantAlgo.INT8) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_moe_scalars.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_moe_scalars.py deleted file mode 100644 index 7706883208f5..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_moe_scalars.py +++ /dev/null @@ -1,74 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""qwen4_exp's routed-expert scalars, derived on CPU from the checkpoint's scale rows. - -CuteDSL requantizes FC1's output with ONE global scale per layer, where trtllm-gen folds -each expert's `down_proj.input_scale` into its own FC1 scalar. The CuteDSL scalars are -exact for any spread of `down_proj.input_scale` across experts: they are the engine's -`max(w2 input_scale)` convention. -""" - -from __future__ import annotations - -import pytest -import torch - -__extra_import_path__ = ["../../.."] # noqa: F841 -- repo's file-scoped import hook - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - modeling as M, -) - -pytestmark = pytest.mark.cpu_only - -EXPERTS = 8 - - -def _rows(isc2: torch.Tensor, seed: int = 0) -> dict: - gen = torch.Generator().manual_seed(seed) - isc1 = torch.full((EXPERTS,), 0.0123) - ws2_1 = torch.rand(EXPERTS, generator=gen) * 1e-3 + 1e-4 - return { - "l0_e_isc1": isc1, - "l0_e_isc1_up": isc1.clone(), - "l0_e_ws2_1": ws2_1, - "l0_e_ws2_1_up": ws2_1.clone(), - "l0_e_isc2": isc2, - "l0_e_ws2_2": torch.rand(EXPERTS, generator=gen) * 1e-3 + 1e-4, - } - - -def test_uniform_isc2_keeps_the_per_expert_products_bitwise() -> None: - w = _rows(torch.full((EXPERTS,), 0.0371)) - got = M._moe_layer_scalars(w, 0, "CUTEDSL") - assert torch.equal(got["alpha2"], w["l0_e_isc2"] * w["l0_e_ws2_2"]) - assert torch.equal(got["global_sf"], 1.0 / w["l0_e_isc2"][:1]) - assert got["global_sf"].shape == (1,) - assert torch.equal(got["alpha1"], w["l0_e_isc1"] * w["l0_e_ws2_1"]) - - -def test_isc2_spread_is_folded_into_alpha2() -> None: - """Whatever the spread (here 40x), `global_sf * alpha2[e]` is expert e's own `ws2_2[e]`, - the product trtllm-gen forms from `(1 / isc2[e]) * (isc2[e] * ws2_2[e])`: the FC2 value an - expert produces is the same whether its input was quantized at its own `1 / isc2[e]` or - at the layer's single scale.""" - isc2 = 0.01 * torch.linspace(1.0, 40.0, EXPERTS) - w = _rows(isc2.flip(0)) - cute = M._moe_layer_scalars(w, 0, "CUTEDSL") - gen = M._moe_layer_scalars(w, 0, "TRTLLM") - isc2 = w["l0_e_isc2"] - assert torch.equal(cute["global_sf"], 1.0 / isc2.max().reshape(1)) - # The CuteDSL FC1 quantizes at the smallest per-expert global scale: no expert's FC2 - # input is scaled above what the checkpoint calibrated for it. - assert (cute["global_sf"] <= 1.0 / isc2).all() - assert torch.allclose( - cute["global_sf"] * cute["alpha2"], (1.0 / isc2) * gen["output2_scale_scalar"], rtol=1e-6 - ) - assert torch.allclose(cute["global_sf"] * cute["alpha2"], w["l0_e_ws2_2"], rtol=1e-6) - - -def test_gate_up_input_scale_mismatch_is_refused() -> None: - w = _rows(torch.full((EXPERTS,), 0.0371)) - w["l0_e_isc1_up"] = w["l0_e_isc1_up"] * 1.5 - with pytest.raises(AssertionError, match="gate/up input_scale differ"): - M._moe_layer_scalars(w, 0, "CUTEDSL") diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_multimodal_embeds.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_multimodal_embeds.py deleted file mode 100644 index d7c71b3b43b9..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_multimodal_embeds.py +++ /dev/null @@ -1,274 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""qwen4_exp's multimodal splice runs the vision tower once per modality a request -carries and puts the rows back in prompt order; the tower itself runs a request's frames -in calls of whole frames under a row ceiling. - -The tower is a stub whose rows name the item they came from, so the bookkeeping -- -modality dispatch, item metadata, the chunk slice, the frame grouping and the order of -the concatenated calls -- runs on the CPU without the vision ops or an engine. -""" - -from __future__ import annotations - -import types -from types import SimpleNamespace - -import pytest -import torch - -__extra_import_path__ = ["../../.."] # noqa: F841 -- repo's file-scoped import hook - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - modeling as M, -) -from tensorrt_llm.inputs.multimodal import ( - MULTIMODAL_ENCODER_ITEM_METADATA_KEY, - MultimodalItemMetadata, - MultimodalParams, -) - -pytestmark = pytest.mark.cpu_only - -Shell = M.ModelingV2Qwen38FlashNextNvfp4Sm103Tp1 - -HIDDEN = 3 -MERGE = 2 -IMAGE, VIDEO = 1, 2 - - -class _Tower: - """Stands in for `VisionTower`: every row of item `i` of modality `m` holds - `100 * m + i`, so a spliced row says where it came from.""" - - def __init__(self): - self.calls = 0 - - def __call__(self, pixel_values, grid_thw): - self.calls += 1 - modality = int(pixel_values[0, 0]) - return torch.cat( - [ - torch.full((t * h * w // (MERGE * MERGE), HIDDEN), float(100 * modality + item)) - for item, (t, h, w) in enumerate(grid_thw.tolist()) - ] - ) - - -def _shell(tower): - shell = SimpleNamespace( - model=SimpleNamespace( - v_merge=MERGE, w={"embed": torch.zeros(16, HIDDEN)}, vision_tower=tower - ) - ) - shell._run_tower = types.MethodType(Shell._run_tower, shell) - shell._in_prompt_order = types.MethodType(Shell._in_prompt_order, shell) - return shell - - -def _media(modality, grids): - """One `multimodal_data` entry; `pixel_values` only carries the stub's modality.""" - patches = sum(t * h * w for t, h, w in grids) - key = "pixel_values" if modality == IMAGE else "pixel_values_videos" - grid_key = "image_grid_thw" if modality == IMAGE else "video_grid_thw" - return {key: torch.full((patches, 1), float(modality)), grid_key: torch.tensor(grids)} - - -def _params(image=None, video=None, order=None, runtime=None, mrope=True): - data = {} - if mrope: - data["mrope_config"] = {"mrope_position_ids": torch.zeros(3, 1, 8, dtype=torch.int32)} - if image is not None: - data["image"] = _media(IMAGE, image) - if video is not None: - data["video"] = _media(VIDEO, video) - if order is not None: - data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata(item_refs=order) - return MultimodalParams(multimodal_data=data, multimodal_runtime=runtime) - - -def _splice(tower, params, num_rows, padding=0): - """The tower's rows spliced, as the prefill target's break does, over a prompt whose - tokens 1..num_rows are the placeholders and which `padding` rows follow; returns the - rows that landed there, one value per row, or None when the step carries no media.""" - kwargs = { - "multimodal_params": [params], - "mm_token_indices": torch.arange(1, 1 + num_rows, dtype=torch.int64), - } - vision = Shell._multimodal_embeds(_shell(tower), kwargs) - if vision is None: - return None - total = num_rows + 2 - h = torch.full((total + padding, HIDDEN), -1.0) - target = SimpleNamespace(_step=SimpleNamespace(vision=vision, num_tokens=total)) - M.PrefillTarget._splice(target, h) - assert torch.count_nonzero(h[total:]) == 0, "padding rows must be zeroed" - assert torch.all(h[0] == -1.0) and torch.all(h[total - 1] == -1.0) - return h[1 : 1 + num_rows, 0].tolist() - - -def _item(modality, index): - return modality, index - - -# Two videos of 4 and 8 rows, two images of 4 and 1. -VIDEOS = [(2, 2, 4), (2, 4, 4)] -IMAGES = [(1, 4, 4), (1, 2, 2)] -MIXED_ORDER = [_item("video", 0), _item("image", 0), _item("video", 1), _item("image", 1)] -MIXED_ROWS = [200.0] * 4 + [100.0] * 4 + [201.0] * 8 + [101.0] - - -def test_the_splice_fills_the_placeholders_and_zeroes_the_padding(): - """An image request's rows land on its placeholders in one tower call and the padding - rows after the step are zeroed (`_splice` checks); a text step writes only the padding.""" - tower = _Tower() - assert _splice(tower, _params(image=IMAGES), 5, padding=3) == [100.0] * 4 + [101.0] - assert tower.calls == 1 - target = SimpleNamespace(_step=SimpleNamespace(vision=None, num_tokens=2)) - h = torch.full((4, HIDDEN), -1.0) - M.PrefillTarget._splice(target, h) - assert torch.all(h[:2] == -1.0) and torch.count_nonzero(h[2:]) == 0 - - -def test_video_rows_reach_the_placeholders_item_by_item(): - tower = _Tower() - assert _splice(tower, _params(video=VIDEOS), 12) == [200.0] * 4 + [201.0] * 8 - assert tower.calls == 1 - - -def test_a_mixed_request_is_reassembled_in_prompt_order(): - tower = _Tower() - params = _params(image=IMAGES, video=VIDEOS, order=MIXED_ORDER) - assert _splice(tower, params, len(MIXED_ROWS)) == MIXED_ROWS - assert tower.calls == 2 - - -def test_a_mixed_request_needs_its_manifest(): - with pytest.raises(AssertionError, match="needs item metadata"): - _splice(_Tower(), _params(image=IMAGES, video=VIDEOS), len(MIXED_ROWS)) - - -def test_a_chunk_takes_its_slice_of_the_prompt_order_rows(): - runtime = SimpleNamespace( - num_cached_mm_tokens=3, num_mm_tokens_in_chunk=6, total_embeds_in_request=len(MIXED_ROWS) - ) - params = _params(image=IMAGES, video=VIDEOS, order=MIXED_ORDER, runtime=runtime) - assert _splice(_Tower(), params, 6) == MIXED_ROWS[3:9] - - -def test_a_chunk_without_placeholders_skips_the_tower(): - runtime = SimpleNamespace( - num_cached_mm_tokens=len(MIXED_ROWS), - num_mm_tokens_in_chunk=0, - total_embeds_in_request=len(MIXED_ROWS), - ) - tower = _Tower() - params = _params(image=IMAGES, video=VIDEOS, order=MIXED_ORDER, runtime=runtime) - assert _splice(tower, params, 0) is None - assert tower.calls == 0 - - -def _frames(grid): - """The item-by-item frame list a grid describes: one `(h, w)` per frame.""" - return [(h, w) for t, h, w in grid for _ in range(t)] - - -def _random_grid(gen): - return [ - (int(torch.randint(1, 6, (1,), generator=gen)), 2 * h, 2 * w) - for h, w in torch.randint( - 1, 5, (int(torch.randint(1, 6, (1,), generator=gen)), 2), generator=gen - ).tolist() - ] - - -def test_frames_are_grouped_whole_and_in_order(): - # 16-row frames, 32 rows to a call: one item's frames split across calls. - assert M._frame_calls([(5, 4, 4)], 32) == [[(2, 4, 4)], [(2, 4, 4)], [(1, 4, 4)]] - # A call takes the end of one item and the start of the next. - grid = [(1, 8, 4), (3, 4, 4), (1, 2, 2)] - assert M._frame_calls(grid, 48) == [[(1, 8, 4), (1, 4, 4)], [(2, 4, 4), (1, 2, 2)]] - assert M._frame_calls(grid, 1 << 20) == [grid] - - -def test_a_call_is_filled_to_the_ceiling_and_never_past_it(): - gen = torch.Generator().manual_seed(0) - for _ in range(200): - grid = _random_grid(gen) - cap = int(torch.randint(64, 300, (1,), generator=gen)) - calls = M._frame_calls(grid, cap) - assert [f for call in calls for f in _frames(call)] == _frames(grid) - sizes = [sum(t * h * w for t, h, w in call) for call in calls] - assert all(0 < size <= cap for size in sizes) - # Greedy: the next call's first frame would not have fit. - for size, call in zip(sizes, calls[1:]): - h, w = _frames(call)[0] - assert size + h * w > cap - - -def test_one_frame_past_the_ceiling_is_refused_and_one_at_it_is_not(): - assert M._frame_calls([(2, 8, 8)], 64) == [[(1, 8, 8)], [(1, 8, 8)]] - with pytest.raises(AssertionError, match="one frame is 8x10 = 80 patch rows"): - M._frame_calls([(1, 4, 4), (3, 8, 10)], 64) - - -PIXEL_WIDTH = 4 - - -def _tower(monkeypatch, cap): - """A `VisionTower` whose per-call body is a stub: merged row `j` of a call carries the - pixel value of the call's patch row `4 * j`, which `_pixels` set to the patch's - index in the request.""" - monkeypatch.setattr(M, "_VISION_CALL_MAX_PATCHES", cap) - tower = object.__new__(M.VisionTower) - tower.core = SimpleNamespace( - w={"embed": torch.zeros(1)}, v_pixel_width=PIXEL_WIDTH, v_merge=MERGE - ) - tower.calls = [] - - def run(pixel_values, grid, grid_thw): - tower.calls.append((pixel_values, grid, grid_thw)) - return pixel_values[:: MERGE * MERGE, :1].expand(-1, HIDDEN).clone() - - tower._run = run - return tower - - -def _pixels(grid): - patches = sum(t * h * w for t, h, w in grid) - return torch.arange(patches, dtype=torch.float32)[:, None].expand(-1, PIXEL_WIDTH).clone() - - -def test_the_tower_calls_whole_frames_and_concatenates_their_rows_in_order(monkeypatch): - grid = [(1, 8, 4), (3, 4, 4), (2, 2, 6), (1, 2, 2)] - tower = _tower(monkeypatch, 48) - out = tower(_pixels(grid), torch.tensor(grid)) - patches = sum(t * h * w for t, h, w in grid) - assert out[:, 0].tolist() == [float(4 * j) for j in range(patches // 4)] - assert [call[1] for call in tower.calls] == M._frame_calls(grid, 48) - start = 0 - for pixel_values, call, grid_thw in tower.calls: - rows = sum(t * h * w for t, h, w in call) - assert pixel_values[:, 0].tolist() == [float(i) for i in range(start, start + rows)] - assert rows <= 48 - assert grid_thw.tolist() == [list(row) for row in call] - assert grid_thw.dtype == torch.int64 - start += rows - assert start == patches - - -def test_a_request_past_the_old_whole_request_ceiling_is_served(monkeypatch): - # 33 temporal groups of 32x32 patches (66 frames of 512x512 pixels): 33792 rows, 1024 - # rows past what a single call may carry. - ceiling = M._VISION_CALL_MAX_PATCHES - grid = [(33, 32, 32)] - tower = _tower(monkeypatch, ceiling) - out = tower(_pixels(grid), torch.tensor(grid)) - assert [call[1] for call in tower.calls] == [[(32, 32, 32)], [(1, 32, 32)]] - assert out.shape[0] == 33 * 32 * 32 // 4 - assert out[:, 0].tolist() == [float(4 * j) for j in range(out.shape[0])] - - -def test_a_media_request_needs_the_processors_position_ids(): - with pytest.raises(AssertionError, match="mrope_position_ids"): - _splice(_Tower(), _params(video=VIDEOS, mrope=False), 12) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_validate_request.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_validate_request.py deleted file mode 100644 index 80688e4b3cf9..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_validate_request.py +++ /dev/null @@ -1,209 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""qwen4_exp's admission check fails a bad media request before the forward would assert. - -The check reads only the request's host-side payload, so these run on fake request -objects with meta pixel tensors: no GPU, vision ops or engine. Every refusal is a -`ValueError`, which the executor turns into that one request's error. -""" - -from __future__ import annotations - -import types -from types import SimpleNamespace - -import pytest -import torch - -__extra_import_path__ = ["../../.."] # noqa: F841 -- repo's file-scoped import hook - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - modeling as M, -) -from tensorrt_llm.inputs.multimodal import ( - MULTIMODAL_ENCODER_ITEM_METADATA_KEY, - MultimodalItemMetadata, -) - -pytestmark = pytest.mark.cpu_only - -Shell = M.ModelingV2Qwen38FlashNextNvfp4Sm103Tp1 - -MERGE = 2 -PATCH = 16 -PIXEL_WIDTH = 4 -TEXT_TOKENS = 5 - - -def _shell(): - shell = SimpleNamespace( - model=SimpleNamespace(v_merge=MERGE, v_patch=PATCH, v_pixel_width=PIXEL_WIDTH) - ) - shell.validate_request = types.MethodType(Shell.validate_request, shell) - shell._media_grid = types.MethodType(Shell._media_grid, shell) - shell._check_item_order = Shell._check_item_order - return shell - - -def _media(modality, grids): - """A payload whose pixel tensor has the right shape and no storage.""" - patches = sum(t * h * w for t, h, w in grids) - key = "pixel_values" if modality == "image" else "pixel_values_videos" - grid_key = "image_grid_thw" if modality == "image" else "video_grid_thw" - return { - key: torch.empty(patches, PIXEL_WIDTH, device="meta"), - grid_key: torch.tensor(grids), - } - - -def _request(image=None, video=None, order=None, **overrides): - """A request whose prompt is `TEXT_TOKENS` text tokens around the placeholders the - media produce, with the position ids and placeholder cumsum the processor emits.""" - rows = sum(t * h * w for t, h, w in (image or []) + (video or [])) // (MERGE * MERGE) - prompt_len = rows + TEXT_TOKENS - data = { - "mrope_config": {"mrope_position_ids": torch.zeros(3, 1, prompt_len, dtype=torch.int32)}, - "multimodal_embed_mask_cumsum": torch.cat( - [torch.arange(1, rows + 1), torch.full((TEXT_TOKENS,), rows)] - ), - } - if image is not None: - data["image"] = _media("image", image) - if video is not None: - data["video"] = _media("video", video) - if order is not None: - data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata(item_refs=order) - fields = {"py_multimodal_data": data, "prompt_len": prompt_len} - fields.update(overrides) - return SimpleNamespace(**fields) - - -def _admit(request): - _shell().validate_request(request) - - -def _item(modality, index): - return modality, index - - -IMAGES = [(1, 4, 4), (1, 2, 2)] -VIDEOS = [(2, 2, 4), (2, 4, 4)] -MIXED_ORDER = [_item("video", 0), _item("image", 0), _item("video", 1), _item("image", 1)] - - -def test_clean_requests_are_admitted(): - # A single-modality request needs no manifest. - _admit(_request(image=IMAGES)) - _admit(_request(video=VIDEOS)) - _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER)) - - -def test_a_request_without_media_costs_nothing(): - # Not even the other request attributes are read. - _admit(SimpleNamespace(py_multimodal_data=None)) - _admit(SimpleNamespace(py_multimodal_data={})) - # A warmup request carries position ids and no media. - _admit( - SimpleNamespace( - py_multimodal_data={"mrope_config": {"mrope_position_ids": torch.zeros(3, 1, 8)}} - ) - ) - # Precomputed embeddings arrive without pixels; the splice skips them. - _admit(SimpleNamespace(py_multimodal_data={"multimodal_embedding": torch.zeros(4, 3)})) - - -def test_the_longest_admitted_frame_is_the_calls_ceiling(monkeypatch): - monkeypatch.setattr(M, "_VISION_CALL_MAX_PATCHES", 64) - _admit(_request(image=[(1, 8, 8)])) - _admit(_request(video=[(3, 4, 16)])) - with pytest.raises(ValueError, match="frames of 8x10 = 80 patches"): - _admit(_request(image=[(1, 8, 8), (1, 8, 10)])) - - -def test_a_frame_past_the_ceiling_fails_the_request(): - # 3072x3072 pixels: a 192x192 grid, 36864 patches in one frame. - request = _request(video=[(1, 192, 192)]) - with pytest.raises(ValueError, match=r"video 0 has frames of 192x192 = 36864 patches"): - _admit(request) - # The message names the pixels and the bound. - with pytest.raises(ValueError, match=r"\(3072x3072 pixels\).*at most 32768 patches"): - _admit(request) - - -def test_the_frame_bound_is_per_frame_not_per_item(): - # 33 temporal groups of 32x32 patches: 33792 rows in all, 1024 past one call. - _admit(_request(video=[(33, 32, 32)])) - # 4 frames of 128x128 patches (16384 each) are fine; 192x192 in the second item is not. - _admit(_request(video=[(4, 128, 128)])) - with pytest.raises(ValueError, match="video 1 has frames of 192x192"): - _admit(_request(video=[(4, 128, 128), (1, 192, 192)])) - - -def test_a_malformed_grid_is_refused(): - request = _request(image=[(1, 4, 4)]) - request.py_multimodal_data["image"]["image_grid_thw"] = torch.tensor([(1, 4, 5)]) - with pytest.raises(ValueError, match="spatial_merge_size 2"): - _admit(request) - - -def test_a_grid_that_is_not_one_row_per_item_is_refused(): - request = _request(image=IMAGES) - request.py_multimodal_data["image"]["image_grid_thw"] = torch.tensor([1, 4, 4]) - with pytest.raises(ValueError, match="one \\(t, h, w\\) row per image"): - _admit(request) - - -def test_pixels_that_disagree_with_the_grid_are_refused(): - request = _request(video=VIDEOS) - video = request.py_multimodal_data["video"] - video["pixel_values_videos"] = torch.empty(5, PIXEL_WIDTH, device="meta") - with pytest.raises(ValueError, match="pixel_values_videos is .*describe 48 patch rows"): - _admit(request) - - -def test_a_payload_without_its_pixels_or_grid_is_refused(): - request = _request(video=VIDEOS) - del request.py_multimodal_data["video"]["video_grid_thw"] - with pytest.raises(ValueError, match="this target reads 'pixel_values_videos'"): - _admit(request) - - -def test_missing_or_mis_sized_position_ids_fail_the_request(): - request = _request(video=VIDEOS) - request.py_multimodal_data["mrope_config"]["mrope_position_ids"] = torch.zeros(3, 1, 40) - with pytest.raises(ValueError, match="one position per prompt token"): - _admit(request) - del request.py_multimodal_data["mrope_config"] - with pytest.raises(ValueError, match="mrope_position_ids.*got None"): - _admit(request) - - -def test_a_mixed_request_needs_its_manifest(): - with pytest.raises(ValueError, match="needs item metadata"): - _admit(_request(image=IMAGES, video=VIDEOS)) - with pytest.raises(ValueError, match="needs item metadata"): - _admit(_request(image=IMAGES, video=VIDEOS, order=[])) - - -def test_a_manifest_that_disagrees_with_the_payload_is_refused(): - with pytest.raises(ValueError, match="lists image items \\[0\\]; the payload carries 2"): - _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER[:3])) - with pytest.raises(ValueError, match="lists video items \\[0, 1, 0\\]"): - _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER + [_item("video", 0)])) - with pytest.raises(ValueError, match="is not an item of this request"): - _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER + [_item("audio", 0)])) - with pytest.raises(ValueError, match="is not an item of this request"): - _admit(_request(image=IMAGES, video=VIDEOS, order=[("image",)])) - - -def test_the_placeholder_count_is_the_sum_over_both_modalities(): - request = _request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER) - request.py_multimodal_data["multimodal_embed_mask_cumsum"] = torch.arange(1, 14) - with pytest.raises(ValueError, match="holds 13 image/video placeholder tokens.*17"): - _admit(request) - - -def test_a_request_without_a_cumsum_skips_the_count(): - request = _request(video=VIDEOS) - del request.py_multimodal_data["multimodal_embed_mask_cumsum"] - _admit(request) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_weight_table.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_weight_table.py deleted file mode 100644 index a3bba2823db9..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_weight_table.py +++ /dev/null @@ -1,401 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""qwen4_exp's weight table declares what the hand-written code did. - -The declaration reference below is a transcription of the `__init__` block the table -replaced, keyed the way the table keys it: vision block `b`'s weights are `l{b}_vis_*` -and the MTP module's are `l{num_hidden_layers}_mtp_*`, its layer index in the checkpoint -and the KV pool. The n-gram table is not in either: its height is a checkpoint fact, -declared at load. - -No GPU: nothing here allocates, and `dims` reads the configs alone. -""" - -from __future__ import annotations - -from types import SimpleNamespace - -import pytest -import torch - -__extra_import_path__ = ["../../.."] # noqa: F841 -- repo's file-scoped import hook - -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - modeling as M, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 - weights as W, -) - -pytestmark = pytest.mark.cpu_only - -# Qwen3.8-Flash-Next's own numbers, so the comparison is against the shapes this target -# actually ships rather than a toy configuration. -_TEXT = dict( - num_hidden_layers=48, - layer_types=["full_attention" if i % 4 == 3 else "linear_attention" for i in range(48)], - hidden_size=2560, - hc_count=4, - hc_lowrank=320, - head_dim=256, - num_attention_heads=24, - num_key_value_heads=2, - indexer_n_heads=4, - indexer_kv_heads=1, - indexer_head_dim=128, - linear_num_key_heads=16, - linear_num_value_heads=48, - linear_key_head_dim=128, - linear_value_head_dim=128, - linear_conv_kernel_dim=4, - num_experts=512, - moe_intermediate_size=640, - shared_expert_intermediate_size=640, - vocab_size=248320, - ple_layer_ids=[2], - ple_embed_dim=2560, - ple_conv_kernel_size=4, - ngram_size=3, - heads_per_ngram=8, -) -_VISION = dict( - depth=27, - hidden_size=1152, - intermediate_size=4304, - in_channels=3, - temporal_patch_size=2, - patch_size=16, - num_position_embeddings=2304, - spatial_merge_size=2, -) - - -def _core(mtp: bool, backend: str = "TRTLLM", prefill_copy: bool | None = None): - """The configuration surface the table reads; the second expert copy defaults to kept - under trtllm-gen.""" - return SimpleNamespace( - model_config=SimpleNamespace( - pretrained_config=SimpleNamespace(**_TEXT), - torch_dtype=torch.bfloat16, - spec_config=object() if mtp else None, - ), - vision_config=SimpleNamespace(**_VISION), - mtp_enabled=mtp, - trunk_moe_backend=backend, - moe_prefill_copy=backend == "TRTLLM" if prefill_copy is None else prefill_copy, - ) - - -def test_every_declared_weight_has_a_source(): - """A declared weight with no source allocates memory nothing ever fills -- silently - zero rather than an error, which is what `load`'s coverage assert catches from the - other side. With MTP live: the rows without it are a subset - (`test_mtp_rows_appear_only_when_mtp_is_live`).""" - core = _core(True) - manifest = W.MODEL_WEIGHTS.manifest(core) - for key, sources in manifest.items(): - assert sources, f"{key} is declared with no checkpoint source" - with torch.device("meta"): - declared = W.MODEL_WEIGHTS.declare(W.MODEL_WEIGHTS.dims(core)) - assert set(manifest) == set(declared) - - -def test_mtp_rows_appear_only_when_mtp_is_live(): - off = W.MODEL_WEIGHTS.manifest(_core(False)) - on = W.MODEL_WEIGHTS.manifest(_core(True)) - added = set(on) - set(off) - assert set(off) <= set(on) - m = f"l{_TEXT['num_hidden_layers']}_mtp_" - assert added and all(k.startswith(m) for k in added) - # The draft layer runs QSA over its own indexer, read from the module's keys. - indexer = {f"{m}idx_w", f"{m}idx_qn", f"{m}idx_kn"} - assert indexer <= added - assert all( - str(src).startswith("mtp.layers.0.self_attn.indexer.") - for k in indexer - for src, _, _ in on[k] - ) - - -def _reference_declaration(mtp: bool) -> dict: - """The declaration block from modeling.py, transcribed.""" - t, v = SimpleNamespace(**_TEXT), SimpleNamespace(**_VISION) - bf16, u8, f8, f32, i64 = ( - torch.bfloat16, - torch.uint8, - torch.float8_e4m3fn, - torch.float32, - torch.int64, - ) - hidden, stream, rank = t.hidden_size, t.hc_count * t.hidden_size, t.hc_lowrank - q_width = t.num_attention_heads * t.head_dim - kv_width = t.num_key_value_heads * t.head_dim - index_width = (t.indexer_n_heads + t.indexer_kv_heads) * t.indexer_head_dim - key_dim = t.linear_num_key_heads * t.linear_key_head_dim - value_dim = t.linear_num_value_heads * t.linear_value_head_dim - conv_dim = 2 * key_dim + value_dim - experts, inter, shared = ( - t.num_experts, - t.moe_intermediate_size, - t.shared_expert_intermediate_size, - ) - ngram_heads = t.heads_per_ngram * (t.ngram_size - 1) - vh, vi = v.hidden_size, v.intermediate_size - vm = vh * v.spatial_merge_size**2 - nl = t.num_hidden_layers - w = {} - - def P(name, *shape, dtype=bf16): - w[name] = (tuple(shape), dtype) - - def hyper_connection(prefix): - P(f"{prefix}_norm", stream) - P(f"{prefix}_down", rank, stream) - P(f"{prefix}_up", stream, rank) - P(f"{prefix}_inject", t.hc_count, stream) - - def gated_attention(prefix, fused_in=False): - if fused_in: - # [q | k | v | gate | index q, k | one index head of zero pad] - rows = 2 * q_width + 2 * kv_width + index_width + t.indexer_head_dim - P(f"{prefix}attn_in_w", rows, hidden) - else: - P(f"{prefix}q_w", 2 * q_width, hidden) - P(f"{prefix}k_w", kv_width, hidden) - P(f"{prefix}v_w", kv_width, hidden) - P(f"{prefix}o_w", hidden, q_width) - P(f"{prefix}q_norm", t.head_dim) - P(f"{prefix}k_norm", t.head_dim) - - def shared_expert(prefix): - # gate|up, then the per-token gate's row and zero rows (`weights._cat_rows`). - P(f"{prefix}sh_gu", 2 * shared + 8, hidden) - P(f"{prefix}sh_dn", hidden, shared) - - for i in range(nl): - hyper_connection(f"l{i}_ahc") - hyper_connection(f"l{i}_mhc") - if t.layer_types[i] == "full_attention": - gated_attention(f"l{i}_", fused_in=True) - P(f"l{i}_idx_qn", t.indexer_head_dim) - P(f"l{i}_idx_kn", t.indexer_head_dim) - else: - # [in_proj_qkv | in_proj_z | in_proj_a | in_proj_b], one row block each. - P(f"l{i}_in_proj_w", conv_dim + value_dim + 2 * t.linear_num_value_heads, hidden) - P(f"l{i}_A_log", t.linear_num_value_heads) - P(f"l{i}_dt_bias", t.linear_num_value_heads) - P(f"l{i}_conv_w", conv_dim, 1, t.linear_conv_kernel_dim) - P(f"l{i}_gnorm", t.linear_value_head_dim) - P(f"l{i}_out_w", hidden, value_dim) - P(f"l{i}_router", experts, hidden) - P(f"l{i}_fc1_w", experts, 2 * inter, hidden // 2, dtype=u8) - P(f"l{i}_fc1_s", experts, 2 * inter, hidden // 16, dtype=u8) - P(f"l{i}_fc2_w", experts, hidden, inter // 2, dtype=u8) - P(f"l{i}_fc2_s", experts, hidden, inter // 16, dtype=u8) - # The trtllm-gen backend's second, CuteDSL-layout copy of the experts. - P(f"l{i}_fc1_w_cd", experts, 2 * inter, hidden // 2, dtype=u8) - P(f"l{i}_fc1_s_cd", experts, 2 * inter, hidden // 16, dtype=u8) - P(f"l{i}_fc2_w_cd", experts, hidden, inter // 2, dtype=u8) - P(f"l{i}_fc2_s_cd", experts, hidden, inter // 16, dtype=u8) - for name in ("e_isc1", "e_isc1_up", "e_ws2_1", "e_ws2_1_up", "e_isc2", "e_ws2_2"): - P(f"l{i}_{name}", experts, dtype=f32) - shared_expert(f"l{i}_") - P("fhc_norm", stream) - P("fhc_down", rank, stream) - P("fhc_up", stream, rank) - P("embed", t.vocab_size, hidden) - P("ple_key_w", stream, hidden) - P("ple_val_w", t.ple_embed_dim, hidden) - P("ple_nk", stream) - P("ple_nq", stream) - P("ple_nc", stream) - P("ple_conv_w", stream, 1, t.ple_conv_kernel_size) - P("ple_mult", t.ngram_size, dtype=i64) - P("ple_vocab", ngram_heads, dtype=i64) - P("ple_offs", ngram_heads, dtype=i64) - P("ple_tscale", 1) - P("vis_patch_w", vh, v.in_channels, v.temporal_patch_size, v.patch_size, v.patch_size) - P("vis_patch_b", vh) - P("vis_pos", v.num_position_embeddings, vh) - for b in range(v.depth): - P(f"l{b}_vis_n1_w", vh) - P(f"l{b}_vis_n1_b", vh) - P(f"l{b}_vis_qkv_w", 3 * vh, vh) - P(f"l{b}_vis_qkv_b", 3 * vh) - P(f"l{b}_vis_proj_w", vh, vh) - P(f"l{b}_vis_proj_b", vh) - P(f"l{b}_vis_n2_w", vh) - P(f"l{b}_vis_n2_b", vh) - P(f"l{b}_vis_fc1_w", vi, vh) - P(f"l{b}_vis_fc1_b", vi) - P(f"l{b}_vis_fc2_w", vh, vi) - P(f"l{b}_vis_fc2_b", vh) - P("vis_m_n_w", vh) - P("vis_m_n_b", vh) - P("vis_m_fc1_w", vm, vm) - P("vis_m_fc1_b", vm) - P("vis_m_fc2_w", hidden, vm) - P("vis_m_fc2_b", hidden) - if mtp: - m = f"l{nl}_mtp_" - P(f"{m}pre_e_norm", hidden) - P(f"{m}pre_h_norm", stream) - P(f"{m}fc_e", hidden, hidden) - P(f"{m}fc_h", hidden, hidden) - hyper_connection(f"{m}ahc") - hyper_connection(f"{m}mhc") - gated_attention(m) - P(f"{m}idx_w", index_width, hidden) - P(f"{m}idx_qn", t.indexer_head_dim) - P(f"{m}idx_kn", t.indexer_head_dim) - P(f"{m}router", experts, hidden) - P(f"{m}fc1_w", experts, 2 * inter, hidden, dtype=f8) - P(f"{m}fc1_s", experts, (2 * inter) // 128, hidden // 128, dtype=f32) - P(f"{m}fc2_w", experts, hidden, inter, dtype=f8) - P(f"{m}fc2_s", experts, hidden // 128, inter // 128, dtype=f32) - shared_expert(m) - P(f"{m}fhc_norm", stream) - P(f"{m}fhc_down", rank, stream) - P(f"{m}fhc_up", stream, rank) - return w - - -def test_the_table_declares_what_the_init_block_did(): - """With MTP live; without it the table drops exactly the MTP module's rows - (`test_mtp_rows_appear_only_when_mtp_is_live`).""" - core = _core(True) - with torch.device("meta"): - declared = W.MODEL_WEIGHTS.declare(W.MODEL_WEIGHTS.dims(core)) - got = {k: (tuple(p.shape), p.dtype) for k, p in declared.items()} - want = _reference_declaration(True) - assert set(got) == set(want), ( - sorted(set(got) - set(want))[:8], - sorted(set(want) - set(got))[:8], - ) - wrong = {k: (got[k], want[k]) for k in got if got[k] != want[k]} - assert not wrong, list(wrong.items())[:8] - - -def test_gdn_input_projections_fill_adjacent_row_blocks(): - """The concatenated GDN input projection: each checkpoint tensor lands in its own - row block, in the column order the forward splits the GEMM output by, and the - blocks tile the declared rows exactly.""" - t = SimpleNamespace(**_TEXT) - core = _core(False) - manifest = W.MODEL_WEIGHTS.manifest(core) - key_dim = t.linear_num_key_heads * t.linear_key_head_dim - value_dim = t.linear_num_value_heads * t.linear_value_head_dim - widths = ( - 2 * key_dim + value_dim, - value_dim, - t.linear_num_value_heads, - t.linear_num_value_heads, - ) - gdn = [i for i, kind in enumerate(t.layer_types) if kind == "linear_attention"] - assert len(gdn) == 36 - for i in gdn: - rows = manifest[f"l{i}_in_proj_w"] - names = [ckpt.rsplit(".", 2)[-2] for ckpt, _, _ in rows] - assert names == ["in_proj_qkv", "in_proj_z", "in_proj_a", "in_proj_b"], names - assert all(ckpt.startswith(f"{W._LM}.layers.{i}.linear_attn.") for ckpt, _, _ in rows) - assert all(transform is None for _, _, transform in rows) - spans = [(index[0].start, index[0].stop) for _, index, _ in rows] - starts = [0, widths[0], widths[0] + widths[1], widths[0] + widths[1] + widths[2]] - assert spans == [(s, s + w) for s, w in zip(starts, widths)], spans - assert spans[-1][1] == sum(widths) == 16480 - - -def test_attention_input_projection_relayout(): - """`_attn_in_w` puts q, k, v, the output gate and the index projection side by side - so that one GEMM's columns are exactly what the four separate projections gave: - q and gate de-interleaved from q_proj's per-head `[q | gate]` rows, then a zero pad - of one index head. Checked at the checkpoint's head geometry, a narrow hidden.""" - t = SimpleNamespace(**_TEXT) - heads, d, kvh, hidden = t.num_attention_heads, t.head_dim, t.num_key_value_heads, 64 - idx_rows = (t.indexer_n_heads + t.indexer_kv_heads) * t.indexer_head_dim - g = torch.Generator().manual_seed(0) - q = torch.randn(2 * heads * d, hidden, generator=g, dtype=torch.float64) - k = torch.randn(kvh * d, hidden, generator=g, dtype=torch.float64) - v = torch.randn(kvh * d, hidden, generator=g, dtype=torch.float64) - idx = torch.randn(idx_rows, hidden, generator=g, dtype=torch.float64) - core = SimpleNamespace(heads_q=heads, head_dim=d, index_head_dim=t.indexer_head_dim) - fused = W._attn_in_w((q, k, v, idx), core) - rows = 2 * heads * d + 2 * kvh * d + idx_rows + t.indexer_head_dim - assert fused.shape == (rows, hidden) and fused.is_contiguous() - - x = torch.randn(3, hidden, generator=g, dtype=torch.float64) - y = x @ fused.t() - qg = (x @ q.t()).reshape(3, heads, 2 * d) - want = torch.cat( - [ - qg[..., :d].reshape(3, -1), - x @ k.t(), - x @ v.t(), - qg[..., d:].reshape(3, -1), - x @ idx.t(), - torch.zeros(3, t.indexer_head_dim, dtype=torch.float64), - ], - 1, - ) - torch.testing.assert_close(y, want, rtol=0, atol=1e-12) - # The row is a whole number of head_dim heads: the main norm+rope addresses it so. - assert rows % d == 0 - - -def test_cutedsl_fc1_pairs_up_and_gate_rows_in_64_row_groups(): - """Destination rows `[128k, 128k + 64)` hold up rows `[64k, 64k + 64)` and the next 64 - hold the matching gate rows, which sit `I` rows later in the `[up ; gate]` stack.""" - inter = _TEXT["moe_intermediate_size"] - core = SimpleNamespace(trunk_moe_backend="CUTEDSL") - perm = W._fc1_perm(2 * inter, "cpu", core) - assert sorted(perm.tolist()) == list(range(2 * inter)) - for k in range(inter // 64): - assert perm[128 * k : 128 * k + 64].tolist() == list(range(64 * k, 64 * k + 64)) - assert perm[128 * k + 64 : 128 * k + 128].tolist() == list( - range(inter + 64 * k, inter + 64 * k + 64) - ) - - -def test_cutedsl_fc2_keeps_the_checkpoint_row_order_and_trtllm_gen_shuffles_it(): - """trtllm-gen's 32-row block shuffle moves source row `4u + v` to destination row - `8v + u`; CuteDSL reads `down_proj` as the checkpoint stores it.""" - rows = torch.arange(64).unsqueeze(1).expand(64, 8).contiguous() - kept = W._expert_fc2_w(rows, SimpleNamespace(trunk_moe_backend="CUTEDSL")) - assert torch.equal(kept, rows) - shuffled = W._expert_fc2_w(rows, SimpleNamespace(trunk_moe_backend="TRTLLM")) - for u in range(8): - for v in range(4): - assert int(shuffled[8 * v + u, 0]) == 4 * u + v - assert int(shuffled[32 + 8 * v + u, 0]) == 32 + 4 * u + v - - -@pytest.mark.parametrize("backend, prefill_copy", [("TRTLLM", True), ("CUTEDSL", False)]) -def test_cute_dsl_expert_copy_rows_only_when_kept(backend, prefill_copy): - """With the second expert copy (trtllm-gen, a step able to reach the CuteDSL dispatch - threshold) every layer also carries its experts in the CuteDSL layout (`fc*_cd`); - without it, and under CuteDSL whose one copy is that layout, no `_cd` row is declared.""" - manifest = W.MODEL_WEIGHTS.manifest(_core(False, backend, prefill_copy)) - rows = [k for k in manifest if k.endswith(("_fc1_w_cd", "_fc1_s_cd", "_fc2_w_cd", "_fc2_s_cd"))] - layers = _TEXT["num_hidden_layers"] - assert len(rows) == (4 * layers if prefill_copy else 0), len(rows) - - -@pytest.mark.parametrize( - "backend, max_num_tokens, env, kept", - [ - ("TRTLLM", 4096, None, True), - ("TRTLLM", 4095, None, False), - ("TRTLLM", 16384, "0", False), - ("CUTEDSL", 16384, None, False), - ], -) -def test_second_expert_copy_kept_only_where_its_launches_can_happen( - monkeypatch, backend, max_num_tokens, env, kept -): - """The copy costs about 1.3 GiB a layer: it is kept only under trtllm-gen, only when a - step's rows can reach the CuteDSL dispatch threshold, and not when the env var is off.""" - if env is None: - monkeypatch.delenv(M._MOE_PREFILL_COPY_ENV, raising=False) - else: - monkeypatch.setenv(M._MOE_PREFILL_COPY_ENV, env) - assert M._moe_prefill_copy(backend, max_num_tokens) is kept diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/vision/test_modeling_v2_patch_embed_conv3d.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/vision/test_modeling_v2_patch_embed_conv3d.py new file mode 100644 index 000000000000..62f900fd6ed8 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/vision/test_modeling_v2_patch_embed_conv3d.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp default target's patch_embed_conv3d entry. + +The reference is a true-fp32 matmul of the flattened patch against the +flattened filter. The op is the same contraction made by cuBLASLt, so each +cell also holds the reference to an fp32 `F.conv3d` of the 5-D patches: that +pins the flatten order the two share to `nn.Conv3d`'s rather than letting one +restate the other. +""" + +import contextlib +import zlib + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.ops.vision.patch_embed_conv3d import ( # noqa: E501 + patch_embed_conv3d as op, +) + +__extra_import_path__ = ["../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "patch_embed_conv3d requires a CUDA device" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """torch's default band for the dtype. + + One fp32 accumulation and one rounding on both sides, so the difference is + accumulation order alone: nearly every output is bitwise equal and the rest + are one-ulp flips of the element, inside the default bf16 band. Wrong + variants (bias dropped, kh/kw transposed, channel-fastest rows, kt swapped) + land thousands of times outside it. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + torch.testing.assert_close(out, ref) + + +@contextlib.contextmanager +def _cudnn_true_fp32(): + prev = torch.backends.cudnn.allow_tf32 + torch.backends.cudnn.allow_tf32 = False + try: + yield + finally: + torch.backends.cudnn.allow_tf32 = prev + + +def _conv3d_fp32(pixel_values, weight, bias): + """`nn.Conv3d` semantics, independently: an fp32 cuDNN convolution of the + 5-D patches with the stride equal to the kernel.""" + out_channels, in_channels, kt, kh, kw = weight.shape + with _cudnn_true_fp32(): + out = F.conv3d( + pixel_values.float().view(-1, in_channels, kt, kh, kw), + weight.float(), + None if bias is None else bias.float(), + stride=(kt, kh, kw), + ) + return out.view(-1, out_channels) + + +IN_CHANNELS = 3 +KT, KH, KW = 2, 16, 16 +EMBED_DIM = 1152 +PATCH_NUMEL = IN_CHANNELS * KT * KH * KW # 1536 + + +def _operands( + num_patches, + dtype=torch.bfloat16, + seed=0, + in_channels=IN_CHANNELS, + kernel=(KT, KH, KW), + embed_dim=EMBED_DIM, + bias=True, +): + """Pixels in [-1, 1] (processor mean/std 0.5), weights at the tower's + initializer_range 0.02.""" + g = torch.Generator(device="cuda").manual_seed(seed) + kt, kh, kw = kernel + pixel_values = torch.empty(num_patches, in_channels * kt * kh * kw, device="cuda") + pixel_values = pixel_values.uniform_(-1.0, 1.0, generator=g).to(dtype) + weight = ( + 0.02 * torch.randn(embed_dim, in_channels, kt, kh, kw, generator=g, device="cuda") + ).to(dtype) + bias_t = (0.02 * torch.randn(embed_dim, generator=g, device="cuda")).to(dtype) if bias else None + return pixel_values, weight, bias_t + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + spec = cell.spec + pixel_values, weight, bias = _operands( + spec["num_patches"], dtype=spec["dtype"], seed=zlib.crc32(cell.why.encode()) + ) + before = pixel_values.clone() + out = op(pixel_values, weight, bias) + assert out.shape == (pixel_values.shape[0], weight.shape[0]) + assert out.dtype == pixel_values.dtype and out.is_contiguous() + assert out.data_ptr() != pixel_values.data_ptr() + assert torch.equal(pixel_values, before), "the pixels are read-only" + + # The reference is the convolution: fp32 against an fp32 F.conv3d. + ref32 = op.reference(pixel_values.float(), weight.float(), bias.float()) + torch.testing.assert_close(ref32, _conv3d_fp32(pixel_values, weight, bias)) + + _compare(out, op.reference(pixel_values, weight, bias)) diff --git a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py index a7a802d8c6d9..13dee2f51c6d 100644 --- a/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py +++ b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py @@ -1925,7 +1925,7 @@ def _cell_layer(spec: dict, layer_factory: Callable[[], tuple[dict, dict]]) -> t ), ), ), - "qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1": ( + "qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default": ( Cell( why="qwen4_exp decode: one token over 512 experts, top-10 renormalized, H 2560 / I 640", spec=dict(tokens=1, **_QWEN4_EXP), diff --git a/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py b/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py index 105296323d6c..cb794c987b0a 100644 --- a/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py +++ b/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py @@ -50,7 +50,7 @@ def _unpack(packed): spec=dict(shape=(16, 256), dtype=torch.bfloat16, vec=16, swizzled=True), ), ), - "qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1": ( + "qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default": ( Cell( why="qwen3.8 decode: 1 token, hidden 2560, linear -- its routed-expert input", spec=dict(shape=(1, 2560), dtype=torch.bfloat16, vec=16, swizzled=False), diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_reference_precision.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_reference_precision.py index f070ddd2f96e..4218ce9bb5df 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_reference_precision.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_reference_precision.py @@ -151,12 +151,9 @@ def test_moe_reference_does_not_override_precision(relative_path: str) -> None: "moe/test_modeling_v2_fp4_block_scale_moe_runner.py", "moe/test_modeling_v2_mxe4m3_mxe2m1_block_scale_moe_runner.py", "models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/moe/test_modeling_v2_fused_moe.py", - "models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/" - "test_modeling_v2_fp4_block_scale_moe_renorm_routed.py", - "models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/" - "test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py", - "models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/" - "test_modeling_v2_fp4_moe_renorm_cute_dsl.py", + "models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py", + "models/qwen4_exp/moe/test_modeling_v2_fp4_block_scale_moe_renorm_permuted.py", + "models/qwen4_exp/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py", ], ) def test_moe_test_reference_calls_scope_precision(relative_path: str) -> 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 cb9e43bb2c6b..73b430a3f50a 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -1058,7 +1058,7 @@ def test_kimi_k3_on_another_sm_does_not_match(monkeypatch): # --- qwen4_exp: one-model MTP-Eagle over a linear chain, one rank --------------- -_QWEN4_EXP_TARGET = "models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.modeling." +_QWEN4_EXP_TARGET = "models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.modeling." @pytest.fixture @@ -1176,7 +1176,7 @@ def test_qwen4_exp_other_kv_cache_algorithms_are_declined(algo): _model_config(_qwen4_exp_config()), _llm_args("auto") ) with pytest.raises(NoTarget, match="KV cache"): - qwen4_exp_routing._check_engine(None, algo, choices) + qwen4_exp_routing._check_engine(None, algo, choices, "tp1") @pytest.mark.usefixtures("_qwen4_exp_sm") @@ -1225,10 +1225,10 @@ def test_qwen4_exp_runtime_draft_length_falls_back(spec_kwargs): @pytest.mark.usefixtures("_qwen4_exp_sm") -@pytest.mark.parametrize("max_draft_len", [4, 7]) -def test_qwen4_exp_draft_beyond_three_needs_the_cached_replay(monkeypatch, max_draft_len): - """The legacy per-step recurrent states are certified at up to three drafted tokens and - hold one fp32 state per verify row.""" +@pytest.mark.parametrize("max_draft_len", [1, 3, 7]) +def test_qwen4_exp_mtp_needs_the_cached_replay(monkeypatch, max_draft_len): + """The target verifies on the GDN cached replay only, never on the legacy per-step + recurrent states.""" monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "0") config = _model_config( _qwen4_exp_config(), spec_config=MTPDecodingConfig(max_draft_len=max_draft_len) @@ -1239,13 +1239,6 @@ def test_qwen4_exp_draft_beyond_three_needs_the_cached_replay(monkeypatch, max_d _claimed(config, dict(kv_cache_config=kv_cache_config)) -@pytest.mark.usefixtures("_qwen4_exp_sm") -def test_qwen4_exp_legacy_verify_matches_up_to_three_drafted_tokens(monkeypatch): - monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "0") - config = _model_config(_qwen4_exp_config(), spec_config=MTPDecodingConfig(max_draft_len=3)) - _claimed(config, dict(kv_cache_config=KvCacheConfig(mamba_ssm_cache_dtype="float32"))) - - @pytest.mark.usefixtures("_qwen4_exp_sm") @pytest.mark.parametrize("max_draft_len", [4, 7]) def test_qwen4_exp_cutedsl_declines_drafts_beyond_three(monkeypatch, max_draft_len): @@ -1302,10 +1295,9 @@ def test_qwen4_exp_mtp_on_a_float16_state_cache_declines_with_replay_off(monkeyp @pytest.mark.usefixtures("_qwen4_exp_sm") -def test_qwen4_exp_mtp_on_a_bf16_state_cache_matches_with_replay_off(monkeypatch): +def test_qwen4_exp_without_speculation_matches_with_replay_off(monkeypatch): monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "0") - config = _model_config(_qwen4_exp_config(), spec_config=MTPDecodingConfig(max_draft_len=3)) - _claimed(config) + _claimed(_model_config(_qwen4_exp_config())) @pytest.mark.usefixtures("_qwen4_exp_sm")