diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 7fe0694b9b36..b1089b34134f 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -1355,7 +1355,11 @@ void attention(torch::Tensor q, std::optional k, std::optionalmRotaryEmbeddingLongMscale = rope_long_m_scale; op->mRotaryEmbeddingMaxPositions = rope_max_positions; op->mRotaryEmbeddingOriginalMaxPositions = rope_original_max_positions; + // Only the FP8 context FMHA kernels read an NVFP4 KV cache, whatever the output type; they do so on SM100-family + // GPUs only, so elsewhere an NVFP4 KV cache still fails the FP4 KV check in AttentionOp. MLA keeps its own NVFP4 KV + // handling (mUseNvfp4MlaKvCache, mFP8ContextMLA below) and must not pick this up. op->mFP8ContextFMHA = is_fp8_out || is_fp4_out || (op->mKVCacheQuantMode.hasFp8KvCache() && use_paged_context_fmha) + || (!is_mla_enable && op->mKVCacheQuantMode.hasFp4KvCache() && tensorrt_llm::common::isSM100Family()) || use_sage_attn; // SageAttention block sizes and quantization mode. op->mSageAttnNumEltsPerBlkQ = static_cast(sage_attn_num_elts_per_blk_q); diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md index 5e34934827bc..4e86b1664588 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md @@ -36,10 +36,15 @@ exists to catch. The decision is one function, `route`, in the family's `routing.py`. It reads the facts that fix which code runs -- the checkpoint's config shape, the -parallel mapping, the GPU's compute capability and the speculative-decoding -config -- and either returns a target class or raises `NoTarget` with the -reason in words. Tuning knobs such as `max_batch_size`, `max_num_tokens` and CUDA graph -settings are not routing inputs: they do not change which code runs, so a +parallel mapping, the GPU's compute capability, the speculative-decoding +config, the KV cache algorithm the engine will run with, and the engine choices +it derives from the `ModelConfig` and the LLM arguments it is also given +(qwen4_exp: MoE backend, recurrent-state cache dtype, prefill CUDA graph +backend) -- and either returns a target class or raises `NoTarget` with the +reason in words. A target declines what its +constructor would refuse, so that `auto` falls back instead of failing mid-build. +Tuning knobs such as `max_batch_size`, `max_num_tokens` and the CUDA graph +bucket lists are not routing inputs: they do not change which code runs, so a target claims a configuration regardless of them. `require` therefore promises that the target ran, not that a particular knob setting was measured. The model loader decides when it is about to build the model, after the built-in model's @@ -67,7 +72,8 @@ LLM(model=..., modeling_v2=...) -> ModelLoader | _router_index.py +-- architectures[0] -> routing module models//routing.py +-- route(config, mapping, compute_capability, - | spec_config, kv_cache_quant_algo) + | spec_config, kv_cache_quant_algo, + | model_config, llm_args) | -> target class, or NoTarget | config.modeling_v2_target = class diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index e8228918397e..cd2a1169d7bd 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -17,10 +17,12 @@ architecture family does. * The routing module's ``route`` owns the rest of the decision as one forward-reading tree over the facts that fix which code runs: the - checkpoint's config, the parallel mapping, the GPU's compute capability and - the speculative-decoding config. It returns the target class, or raises - :class:`NoTarget` with the reason in words. Reading that one file tells you - where any configuration lands. + checkpoint's config, the parallel mapping, the GPU's compute capability, the + speculative-decoding config and the KV-cache algorithm the engine will run + with. It also receives the ``ModelConfig`` and the LLM arguments, from which a + family derives any other engine choice its targets depend on. It returns the + target class, or raises :class:`NoTarget` with the reason in words. Reading + that one file tells you where any configuration lands. Tuning knobs -- ``max_batch_size``, ``max_num_tokens``, CUDA graph settings -- are deliberately not routing inputs. They do not change which code runs, so a @@ -195,6 +197,8 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti _device_compute_capability(), config.spec_config, _kv_cache_quant_algo(config, llm_args), + model_config=config, + llm_args=llm_args, ) except NoTarget as why: if mode is ModelingV2Mode.REQUIRE: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py index 73c6f700c652..e6dbc41e29ba 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py @@ -103,6 +103,21 @@ class FusedQkNormRope(OpWrapper): "GEMMA_PREFILL_ATOL", spec=dict(tokens=8192, gemma_prefill_band=True, **_QWEN_MAIN), ), + # The main rope run in place on qwen4_exp's fused input projection, whose + # row is [q | k | v | gate | index | pad] = 55 heads of 256: everything + # after k is passed as v heads, which only sets the row stride. + Cell( + why="qwen4_exp fused-row main rope, 1 token: 29 v heads only set the row stride", + spec=dict(tokens=1, **{**_QWEN_MAIN, "num_heads_v": 29}), + ), + Cell( + why="qwen4_exp fused-row main rope, 128 tokens: decode at max concurrency", + spec=dict(tokens=128, **{**_QWEN_MAIN, "num_heads_v": 29}), + ), + Cell( + why="qwen4_exp fused-row main rope, 8192 tokens: prefill, under GEMMA_PREFILL_ATOL", + spec=dict(tokens=8192, gemma_prefill_band=True, **{**_QWEN_MAIN, "num_heads_v": 29}), + ), Cell( why="qwen4_exp index rope (mRoPE), 1 token: one k and one v head, the " "narrowest layout; 5 warps of an 8-warp block", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py index fed8d5fa4755..80a2ad211f21 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.py @@ -26,6 +26,10 @@ #: quant_mode bit for an fp8 KV cache. QM_FP8_KV_CACHE = 128 +#: quant_mode bit for an NVFP4 KV cache: e2m1 values in the pool, one e4m3 scale +#: per 16 values in a second pool, fp32 [3] kv scales (q, K, V). +QM_NVFP4_KV_CACHE = 8192 + #: qwen4_exp's full-attention layer -- the twelve trunk layers and the MTP #: module alike: packed-QKV causal GQA over a bf16 pool, rope applied outside, #: the paged-context path on (the target serves with block reuse and chunked @@ -63,6 +67,31 @@ mask_type=0, ) +#: The qwen4_exp full-attention call over an fp8-e4m3 pool at the production kv +#: scale 1.0 (fp32 [1] scale tensors, quant_mode 128): what the trunk layers and +#: the MTP draft layer pass when the engine serves an fp8 KV cache. Same chosen +#: pool sizes as `_QWEN4_DENSE`. +_QWEN4_FP8 = dict(_QWEN4_DENSE, quant_mode=QM_FP8_KV_CACHE, kv_scale=1.0) + +#: The same call over an NVFP4 pool: e2m1 data and a second pool of e4m3 block +#: scales, fp32 [3] unit kv scales, `[1, 2, 2]` pool pointers, bf16 output -- +#: what the trunk layers and the MTP draft layer pass when the engine serves an +#: NVFP4 KV cache. Same chosen pool sizes as `_QWEN4_DENSE`. +_QWEN4_NVFP4 = dict(_QWEN4_DENSE, quant_mode=QM_NVFP4_KV_CACHE, kv_scale=1.0) + +#: `_QWEN4_DENSE` and `_QWEN4_FP8` for a verify step of up to 32 requests, the captured-batch +#: ceiling: one block-offset row per request and pages for 32 histories (the largest cell +#: needs 133 of the 144). +_QWEN4_DENSE_G32 = dict(_QWEN4_DENSE, max_num_requests=32, num_blocks=144) +_QWEN4_FP8_G32 = dict(_QWEN4_FP8, max_num_requests=32, num_blocks=144) + + +def _spread(batch: int, longest: int) -> tuple[int, ...]: + """`batch` histories rising evenly from 1 to `longest` tokens: one-page, page-boundary + and multi-page sequences in one verify step.""" + return tuple(1 + (longest - 1) * i // (batch - 1) for i in range(batch)) + + #: A GQA 24/4/256 full-attention layer (Qwen3.8-27B) over an fp8-e4m3 pool at #: the production kv scale 1.0. Same chosen pool sizes as `_QWEN4_DENSE`. _FP8_24_4_256 = dict( @@ -90,8 +119,9 @@ class ThopAttention(OpWrapper): it, by different paths: gpt-oss on the standard path with attention sinks and an alternating sliding window, deepseek on MLA's two phases, and qwen4_exp on the standard path twice over -- its full-attention layers and - MTP module (GQA 24/2/256, a P-token draft block per generation sequence, - a two-pool hybrid mapping), and its vision tower with no KV cache at all. + MTP module (GQA 24/2/256 over a bf16 or an fp8-e4m3 pool, a P-token draft + block per generation sequence, a two-pool hybrid mapping), and its vision + tower with no KV cache at all. Writes rows `[:num_tokens]` of `output` and returns `output` -- the same buffer, mutated and handed back so a call site can show what it touched @@ -277,7 +307,7 @@ class ThopAttention(OpWrapper): ), ), Cell( - why="qwen4_exp MTP verify P=4 (max_draft_len 3): the target's ceiling", + why="qwen4_exp MTP verify P=4 (max_draft_len 3)", spec=dict( _QWEN4_DENSE, cached=(70, 35), @@ -286,6 +316,61 @@ class ThopAttention(OpWrapper): predicted_tokens_per_seq=4, ), ), + *( + Cell( + why=f"qwen4_exp MTP verify P={p} (max_draft_len {p - 1})" + + (": the target's ceiling" if p == 8 else ""), + spec=dict( + _QWEN4_DENSE, + cached=(70, 35), + new=(p, p), + num_contexts=0, + predicted_tokens_per_seq=p, + ), + ) + for p in (5, 6, 7, 8) + ), + Cell( + why="qwen4_exp P=8 G=1: one draft block from row 95, across the page edge at 96", + spec=dict( + _QWEN4_DENSE, + cached=(95,), + new=(8,), + num_contexts=0, + predicted_tokens_per_seq=8, + ), + ), + # The trtllm-gen launch picks its generation kernel by batch size at P 5..8, so each P + # is also driven at 32 requests and at one batch in 17..20, with histories of every + # page shape in the step. + *( + Cell( + why=f"qwen4_exp bf16 batch P={p} G=32: 32 blocks over histories 1..218, " + f"{32 * p} q rows", + spec=dict( + _QWEN4_DENSE_G32, + cached=_spread(32, 218), + new=(p,) * 32, + num_contexts=0, + predicted_tokens_per_seq=p, + ), + ) + for p in (5, 6, 7, 8) + ), + *( + Cell( + why=f"qwen4_exp bf16 batch P={p} G={g}: {g} blocks over histories 200..1, " + f"{g * p} q rows", + spec=dict( + _QWEN4_DENSE_G32, + cached=tuple(reversed(_spread(g, 200))), + new=(p,) * g, + num_contexts=0, + predicted_tokens_per_seq=p, + ), + ) + for p, g in ((5, 17), (6, 18), (7, 19), (8, 20)) + ), Cell( why="qwen4_exp mixed step: a fresh 20-row context ahead of two P=2 " "draft blocks, generation rows token-major after the context rows", @@ -297,6 +382,17 @@ class ThopAttention(OpWrapper): predicted_tokens_per_seq=2, ), ), + Cell( + why="qwen4_exp mixed step P=8: a fresh 20-row context ahead of two P=8 " + "draft blocks over histories 70 and 35", + spec=dict( + _QWEN4_DENSE, + cached=(0, 70, 35), + new=(20, 8, 8), + num_contexts=1, + predicted_tokens_per_seq=8, + ), + ), Cell( why="qwen4_exp hybrid pool: row 47 of the 48-row two-pool mapping " "(pool 0, layer-in-pool 11), P=2 verify", @@ -310,6 +406,19 @@ class ThopAttention(OpWrapper): predicted_tokens_per_seq=2, ), ), + Cell( + why="qwen4_exp P=8 hybrid pool: row 47 of the 48-row two-pool mapping " + "(pool 0, layer-in-pool 11), P=8 verify", + spec=dict( + _QWEN4_DENSE, + pool_layout="hybrid48", + local_layer_idx=47, + cached=(64, 25), + new=(8, 8), + num_contexts=0, + predicted_tokens_per_seq=8, + ), + ), Cell( why="48-layer single pool, identity mapping: row 47 = layer-in-pool 47, " "24 rows over a 40-token cached prefix", @@ -371,6 +480,24 @@ class ThopAttention(OpWrapper): "12544-row frame is the longest segment", spec=dict(_QWEN4_VISION, seg_lens=(12544, 9600)), ), + # A vision-tower call is one frame per frame group, at most 32768 rows per call. + Cell( + why="video: 384 frame groups of 64 rows (8x8 patches), the untiled " + "kernel's longest frame", + spec=dict(_QWEN4_VISION, seg_lens=(64,) * 384), + ), + Cell( + why="video: 256 frame groups of 128 rows, the 32768-row ceiling", + spec=dict(_QWEN4_VISION, seg_lens=(128,) * 256), + ), + Cell( + why="video: 2 frame groups of 16384 rows (128x128), the longest frame at the ceiling", + spec=dict(_QWEN4_VISION, seg_lens=(16384, 16384)), + ), + Cell( + why="video: 8192 frame groups of 2x2 patches, the most a 32768-row call can carry", + spec=dict(_QWEN4_VISION, seg_lens=(4,) * 8192), + ), Cell( why="fp8 pool packed prefill, GQA 24q/4kv d256: the FMHA reads the " "bf16 q rows, not the pool", @@ -479,6 +606,165 @@ class ThopAttention(OpWrapper): predicted_tokens_per_seq=3, ), ), + # qwen4_exp's own geometry over an fp8 pool: the kv head count changes the compiled + # decode kernel, so the 24/4/256 cells above do not cover it. + Cell( + why="qwen4_exp fp8 pool paged-context prefill, 24/2/256: a fresh 48-row " + "sequence beside 33 rows over a 48-token e4m3 prefix", + spec=dict( + _QWEN4_FP8, + cached=(0, 48), + new=(48, 33), + num_contexts=2, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="qwen4_exp fp8 pool decode P=1, 24/2/256: histories 200 and 48 span " + "pages and KV tiles", + spec=dict( + _QWEN4_FP8, + cached=(200, 48), + new=(1, 1), + num_contexts=0, + predicted_tokens_per_seq=1, + ), + ), + *( + Cell( + why=f"qwen4_exp fp8 pool MTP verify P={p}, 24/2/256 (max_draft_len {p - 1})", + spec=dict( + _QWEN4_FP8, + cached=(70, 35), + new=(p, p), + num_contexts=0, + predicted_tokens_per_seq=p, + ), + ) + for p in (2, 4, 5, 6, 7, 8) + ), + Cell( + why="qwen4_exp fp8 pool mixed step P=3, 24/2/256: a 20-row context ahead " + "of two P=3 blocks over histories 31 and 64", + spec=dict( + _QWEN4_FP8, + cached=(0, 31, 64), + new=(20, 3, 3), + num_contexts=1, + predicted_tokens_per_seq=3, + ), + ), + Cell( + why="qwen4_exp fp8 pool mixed step P=8, 24/2/256: a 20-row context ahead " + "of two P=8 blocks over histories 31 and 64", + spec=dict( + _QWEN4_FP8, + cached=(0, 31, 64), + new=(20, 8, 8), + num_contexts=1, + predicted_tokens_per_seq=8, + ), + ), + # KV stays within one 128-key tile (histories to 119 plus 8 rows), where the e4m3 + # probability model of `reference` is the kernel's own at any batch. + Cell( + why="qwen4_exp fp8 batch P=8 G=32: 32 blocks over histories 1..119, 256 q rows", + spec=dict( + _QWEN4_FP8_G32, + cached=_spread(32, 119), + new=(8,) * 32, + num_contexts=0, + predicted_tokens_per_seq=8, + ), + ), + Cell( + why="qwen4_exp fp8 hybrid pool P=2: row 47 of the 48-row two-pool mapping " + "(pool 0, layer-in-pool 11) over e4m3 slabs", + spec=dict( + _QWEN4_FP8, + pool_layout="hybrid48", + local_layer_idx=47, + cached=(64, 25), + new=(2, 2), + num_contexts=0, + predicted_tokens_per_seq=2, + ), + ), + Cell( + why="qwen4_exp fp8 hybrid pool P=8: row 47 of the 48-row two-pool mapping " + "(pool 0, layer-in-pool 11) over e4m3 slabs", + spec=dict( + _QWEN4_FP8, + pool_layout="hybrid48", + local_layer_idx=47, + cached=(64, 25), + new=(8, 8), + num_contexts=0, + predicted_tokens_per_seq=8, + ), + ), + Cell( + why="qwen4_exp nvfp4 pool paged-context prefill, 24/2/256: a fresh 48-row " + "sequence beside 33 rows over a 48-token e2m1 prefix", + spec=dict( + _QWEN4_NVFP4, + cached=(0, 48), + new=(48, 33), + num_contexts=2, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="qwen4_exp nvfp4 pool decode P=1, 24/2/256: histories 200 and 48", + spec=dict( + _QWEN4_NVFP4, + cached=(200, 48), + new=(1, 1), + num_contexts=0, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="qwen4_exp nvfp4 pool MTP verify P=2, 24/2/256 (max_draft_len 1)", + spec=dict( + _QWEN4_NVFP4, + cached=(70, 35), + new=(2, 2), + num_contexts=0, + predicted_tokens_per_seq=2, + ), + ), + Cell( + why="qwen4_exp nvfp4 pool MTP verify P=4, 24/2/256 (max_draft_len 3)", + spec=dict( + _QWEN4_NVFP4, + cached=(70, 35), + new=(4, 4), + num_contexts=0, + predicted_tokens_per_seq=4, + ), + ), + Cell( + why="qwen4_exp nvfp4 pool mixed step P=1: a 20-row context ahead of two " + "1-row sequences over histories 31 and 64", + spec=dict( + _QWEN4_NVFP4, + cached=(0, 31, 64), + new=(20, 1, 1), + num_contexts=1, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="qwen4_exp nvfp4 pool mixed step P=4: a 20-row context ahead of two 4-row blocks", + spec=dict( + _QWEN4_NVFP4, + cached=(0, 31, 64), + new=(20, 4, 4), + num_contexts=1, + predicted_tokens_per_seq=4, + ), + ), ) note = """ @@ -510,24 +796,51 @@ class its own layer group, so its mapping carries two pool ids; the op reads `is_spec_decoding_enabled` False, which is what sm_103 passes for a linear draft, and mask_type 0 and 1 are bitwise identical at P > 1. Nothing checks q's row count against P. P reaches the decode kernel as `maxSeqLenQ`: one - JIT compile per P. + JIT compile per P. P is driven at 1..8 (max_draft_len 1..7) at 24/2/256 over + bf16 and fp8 pools. From P = 5 the launch picks its generation kernel by + batch size, so P 5..8 is also driven at 32 requests and at one batch in + 17..20 (bf16), and P = 8 at 32 requests over the fp8 pool. No KV cache (an encoder): pass `kv_cache_block_offsets`, `host_kv_cache_pool_pointers` and `host_kv_cache_pool_mapping` all None, with `update_kv_cache=True` -- False is rejected ("KV cache update cannot be disabled now"). One context sequence per frame; mask_type 0 confines a frame to itself. `tokens_per_block` and `use_paged_context_fmha` are - bitwise inert there. Passing only some of the three None is unmeasured. + bitwise inert there. Passing only some of the three None is unmeasured. An + image is one frame and a video one frame per frame group; the cells run up to + 8192 frames and 32768 rows per call at 16/16/72 under the padding mask. + `FusedMHARunnerV2` picks a different fmha_v2 kernel for a call whose longest + frame is at most 64 rows than for a longer one, so a frame's last bits can + move with the other frames it shares a call with, within `reference`'s band. head_size 72 exists in the context phase only: a generation call raises "unsupported head_size 72", and the process stays usable. Over an fp8-e4m3 pool the packed context path reads the bf16 q rows and never the pool, while `use_paged_context_fmha=True` reads the pool, so its operands are e4m3: the flag is output-inert over a bf16 pool only. Every - fp8 cell runs kv scale 1.0; at 1.5 `reference` does not reproduce the - kernel's probability rounding. + fp8 cell runs kv scale 1.0 through fp32 [1] scale tensors; at 1.5 `reference` + does not reproduce the kernel's probability rounding. quant_mode 384, 1152, 4224 and 4480 are bit-identical to 128 at 24/4/256; - 4352 (the same bits without FP8_KV_CACHE) reads e4m3 slabs as bf16. + 4352 (the same bits without FP8_KV_CACHE) reads e4m3 slabs as bf16. At + 24/2/256 only 128 was driven, with KV up to 200 keys at P = 1 and up to + 127 at P 2..8: the fp8 kernels round the softmax probabilities against a + per-tile max, which the row-max model of + `reference(e4m3_probabilities=True)` reproduces within one 128-key tile + and approximates up to 200 keys. + + NVFP4 pool (quant_mode 8192): K and V are stored as e2m1 values, two per + byte (low nibble = even element), per head page `[tpb, D/2]`, with one e4m3 + scale per 16 values in a second pool whose pointer is the third axis of + `host_kv_cache_pool_pointers` (`[pools, 2, 2]`: primary/secondary x + data/scales); K's scales are token-major per head page, V's in four-token + groups `[tpb/4, D/16, 4]`. The append quantizes with `kv_scale_orig_quant[1]` + / `[2]` as second-level scales (fp32 `[3]`, `[0]` is q's e4m3 scale). One + block-offset table addresses both pools, so they must hold the same pages + per slot. The kernels are trtllm-gen cubins with e4m3 q and e2m1 K/V and + exist for the paged-context path only, which needs the op to enable FP8 + context FMHA for an FP4 KV cache. Every NVFP4 cell runs the paged-context + path at kv scale 1.0 and reads K and V through e4m3, so `reference` runs + over the NVFP4 values with `e4m3_inputs` and `e4m3_probabilities`. The sparse / skip-softmax arguments are bitwise inert in the standard configuration (nine off-default settings at 24/2/256): passing them does diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/graph/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/graph/__init__.py new file mode 100644 index 000000000000..c976f03c8914 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/graph/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Graph-capture entries.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/graph/eager_on_graph.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/graph/eager_on_graph.py new file mode 100644 index 000000000000..43f18a5ed022 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/graph/eager_on_graph.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One break of a breakable CUDA graph: a callable run eagerly between captured segments. + +Wraps the engine's `eager_on_graph` (`pyexecutor/breakable_cuda_graph`). It launches no +kernel of its own; what it certifies is the control flow around the callable it is handed. +""" + +from __future__ import annotations + +from typing import Any, Callable + +import torch + +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import eager_on_graph as _eager_on_graph + +from .._op import Arch, Cell, OpWrapper + + +class EagerOnGraph(OpWrapper): + """`inner(*args)`, eagerly; under breakable CUDA-graph capture, a break. + + Outside a `BreakableCUDAGraphCapture` the call is a passthrough. Inside one it ends the + captured segment, runs `inner(*args)` once, records a replay function holding `inner` + and the arguments, and opens the next segment. `BreakableCUDAGraph.replay` then runs + each segment and, between two of them, re-runs `inner` with the **capture-time** + arguments: a CUDA tensor by its capture-time address, a CPU tensor as the same object, + an `int` / `float` / `bool` / `str` / `None` frozen at its capture-time value. Whatever + varies per step therefore reaches `inner` through state it reads itself, which the + caller refreshes before each replay, never through an argument. + + Fusion boundary: none. The break runs exactly what `inner` runs, on the current stream. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # Spec axes: `capture` (inside a breakable capture, or plain eager), `breaks` (breaks in + # one body), `tokens` (the bucket `T` the body is captured at), `width` (row width), + # `live_rows` (the rows `n` each replay carries, set together with a fresh per-step scale + # and fresh input rows). Each break's `inner` reads `n` and the scale from live state, + # mixes rows `[:n]` across the batch, writes them into a fixed in-place output, zeroes + # rows `[n:]` and returns None: the shape of every break a prefill body makes. + CELLS: tuple[Cell, ...] = ( + Cell( + why="outside capture: a passthrough, bitwise the inner call", + spec=dict(capture=False, breaks=2, tokens=256, width=128, live_rows=(256, 17)), + ), + Cell( + why="one break, bucket 256: replays at new live rows equal eager", + spec=dict(capture=True, breaks=1, tokens=256, width=128, live_rows=(17, 256, 1)), + ), + Cell( + why="50 breaks in one body, bucket 1024", + spec=dict(capture=True, breaks=50, tokens=1024, width=128, live_rows=(1000, 255, 1024)), + ), + Cell( + why="bucket 8192, live rows across the 2048 boundary and the full bucket", + spec=dict(capture=True, breaks=2, tokens=8192, width=64, live_rows=(6144, 8192, 2047)), + ), + ) + + note = """ + Preconditions and hazards, each silent unless stated otherwise: + + * **Scalars are frozen.** An `int`, `float`, `bool`, `str` or `None` argument replays at + its capture-time value. A static per-site constant (a layer index) is right; a + per-step count is wrong at every replay. + * **CUDA tensors are re-read by address.** An argument must live at a stable address for + the graph's life: allocated inside the captured body (the graph pool) or an + engine-owned persistent buffer. A tensor allocated per step outside the graph replays + the capture-time memory. + * **Any other argument raises.** A dataclass, a namespace or a bound object raises + TypeError at capture, after `inner` already ran once. Loud, but it is a break raising + under capture, so the next hazard follows it. + * **A break that raises under capture leaves a CUDA error pending.** The raise comes + after the break ended the captured segment and before it opened the next, and the + capture's exit ends the ended segment again: it raises RuntimeError over the break's + own error and leaves cudaErrorIllegalState pending in the CUDA runtime, so the + process's next CUDA launch fails once with "the operation cannot be performed in the + present state". A caller that recovers must take the pending error + (`cudaGetLastError`) before launching anything else. + * **Side-stream waits are dropped.** While a capture is open, including inside `inner`'s + own capture-time run, a wait of the capture stream on a stream that is not capturing + is dropped with a warning, and replay does not preserve it. + * **The return value is copied each replay** onto the object `inner` returned at + capture. `inner` writing a fixed output in place returns None. + * **A fixed output must be fully written.** Rows `inner` leaves alone keep whatever an + earlier replay or another bucket left there. + + Hard limits + none: no index-typed argument + """ + + def __call__(self, *args: Any, inner: Callable[..., Any]) -> Any: + return _eager_on_graph(inner)(*args) + + def reference(self, *args: Any, inner: Callable[..., Any]) -> Any: + """The inner call, eagerly: what every replay of the break must reproduce.""" + return inner(*args) + + def compare(self, out: torch.Tensor, ref: torch.Tensor) -> None: + """Bitwise: a break adds no arithmetic, so any difference is a frozen or stale + operand, not rounding.""" + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + assert torch.equal(out, ref), ( + f"{int((out != ref).sum())} of {out.numel()} elements differ from the eager body" + ) + + +eager_on_graph = EagerOnGraph() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/routing.py index 6c2ebcaf4a15..6f4dd9fd6c2f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/routing.py @@ -11,7 +11,7 @@ from __future__ import annotations -from typing import Any, Optional, Tuple +from typing import TYPE_CHECKING, Any, Optional, Tuple from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig from tensorrt_llm.quantization.mode import QuantAlgo @@ -69,6 +69,10 @@ } ) +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + def route( pretrained_config: Any, @@ -76,6 +80,8 @@ def route( compute_capability: Tuple[int, int], spec_config: Optional[Any], kv_cache_quant_algo: Optional[Any], + model_config: Optional[ModelConfig] = None, + llm_args: Optional[TorchLlmArgs] = None, ) -> type: """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits. @@ -83,7 +89,8 @@ def route( (see ``_router_index._kv_cache_quant_algo``), not the checkpoint's alone. Each target is imported where it is returned, so declining costs no import - of modeling code. + of modeling code. No engine choice in ``model_config`` or ``llm_args`` changes + which of these targets runs, so neither is read. """ # The checkpoint is recognized by its config shape; see the gpt_oss routing # module for what shape-sniffing does and does not pin. q_lora_rank is in diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/routing.py index 3582cc17a695..7bd51d78e37a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/routing.py @@ -10,7 +10,7 @@ from __future__ import annotations -from typing import Any, Optional, Tuple +from typing import TYPE_CHECKING, Any, Optional, Tuple from ..._router_index import NoTarget, tp_reaches_across_nodes @@ -42,6 +42,10 @@ } ) +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + def route( pretrained_config: Any, @@ -49,6 +53,8 @@ def route( compute_capability: Tuple[int, int], spec_config: Optional[Any], kv_cache_quant_algo: Optional[Any], + model_config: Optional[ModelConfig] = None, + llm_args: Optional[TorchLlmArgs] = None, ) -> type: """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits. @@ -56,7 +62,8 @@ def route( (see ``_router_index._kv_cache_quant_algo``), not the checkpoint's alone. Each target is imported where it is returned, so declining costs no import - of modeling code. + of modeling code. No engine choice in ``model_config`` or ``llm_args`` changes + which of these targets runs, so neither is read. """ # The checkpoint is recognized by its config shape, the upstream idiom # (``is_mla``, ``is_nemotron_hybrid`` do the same). A fine-tune with the 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 index f71b25463bdd..14ace1194c5a 100644 --- 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 @@ -11,19 +11,45 @@ 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. +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 and the prefill PLE block keep the unfused + ``qwen4_exp_gated_residual`` / ``qwen4_exp_post_inject`` / ``_grouped_norm``. * **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. + 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. +* **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: @@ -40,6 +66,37 @@ ``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. A bf16 pool makes the decode step call its recurrence without ``cu_seqlens``; +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 @@ -64,45 +121,86 @@ 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.moe.fp4_block_scale_moe_runner import ( - fp4_block_scale_moe_runner, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( - fp4_quantize, +from tensorrt_llm._torch._experimental.modeling_v2.catalog.graph.eager_on_graph import ( + eager_on_graph, ) 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.speculative import get_spec_worker, should_use_separate_draft_kv_cache -from tensorrt_llm._torch.utils import is_gdn_replay_enabled +from tensorrt_llm._torch.utils import ActivationType +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_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_prefill_compress import qsa_prefill_compress +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 qwen4_exp_attn_decode_prologue +from .ops.attention.qwen4_exp_qk_norm_rope import qwen4_exp_qk_norm_rope from .ops.attention.vision_rope_2d import vision_rope_2d from .ops.embedding.ngram_hash import ngram_hash +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.low_m_bf16_mm import low_m_bf16_mm from .ops.hc.qwen4_exp_gated_residual import qwen4_exp_gated_residual +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_tc_first import qwen4_exp_hc_decode_tc_first +from .ops.hc.qwen4_exp_hc_inject_norm import qwen4_exp_hc_inject_norm +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_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.gated_ngram_injection import gated_ngram_injection from .ops.ple.ngram_token_history import ngram_token_history +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.dequantize_e4m3_per_tensor import dequantize_e4m3_per_tensor from .ops.quantization.fp8_quantize_1x128 import fp8_quantize_1x128 from .ops.ssm.causal_conv1d_fwd import causal_conv1d_fwd @@ -112,10 +210,20 @@ from .ops.ssm.cu_seqlens_from_seq_lens import cu_seqlens_from_seq_lens from .ops.ssm.dilated_short_conv_paged import dilated_short_conv_paged from .ops.ssm.extract_transpose_prefill_slice import extract_transpose_prefill_slice +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_gdn_post_conv import fused_gdn_post_conv +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.fused_sigmoid_gating_delta_rule_update import fused_sigmoid_gating_delta_rule_update +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.ssm.reset_recurrent_state_rows import reset_recurrent_state_rows from .ops.vision.patch_embed_conv3d import patch_embed_conv3d @@ -147,11 +255,38 @@ ) +# 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 + +#: 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 + + +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. More rows 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`. + """ + if mat_a.shape[0] == 1 and mat_b.shape[0] in _GEMV_TILED_K: + 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 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` and `predicted_tokens_per_seq` are passed -# beside these. +# 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, @@ -163,9 +298,6 @@ mask_type=1, q_scaling=1.0, attention_sinks=None, - # bf16 pool: no KV scale. - kv_scale_orig_quant=None, - kv_scale_quant_orig=None, # In-kernel RoPE disabled; rope_* values are inert placeholders. position_embedding_type=0, rotary_inv_freq=None, @@ -310,12 +442,14 @@ #: 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 -# The trtllm-gen MoE runner reads activation scales as a *linear* buffer, not the -# 128x4-swizzled one of a dense NVFP4 GEMM. Byte counts match when num_tokens is a -# multiple of 128 (every decode graph batch), so the wrong one passes silently. -_SF_LINEAR = False # 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 @@ -323,6 +457,18 @@ _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 @@ -331,7 +477,7 @@ _PAD_SLOT_ID = -1 # Gated-DeltaNet gating form: the only pair `fused_sigmoid_gating_delta_rule_update` -# is certified at; matches `fused_gdn_post_conv`'s default so prefill and decode +# is certified at; matches `gdn_prefill_conv_post`'s default so prefill and decode # compute the same `g`. _SOFTPLUS_BETA = 1.0 _SOFTPLUS_THRESHOLD = 20.0 @@ -359,6 +505,24 @@ def _ple_host_offload_enabled() -> bool: 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 @@ -370,8 +534,9 @@ def _ple_host_offload_enabled() -> bool: # `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 batch the index plane sees a frame -# at one position, changing block selection only above `dense_seq_len_threshold`. +# 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 @@ -388,13 +553,54 @@ def _ple_host_offload_enabled() -> bool: #: draft rejected; only the acceptance gate can tell. _MTP_EMBED_BLOCK_ALL_STREAMS = True +#: 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 +#: (one 128-row MMA tile); larger batches keep cuBLAS. +_HC_TC_FIRST_MAX_ROWS = ( + 128 # capacity: fixed: dispatch threshold (one MMA tile), not a buffer bound +) + #: 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 +#: 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 --- # @@ -402,19 +608,23 @@ def _ple_host_offload_enabled() -> bool: # 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. +# (`_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, replaced -# atomically and keyed on manager identity, (2) pure contract checks, (3) refusing -# a re-bind during CUDA graph capture. 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. +# 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): @@ -427,6 +637,7 @@ class _PoolGeometry(NamedTuple): 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 @@ -434,6 +645,29 @@ class _PoolGeometry(NamedTuple): 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 class _Bound(NamedTuple): @@ -443,14 +677,27 @@ class _Bound(NamedTuple): 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 -> (K, V) views + 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 runtime's per-step rollback buffers. + # 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: @@ -480,15 +727,47 @@ def check_attention_pools(g: _PoolGeometry, layer: int, index_cache, pool) -> No 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.head_dim) - _check_pool(pool, f"layer {layer}: paged K/V", 2, kv_shape) + 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: @@ -501,22 +780,165 @@ def check_ple_pools(g: _PoolGeometry, ple_conv, ple_ngram, conv) -> None: ) -def check_step_buffers(g: _PoolGeometry, int_conv, int_ssm, conv, ssm, int_idx) -> None: - """MTP verify rollback buffers: `[slots, steps, ...]`, at the pools' dtypes - (the promotion kernel copies between them without converting).""" - assert int_conv is not None and int_ssm is not None, ( - "the runtime allocated no per-step conv/recurrent state buffers; a " - "partially accepted draft cannot be rolled back" +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) - 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) - # `update_mamba_states` reads its source rows from this vector. 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 @@ -621,6 +1043,13 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig 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 @@ -769,10 +1198,27 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig 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 ---------------------------------------------- @@ -790,6 +1236,8 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig 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", ( @@ -810,7 +1258,9 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig # 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. + # 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, ( @@ -818,6 +1268,17 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig ) # 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 + ) # ---- PLE ----------------------------------------------------------- ple_ids = list(getattr(cfg, "ple_layer_ids", None) or []) @@ -931,16 +1392,32 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig "must publish them as `multimodal_token_ids`" ) - # ---- KV pool dtype -------------------------------------------------- - # An unsupported KV dtype makes the QSA backend silently fall back to dense - # (`QSA_SPARSE_KV_CACHE_DTYPES`). + # ---- 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 is None, ( - f"this target runs the checkpoint's own bf16 KV pool, which is what " - f"the accuracy reference was measured at and what the QSA sparse " - f"path requires; the engine resolved kv_cache_quant_algo {kv_algo!r}" - ) - self.quant_mode = 0 + 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 @@ -961,14 +1438,18 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig f"the checkpoint; this target carries the weights of " f"{_MTP_NUM_LAYERS} MTP module" ) - # `max_draft_len` becomes `predicted_tokens_per_seq` on every attention - # call, certified at 1..4, so attention binds. + # `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 <= 3, ( - f"max_draft_len={max_draft} needs predicted_tokens_per_seq " - f"{max_draft + 1}, and attention/thop_attention certifies that " - f"argument at 1..4. A longer draft needs a kernel-verify " - f"dispatch, not a bigger number" + 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), ( @@ -983,19 +1464,10 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig f"{spec_config.max_total_draft_tokens} against {max_draft}" ) - # The engine picks the replay (default) or legacy verify path at - # cache-manager construction; this target is closed over the legacy one - # (`_gdn_verify` + `update_mamba_states`). Checked again against the - # manager at the first forward. - assert not is_gdn_replay_enabled(), ( - "the GDN cached-replay verify path is on and this target does " - "not implement it: its commit kernel " - "(fused_recurrent_gated_delta_rule_cached_replay_update) has no " - "catalog entry, and with replay on the runtime allocates no " - "intermediate SSM buffer for the legacy path to write. Export " - "TRTLLM_USE_GDN_REPLAY=0 before the ranks start. This is a " - "requirement of the MTP path, not a tuning knob" - ) + # 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 " @@ -1021,17 +1493,41 @@ def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig # 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 + 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 + # 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.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 --- @@ -1039,47 +1535,59 @@ 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"]) + 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 = [] - for i in range(self.num_layers): - # The checkpoint stores modelopt reciprocals (`input_scale = 1 / g_act`, - # `weight_scale_2 = 1 / g_w`), so the 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 = w[f"l{i}_e_isc2"] - gate1 = (isc1 * w[f"l{i}_e_ws2_1"]).contiguous() - moe_scalars.append( - dict( - # [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 * w[f"l{i}_e_ws2_2"]).contiguous(), - # Routed FC1 activation global scale (one per layer). - global_scale=(1.0 / isc1[:1]).contiguous(), - ) - ) + 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 = { @@ -1107,23 +1615,44 @@ def post_load_weights(self) -> None: f"do not concatenate to the padded id space this checkpoint " f"declares" ) - # The index branch's rope table 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 / (self.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) - self._cos_sin = torch.cat([torch.cos(angles), torch.sin(angles)], dim=1).contiguous() - assert self._cos_sin.shape == (self.max_positions, 2 * pairs), self._cos_sin.shape + 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) + self._fuse_hc_first_projections() + + # 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 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. @@ -1172,6 +1701,7 @@ def _pool_geometry(self, tokens_per_block: int) -> _PoolGeometry: 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, @@ -1179,7 +1709,21 @@ def _pool_geometry(self, tokens_per_block: int) -> _PoolGeometry: 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 @@ -1239,6 +1783,13 @@ def _bind_cache_manager(self, md, position_ids) -> None: "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, @@ -1257,9 +1808,7 @@ def _bind_cache_manager(self, md, position_ids) -> None: index_cache, kv = {}, {} for i in self.attn_layers: index_cache[i] = cm.get_index_k_buffer(i) - pool = cm.get_buffers(i, kv_layout="HND") - check_attention_pools(g, i, index_cache[i], pool) - kv[i] = (pool[:, 0], pool[:, 1]) + 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. @@ -1305,33 +1854,113 @@ def _bind_cache_manager(self, md, position_ids) -> None: cm, pool_idx, index_cache, position_cache, kv, conv, ssm, (ple_conv, ple_ngram) ) if self.mtp_enabled: - bound = bound._replace(**self._bind_step_buffers(cm, g, conv0, ssm0)) + bound = bound._replace(**self._bind_step_buffers(cm, g, mm, conv0, ssm0)) self._bound = bound - def _bind_step_buffers(self, cm, g: _PoolGeometry, conv0, ssm0) -> dict: + 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 per-step buffers that exist only under the legacy GDN path. - Returns the `_Bound` fields it adds.""" + 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" ) - assert not getattr(cm, "use_replay_state_update", False), ( - "the GDN cached-replay verify path is on. Export TRTLLM_USE_GDN_REPLAY=0 " - "before the ranks start; see this target's __init__ for why" - ) _require( cm, ("get_intermediate_conv_states", "get_intermediate_ssm_states"), - "reads both per-step state buffers", + "binds its per-step state buffers", ) # `get_intermediate_*` builds a fresh view per call, so bind them here. int_conv = {i: cm.get_intermediate_conv_states(i) for i in self.gdn_layers} - int_ssm = {i: cm.get_intermediate_ssm_states(i) for i in self.gdn_layers} int_idx = getattr(cm, "intermediate_state_indices", None) g0 = self.gdn_layers[0] - check_step_buffers(g, int_conv[g0], int_ssm[g0], conv0, ssm0, int_idx) - return dict(int_idx=int_idx, int_conv=int_conv, int_ssm=int_ssm) + check_step_buffers(g, int_conv[g0], conv0, int_idx) + int_ssm0 = cm.get_intermediate_ssm_states(g0) + if not getattr(cm, "use_replay_state_update", False): + check_legacy_state(g, int_ssm0, ssm0) + int_ssm = {i: cm.get_intermediate_ssm_states(i) for i in self.gdn_layers} + 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, ( + "use_replay_state_update is True but the manager returns no 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)) + 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 + ) + + 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 --- @@ -1342,16 +1971,26 @@ def _select_target(self, attn_metadata) -> "PrefillTarget | DecodeTarget": draft step routes through here per call.""" return self._decode if attn_metadata.num_contexts == 0 else self._prefill - def forward( + def begin_step( self, attn_metadata: AttentionMetadata, - input_ids: torch.IntTensor | None = None, - position_ids: torch.IntTensor | None = None, + input_ids: torch.IntTensor | None, + position_ids: torch.IntTensor | None, inputs_embeds: torch.FloatTensor | None = None, lora_params: dict | None = None, - **kwargs, - ): - """Check what holds for every step, bind the cache manager, then route.""" + 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" ) @@ -1359,59 +1998,92 @@ def forward( 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 even with `inputs_embeds` (upstream passes - # `input_ids=None` after splicing): the PLE block hashes token ids every step, - # and text gates would not catch a mis-hash. + # `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 whose embeddings were spliced upstream" + "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. - spec_metadata = kwargs.get("spec_metadata") 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`). + # 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, - inputs_embeds=inputs_embeds, - **kwargs, + 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: tuple - ahc_down: torch.Tensor + ahc_norm: torch.Tensor + ahc_first: torch.Tensor ahc_up: torch.Tensor - ahc_inject: torch.Tensor - mhc_norm: tuple - mhc_down: torch.Tensor + mhc_norm: torch.Tensor + mhc_first: torch.Tensor mhc_up: torch.Tensor - mhc_inject: torch.Tensor - q: torch.Tensor | None - k: torch.Tensor | None - v: torch.Tensor | None + attn_in: torch.Tensor | None o: torch.Tensor | None q_norm: torch.Tensor | None k_norm: torch.Tensor | None - idx: torch.Tensor | None idx_qn: torch.Tensor | None idx_kn: torch.Tensor | None - in_qkv: torch.Tensor | None - in_z: torch.Tensor | None - in_a: torch.Tensor | None - in_b: torch.Tensor | None + in_proj: torch.Tensor | None out: torch.Tensor | None conv_w: torch.Tensor | None a_log: torch.Tensor | None @@ -1423,35 +2095,53 @@ class _LayerWeights(NamedTuple): 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 + # 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 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( - q=None, - k=None, - v=None, + attn_in=None, o=None, q_norm=None, k_norm=None, - idx=None, idx_qn=None, idx_kn=None, ) gdn: dict = dict( - in_qkv=None, - in_z=None, - in_a=None, - in_b=None, + in_proj=None, out=None, conv_w=None, a_log=None, @@ -1460,23 +2150,17 @@ def _layer_views(core: Qwen4ExpTp1ModelingV2) -> list[_LayerWeights]: ) if core.is_attn[i]: attn = dict( - q=w[f"l{i}_q_w"].t(), - k=w[f"l{i}_k_w"].t(), - v=w[f"l{i}_v_w"].t(), + 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=w[f"l{i}_idx_w"].t(), 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_qkv=w[f"l{i}_in_qkv_w"].t(), - in_z=w[f"l{i}_in_z_w"].t(), - in_a=w[f"l{i}_in_a_w"].t(), - in_b=w[f"l{i}_in_b_w"].t(), + 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), @@ -1485,30 +2169,53 @@ def _layer_views(core: Qwen4ExpTp1ModelingV2) -> list[_LayerWeights]: 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=core._norm_slices(f"l{i}_ahc_norm"), - ahc_down=w[f"l{i}_ahc_down"], - ahc_up=w[f"l{i}_ahc_up"], - ahc_inject=w[f"l{i}_ahc_inject"], - mhc_norm=core._norm_slices(f"l{i}_mhc_norm"), - mhc_down=w[f"l{i}_mhc_down"], - mhc_up=w[f"l{i}_mhc_up"], - mhc_inject=w[f"l{i}_mhc_inject"], + 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"], - fc1_s=w[f"l{i}_fc1_s"].view(torch.float8_e4m3fn), - fc2_w=w[f"l{i}_fc2_w"], - fc2_s=w[f"l{i}_fc2_s"].view(torch.float8_e4m3fn), - out1_scale=scalars["output1_scale_scalar"], - out1_gate_scale=scalars["output1_scale_gate_scalar"], - out2_scale=scalars["output2_scale_scalar"], - sh_gu=w[f"l{i}_sh_gu"].t(), + 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_gate"].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(), ) ) return layers @@ -1518,14 +2225,17 @@ 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: tuple - fhc_down: torch.Tensor + fhc_norm: torch.Tensor + fhc_first: torch.Tensor fhc_up: torch.Tensor ple_key: torch.Tensor ple_val: torch.Tensor ple_nk: tuple ple_nq: tuple ple_nc: tuple + 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 @@ -1536,14 +2246,19 @@ class _StackWeights(NamedTuple): def _stack_views(core: Qwen4ExpTp1ModelingV2) -> _StackWeights: w = core.w return _StackWeights( - fhc_norm=core._norm_slices("fhc_norm"), - fhc_down=w["fhc_down"], - fhc_up=w["fhc_up"], + 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(), ple_nk=core._norm_slices("ple_nk"), ple_nq=core._norm_slices("ple_nq"), ple_nc=core._norm_slices("ple_nc"), + # The same three weights whole, for the decode target's fused tail, which + # 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"], @@ -1574,6 +2289,9 @@ class _MtpWeights(NamedTuple): 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 @@ -1609,20 +2327,321 @@ def _mtp_views(core: Qwen4ExpTp1ModelingV2) -> _MtpWeights: 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"].t(), + 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_gate"].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 _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, +): + """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.""" + if kv.k_scale is None: + 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, + ) + 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 _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) -> 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). + `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 conv before it does not signal + dependents, so 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) + out = fused_recurrent_gated_delta_rule_cached_replay_update( + 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=False, + 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. 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 + + 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.""" @@ -1633,8 +2652,16 @@ def __init__(self, core: Qwen4ExpTp1ModelingV2) -> None: # 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 + ) self._layers = _layer_views(core) self._stack = _stack_views(core) + # 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 @@ -1696,6 +2723,38 @@ def step_args(self, md: TrtllmAttentionMetadata) -> dict: # --------------------------------------------------------------- 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 _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 @@ -1706,61 +2765,157 @@ def _grouped_norm(self, x: torch.Tensor, weights: tuple) -> torch.Tensor: [flashinfer_gemma_rmsnorm(part, w, core.eps) for w, part in zip(weights, parts)], 1 ) - def _moe_block(self, x: torch.Tensor, L: _LayerWeights, dt: torch.dtype) -> torch.Tensor: - """NVFP4 routed experts (quantized once in the linear block-scale layout, with - finished top-k ids and weights so the runner's routing arguments are inert) - plus a bf16 shared expert folded in by a per-token sigmoid gate.""" + 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 - # The routing op reads logits from data_ptr() ignoring strides; cublas_mm - # allocates them contiguous. + 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) - topk_ids, topk_w = renorm_moe_routing_op(logits, topk=core.topk, output_dtype=dt) - xq, xsf = fp4_quantize(x, L.moe_gscale, _SF_VEC, False, _SF_LINEAR) - routed = fp4_block_scale_moe_runner( - routing_logits=None, - routing_bias=None, - hidden_states=xq, - hidden_states_scale=xsf.view(torch.float8_e4m3fn), - gemm1_weights=L.fc1_w, - gemm1_weights_scale=L.fc1_s, - gemm1_bias=None, - gemm1_alpha=None, - gemm1_beta=None, - gemm1_clamp_limit=None, - gemm2_weights=L.fc2_w, - gemm2_weights_scale=L.fc2_s, - gemm2_bias=None, - output1_scale_scalar=L.out1_scale, - output1_scale_gate_scalar=L.out1_gate_scale, - output2_scale_scalar=L.out2_scale, - num_experts=core.num_experts, - top_k=core.topk, - 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, - do_finalize=True, - act_type=_ACT_TYPE_SWIGLU, - topk_weights=topk_w, - topk_ids=topk_ids, - )[0] - shared = cublas_mm(flashinfer_silu_and_mul(cublas_mm(x, L.sh_gu)), L.sh_dn) - gate_logits = cublas_mm(x, L.sh_gate) - return fused_sigmoid_gate_mul_add(routed, gate_logits, shared) + 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, + ) + 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, qkv, index_scratch, attn_out, qsa, step): + def _attention_block(self, x, i, L, proj_qkv, index_scratch, attn_out, pos3): """One gated GQA layer plus its QSA index branch. - The q projection is twice as wide as q (q and output gate interleave per - head, as HF chunks the last dim); only the q half is packed into the QKV - buffer and the gate half is read later through its own strides. 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 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 = cublas_mm(x, L.attn_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 cublas_mm(attn_out, L.o) + + 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, @@ -1770,63 +2925,28 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step): topk_row_starts, block_table, tokens_per_block, - pos, - pos3, coords, sparse, - ) = qsa - num_tokens = x.shape[0] - - qg = cublas_mm(x, L.q) - k = cublas_mm(x, L.k) - v = cublas_mm(x, L.v) - - # Destination views into the packed QKV buffer. - qkv_h = torch.reshape(qkv, [num_tokens, 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, [num_tokens, core.heads_q, 2 * core.head_dim]), - [core.head_dim, core.head_dim], - 2, - ) - q_dst.copy_(q_src) - k_dst.copy_(torch.reshape(k, [num_tokens, core.heads_kv, core.head_dim])) - v_dst.copy_(torch.reshape(v, [num_tokens, core.heads_kv, core.head_dim])) - 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=L.q_norm, - k_weight=L.k_norm, - base=core.theta, - is_neox=True, - position_ids=pos3, - is_qk_norm=True, - # The gemma QK norm, `rmsnorm(x) * (1 + w)`. - use_gemma=True, - # Real mRoPE, unconditionally. A wrong wiring is invisible on text (equal - # rows) and ~500x out on an image batch. - use_mrope=True, - mrope_section1=core.mrope_section1, - mrope_section2=core.mrope_section2, + 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) - # ---- the QSA index branch ------------------------------------------ - # One replicated projection produces [index_heads q | 1 k]. The raw 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). - qk_index = cublas_mm(x, L.idx) index_cache = core._bound.index_cache[i] token_k = torch.split( - torch.reshape(qk_index, [num_tokens, core.index_heads + core.index_kv_heads, -1]), - [core.index_heads, core.index_kv_heads], + 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, @@ -1837,33 +2957,9 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step): block_table, tokens_per_block, ) - # Query norm+rotation in the fused call's packed layout; the k slot is - # computed and discarded, the v slot exists only because the layout has one. - # Coordinate 0 alone, matching `_QSA_MROPE_SECTION` on the key side. - scratch_h = torch.reshape(index_scratch, [num_tokens, core.index_scratch_width]) - scratch_h[:, : 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=L.idx_qn, - k_weight=L.idx_kn, - base=core.theta, - is_neox=True, - position_ids=pos3 if _QSA_INDEX_USE_MROPE else 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( - scratch_h, + 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], @@ -1884,82 +2980,151 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step): ) if sparse: - k_cache, v_cache = core._bound.kv[i] + 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. - qsa_paged_kv_store( - k_dst, v_dst, k_cache, v_cache, req_idx, logical, block_table, tokens_per_block + # 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, ) - 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=True, - ) - # `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. - qsa_topk_select( - scores, topk_row_starts, visible_blocks, topk_indices, index_topk=core.block_topk - ) - selected = qsa_expand_block_indices( - topk_indices, - logical, - seq_lens_q, - compress_ratio=core.compress_ratio, - token_topk=core.token_topk, - ) - out = torch.reshape( - qsa_paged_sparse_gqa( + 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, + ) + ) + # 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. + if num_tokens >= _ROWTILE_MIN_ROWS and kv.k.dtype == q_dst.dtype: + out = qsa_sparse_gqa_rowtile( q=q_dst, - k_cache=k_cache, - v_cache=v_cache, + 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, + ) + 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, - ), - [num_tokens, core.q_width], - ) + ) + attn_out[:num_tokens].copy_(torch.reshape(out, [num_tokens, core.q_width])) else: - out = thop_attention( + # 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, + 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, - **step, + kv_scale_orig_quant=core.kv_scale_orig_quant, + kv_scale_quant_orig=core.kv_scale_quant_orig, + **step.attn, **_CALL_CONSTANTS, ) - fused_sigmoid_mul_inplace(out, gate) - return cublas_mm(out, L.o) + 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) -> torch.Tensor: + 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, with per-step conv window and recurrent - state written to the runtime's rollback buffers, not the pools. - - The pools must not move: which prefix is real is unknown until the sampler - accepts, and `update_mamba_states` then promotes step `k - 1`. Advancing them - would leave rejected tokens' state in every recurrent layer, invisible to the - acceptance rate. `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); `pack_gdn_decode_qkv` does no Q/K L2 norm, - hence `use_qk_l2norm_in_kernel=True` below.""" + `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 (the - # transposed view is certified bitwise equal to the contiguous form). + # [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, @@ -1977,6 +3142,8 @@ def _gdn_verify(self, mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, n, p) -> ) # 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 ) @@ -2017,10 +3184,29 @@ def _gdn_verify(self, mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, n, p) -> ) return torch.reshape(out, [1, n * p, core.num_v_heads, core.head_v_dim]) - def _gdn_block(self, x, i, L, gdn_state) -> torch.Tensor: - """One Gated-DeltaNet layer over a batch with a context prefix. Every branch - is on a host int the engine fixes per capture, never a tensor value.""" + 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 = cublas_mm(x, L.in_proj) + 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. + + Over an fp32 recurrent pool outside a verify 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 FlashInfer's single-token decode kernel; + otherwise `_gdn_core_generic`.""" + core = self.core + L = self._layers[i] + step = self._step ( num_prefill_tokens, num_decode_tokens, @@ -2037,13 +3223,123 @@ def _gdn_block(self, x, i, L, gdn_state) -> torch.Tensor: gen_p, int_idx, qsl_long_p, - ) = gdn_state - num_tokens = x.shape[0] + replay, + ) = step.gdn + num_tokens = step.num_tokens + ssm_states = core._bound.ssm[i] + if is_verify or ssm_states.dtype != torch.float32: + self._gdn_core_generic(proj, out, i) + return + mixed_qkv, _, a, b = torch.split(proj[:num_tokens], core.in_proj_split, 1) + conv_states = core._bound.conv[i] - mixed_qkv = cublas_mm(x, L.in_qkv) - z = cublas_mm(x, L.in_z) - a = cublas_mm(x, L.in_a) - b = cublas_mm(x, L.in_b) + # A context sequence carrying nothing starts from a zeroed slot; only the + # context prefix is a candidate (a generation row's state must survive). + reset_recurrent_state_rows(ssm_states, si_p, hi_p, conv_states) + 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, + ) + # 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' reset and recurrence touch 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, + ) + out[:, num_tokens:].zero_() + + def _gdn_core_generic(self, proj, out, i) -> None: + """`_gdn_core` on a verify step or over a bf16 recurrent pool: the conv, the + post-conv gating and the chunked delta rule over the step's `n` rows, or under + speculation the verify path for the generation rows. 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() conv_states = core._bound.conv[i] ssm_states = core._bound.ssm[i] @@ -2088,7 +3384,7 @@ def _gdn_block(self, x, i, L, gdn_state) -> torch.Tensor: softplus_threshold=_SOFTPLUS_THRESHOLD, ) # q and k are already L2-normalized by the post-conv call. - out_p, _ = chunk_gated_delta_rule( + chunk_gated_delta_rule( q, k, v, @@ -2101,12 +3397,14 @@ def _gdn_block(self, x, i, L, gdn_state) -> torch.Tensor: output_final_state=False, cu_seqlens=qsl_long_p, use_qk_l2norm_in_kernel=False, - output=None, + output=out[:, :num_prefill_tokens], ) out_d = self._gdn_verify( - mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, num_decodes, gen_p + mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, num_decodes, gen_p, replay ) - return self._gdn_epilogue(torch.cat([out_p, out_d], 1), z, L, num_tokens) + out[:, num_prefill_tokens:num_tokens].copy_(out_d) + out[:, num_tokens:].zero_() + return causal_conv1d_update( x=mq_d, conv_state=conv_states, @@ -2160,7 +3458,7 @@ def _gdn_block(self, x, i, L, gdn_state) -> torch.Tensor: softplus_threshold=_SOFTPLUS_THRESHOLD, ) # q and k are already L2-normalized; the recurrence must not repeat it. - out, _ = chunk_gated_delta_rule( + chunk_gated_delta_rule( q, k, v, @@ -2173,9 +3471,35 @@ def _gdn_block(self, x, i, L, gdn_state) -> torch.Tensor: output_final_state=False, cu_seqlens=qsl_long, use_qk_l2norm_in_kernel=False, - output=None, + output=out[:, :num_tokens], + ) + 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 -- 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, ) - return self._gdn_epilogue(out, z, L, num_tokens) def _gdn_epilogue(self, out, z, L, num_tokens) -> torch.Tensor: """Gated RMS norm and output projection shared by every recurrence.""" @@ -2191,9 +3515,11 @@ def _gdn_epilogue(self, out, z, L, num_tokens) -> torch.Tensor: ) return cublas_mm(torch.reshape(o, [num_tokens, core.value_dim]), L.out) - def _ple_block(self, hc, input_ids, ple_state) -> torch.Tensor: - """The n-gram side path on one decoder layer; the caller adds the result into - the hyper-connection stream. + 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 @@ -2201,13 +3527,79 @@ def _ple_block(self, hc, input_ids, ple_state) -> torch.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 hash cannot distinguish. The id history and the conv state publish the + state after the whole chunk. + + Outside a verify step 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`. On a verify step (`stage`) the + generation rows' pre-step pool rows, id windows and conv input go to + `SpeculativeRollback`, which rewrites both pools to the accepted prefix once the + sampler has decided it, so that step takes `_ple_delta`, which exposes them.""" core = self.core S = self._stack - conv_states, ngram_pool, si_all, carry_all, qsl = ple_state - + 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] + if step.stage: + stream = torch.add(hc, self._ple_delta(hc)) + stream, normed = qwen4_exp_hc_inject_norm( + stream, None, None, next_norm, core.eps, core.hc_count + ) + else: + 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, + ) + stream, normed = 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, + ) + hc_out[:num_tokens].copy_(stream) + normed_out[:num_tokens].copy_(normed) + hc_out[num_tokens:].zero_() + normed_out[num_tokens:].zero_() + + def _ple_delta(self, hc) -> torch.Tensor: + """A verify step's PLE output for the batch's `n` stream rows `hc`, before it is + added into the stream, through the unfused ops whose intermediates + `SpeculativeRollback` stages: the pre-step pool rows, the id windows and the conv + input.""" + core = self.core + S = self._stack + step = self._step + conv_states, ngram_pool, si_all, carry_all, qsl = step.ple + core.spec_rollback.stage_ple_pools(ngram_pool, conv_states) window = ngram_token_history( - input_ids, ngram_pool, si_all, carry_all, qsl, core.eos_token_id, core.ngram_size + step.input_ids, ngram_pool, si_all, carry_all, qsl, core.eos_token_id, core.ngram_size ) rows = torch.reshape( ngram_hash( @@ -2231,9 +3623,11 @@ def _ple_block(self, hc, input_ids, ple_state) -> torch.Tensor: gated = gated_ngram_injection( key_normed, query_normed, value, core.hc_count, _PLE_CLAMP_MIN ) + conv_in = self._grouped_norm(gated, S.ple_nc) + core.spec_rollback.stage_ple_rows(window, conv_in) # The short conv reads channel-major; a bare transpose is the certified form. convolved = dilated_short_conv_paged( - x=torch.transpose(self._grouped_norm(gated, S.ple_nc), 0, 1), + x=torch.transpose(conv_in, 0, 1), weight=S.ple_conv_w, bias=None, conv_states=conv_states, @@ -2246,70 +3640,64 @@ def _ple_block(self, hc, input_ids, ple_state) -> torch.Tensor: ) return torch.add(gated, torch.transpose(convolved, 0, 1)) + 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 forward( + def prepare( self, attn_metadata: AttentionMetadata, - input_ids: torch.IntTensor | None = None, - position_ids: torch.IntTensor | None = None, - inputs_embeds: torch.FloatTensor | None = None, - **kwargs, - ): + 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) - step = self.step_args(md) - spec_metadata = kwargs.get("spec_metadata") + attn_args = self.step_args(md) mm = md.mamba_metadata - # Three views of one `[3, 1, T]` buffer: `pos` (row 0, flat `[T]`: QSA index - # rope), `pos3` (contiguous `[3, T]` for `use_mrope=True`; the engine hands a - # slice of a larger buffer and the op rejects the row stride, hence the - # copy), `coords` (`[T, 3]`, stored by the QSA position side cache). - pos = torch.reshape(torch.split(position_ids, [1, _QSA_POSITION_AXES - 1], 0)[0], [-1]) - coords = torch.transpose(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]), 0, 1) - pos3 = torch.empty( - [_QSA_POSITION_AXES, position_ids.shape[-1]], - dtype=torch.int32, - device=position_ids.device, + # **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})" ) - pos3.copy_(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1])) - - # Spliced vision rows are already the embedding; a lookup would overwrite them. - if inputs_embeds is None: - h = nn.functional.embedding(input_ids, core.w["embed"]) - else: - assert inputs_embeds.shape[0] == input_ids.shape[0], ( - f"the spliced embedding carries {inputs_embeds.shape[0]} rows for " - f"{input_ids.shape[0]} token ids" - ) - assert inputs_embeds.shape[1] == core.hidden, ( - f"the spliced embedding is {inputs_embeds.shape[1]} wide, not {core.hidden}" - ) - h = inputs_embeds - num_tokens = h.shape[0] - dt = h.dtype - dev = h.device # 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 - # Host-int check (graph-safe): an activation buffer longer than the batch - # would push padding rows through the recurrences into the state pools. - 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`; read from the metadata because the worker rewrites # it between draft steps. @@ -2342,30 +3730,35 @@ def forward( "sized from the second and written from the first" ) # Per-iteration counterpart of the core's `max_draft_len` bound. - assert 1 <= gen_p <= 4, ( - f"predicted_tokens_per_seq {gen_p} is outside the 1..4 range " - f"attention/thop_attention certifies" - ) + 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 - step["predicted_tokens_per_seq"] = gen_p + 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. + # `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 - qsl_long_p = 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] - # Context rows keep an ordinary chunked recurrence over their prefix. - qsl_long_p = qsl_long[: num_prefills + 1] + if core._bound.replay is not None: + replay = _replay_work(core, mm, num_decodes) gdn_state = ( num_prefill_tokens, num_decode_tokens, @@ -2382,6 +3775,7 @@ def forward( gen_p, int_idx, qsl_long_p, + replay, ) # PLE carry flags must cover the whole batch, but `has_initial_states` is @@ -2403,6 +3797,9 @@ def forward( # 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], @@ -2412,47 +3809,148 @@ def forward( md.qsa_topk_row_starts[:num_tokens], md.qsa_block_table, md.tokens_per_block, - pos, - pos3, - coords, + 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, + ) + + 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, ) - # Scratch shared by all attention layers (each fully rewrites them). - qkv = torch.empty([num_tokens, core.qkv_width], dtype=dt, device=dev) + 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`).""" + 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) - for i, L in enumerate(self._layers): + layers, S = self._layers, self._stack + _, 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: - hc = torch.add(hc, self._ple_block(hc, input_ids, ple_state)) - mixed, inject = qwen4_exp_gated_residual( - self._grouped_norm(hc, L.ahc_norm), L.ahc_down, L.ahc_up, L.ahc_inject, hc_count - ) + eager_on_graph(hc, ple_hc, ple_normed, inner=self._ple_block) + hc, normed = ple_hc, ple_normed + 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, qkv, index_scratch, attn_out, qsa, step + mixed, i, L, proj_qkv, index_scratch, attn_out, pos3 ) else: - block_out = self._gdn_block(mixed, i, L, gdn_state) - hc = qwen4_exp_post_inject(hc, block_out, inject) - mixed, inject = qwen4_exp_gated_residual( - self._grouped_norm(hc, L.mhc_norm), L.mhc_down, L.mhc_up, L.mhc_inject, hc_count + block_out = self._gdn_block(mixed, i, L, gdn_out) + hc, normed = qwen4_exp_hc_inject_norm( + hc, block_out, inject, L.mhc_norm, core.eps, hc_count ) - hc = qwen4_exp_post_inject(hc, self._moe_block(mixed, L, dt), inject) + 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) + 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. - S = self._stack - final, _ = qwen4_exp_gated_residual( - self._grouped_norm(hc, S.fhc_norm), S.fhc_down, S.fhc_up, None, hc_count - ) - if spec_metadata is None: + 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. - return final, hc + # `[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 --- @@ -2497,6 +3995,36 @@ def _mtp_moe_block(self, x, dt) -> torch.Tensor: 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) -> 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.""" + 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 + 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, @@ -2567,6 +4095,43 @@ def mtp_forward( "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, ( @@ -2579,6 +4144,10 @@ def mtp_forward( 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) @@ -2637,19 +4206,139 @@ def mtp_forward( mrope_section1=core.mrope_section1, mrope_section2=core.mrope_section2, ) - # Dense attention; see `MTPLayer` for why the module's QSA indexer is unused. - 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, - **step, - **_CALL_CONSTANTS, + # 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([rows, core.index_scratch_width], dtype=dt, device=dev) + 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=pos3 if _QSA_INDEX_USE_MROPE else 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, + [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) + 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=True, + ) + topk_indices = md.qsa_topk_indices[:rows] + visible_blocks = md.qsa_visible_blocks[:rows] + _select( + core, + scores, + md.qsa_topk_row_starts[:rows], + visible_blocks, + topk_indices, + _live_score_columns(core, md, batch), + ) + if share: + self._capture_selection(md, topk_indices, visible_blocks) + selected = qsa_expand_block_indices( + topk_indices, + logical, + md.qsa_sequence_lengths[:rows], + compress_ratio=core.compress_ratio, + token_topk=core.token_topk, + ) + attn_out = torch.reshape( + _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=logical, + compress_ratio=core.compress_ratio, + ), + [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) x = qwen4_exp_post_inject(x, cublas_mm(attn_out, M.o), inject) @@ -2668,6 +4357,11 @@ def __init__(self, core: Qwen4ExpTp1ModelingV2) -> None: 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) @@ -2729,6 +4423,113 @@ def step_args(self, md: TrtllmAttentionMetadata) -> dict: # --------------------------------------------------------------- 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 either 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.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.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 @@ -2737,176 +4538,262 @@ def _grouped_norm(self, x: torch.Tensor, weights: tuple) -> torch.Tensor: [flashinfer_gemma_rmsnorm(part, w, core.eps) for w, part in zip(weights, parts)], 1 ) - def _moe_block(self, x: torch.Tensor, L: _LayerWeights, dt: torch.dtype) -> torch.Tensor: - """See `PrefillTarget._moe_block`.""" + 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` 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. `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 - # The routing op reads logits from data_ptr() ignoring strides. - logits = cublas_mm(x, L.router) - topk_ids, topk_w = renorm_moe_routing_op(logits, topk=core.topk, output_dtype=dt) - xq, xsf = fp4_quantize(x, L.moe_gscale, _SF_VEC, False, _SF_LINEAR) - routed = fp4_block_scale_moe_runner( - routing_logits=None, - routing_bias=None, - hidden_states=xq, - hidden_states_scale=xsf.view(torch.float8_e4m3fn), - gemm1_weights=L.fc1_w, - gemm1_weights_scale=L.fc1_s, - gemm1_bias=None, - gemm1_alpha=None, - gemm1_beta=None, - gemm1_clamp_limit=None, - gemm2_weights=L.fc2_w, - gemm2_weights_scale=L.fc2_s, - gemm2_bias=None, - output1_scale_scalar=L.out1_scale, - output1_scale_gate_scalar=L.out1_gate_scale, - output2_scale_scalar=L.out2_scale, - num_experts=core.num_experts, - top_k=core.topk, - 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, - do_finalize=True, - act_type=_ACT_TYPE_SWIGLU, - topk_weights=topk_w, - topk_ids=topk_ids, - )[0] - shared = cublas_mm(flashinfer_silu_and_mul(cublas_mm(x, L.sh_gu)), L.sh_dn) - gate_logits = cublas_mm(x, L.sh_gate) - return fused_sigmoid_gate_mul_add(routed, gate_logits, shared) - - def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step): - """One gated GQA layer plus its QSA index branch for a generation-only batch; - see `PrefillTarget._attention_block`.""" + 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): + shared = gemm(flashinfer_silu_and_mul(gemm(x, L.sh_gu)), L.sh_dn) + gate_logits = gemm(x, L.sh_gate) + # The routing reads its logits as a dense row-major buffer from data_ptr() and + # ignores strides; both projection entries allocate them contiguous. + logits = gemm(x, L.router) + 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) + moe_op = fp4_block_scale_moe_renorm_permuted if fold else 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 - ( - req_idx, - logical, - seq_lens_q, - visible_blocks, - topk_indices, - topk_row_starts, - block_table, - tokens_per_block, - pos, + 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, - coords, - sparse, - ) = qsa - num_tokens = x.shape[0] - - qg = cublas_mm(x, L.q) - k = cublas_mm(x, L.k) - v = cublas_mm(x, L.v) - - # Destination views into the packed QKV buffer. - qkv_h = torch.reshape(qkv, [num_tokens, 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, [num_tokens, core.heads_q, 2 * core.head_dim]), - [core.head_dim, core.head_dim], - 2, - ) - q_dst.copy_(q_src) - k_dst.copy_(torch.reshape(k, [num_tokens, core.heads_kv, core.head_dim])) - v_dst.copy_(torch.reshape(v, [num_tokens, core.heads_kv, core.head_dim])) - fused_qk_norm_rope( - qkv, + core._cos_sin, + L.q_norm, + L.k_norm, + L.idx_qn, num_heads_q=core.heads_q, num_heads_k=core.heads_kv, - num_heads_v=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, - q_weight=L.q_norm, - k_weight=L.k_norm, - base=core.theta, - is_neox=True, - position_ids=pos3, - is_qk_norm=True, - # The gemma QK norm, `rmsnorm(x) * (1 + w)`. - use_gemma=True, - # Real mRoPE, unconditionally (see PrefillTarget). - use_mrope=True, 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 (see PrefillTarget) ------------------------ - # The raw key is stored before any norm/rotation; the compression step - # applies both itself. - qk_index = cublas_mm(x, L.idx) - index_cache = core._bound.index_cache[i] - token_k = torch.split( - torch.reshape(qk_index, [num_tokens, 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, - core._bound.position_cache, + # ---- 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 pool, + on a step with one row per request, everything up to the index scores is one + launch, `qwen4_exp_attn_decode_prologue`. A verify step stays on the composed + prologue: at a group's last token that launch reads the group's earlier keys + from the index plane, and in a verify group the rows before it are written by + the same 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] + + proj = gemm(x, L.attn_in) + main = torch.cuda.current_stream() + side = core._side_stream + qkv_view, gate, index_view = torch.split( + proj, [core.qkv_width, core.q_width, core.index_scratch_width], 1 ) - # Query norm+rotation; coordinate 0 alone, matching `_QSA_MROPE_SECTION`. - scratch_h = torch.reshape(index_scratch, [num_tokens, core.index_scratch_width]) - scratch_h[:, : 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=L.idx_qn, - k_weight=L.idx_kn, - base=core.theta, - is_neox=True, - position_ids=pos3 if _QSA_INDEX_USE_MROPE else 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_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( - scratch_h, + 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] - 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, - ) + index_section = (_QSA_MROPE_SECTION or (0, 0, 0)) if _QSA_INDEX_USE_MROPE else (0, 0, 0) + index_cache = core._bound.index_cache[i] + kv = core._bound.kv[i] + fused_prologue = sparse and not stage and kv.k_scale is None and kv.k.dtype == proj.dtype + if fused_prologue: + # The norm + rope of both planes, the K/V append the dense op would have + # owned, 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, + ) + 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: - k_cache, v_cache = core._bound.kv[i] - # Append the fused-norm+rope'd K/V; the selection reads the current - # token from the pool. - qsa_paged_kv_store( - k_dst, v_dst, k_cache, v_cache, req_idx, logical, block_table, tokens_per_block - ) + 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) scores = qsa_paged_index_scores( q=q_index, index_cache=index_cache, @@ -2917,35 +4804,89 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step): 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. + use_pdl=fused_prologue, + prefetch=fused_prologue and num_tokens <= _SPARSE_DECODE_MAX_ROWS, ) - # `index_topk` is the block budget (see PrefillTarget). - qsa_topk_select( - scores, topk_row_starts, visible_blocks, topk_indices, index_topk=core.block_topk - ) - selected = qsa_expand_block_indices( - topk_indices, - logical, - seq_lens_q, - compress_ratio=core.compress_ratio, - token_topk=core.token_topk, - ) - out = torch.reshape( - qsa_paged_sparse_gqa( + 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: + attn = _sparse_gqa( + core, + kv, q=q_dst, - k_cache=k_cache, - v_cache=v_cache, block_table=block_table, selected_tokens=selected, request_indices=req_idx, tokens_per_block=tokens_per_block, - softmax_scale=core.softmax_scale, # The prefill pair bounds a context row's scan; none here. query_positions=None, compress_ratio=None, - ), - [num_tokens, core.q_width], - ) + ) + 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, @@ -2954,13 +4895,18 @@ def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step): 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 cublas_mm(out, L.o) + if not gated: + fused_sigmoid_mul_inplace(out, gate) + return gemm(out, L.o) - def _gdn_verify(self, mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, n, p) -> torch.Tensor: + def _gdn_verify( + self, mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, n, p, replay + ) -> torch.Tensor: """Recurrent half of a speculative verify step; see `PrefillTarget._gdn_verify`.""" core = self.core @@ -2982,6 +4928,8 @@ def _gdn_verify(self, mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, n, p) -> ) # 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 ) @@ -3021,24 +4969,74 @@ def _gdn_verify(self, mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, n, p) -> ) return torch.reshape(out, [1, n * p, core.num_v_heads, core.head_v_dim]) - def _gdn_block(self, x, i, L, gdn_state) -> torch.Tensor: + 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 via the verify path.""" + 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, after the conv step, as one streaming launch (`gdn_decode_rows_gated_norm`). + Over a bf16 pool the recurrence is `fused_sigmoid_gating_delta_rule_update`'s rows + form, which is batch-invariant.""" core = self.core - si_all, qsl_long, is_verify, num_decodes, gen_p, int_idx = gdn_state + si_all, is_verify, num_decodes, gen_p, int_idx, replay = gdn_state num_tokens = x.shape[0] - - mixed_qkv = cublas_mm(x, L.in_qkv) - z = cublas_mm(x, L.in_z) - a = cublas_mm(x, L.in_a) - b = cublas_mm(x, L.in_b) - ssm_states = core._bound.ssm[i] + fp32_state = ssm_states.dtype == torch.float32 + + conv_fused = fused_conv and not is_verify and num_tokens <= _GDN_CONV_FUSED_MAX_ROWS + if 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, + ) + else: + proj = gemm(x, L.in_proj) + mixed_qkv, z, a, b = torch.split(proj, core.in_proj_split, 1) + + if fp32_state 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 gemm(o, L.out) if is_verify: out = self._gdn_verify( - mixed_qkv, a, b, ssm_states, i, L, si_all, int_idx, num_decodes, gen_p + mixed_qkv, + a.contiguous(), + b.contiguous(), + ssm_states, + i, + L, + si_all, + int_idx, + num_decodes, + gen_p, + replay, ) - else: + return self._gdn_epilogue(out, z, L, num_tokens, gemm) + if not conv_fused: causal_conv1d_update( x=mixed_qkv, conv_state=core._bound.conv[i], @@ -3049,31 +5047,56 @@ def _gdn_block(self, x, i, L, gdn_state) -> torch.Tensor: conv_state_indices=si_all, pad_slot_id=_PAD_SLOT_ID, ) - # q/k/v stay views over the conv output (native strides, no copy). - qv, kv, vv = torch.split(mixed_qkv, [core.key_dim, core.key_dim, core.value_dim], 1) + if fp32_state: + # 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 gemm(o, L.out) + else: + # The rows form, `cu_seqlens` None: it finds row n's tokens at + # `n * q.stride(1)`, so q/k/v are split from a contiguous `[T, 10240]` copy + # reshaped to `[T, 1, 80, 128]` (K and V heads are both 128 wide), where the + # size-1 token axis carries the row stride; a/b are contiguous `[T, 48]`. + qkv_rows = torch.reshape(mixed_qkv.contiguous(), [num_tokens, 1, -1, core.head_k_dim]) + q, k, v = torch.split( + qkv_rows, [core.num_k_heads, core.num_k_heads, core.num_v_heads], 2 + ) out = fused_sigmoid_gating_delta_rule_update( A_log=L.a_log, - a=a, + a=torch.reshape(a.contiguous(), [num_tokens, 1, core.num_v_heads]), dt_bias=L.dt_bias, softplus_beta=_SOFTPLUS_BETA, softplus_threshold=_SOFTPLUS_THRESHOLD, - q=torch.reshape(qv, [1, num_tokens, core.num_k_heads, core.head_k_dim]), - k=torch.reshape(kv, [1, num_tokens, core.num_k_heads, core.head_k_dim]), - v=torch.reshape(vv, [1, num_tokens, core.num_v_heads, core.head_v_dim]), - b=b, + q=q, + k=k, + v=v, + b=torch.reshape(b.contiguous(), [num_tokens, 1, core.num_v_heads]), initial_state_source=ssm_states, initial_state_indices=si_all, scale=None, use_qk_l2norm_in_kernel=True, - cu_seqlens=qsl_long, + cu_seqlens=None, output=None, ) - return self._gdn_epilogue(out, z, L, num_tokens) + return self._gdn_epilogue(out, z, L, num_tokens, gemm) - def _gdn_epilogue(self, out, z, L, num_tokens) -> torch.Tensor: - """Gated RMS norm and output projection shared by every recurrence.""" + 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 - # Sigmoid gate, plain (non-gemma) weight. 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]), @@ -3082,53 +5105,109 @@ def _gdn_epilogue(self, out, z, L, num_tokens) -> torch.Tensor: fp8_scale=None, gate_activation="sigmoid", ) - return cublas_mm(torch.reshape(o, [num_tokens, core.value_dim]), L.out) - - def _ple_block(self, hc, input_ids, ple_state) -> torch.Tensor: - """The n-gram side path on one decoder layer; see `PrefillTarget._ple_block`.""" + return gemm(torch.reshape(o, [num_tokens, core.value_dim]), L.out) + + 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`) takes the unfused ops instead, whose + intermediates `SpeculativeRollback` stages (see `PrefillTarget._ple_delta`).""" core = self.core S = self._stack - conv_states, ngram_pool, si_all, carry_all, qsl = ple_state - - window = ngram_token_history( - input_ids, ngram_pool, si_all, carry_all, qsl, core.eos_token_id, core.ngram_size - ) - rows = torch.reshape( - ngram_hash( - window, - layer_multipliers=S.ple_mult, - head_vocab_sizes=S.ple_vocab, - head_offsets=S.ple_offs, - eos_token_id=core.eos_token_id, - num_query_tokens=1, - ), - [-1, core.ngram_heads], - ) - gathered = nn.functional.embedding(rows, core.w["ple_table"]) - ngram = torch.reshape( - dequantize_e4m3_per_tensor(gathered, S.ple_tscale), [-1, core.ple_embed_dim] - ) - - key_normed = self._grouped_norm(cublas_mm(ngram, S.ple_key), S.ple_nk) - query_normed = self._grouped_norm(hc, S.ple_nq) - value = cublas_mm(ngram, S.ple_val) - gated = gated_ngram_injection( - key_normed, query_normed, value, core.hc_count, _PLE_CLAMP_MIN + conv_states, ngram_pool, si_all, carry_all, qsl, gen_p = ple_state + if stage: + core.spec_rollback.stage_ple_pools(ngram_pool, conv_states) + window = ngram_token_history( + input_ids, ngram_pool, si_all, carry_all, qsl, core.eos_token_id, core.ngram_size + ) + rows = torch.reshape( + ngram_hash( + window, + layer_multipliers=S.ple_mult, + head_vocab_sizes=S.ple_vocab, + head_offsets=S.ple_offs, + eos_token_id=core.eos_token_id, + num_query_tokens=1, + ), + [-1, core.ngram_heads], + ) + gathered = nn.functional.embedding(rows, core.w["ple_table"]) + ngram = torch.reshape( + dequantize_e4m3_per_tensor(gathered, S.ple_tscale), [-1, core.ple_embed_dim] + ) + key_normed = self._grouped_norm(cublas_mm(ngram, S.ple_key), S.ple_nk) + query_normed = self._grouped_norm(hc, S.ple_nq) + value = cublas_mm(ngram, S.ple_val) + gated = gated_ngram_injection( + key_normed, query_normed, value, core.hc_count, _PLE_CLAMP_MIN + ) + conv_in = self._grouped_norm(gated, S.ple_nc) + core.spec_rollback.stage_ple_rows(window, conv_in) + # The short conv reads channel-major; a bare transpose is the certified form. + convolved = dilated_short_conv_paged( + x=torch.transpose(conv_in, 0, 1), + weight=S.ple_conv_w, + bias=None, + conv_states=conv_states, + query_start_loc=qsl, + cache_indices=si_all, + has_initial_state=carry_all, + dilation=core.ple_dilation, + silu_activation=True, + pad_slot_id=_PAD_SLOT_ID, + ) + hc = torch.add(hc, torch.add(gated, torch.transpose(convolved, 0, 1))) + return qwen4_exp_hc_inject_norm(hc, None, None, next_norm, core.eps, core.hc_count) + 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, ) - # The short conv reads channel-major; a bare transpose is the certified form. - convolved = dilated_short_conv_paged( - x=torch.transpose(self._grouped_norm(gated, S.ple_nc), 0, 1), - weight=S.ple_conv_w, - bias=None, - conv_states=conv_states, - query_start_loc=qsl, - cache_indices=si_all, - has_initial_state=carry_all, - dilation=core.ple_dilation, - silu_activation=True, - pad_slot_id=_PAD_SLOT_ID, + return 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, ) - return torch.add(gated, torch.transpose(convolved, 0, 1)) # -------------------------------------------------------------- forward --- @@ -3137,7 +5216,6 @@ def forward( attn_metadata: AttentionMetadata, input_ids: torch.IntTensor | None = None, position_ids: torch.IntTensor | None = None, - inputs_embeds: torch.FloatTensor | None = None, **kwargs, ): core = self.core @@ -3146,10 +5224,12 @@ def forward( 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 - # Three views of one `[3, 1, T]` buffer (see PrefillTarget.forward): `pos`, - # contiguous `pos3`, and `coords`. - pos = torch.reshape(torch.split(position_ids, [1, _QSA_POSITION_AXES - 1], 0)[0], [-1]) + # 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 = torch.empty( [_QSA_POSITION_AXES, position_ids.shape[-1]], @@ -3158,25 +5238,14 @@ def forward( ) pos3.copy_(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1])) - # Spliced vision rows are already the embedding; a lookup would overwrite them. - if inputs_embeds is None: - h = nn.functional.embedding(input_ids, core.w["embed"]) - else: - assert inputs_embeds.shape[0] == input_ids.shape[0], ( - f"the spliced embedding carries {inputs_embeds.shape[0]} rows for " - f"{input_ids.shape[0]} token ids" - ) - assert inputs_embeds.shape[1] == core.hidden, ( - f"the spliced embedding is {inputs_embeds.shape[1]} wide, not {core.hidden}" - ) - h = inputs_embeds + 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.forward. + # 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 " @@ -3200,7 +5269,7 @@ def forward( "speculation off; the base path implements one query token per " "generation sequence" ) - # `gen_p > 1` is load-bearing for verify (see PrefillTarget.forward). + # `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, ( @@ -3209,21 +5278,22 @@ def forward( "per-step conv snapshots and the recurrent rollback buffer are " "sized from the second and written from the first" ) - assert 1 <= gen_p <= 4, ( - f"predicted_tokens_per_seq {gen_p} is outside the 1..4 range " - f"attention/thop_attention certifies" - ) + 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] - gdn_state = (si_all, qsl_long, is_verify, batch, gen_p, int_idx) + 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) # Every generation sequence carries its n-gram history; the engine's # `has_initial_states` covers the context prefix only (empty here). @@ -3231,12 +5301,19 @@ def forward( [batch], dtype=mm.has_initial_states.dtype, device=mm.has_initial_states.device ) # PLE segmentation vector: rebuilt from `seq_lens_cuda` when `gen_p > 1` (see - # PrefillTarget.forward). + # PrefillTarget.prepare). 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) + ple_state = ( + core._bound.ple_state[0], + core._bound.ple_state[1], + si_all, + ple_carry, + ple_qsl, + gen_p, + ) # 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. @@ -3251,41 +5328,97 @@ def forward( md.qsa_topk_row_starts[:num_tokens], md.qsa_block_table, md.tokens_per_block, - pos, pos3, coords, sparse, ) - # Scratch shared by all attention layers (each fully rewrites them). + # Staging for the speculative rollback, as on the prefill target (see + # PrefillTarget.prepare); every row here is a generation row. + stage = is_verify + rollback = core.spec_rollback + if rollback is not None: + rollback.discard() + if 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, + ) + + # 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) - for i, L in enumerate(self._layers): + 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: - hc = torch.add(hc, self._ple_block(hc, input_ids, ple_state)) - mixed, inject = qwen4_exp_gated_residual( - self._grouped_norm(hc, L.ahc_norm), L.ahc_down, L.ahc_up, L.ahc_inject, hc_count + 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 + mixed, i, L, qkv, index_scratch, attn_out, qsa, step, stage, gemm ) else: - block_out = self._gdn_block(mixed, i, L, gdn_state) - hc = qwen4_exp_post_inject(hc, block_out, inject) - mixed, inject = qwen4_exp_gated_residual( - self._grouped_norm(hc, L.mhc_norm), L.mhc_down, L.mhc_up, L.mhc_inject, hc_count + 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 ) - hc = qwen4_exp_post_inject(hc, self._moe_block(mixed, L, dt), inject) - # The stack-final mixer with no injection weight is the final norm. - S = self._stack - final, _ = qwen4_exp_gated_residual( - self._grouped_norm(hc, S.fhc_norm), S.fhc_down, S.fhc_up, None, hc_count + # 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 @@ -3330,6 +5463,26 @@ def _mtp_moe_block(self, x, dt) -> torch.Tensor: 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, 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, @@ -3396,6 +5549,44 @@ def mtp_forward( "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, ( @@ -3408,6 +5599,10 @@ def mtp_forward( 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) @@ -3466,19 +5661,150 @@ def mtp_forward( mrope_section1=core.mrope_section1, mrope_section2=core.mrope_section2, ) - # Dense attention; see `MTPLayer` for why the module's QSA indexer is unused. - 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, - **step, - **_CALL_CONSTANTS, - ) + 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 + # The QSA index branch, as in the trunk's `_attention_block` (see the prefill + # target's draft step). Under IndexShare a draft decode step runs none of it: what + # it would write (its row's raw key and coordinates, and the anchor of a group that + # row completes) sits at or past the newly committed length, and nothing reads it + # before the next engine step's draft-extend rewrites those positions, storing + # before it compresses and compressing before it scores. + if not reuse: + 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([rows, core.index_scratch_width], dtype=dt, device=dev) + 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=pos3 if _QSA_INDEX_USE_MROPE else 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, + [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 + _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] + if reuse: + # The capture this loop's draft-extend step made, with the groups completed + # since appended; the expansion then adds the incomplete tail, so the + # positions drafted since the capture stay visible. + 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, + ) + else: + 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, + ) + _select(core, scores, md.qsa_topk_row_starts[:rows], visible_blocks, topk_indices) + if share: + self._capture_selection(md, topk_indices, visible_blocks, gen_seqs, p) + selected = qsa_expand_block_indices( + topk_indices, + logical, + md.qsa_sequence_lengths[:rows], + compress_ratio=core.compress_ratio, + token_topk=core.token_topk, + ) + attn_out = torch.reshape( + _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, + ), + [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) x = qwen4_exp_post_inject(x, cublas_mm(attn_out, M.o), inject) @@ -3555,15 +5881,48 @@ def _vision_views(core: Qwen4ExpTp1ModelingV2) -> _VisionWeights: ) +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 through the 27-block ViT and its patch merger. + """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, 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.""" + 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 @@ -3575,10 +5934,11 @@ def __init__(self, core: Qwen4ExpTp1ModelingV2) -> 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 and stays on the host (no synchronization). Returns - `[sum(t*h*w) / merge**2, out_hidden_size]`, one row per placeholder.""" + 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, dev = pixel_values.dtype, pixel_values.device + 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` " @@ -3588,29 +5948,53 @@ def __call__(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch. f"pixel_values must be [patches, {core.v_pixel_width}] flat patch rows; " f"got {tuple(pixel_values.shape)}" ) - 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. - grid = grid_thw.tolist() - seg_lens: list[int] = [] - for row in grid: + 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(row[0]), int(row[1]), int(row[2]) + 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] - assert sum(seg_lens) == num_patches, ( - f"grid_thw describes {sum(seg_lens)} patch rows and pixel_values carries {num_patches}" - ) num_segments = len(seg_lens) max_seg = max(seg_lens) - # Batch-state tensors for the encoder-only call; they describe the images. + # 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) @@ -3681,12 +6065,26 @@ class MTPLayer: 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 dense, the one deliberate departure from the checkpoint (its QSA - indexer stays unloaded): with a separate draft KV manager the worker swaps - `attn_metadata.kv_cache_manager`, `QSAAttentionMetadata` implements no - `prepare_for_draft_forward`, so `qsa_block_table` would decode the target - manager's page offsets against the draft pool. Below `dense_seq_len_threshold` - dense is bit-for-bit the trunk's. + 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.""" @@ -3758,6 +6156,197 @@ 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_pools` / + `stage_ple_rows` 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.""" + + def __init__(self) -> None: + self._pending: dict | None = None + + # ------------------------------------------------------------- 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, + ) -> 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.""" + 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, + ) + + 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_rows(self, window: torch.Tensor, conv_input: torch.Tensor) -> None: + """After the id history and before the conv: the generation rows' id windows + (`[T, ngram_size]`, newest id last) and conv input (`[T, C]`, token-major). Copied + rather than viewed, so a long context prefix in a mixed batch is not kept alive + through the draft loop.""" + pending = self._pending + prefix, n, p = pending["prefix"], pending["n"], pending["p"] + # Row t of the window without its oldest id is the history after token t. + pending["ngram_candidates"] = torch.reshape(window[prefix:, 1:], [n, p, -1]).clone() + pending["conv_x"] = torch.reshape(conv_input[prefix:], [n, p, -1]).clone() + + 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]) + 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" + ) + 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"] + 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 _WideningLogitsProcessor(LogitsProcessor): + """The shell's logits processor (its row gather, its fp32 logits) over `_WidenedHead`.""" + + def forward( + self, hidden_states, lm_head, attn_metadata, return_context_logits: bool = False + ) -> torch.Tensor: + return super().forward( + hidden_states, _WidenedHead(lm_head), attn_metadata, return_context_logits + ) + + class ModelingV2Qwen38FlashNextNvfp4Sm103Tp1( DecoderModelForCausalLM[Qwen4ExpTp1ModelingV2, PretrainedConfig] ): @@ -3799,6 +6388,8 @@ def __init__(self, model_config: ModelConfig): 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) @@ -3844,6 +6435,10 @@ def __init__(self, model_config: ModelConfig): ) # 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 @@ -3852,10 +6447,11 @@ def __init__(self, model_config: ModelConfig): @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` on the device where + 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", ] @@ -3871,43 +6467,164 @@ def multimodal_token_ids(self) -> torch.Tensor: self._mm_token_ids_cache = torch.tensor(self._mm_token_ids_value, dtype=torch.int32) return self._mm_token_ids_cache - def _multimodal_embeds(self, input_ids, kwargs): - """Run the vision tower for this batch and splice its rows in; returns the - `[T, hidden]` embedding buffer, or None when the step carries no media. + 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 `mm_item_order` 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: + self._check_item_order(grids, getattr(request, "py_mm_item_order", 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 mm_item_order 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: + listed[entry["modality"]].append(int(entry["index"])) + except (KeyError, TypeError, ValueError) as error: + raise ValueError( + f"mm_item_order entry {entry!r} is not an item of this request, which " + f"carries {sorted(grids)}; entries are {{'modality', 'index'}}" + ) from error + for modality, grid in grids.items(): + if sorted(listed[modality]) != list(range(len(grid))): + raise ValueError( + f"mm_item_order 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 cached and - in-chunk row counts; `mm_token_indices` names every placeholder row.""" + `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 - core = self.model rows = [] for params in mm_params: data = getattr(params, "multimodal_data", None) or {} - image = data.get("image") - assert data.get("video") is None, ( - "this target implements the image path only; a video request " - "reached it. The tower itself is per-frame and would run, but " - "the temporal half of the position bookkeeping is not derived " - "here and a wrong one is invisible to every gate" - ) - if image is None: + image, video = data.get("image"), data.get("video") + if image is None and video is None: continue - pixel_values = image.get("pixel_values") - grid_thw = image.get("image_grid_thw") - assert pixel_values is not None and grid_thw is not None, ( - f"the multimodal payload carries {sorted(image)}; this target " - "reads 'pixel_values' and 'image_grid_thw' from it" + 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)}" ) - assert not grid_thw.is_cuda, ( - "image_grid_thw must stay on the host; see multimodal_data_device_paths" - ) - embeds = core.vision_tower(pixel_values, grid_thw) 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) @@ -3916,12 +6633,10 @@ def _multimodal_embeds(self, input_ids, kwargs): 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}; the two are derived from the same grid and a " - "disagreement means the merger's row count is wrong" + 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) - if count == 0: - continue embeds = embeds[start : start + count] rows.append(embeds) if not rows: @@ -3937,10 +6652,59 @@ def _multimodal_embeds(self, input_ids, kwargs): f"the batch has {mm_token_indices.shape[0]} placeholder rows and the " f"tower produced {vision.shape[0]} embedding rows" ) - h = nn.functional.embedding(input_ids, core.w["embed"]) - # Indexed write: placeholder rows form several disjoint, data-dependent runs. - h.index_copy_(0, mm_token_indices, vision) - return h + 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 `mm_item_order` + (`{"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] + order = getattr(params, "mm_item_order", None) + assert order, ( + f"a request carrying {sorted(towers)} needs mm_item_order 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 entry in order: + assert entry["modality"] in items, ( + f"mm_item_order lists a {entry['modality']!r} item; this request " + f"carries {sorted(items)}" + ) + listed[entry["modality"]].append(int(entry["index"])) + for modality, parts in items.items(): + assert sorted(listed[modality]) == list(range(len(parts))), ( + f"mm_item_order lists {modality} items {listed[modality]}; the " + f"payload carries {len(parts)}" + ) + return torch.cat([items[entry["modality"]][int(entry["index"])] for entry in order], 0) def forward( self, @@ -3960,16 +6724,18 @@ def forward( 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 splice happens - here because the base forward would not pass `multimodal_params` through.""" - spliced = self._multimodal_embeds(input_ids, kwargs) - if spliced is not None: - assert inputs_embeds is None, ( - "the engine supplied inputs_embeds and this step also carries " - "multimodal params; this target owns the splice and cannot " - "reconcile two embedding sources" - ) - inputs_embeds = spliced + (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 " @@ -3989,14 +6755,16 @@ def forward( "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( - 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, + 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. @@ -4006,6 +6774,12 @@ def forward( 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, 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 index 4982ae3f00fe..7987c692ddfd 100644 --- 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 @@ -34,67 +34,16 @@ class FusedSigmoidMulInplace(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell( - why="decode: 1 token, 24x256 gate as the strided half of q_proj", - spec=dict(tokens=1, **_CKPT), - ), - Cell( - why="17 tokens: several rows, 1024-element blocks over 6144", - spec=dict(tokens=17, **_CKPT), - ), - Cell( - why="prefill 2048 tokens: >= 1024 rows selects 2048-element blocks", - spec=dict(tokens=2048, **_CKPT), - ), - Cell( - why="attention_output a [33, 6144] window of a [35, 8192] buffer", - spec=dict(tokens=33, out="window", **_CKPT), - ), - Cell( - why="attention_output every other row: row stride 2 * hidden", - spec=dict(tokens=32, out="every_other_row", **_CKPT), - ), - Cell( - why="2-D gate: one head spanning the whole row", - spec=dict(tokens=9, heads=1, head_dim=6144, dtype=torch.bfloat16, gate="flat"), - ), - Cell( - why="gate fp32 under a bf16 output: widened, not rejected", - spec=dict( - tokens=9, - heads=1, - head_dim=512, - dtype=torch.bfloat16, - gate="flat", - gate_dtype=torch.float32, - ), + 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="fp16 operands", - spec=dict(tokens=33, heads=4, head_dim=128, dtype=torch.float16, gate="packed"), + 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="fp32 operands", - spec=dict(tokens=33, heads=4, head_dim=128, dtype=torch.float32, gate="packed"), - ), - Cell( - why="hidden 16: a single block narrower than any vector", - spec=dict(tokens=1, heads=4, head_dim=4, dtype=torch.bfloat16, gate="packed"), - ), - Cell( - why="hidden 100: not a multiple of any vector width", - spec=dict(tokens=3, heads=4, head_dim=25, dtype=torch.bfloat16, gate="packed"), - ), - Cell( - why="hidden 2049 at 1030 rows: 2048-element blocks plus a masked tail", - spec=dict(tokens=1030, heads=3, head_dim=683, dtype=torch.bfloat16, gate="packed"), - ), - Cell( - why="hidden 6145: odd, multi-block, masked tail", - spec=dict(tokens=4, heads=1, head_dim=6145, dtype=torch.bfloat16, gate="packed"), - ), - Cell( - why="0 tokens: returns without launching", - spec=dict(tokens=0, **_CKPT), + why="MTP decode: 1 token, 24x256 gate as the strided half of q_proj", + spec=dict(tokens=1, **_CKPT), ), ) @@ -107,25 +56,30 @@ class FusedSigmoidMulInplace(OpWrapper): (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. Hence no - `is_valid`. Cross-device operands were not probed. + `stride(-1) == 1`. A CPU tensor raises ValueError from Triton (measured). + Hence no `is_valid`. Cross-device operands were not probed. - `gate.dtype` need not equal `attention_output.dtype`: an fp32 gate under a - bf16 output returns the fp32-gate result. The pre-call contents of - `attention_output` are destroyed; writes stay inside its `[T, hidden]` - window (padded columns and trailing rows bitwise unchanged). + The pre-call contents of `attention_output` are destroyed; the gate is + left bitwise unchanged. - Numerics: bf16/fp16 bitwise equal to the fp32 reference except ~1e-6 of + 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], so the per-head indexing is under - test. + 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 2048 rows) + (largest run 8192 rows) - Not certified: CUDA-graph capture, torch.compile. + 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: 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py index 7bfbe1ee84a6..07150279c931 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py @@ -23,11 +23,6 @@ #: `merger.linear_fc1` is `[4608, 4608]`: spatial merge 2 over a 1152-wide tower. MERGER_WIDTH = 4608 -#: Rows the merger is asked for, `sum(t*h*w) / 4`: 64 is a 224x224 image -#: (grid [1, 16, 16]), 140 a 336x448 one ([1, 20, 28]) -- not a multiple of -#: any MoE tile size. -MERGER_ROW_COUNTS = (1, 7, 64, 65, 140, 1024, 4096) - _SQRT_2 = math.sqrt(2.0) @@ -45,91 +40,43 @@ class GeluErf(OpWrapper): # spec: shape, dtype; `dist` "uniform" over [lo, hi] (default) or "merger" # (a tower hidden state scaled by the merger norm's measured gamma, grouped - # 4-to-1, through an fc1 at the checkpoint's measured weight/bias scale); - # `layout` "contiguous" (default), "column_slice" or "transposed". + # 4-to-1, through an fc1 at the checkpoint's measured weight/bias scale). CELLS: tuple[Cell, ...] = ( - tuple( - Cell( - why=f"merger activand, {rows} merged rows x 4608 bf16", - spec=dict(shape=(rows, MERGER_WIDTH), dtype=torch.bfloat16, lo=-16.0, hi=16.0), - ) - for rows in MERGER_ROW_COUNTS - ) - + ( - Cell( - why="activand built the way the merger builds it", - spec=dict(shape=(512, MERGER_WIDTH), dtype=torch.bfloat16, dist="merger"), - ), - Cell( - why="fp16 over [-16, 16]", - spec=dict(shape=(256, MERGER_WIDTH), dtype=torch.float16, lo=-16.0, hi=16.0), - ), - Cell( - why="fp32 over [-16, 16]", - spec=dict(shape=(256, MERGER_WIDTH), dtype=torch.float32, lo=-16.0, hi=16.0), - ), - ) - + tuple( - Cell( - why=f"{str(dtype)[6:]} over [-3, 3], the erf/tanh divergence peak", - spec=dict(shape=(256, MERGER_WIDTH), dtype=dtype, lo=-3.0, hi=3.0), - ) - for dtype in (torch.bfloat16, torch.float16, torch.float32) - ) - + ( - Cell( - why="column slice of a wider buffer: row stride free", - spec=dict( - shape=(64, MERGER_WIDTH), - dtype=torch.bfloat16, - lo=-16.0, - hi=16.0, - layout="column_slice", - ), - ), - Cell( - why="transposed buffer: output keeps the input's strides", - spec=dict( - shape=(64, MERGER_WIDTH), - dtype=torch.bfloat16, - lo=-16.0, - hi=16.0, - layout="transposed", - ), - ), - Cell( - why="1-D input", - spec=dict(shape=(MERGER_WIDTH,), dtype=torch.bfloat16, lo=-16.0, hi=16.0), - ), - Cell( - why="3-D input", - spec=dict(shape=(2, 37, MERGER_WIDTH), dtype=torch.bfloat16, lo=-16.0, hi=16.0), - ), - Cell( - why="zero-row batch", - spec=dict(shape=(0, MERGER_WIDTH), dtype=torch.bfloat16, lo=-16.0, hi=16.0), - ), - ) + Cell( + why="activand built the way the merger builds it", + spec=dict(shape=(512, MERGER_WIDTH), dtype=torch.bfloat16, dist="merger"), + ), + Cell( + why="merger activand, 8192 merged rows x 4608 bf16: the most one tower call carries", + spec=dict(shape=(8192, MERGER_WIDTH), dtype=torch.bfloat16, lo=-16.0, hi=16.0), + ), + Cell( + why="merger activand, 1 merged row x 4608 bf16", + spec=dict(shape=(1, MERGER_WIDTH), dtype=torch.bfloat16, lo=-16.0, hi=16.0), + ), ) note = """ Metadata consumed: none. Out-of-place; `x` is not mutated. - The output inherits `x`'s memory format: a transposed `x` returns a - non-contiguous result (bitwise equal to the call on a contiguous copy). + 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 the fp32 cells only. 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. + 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. - Values: `-inf -> nan` (the closed form does the same); `+inf -> +inf`; - large negatives give `-0.0`; `finfo.max` comes back unchanged. - CUDA-graph capturable; the output is a fresh allocation. + Measured, not certified: `-inf -> nan` (the closed form does the same); + `+inf -> +inf`; large negatives give `-0.0`; `finfo.max` comes back + unchanged; CUDA-graph capturable; the output is a fresh allocation. - Not certified: float8 operands. No index-typed argument; numel past 2^31 - not probed. + 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. """ 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/gelu_tanh.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py index 2f6c5955f1dc..e9cd2ebec23a 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py @@ -40,93 +40,43 @@ class GeluTanh(OpWrapper): # spec: shape, dtype; `dist` "uniform" over [lo, hi] (default) or "fc1" (a # unit-normal hidden state through an fc1 at the checkpoint's measured - # weight std 0.0197 / bias std 1.676); `layout` "contiguous" (default), - # "column_slice" (the first half of a twice-as-wide buffer) or "transposed". + # weight std 0.0197 / bias std 1.676). CELLS: tuple[Cell, ...] = ( - tuple( - Cell( - why=f"block MLP activand, {rows} patch rows x 4304 bf16", - spec=dict(shape=(rows, INTERMEDIATE_SIZE), dtype=torch.bfloat16, lo=-16.0, hi=16.0), - ) - for rows in (1, 2, 7, 64, 1024, 4096) - ) - + ( - Cell( - why="activand built through fc1 at the checkpoint's measured scale", - spec=dict(shape=(512, INTERMEDIATE_SIZE), dtype=torch.bfloat16, dist="fc1"), - ), - Cell( - why="fp16 over [-16, 16]", - spec=dict(shape=(256, INTERMEDIATE_SIZE), dtype=torch.float16, lo=-16.0, hi=16.0), - ), - Cell( - why="fp32 over [-16, 16]", - spec=dict(shape=(256, INTERMEDIATE_SIZE), dtype=torch.float32, lo=-16.0, hi=16.0), - ), - ) - + tuple( - Cell( - why=f"{str(dtype)[6:]} over [-3, 3], the tanh/erf divergence peak", - spec=dict(shape=(256, INTERMEDIATE_SIZE), dtype=dtype, lo=-3.0, hi=3.0), - ) - for dtype in (torch.bfloat16, torch.float16, torch.float32) - ) - + ( - Cell( - why="column slice of a wider buffer: row stride free", - spec=dict( - shape=(64, INTERMEDIATE_SIZE), - dtype=torch.bfloat16, - lo=-16.0, - hi=16.0, - layout="column_slice", - ), - ), - Cell( - why="transposed buffer: output keeps the input's strides", - spec=dict( - shape=(64, INTERMEDIATE_SIZE), - dtype=torch.bfloat16, - lo=-16.0, - hi=16.0, - layout="transposed", - ), - ), - Cell( - why="1-D input", - spec=dict(shape=(INTERMEDIATE_SIZE,), dtype=torch.bfloat16, lo=-16.0, hi=16.0), - ), - Cell( - why="3-D input", - spec=dict( - shape=(2, 37, INTERMEDIATE_SIZE), dtype=torch.bfloat16, lo=-16.0, hi=16.0 - ), - ), - Cell( - why="zero-row batch", - spec=dict(shape=(0, INTERMEDIATE_SIZE), dtype=torch.bfloat16, lo=-16.0, hi=16.0), - ), - ) + Cell( + why="activand built through fc1 at the checkpoint's measured scale", + spec=dict(shape=(512, 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), + ), + Cell( + why="block MLP activand, 1 patch row x 4304 bf16", + spec=dict(shape=(1, INTERMEDIATE_SIZE), dtype=torch.bfloat16, lo=-16.0, hi=16.0), + ), ) note = """ Metadata consumed: none. Out-of-place; `x` is not mutated. - The output inherits `x`'s memory format: a transposed `x` returns a - non-contiguous result (bitwise equal to the call on a contiguous copy). + 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). - Identity is pinned by the fp32 cells only. 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. + 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. - Values: `-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 allocation. + 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 + allocation. - Not certified: float8 operands. No index-typed argument; numel past 2^31 + 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. """ 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/silu_and_mul_rows.py new file mode 100644 index 000000000000..cae68c64e1a9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/silu_and_mul_rows.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""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. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: This checkpoint's shared expert: intermediate 640, the fused GEMM output row 2 * 640 + 8. +_INTER = 640 +_WIDTH = 2 * _INTER + 8 +_BLOCK = 1024 + + +@triton.jit +def _silu_and_mul_rows_kernel(x_ptr, out_ptr, stride_x, INTER: tl.constexpr, BLOCK: tl.constexpr): + row = tl.program_id(0).to(tl.int64) + cols = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + mask = cols < INTER + g = tl.load(x_ptr + row * stride_x + cols, mask=mask, other=0.0).to(tl.float32) + u = tl.load(x_ptr + row * stride_x + INTER + cols, mask=mask, other=0.0).to(tl.float32) + y = g / (1.0 + tl.exp(-g)) * u + tl.store(out_ptr + row * INTER + cols, y.to(out_ptr.dtype.element_ty), mask=mask) + + +class SiluAndMulRows(OpWrapper): + """`out[t, i] = silu(x[t, i]) * x[t, inter + i]` for `i < inter`: the gate half first, as + `flashinfer_silu_and_mul` reads it, in fp32 and rounded once to `x`'s dtype. Columns of + `x` past `2 * inter` are not read. Returns a fresh contiguous `[T, inter]` tensor. + + Fusion boundary: the activation only; the GEMM that forms `x` and the one that reads + `out` are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + 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)), + ) + + 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. + + 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). + + 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). + + 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. + """ + + def __call__(self, x: torch.Tensor, inter: int) -> torch.Tensor: + rows = x.shape[0] + out = torch.empty((rows, inter), dtype=x.dtype, device=x.device) + if rows > 0: + _silu_and_mul_rows_kernel[(rows, triton.cdiv(inter, _BLOCK))]( + x, out, x.stride(0), INTER=inter, BLOCK=_BLOCK, num_warps=4 + ) + return out + + def reference(self, x: torch.Tensor, inter: int) -> torch.Tensor: + """fp32; the gate half first.""" + gate, up = x[:, : 2 * inter].float().chunk(2, dim=-1) + return (F.silu(gate) * up).to(x.dtype) + + def is_valid(self, x: torch.Tensor, inter: int) -> None: + assert x.dim() == 2 and x.stride(1) == 1, "x [T, width] with unit column stride" + assert x.shape[1] >= 2 * inter, "x carries the gate and up halves" + assert x.dtype in (torch.bfloat16, torch.float16), "bf16 / fp16" + + +silu_and_mul_rows = SiluAndMulRows() 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 new file mode 100644 index 000000000000..b2736198b176 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_canonical_selection.py @@ -0,0 +1,214 @@ +# 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_expand_block_indices.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_expand_block_indices.py index 3bf02f3b3eac..ba1cf5acd32e 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_expand_block_indices.py @@ -17,6 +17,10 @@ _CKPT = dict(compress_ratio=4, token_topk=2048) +#: 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)) + + def _spec(batch, **overrides) -> dict: """A cell spec at this checkpoint's geometry and the engine's metadata dtypes. @@ -24,22 +28,18 @@ def _spec(batch, **overrides) -> dict: ("sweep", n) n rows at 4096 + 13 i, length = position + 1 ("serving", requests, context_lens) one packed step: context chunks, then two MTP rows per remaining request - ("rows", positions) explicit positions, length = position + 1 + ("verify", starts, p) one MTP verify block of p rows per start, + every row's length the block's end """ base = dict( batch=batch, position_dtype=torch.int64, # qsa_logical_positions length_dtype=torch.int32, # qsa_sequence_lengths - layout="dense", **_CKPT, ) return {**base, **overrides} -def _geometry(compress_ratio: int, token_topk: int) -> dict: - return _spec(("rows", (4096, 97, 30)), compress_ratio=compress_ratio, token_topk=token_topk) - - class QsaExpandBlockIndices(OpWrapper): """Expand each row's selected compressed groups into token ids, then append the incomplete group at the causal frontier. @@ -70,67 +70,20 @@ class QsaExpandBlockIndices(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - # ---- decode rows: the column-split launch below 128 rows ------------- + # ---- decode rows: the column-split launch below 128 rows, whole rows from 128 Cell(why="decode: one row, columns split across programs", spec=_spec(("sweep", 1))), - Cell(why="decode MTP: one request's two rows", spec=_spec(("sweep", 2))), - Cell(why="decode: one row for each of 32 request slots", spec=_spec(("sweep", 32))), - Cell(why="127 rows: the last row count on the column split", spec=_spec(("sweep", 127))), - Cell(why="128 rows: the first row count on whole-row programs", spec=_spec(("sweep", 128))), - Cell(why="256 rows on whole-row programs", spec=_spec(("sweep", 256))), - # ---- serving steps: sequence_lengths > query_position + 1 inside chunks - Cell( - why="mixed 1334 rows: 4 chunks + 28 MTP decodes, 32 slots", - spec=_spec(("serving", 32, (900, 200, 150, 28))), - ), - Cell( - why="mixed 8190 rows: 4 chunks + 28 MTP decodes, 32 slots", - spec=_spec(("serving", 32, (8000, 100, 30, 4))), - ), + # ---- MTP verify blocks (max_draft_len up to 7): p rows per request, every + # row's length the block's end, so only the last row has length == position + 1. Cell( - why="warm-up capture: 8191 rows of one request from 0", - spec=_spec(("serving", 1, (8191,))), + why="verify MTP7: 32 requests x 8 = 256 rows, whole-row programs", + spec=_spec(("verify", _VERIFY_STARTS, 8)), ), + # ---- context rows: sequence_lengths > query_position + 1 inside chunks; chunk + # 0 starts at position 0, so its rows also carry every tail width 0..3. Cell( why="warm-up capture: 2 x 4096 rows, the [8192, 512] buffer's height", spec=_spec(("serving", 2, (4096, 4096))), ), - Cell( - why="tail widths 0..3 across 4096, and positions 0, 3, 7", - spec=_spec(("rows", (4095, 4096, 4097, 4098, 4099, 7, 3, 0))), - ), - # ---- geometry is a free axis: neither cr nor block_topk need be 2^k --- - Cell(why="geometry cr 2, block_topk 8", spec=_geometry(2, 16)), - Cell(why="geometry cr 2, block_topk 5", spec=_geometry(2, 10)), - Cell(why="geometry cr 2, block_topk 1024", spec=_geometry(2, 2048)), - Cell(why="geometry cr 3, block_topk 7", spec=_geometry(3, 21)), - Cell(why="geometry cr 3, block_topk 8", spec=_geometry(3, 24)), - Cell(why="geometry cr 4, block_topk 3", spec=_geometry(4, 12)), - Cell(why="geometry cr 4, block_topk 511", spec=_geometry(4, 2044)), - Cell(why="geometry cr 5, block_topk 8", spec=_geometry(5, 40)), - Cell(why="geometry cr 8, block_topk 8", spec=_geometry(8, 64)), - Cell(why="geometry cr 16, block_topk 128", spec=_geometry(16, 2048)), - # ---- metadata dtypes and layouts ------------------------------------ - Cell( - why="int32 positions, int32 lengths", - spec=_spec(("rows", (4096, 97)), position_dtype=torch.int32), - ), - Cell( - why="int64 positions, int64 lengths", - spec=_spec(("rows", (4096, 97)), length_dtype=torch.int64), - ), - Cell( - why="int32 positions, int64 lengths", - spec=_spec(("rows", (4096, 97)), position_dtype=torch.int32, length_dtype=torch.int64), - ), - Cell( - why="a row slice of the engine's taller qsa_topk_indices buffer", - spec=_spec(("sweep", 4), layout="row_slice"), - ), - Cell( - why="a column slice of a wider buffer: stride(0) != block_topk", - spec=_spec(("sweep", 4), layout="column_slice"), - ), - Cell(why="zero rows: a shaped empty, no launch", spec=_spec(("sweep", 0))), ) note = """ @@ -140,8 +93,8 @@ class QsaExpandBlockIndices(OpWrapper): request's whole KV length broadcast over its rows, so inside a context chunk it exceeds `position + 1` on every row but the last). - Preconditions that are values, so not `is_valid` guards (both measured to - fail silently): + Preconditions that are values, so not `is_valid` guards (both fail + silently; no test pins them): - `block_indices` is prefix-packed: every group id before every `-1`. The tail's column comes from the *count* of valid groups, so `[-1, 3]` at position 15, length 16, cr 4, topk 8 returns all `-1` -- tokens @@ -152,16 +105,16 @@ class QsaExpandBlockIndices(OpWrapper): inside `sequence_lengths` is expanded verbatim. Duplicate group ids expand to duplicate tokens. - Every metadata violation is rejected loudly by the callable (rank, width - vs geometry, int32-only `block_indices`, int32/int64 metadata, one entry - per row, `stride(1) == 1`, contiguous metadata, CUDA), so `is_valid` adds - nothing. A one-element strided metadata view reports contiguous and is - accepted -- benign at one row. + The callable rejects metadata violations loudly (rank, width vs geometry, + int32-only `block_indices`, int32/int64 metadata, one entry per row, + `stride(1) == 1`, contiguous metadata, CUDA), so `is_valid` adds nothing; + 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; 127 and 128 rows agree on their shared rows. No autotune. - CUDA-graph capturable; a replay recomputes from the live buffers, and the - captured row count fixes the launch shape. + 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 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_mtp_reuse_topk.py new file mode 100644 index 000000000000..582e514682ea --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_mtp_reuse_topk.py @@ -0,0 +1,246 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA MTP IndexShare: draft-decode selections rebuilt from the row captured at draft-extend.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# Tier-2 sourcing: a vendored Triton callable reached as ordinary Python. QSA's +# kernel plane has no torch.ops surface. +from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels + +#: This checkpoint's selection width: `indexer_budget` 2048 tokens over +#: `indexer_compress_ratio` 4. +_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. +_COMBINED_SLOTS = (13, 12) + + +def _spec(contexts, **overrides) -> dict: + """A cell spec at this checkpoint's geometry and the engine's metadata dtypes. + + `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 + ("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. + """ + base = dict( + contexts=contexts, + drafted=(1, 2, 3), + rows_per_request=1, + shared_rows=64, + slots=_COMBINED_SLOTS, + block_topk=_BLOCK_TOPK, + metadata_dtype=torch.int32, + graph_batches=(), + ) + return {**base, **overrides} + + +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. + + Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels. + triton_qsa_mtp_reuse_topk`: one Triton launch, one program per output row. + For row `r`, with `q = request_indices[r]` and `K = output.shape[1]` + (`block_topk`): + + 1. start from the captured row `shared_indices[q]`: valid group ids as a + prefix, then `-1`; + 2. `captured = shared_visible_blocks[q]`, `current = visible_blocks[r]`, + `fresh = clamp(current - captured, 0, K)`: the groups + `captured .. captured + fresh - 1` completed after the capture; + 3. `filled` = the count of non-negative ids in the captured row; the + fresh ids go, ascending, into slots `start .. start + fresh - 1` with + `start = min(filled, K - fresh)` -- right after the prefix when there + is room, otherwise after the captured row's newest `K - fresh` ids: + the oldest `filled + fresh - K` ids are displaced and the survivors + shift to the front; + 4. every other slot is the captured row's survivors in their order, `-1` + padding included. + + Every slot of `output[:rows]` is written and nothing else; returns `output`. + The inputs are read-only. Nothing is sorted, compacted or de-duplicated and + the ranking is not recomputed: the captured row is taken as ascending (the + target makes it canonical), so the displaced ids are the oldest groups, not the + lowest-scored ones, and the result is ascending again. + + Fusion boundary: the copy-and-append only. Capturing each request's row + 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. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + # ---- decode batch: one row per request; the first 16 requests sit at the + # boundaries (empty, short, 511-id, 512-id and full rows), fresh 0 and 1. + Cell( + why="decode MTP: 32 requests (envelope), 256-row shared buffer", + spec=_spec(("mixed", 32), shared_rows=256), + ), + # ---- the riskiest edge: two fresh groups displace a full row's two oldest + Cell( + why="8 requests at fresh 2: full, 511-id and short rows gain two groups", + spec=_spec( + ("explicit", (2050, 402, 2046, 8186, 6, 2047, 1022, 2049)), + drafted=(5, 6, 5, 6, 6, 4, 5, 6), + shared_rows=8, + ), + ), + # ---- CUDA graph ----------------------------------------------------- + Cell( + 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)), + ), + ) + + note = """ + Metadata consumed: none; every input is an argument, all int32 at the call + site: + shared_indices `qsa_shared_topk_indices[slot]`, [max_num_sequences, block_topk] + shared_visible_blocks `qsa_shared_topk_visible_blocks[slot]`, [max_num_sequences] + request_indices `qsa_req_idx_per_token[:rows]` + visible_blocks `qsa_visible_blocks[:rows]` + output `qsa_topk_indices[:rows]`, the [max_num_tokens, block_topk] buffer + `slot` is the layer's row in the cache manager's `qsa_shared_topk_slots`. + + It runs in the MTP draft loop's decode steps after the first: no context + rows, one row per request, so rows = requests, and the request map is the + row index. Draft step `i` sits `i` positions past the captured row, so + `max_draft_len` 1 never reaches this op, 3 reaches 1 and 2 positions, and 7 + (the target's ceiling) reaches 1..6. With compress_ratio 4, `fresh = + ((p + 1) % 4 + i) // 4` for a capture at position `p`: 0 or 1 up to i = 3 + and 0, 1 or 2 from i = 4. In a full row + the fresh ids displace the row's oldest ids (one at fresh 1, two at fresh + 2), and a 511-id row at fresh 2 loses its oldest id. + + Preconditions that are values, so not `is_valid` guards: + - each captured row is prefix-packed (every valid id before every -1), as + `qsa_topk_select` emits it. The append slot comes from the COUNT of + non-negative ids, so a -1 hole puts a fresh id over a valid one and + loses a group silently; + - captured ids are below that row's `shared_visible_blocks`, so a fresh id + never duplicates a captured one; + - every `request_indices[r]` is a row of both shared tensors; + - `shared_visible_blocks[q]` is the visible count of the very row that was + captured, and `visible_blocks[r]` is `(position + 1) // cr` of row `r`. + + Loud rejections (not pinned by a test): `ValueError` when the shared width differs from the output + width, when `request_indices` or `visible_blocks` does not hold one element + per output row, or on CPU tensors; a 1-D `output` or `shared_indices`; a + width that is not a power of two (Triton compilation). + + Accepted and silently wrong, hence `is_valid`: any non-contiguous operand + (the kernel addresses raw pointers densely and ignores strides); a 3-D + `shared_indices` whose middle extent equals `block_topk` (read as slot 0); + 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. + """ + + def __call__( + self, + *, + shared_indices: torch.Tensor, + shared_visible_blocks: torch.Tensor, + request_indices: torch.Tensor, + visible_blocks: torch.Tensor, + output: torch.Tensor, + ) -> torch.Tensor: + return _qsa_kernels.triton_qsa_mtp_reuse_topk( + shared_indices=shared_indices, + shared_visible_blocks=shared_visible_blocks, + request_indices=request_indices, + visible_blocks=visible_blocks, + output=output, + displace_oldest=True, + ) + + def reference( + self, + *, + shared_indices: torch.Tensor, + shared_visible_blocks: torch.Tensor, + request_indices: torch.Tensor, + visible_blocks: torch.Tensor, + output: torch.Tensor, + ) -> torch.Tensor: + """The rebuilt rows as the contract states them, in int64 torch. + + `output` supplies only the shape and dtype; its contents are ignored. + """ + rows, block_topk = output.shape + device = output.device + requests = request_indices.long().reshape(-1) + captured = shared_visible_blocks.long()[requests] + current = visible_blocks.long().reshape(-1) + rebuilt = shared_indices.long()[requests].clone() + + filled = (rebuilt >= 0).sum(dim=1) + fresh = (current - captured).clamp(min=0, max=block_topk) + start = torch.minimum(filled, block_topk - fresh) + slot = torch.arange(block_topk, device=device)[None, :] + # The oldest `filled - start` ids go; the survivors move to the front. + source = slot + (filled - start)[:, None] + rebuilt = torch.where( + source < block_topk, + rebuilt.gather(1, source.clamp(max=block_topk - 1)), + torch.full_like(rebuilt, -1), + ) + rank = slot - start[:, None] + appended = (rank >= 0) & (rank < fresh[:, None]) + rebuilt[appended] = (captured[:, None] + rank)[appended] + return rebuilt.to(output.dtype) + + def is_valid( + self, + *, + shared_indices: torch.Tensor, + shared_visible_blocks: torch.Tensor, + request_indices: torch.Tensor, + visible_blocks: torch.Tensor, + output: torch.Tensor, + ) -> None: + # The kernel takes raw pointers and addresses them densely; each case + # below returns a plausible wrong answer rather than an error. + if shared_indices.dim() != 2 or not shared_indices.is_contiguous(): + raise ValueError( + "shared_indices must be one slot's contiguous [num_seqs, block_topk] " + f"rows; got shape {tuple(shared_indices.shape)}, " + f"stride {shared_indices.stride()}" + ) + if shared_visible_blocks.dim() != 1 or not shared_visible_blocks.is_contiguous(): + raise ValueError( + "shared_visible_blocks must be one slot's contiguous [num_seqs] counts; " + f"got shape {tuple(shared_visible_blocks.shape)}, " + f"stride {shared_visible_blocks.stride()}" + ) + for name, tensor in ( + ("request_indices", request_indices), + ("visible_blocks", visible_blocks), + ("output", output), + ): + if not tensor.is_contiguous(): + raise ValueError( + f"{name} must be contiguous; got shape {tuple(tensor.shape)}, " + f"stride {tensor.stride()}" + ) + + +qsa_mtp_reuse_topk = QsaMtpReuseTopk() 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 index d0dbbfa1c66a..fff162545bc7 100644 --- 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 @@ -35,40 +35,15 @@ def _spec(**overrides) -> dict: return {**_CHECKPOINT, **overrides} -def _small(rows: int, requests: int, **overrides) -> dict: - """The narrow-table cells: random positions over an 8-page table, rows - grouped by request, each request's table row fully mapped, a contiguous q. - Eight table rows, as in the serving dump, whatever the request count.""" - base = dict(layout="random", rows=rows, requests=requests, slots=8, pages=8) - base.update(q_layout="packed", positions_dtype=torch.int32) - base.update(only_visible_blocks=False, context_rows=False) - base["branch"] = "row_wise" - base.update(overrides) - return _spec(**base) - - #: 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 -_FREE_HEADS = tuple( - Cell( - why=f"free axis, 64 rows: {h} index heads x dim {d}", - spec=_small(64, 3, heads=h, head_dim=d), - ) - for d in (64, 128) - for h in (1, 2, 3, 4, 8, 16) - if (h, d) != (4, 128) -) -_FREE_PAGING = tuple( - Cell( - why=f"free axis, 48 rows: compress {c} x {t} tokens per block", - spec=_small(48, 3, compress_ratio=c, tokens_per_block=t), - ) - for t in (16, 32, 64) - for c in (2, 4, 8, 16) - if (c, t) != (4, 32) -) + +#: 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): @@ -92,10 +67,10 @@ class QsaPagedIndexScores(OpWrapper): `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 16 into more than one row (1, 2, 4, 8) it selects a - query-tiled kernel sharing one key gather across rows; otherwise the - row-wise kernel runs. Each cell names the kernel it was measured on. The - two agree to one fp32 ulp, not bitwise. + 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 @@ -106,22 +81,8 @@ class QsaPagedIndexScores(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - # ---- the target's call shapes: [*, 260] table, 2080 columns ------- - Cell( - why="decode: 1 row deep in history, [32, 260] table, visible-only", - spec=_spec( - layout="decode", - gen_positions=(6000,), - rows_per_request=1, - slots=32, - pages=_WIDTH, - only_visible_blocks=True, - context_rows=False, - branch="row_wise", - ), - ), Cell( - why="decode MTP: 4 requests x 2 rows, visible-only", + why="NVFP4 trunk plane, page stride 49344: decode MTP 4 requests x 2 rows", spec=_spec( layout="decode", gen_positions=(7000, 5000, 4200, 2100), @@ -131,27 +92,15 @@ class QsaPagedIndexScores(OpWrapper): only_visible_blocks=True, context_rows=False, branch="row_wise", + **_TRUNK_POOL_NVFP4, ), ), Cell( - why="decode MTP: 4 requests x 2 rows, -inf padded mode", - spec=_spec( - layout="decode", - gen_positions=(7000, 5000, 4200, 2100), - rows_per_request=2, - slots=32, - pages=_WIDTH, - only_visible_blocks=False, - context_rows=False, - branch="row_wise", - ), - ), - Cell( - why="decode MTP: every one of 32 slots x 2 rows = 64 rows", + 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=2, + rows_per_request=8, slots=32, pages=_WIDTH, only_visible_blocks=True, @@ -159,54 +108,6 @@ class QsaPagedIndexScores(OpWrapper): branch="row_wise", ), ), - Cell( - why="capture: 8191 rows, one request from 0, visible-only", - spec=_spec( - layout="capture", - lengths=(8191,), - slots=8, - pages=_WIDTH, - only_visible_blocks=True, - context_rows=True, - branch="query_tile", - ), - ), - Cell( - why="capture: 8191 rows, one request from 0, -inf padded", - spec=_spec( - layout="capture", - lengths=(8191,), - slots=8, - pages=_WIDTH, - only_visible_blocks=False, - context_rows=True, - branch="query_tile", - ), - ), - Cell( - why="capture: 8192 rows = 2 x 4096, visible-only", - spec=_spec( - layout="capture", - lengths=(4096, 4096), - slots=8, - pages=_WIDTH, - only_visible_blocks=True, - context_rows=True, - branch="query_tile", - ), - ), - Cell( - why="capture: 8192 rows = 2 x 4096, -inf padded", - spec=_spec( - layout="capture", - lengths=(4096, 4096), - slots=8, - pages=_WIDTH, - only_visible_blocks=False, - context_rows=True, - branch="query_tile", - ), - ), Cell( why="capture: 8191 rows over 8 ragged requests, visible-only", spec=_spec( @@ -219,155 +120,6 @@ class QsaPagedIndexScores(OpWrapper): branch="query_tile", ), ), - Cell( - why="capture: 8191 rows over 8 ragged requests, -inf padded", - spec=_spec( - layout="capture", - lengths=(1024, 512, 2048, 256, 1500, 700, 1151, 1000), - slots=8, - pages=_WIDTH, - only_visible_blocks=False, - context_rows=True, - branch="query_tile", - ), - ), - Cell( - why="far end: 8191 rows at position 8319, all 2080 columns visible", - spec=_spec( - layout="far", - rows=8191, - slots=8, - pages=_WIDTH, - only_visible_blocks=False, - context_rows=True, - branch="query_tile", - ), - ), - Cell( - why="mixed step: 195-token chunk + 31 MTP decodes = 257 rows", - spec=_spec( - layout="mixed", - requests=32, - context_lens=(195,), - slots=32, - pages=_WIDTH, - only_visible_blocks=True, - context_rows=True, - branch="query_tile", - ), - ), - Cell( - why="mixed step: 4 chunks + 28 MTP decodes = 1334 rows", - spec=_spec( - layout="mixed", - requests=32, - context_lens=(900, 200, 150, 28), - slots=32, - pages=_WIDTH, - only_visible_blocks=True, - context_rows=True, - branch="query_tile", - ), - ), - Cell( - why="mixed step: 4 chunks + 28 MTP decodes = 8190 rows", - spec=_spec( - layout="mixed", - requests=32, - context_lens=(8000, 100, 30, 4), - slots=32, - pages=_WIDTH, - only_visible_blocks=True, - context_rows=True, - branch="query_tile", - ), - ), - Cell( - why="mixed step at 162 rows: context_rows=True stays row-wise", - spec=_spec( - layout="mixed", - requests=32, - context_lens=(100,), - slots=32, - pages=_WIDTH, - only_visible_blocks=True, - context_rows=True, - branch="row_wise", - ), - ), - # ---- the branch rule's edges (8-page table) ------------------------ - Cell( - why="1 row with context_rows=True: still row-wise", spec=_small(1, 4, context_rows=True) - ), - Cell(why="odd 17 rows, generation-sized grid", spec=_small(17, 4)), - Cell(why="256 rows, context_rows=False: last decode-sized", spec=_small(256, 4)), - Cell( - why="256 rows, context_rows=True: at the threshold, row-wise", - spec=_small(256, 4, context_rows=True), - ), - Cell( - why="257 rows, context_rows=True: one past it, query-tiled", - spec=_small(257, 6, context_rows=True, branch="query_tile"), - ), - Cell(why="257 rows, context_rows=False: prefill-sized row-wise", spec=_small(257, 6)), - Cell( - why="1024 rows, context_rows=True: query-tiled", - spec=_small(1024, 6, context_rows=True, branch="query_tile"), - ), - Cell(why="1024 rows, context_rows=False: row-wise", spec=_small(1024, 6)), - Cell( - why="300 rows x 1 head: 16 rows per query tile", - spec=_small(300, 5, heads=1, context_rows=True, branch="query_tile"), - ), - Cell( - why="300 rows x 2 heads: 8 rows per query tile", - spec=_small(300, 5, heads=2, context_rows=True, branch="query_tile"), - ), - Cell( - why="300 rows x 8 heads: 2 rows per query tile", - spec=_small(300, 5, heads=8, context_rows=True, branch="query_tile"), - ), - Cell( - why="300 rows x 16 heads: one row per tile, row-wise", - spec=_small(300, 5, heads=16, context_rows=True), - ), - Cell( - why="300 rows x 3 heads: does not divide 16, row-wise", - spec=_small(300, 5, heads=3, context_rows=True), - ), - Cell( - why="300 rows x 5 heads: does not divide 16, row-wise", - spec=_small(300, 5, heads=5, context_rows=True), - ), - Cell( - why="query tile over rows shuffled across requests", - spec=_small(320, 6, layout="shuffled", context_rows=True, branch="query_tile"), - ), - # ---- axes measured free -------------------------------------------- - Cell(why="fp16 q and index cache", spec=_small(64, 3, dtype=torch.float16)), - Cell( - why="int64 positions and int64 request indices", - spec=_small(48, 3, positions_dtype=torch.int64, request_dtype=torch.int64), - ), - Cell( - why="q as a head-dim slice of a [40, 4, 256] buffer", - spec=_small(40, 3, q_layout="dim_slice"), - ), - Cell( - why="1-page table: the narrowest legal score row", - spec=_small(8, 2, pages=1, tokens_per_block=16), - ), - Cell(why="2-page table", spec=_small(8, 2, pages=2, tokens_per_block=16)), - Cell( - why="positions past the table: every column visible", - spec=_small(32, 2, layout="past_table", pages=4), - ), - Cell( - why="positions 0..7: rows before a complete group see nothing", - spec=_small(8, 2, layout="arange"), - ), - *_FREE_HEADS, - *_FREE_PAGING, ) note = """ @@ -384,27 +136,40 @@ class QsaPagedIndexScores(OpWrapper): 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 (10 `ValueError`s, pinned by the - test), so `is_valid` adds nothing. + 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: 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 + `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, outside - the test's band. No cell covers fp32; the side cache is bf16. - - CUDA graphs: captured at 8191 rows (visible-only), replayed, then replayed - again after `q`, the cache and the table were overwritten in place; both - matched the reference. - - Not certified: more than 8192 rows, tables wider than 260 pages, more than - 32 table rows, and any head geometry other than 4 x 128 at the 2080-column - width. Hard limits: not probed; the largest certified extent is 8192 x - 2080 output elements. + 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__( @@ -418,6 +183,8 @@ def __call__( 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, @@ -429,6 +196,8 @@ def __call__( compress_ratio=compress_ratio, only_visible_blocks=only_visible_blocks, context_rows=context_rows, + use_pdl=use_pdl, + prefetch=prefetch, ) def reference( @@ -442,6 +211,8 @@ def reference( 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 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py index 341287f03b95..0100f49724d9 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_store.py @@ -7,9 +7,114 @@ 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 +_INTEGER_DTYPES = (torch.int32, torch.int64) + + +# The row count and the pool/table extents change from step to step; letting +# Triton specialize on them (== 1, divisible by 16) only multiplies compiles. +@triton.jit(do_not_specialize=["num_rows", "num_requests", "table_width", "num_pages"]) +def _qsa_paged_index_store_kernel( + token_k, + coords, + index_cache, + position_cache, + request_indices, + logical_positions, + block_table, + num_rows, + num_requests, + table_width, + num_pages, + stride_k_row, + stride_k_head, + stride_k_dim, + stride_c_row, + stride_c_axis, + stride_i_page, + stride_i_slot, + stride_i_head, + stride_i_dim, + stride_p_page, + stride_p_slot, + stride_p_axis, + stride_request, + stride_position, + stride_t_request, + stride_t_column, + HKV: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + AXES: tl.constexpr, + BLOCK_A: tl.constexpr, + TOKENS_PER_BLOCK: tl.constexpr, + BLOCK_ROWS: tl.constexpr, +): + # One program stores BLOCK_ROWS rows: the page lookup, then every head of + # the key and the coordinate row, all through the operands' own strides. + rows = (tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)).to(tl.int64) + live = rows < num_rows + request = tl.load(request_indices + rows * stride_request, mask=live, other=-1).to(tl.int64) + position = tl.load(logical_positions + rows * stride_position, mask=live, other=-1).to(tl.int64) + column = position // TOKENS_PER_BLOCK + slot = position % TOKENS_PER_BLOCK + live &= (request >= 0) & (request < num_requests) & (position >= 0) & (column < table_width) + page = tl.load( + block_table + request * stride_t_request + column * stride_t_column, mask=live, other=-1 + ).to(tl.int64) + live &= (page >= 0) & (page < num_pages) + + dims = tl.arange(0, BLOCK_D) + key_mask = live[:, None] & (dims < HEAD_DIM)[None, :] + for head in tl.static_range(HKV): + key = tl.load( + token_k + + rows[:, None] * stride_k_row + + head * stride_k_head + + dims[None, :] * stride_k_dim, + mask=key_mask, + other=0.0, + ) + tl.store( + index_cache + + page[:, None] * stride_i_page + + slot[:, None] * stride_i_slot + + head * stride_i_head + + dims[None, :] * stride_i_dim, + key.to(index_cache.dtype.element_ty), + mask=key_mask, + ) + + axes = tl.arange(0, BLOCK_A) + axis_mask = live[:, None] & (axes < AXES)[None, :] + coordinate = tl.load( + coords + rows[:, None] * stride_c_row + axes[None, :] * stride_c_axis, + mask=axis_mask, + other=0, + ) + tl.store( + position_cache + + page[:, None] * stride_p_page + + slot[:, None] * stride_p_slot + + axes[None, :] * stride_p_axis, + coordinate.to(position_cache.dtype.element_ty), + mask=axis_mask, + ) + + +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.""" + if rows <= 256: + return 1, 1 + return 16, 4 + + #: This checkpoint's index geometry and the engine's QSA metadata. `token_k` is #: the k head of the packed `[T, 4 + 1, 128]` index projection; the coordinates #: are the `[T, 3]` transpose of the engine's `[3, max_num_tokens]` position @@ -26,13 +131,19 @@ request_dtype=torch.int32, position_dtype=torch.int64, key_layout="packed", - coords_layout="transposed", - meta_layout="contiguous", cache_layout="plain", - table_layout="dense", + hole=None, ) -_DECODE = ((0, 1), (1, 1), (31, 1), (32, 1), (33, 1), (63, 1), (64, 1), (95, 1)) +#: The target's serving envelope: ISL 8192 + OSL 1024 puts decode rows up to position +#: 9215 (page column 287), 128 concurrent requests, and `max_num_tokens` 16384 +#: rows in one prefill or mixed step. +_WIDE = dict(slots=160, width=300) +_PREFILL_16K = ((0, 8192), (0, 8192)) + +#: 32 requests of one MTP7 verify block each (8 rows), at starts that put some blocks across a +#: page edge and some inside one. +_MTP7_32 = tuple((29 + 61 * i, 8) for i in range(32)) def _spec(segments, pages, **overrides) -> dict: @@ -40,11 +151,6 @@ def _spec(segments, pages, **overrides) -> dict: return {**_CKPT, "segments": segments, "pages": pages, **overrides} -def _bits(tensor: torch.Tensor) -> torch.Tensor: - as_int = {1: torch.uint8, 2: torch.int16, 4: torch.int32}[tensor.element_size()] - return tensor.view(as_int) - - class QsaPagedIndexStore(OpWrapper): """Write each row's raw index key and position coordinates into its paged slot. @@ -54,16 +160,21 @@ class QsaPagedIndexStore(OpWrapper): index_cache[page, slot] = token_k[t] # cast to the cache dtype position_cache[page, slot, :A] = position_coordinates[t] # A = coordinates' width - The key is stored raw -- un-normed, un-rotated: `qsa_prefill_compress` - pools raw keys and applies the Gemma norm and the rotation itself. Returns - `(index_cache, position_cache)`, both written in place. + A row that resolves to no slot -- negative request, request past the + table, negative position, page column past the table, or a table entry + that is negative (`-1`, unallocated) or past the pool -- is skipped + without error. The key is stored raw -- un-normed, un-rotated: + `qsa_prefill_compress` pools raw keys and applies the Gemma norm and the + rotation itself. Returns `(index_cache, position_cache)`, both written in + place. - Sourcing tier 3, composed from native torch: a block-table gather, a - floor-divide and remainder, two indexed writes (8 launches). 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. + 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). Fusion boundary: page lookup, slot arithmetic, both scatters. The index projection, the query's norm and rotation, group compression, scoring, @@ -74,62 +185,29 @@ class QsaPagedIndexStore(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell( - why="decode: 8 rows at aligned, mid and page-crossing positions", - spec=_spec(_DECODE, pages=32), - ), - Cell( - why="decode MTP: 4 requests x 2 rows, a pair straddling a page", - spec=_spec(((10, 2), (31, 2), (63, 2), (5, 2)), pages=16), - ), - Cell( - why="prefill: 1024 rows, 4 requests x 256 walking 8 pages each", - spec=_spec(((0, 256),) * 4, pages=32), - ), - Cell( - why="mixed: two context chunks and two decode rows, 66 rows", - spec=_spec(((0, 40), (0, 24), (200, 1), (7, 1)), pages=32), - ), - Cell( - why="contiguous token_k and coordinates", - spec=_spec(_DECODE, pages=32, key_layout="contiguous", coords_layout="contiguous"), - ), - Cell( - why="text batch: stride-0 expand of one position column", - spec=_spec(((0, 1), (5, 1), (32, 1), (70, 1)), pages=16, coords_layout="stride0"), - ), - Cell( - why="a 5-axis position plane: axes past the 3rd untouched", - spec=_spec(((2, 1), (33, 1)), pages=8, axes=5), - ), - Cell( - why="metadata as columns of wider buffers (strided)", - spec=_spec(((0, 8),), pages=8, meta_layout="strided"), - ), - Cell( - why="index_cache as a role view of a wider allocation", - spec=_spec(((0, 1), (31, 1), (32, 1), (33, 1)), pages=8, cache_layout="role_view"), - ), - Cell( - why="block table as a column slice of a wider one", - spec=_spec(((0, 1), (1, 1), (32, 1), (33, 1)), pages=8, table_layout="column_slice"), - ), - Cell( - why="fp16 keys: one conversion into the bf16 plane", - spec=_spec(((0, 4),), pages=8, key_dtype=torch.float16), - ), - Cell( - why="fp32 keys: one conversion into the bf16 plane", - spec=_spec(((0, 4),), pages=8, key_dtype=torch.float32), + why="decode MTP7: 32 requests x 8 rows (256), the last one-row-per-program " + "launch, blocks straddling pages", + spec=_spec(_MTP7_32, pages=1200, slots=40), ), Cell( - why="int32 requests and int32 positions", - spec=_spec(_DECODE, pages=32, position_dtype=torch.int32), + why="prefill 16384 rows (max_num_tokens), 16 rows per program, token_k off the " + "fused projection row", + spec=_spec(_PREFILL_16K, pages=600, key_layout="fused_row", **_WIDE), ), Cell( - why="int64 requests and int64 positions", - spec=_spec(_DECODE, pages=32, request_dtype=torch.int64), + why="NVFP4-KV side planes: index and position planes as strided views of one " + "per-slot buffer holding 12 index-K pages and the position page; one row " + "on an unallocated (-1) table entry is skipped", + spec=_spec( + ((0, 1), (31, 1), (32, 1), (33, 40)), + pages=8, + cache_layout="shared_page", + layers=12, + page_stride=49344, + position_strides=(24672, 3, 1), + hole="unallocated_page", + ), ), - Cell(why="zero rows: both planes unchanged", spec=_spec((), pages=8)), ) note = """ @@ -139,31 +217,57 @@ class QsaPagedIndexStore(OpWrapper): `get_index_k_buffer(i)` (`[pages, tpb, 1, 128]` bf16) and the shared `get_qsa_position_buffer()` (`[pages, tpb, 3]` int32, rewritten idempotently by every sparse layer). Every operand is read through its - own strides, unlike the Triton siblings. - - Values the caller owns -- data, not metadata, so not guarded; each - measured, and unlike `qsa_paged_kv_store` nothing is skipped: - - a `-1` (unallocated) table entry wraps to the pool's last page, and a - negative request selects the table's last row; - - a negative position wraps twice (last table column, last slot); - - a page, request, column or slot genuinely out of range is a CUDA - device-side assert that poisons the process -- so none is tested; - - duplicate destinations both store, unordered; - - coordinates wider than the plane's dtype wrap (int64 `2^31 + 5` - stores as -2147483643); floats truncate toward zero. - Loud rejections are torch's own: a 1-D table, a position plane narrower - than the coordinates, mixed devices, non-index metadata dtypes. - `rows == 0` leaves both planes unchanged. - - The 8 launches are launch-bound: ~0.15 ms per call flat from 1 to 1024 - rows. CUDA-graph capturable; a replay follows positions rewritten in - place. Repeating a call is idempotent. - - Hard limits: torch advanced indexing in int64 -- no int32 boundary inside - the call; the coordinate wrap above is the plane's dtype. - - Not certified: an fp8 index plane, `Hkv` above 1, multiple devices, - concurrent streams onto one pool. + own strides, all of them runtime arguments; the cells drive the target's: + `token_k` off the packed index projection (row stride 640) and off the fused + input projection (14080), the `[T, 3]` transpose of the position buffer as + the coordinates, contiguous metadata and table. Under an + NVFP4 KV pool the QSA cache manager packs each page's index-K roles and + then its position page into one per-slot buffer, so both planes are strided + views of one allocation: the trunk's index plane is `[pages, 32, 1, 128]` + at page stride 49344 (12 x 4096 + 192) with the position plane + `[pages, 32, 3]` at (24672, 3, 1); the cell checks that the other index-K + roles of each page survive the call. + + Preconditions: + - destination slots are unique across the batch: two rows naming one + (page, slot) both store, unordered. The engine's mapping cannot + collide. + Loud rejections (`ValueError`, before launching; the test pins the + unequal-length metadata one): metadata that is not + one-dimensional, of unequal length, or not int32/int64; `token_k` or the + coordinates with a row count other than the metadata's; an index plane + whose page size, head count or head width disagrees with + `tokens_per_block` and `token_k`; a position plane whose page size + disagrees or that is narrower than the coordinates; a table that is not a + 2-D integer tensor. `rows == 0` returns before launching. + + Values the caller owns -- data, not metadata (measured, not pinned by a + test): coordinates wider than the plane's dtype wrap (int64 `2^31 + 5` + stores as -2147483643); floats truncate toward zero. + + One launch per call; HEAD_DIM, the head count, the page size and the + coordinate width are constexprs, everything else (strides, pool and table + extents, rows) is a runtime argument and the row count and extents are + not specialized on, so a new batch or pool size does not recompile + (strides still specialize on 1 and 16-divisibility). Up to 256 rows run + one row per program (1 warp), above that 16 rows per program (4 warps). + CUDA-graph capturable (captured at 256 rows, 32 MTP7 verify blocks); a + replay follows positions rewritten in place. + + Hard limits: every index is widened to int64 inside the kernel before + any arithmetic; an int64 position at `2^32 + 5` lands past the table and + is skipped rather than truncated (read from the kernel, not pinned by a + test). Rows are a grid axis: 2^31 programs. + + Skip certified by a cell: an unallocated (`-1`) table entry, one row of the + NVFP4-plane cell. The other skip conditions (a negative request index among + 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. """ def __call__( @@ -177,14 +281,73 @@ def __call__( block_table: torch.Tensor, tokens_per_block: int, ) -> Tuple[torch.Tensor, torch.Tensor]: - positions = logical_positions.to(torch.long) - pages = block_table[request_indices.to(torch.long), positions // tokens_per_block].to( - torch.long - ) - within = positions % tokens_per_block - index_cache[pages, within] = token_k.to(index_cache.dtype) - position_cache[pages, within, : position_coordinates.shape[-1]] = position_coordinates.to( - position_cache.dtype + rows = request_indices.shape[0] + if request_indices.dim() != 1 or logical_positions.shape != request_indices.shape: + raise ValueError( + f"request_indices {tuple(request_indices.shape)} and logical_positions " + f"{tuple(logical_positions.shape)} must be one-dimensional and equally long" + ) + if request_indices.dtype not in _INTEGER_DTYPES or ( + logical_positions.dtype not in _INTEGER_DTYPES + or block_table.dtype not in _INTEGER_DTYPES + ): + raise ValueError("QSA index store metadata and block table must be int32 or int64") + if token_k.dim() != 3 or token_k.shape[0] != rows: + raise ValueError(f"token_k {tuple(token_k.shape)} must be [{rows}, heads, dim]") + if position_coordinates.dim() != 2 or position_coordinates.shape[0] != rows: + raise ValueError( + f"position_coordinates {tuple(position_coordinates.shape)} must be [{rows}, axes]" + ) + if index_cache.dim() != 4 or index_cache.shape[1:] != ( + tokens_per_block, + *token_k.shape[1:], + ): + raise ValueError( + f"index_cache {tuple(index_cache.shape)} must be [pages, {tokens_per_block}, " + f"{token_k.shape[1]}, {token_k.shape[2]}]" + ) + if ( + position_cache.dim() != 3 + or position_cache.shape[1] != tokens_per_block + or position_cache.shape[2] < position_coordinates.shape[1] + ): + raise ValueError( + f"position_cache {tuple(position_cache.shape)} must be [pages, " + f"{tokens_per_block}, >= {position_coordinates.shape[1]}]" + ) + if block_table.dim() != 2: + raise ValueError(f"block_table {tuple(block_table.shape)} must be two-dimensional") + if rows == 0: + return index_cache, position_cache + block_rows, warps = _launch_rows(rows) + axes = position_coordinates.shape[1] + _qsa_paged_index_store_kernel[(triton.cdiv(rows, block_rows),)]( + token_k, + position_coordinates, + index_cache, + position_cache, + request_indices, + logical_positions, + block_table, + rows, + block_table.shape[0], + block_table.shape[1], + index_cache.shape[0], + *token_k.stride(), + *position_coordinates.stride(), + *index_cache.stride(), + *position_cache.stride(), + request_indices.stride(0), + logical_positions.stride(0), + *block_table.stride(), + HKV=token_k.shape[1], + HEAD_DIM=token_k.shape[2], + BLOCK_D=triton.next_power_of_2(token_k.shape[2]), + AXES=axes, + BLOCK_A=triton.next_power_of_2(axes), + TOKENS_PER_BLOCK=tokens_per_block, + BLOCK_ROWS=block_rows, + num_warps=warps, ) return index_cache, position_cache @@ -199,70 +362,39 @@ def reference( block_table: torch.Tensor, tokens_per_block: int, ) -> Tuple[torch.Tensor, torch.Tensor]: - """Both planes after the store, as fresh tensors, built row by row. + """Both planes after the store, as fresh tensors; the inputs are not touched. - Pages and slots are computed in Python and each write is a basic - (integer) index -- none of the advanced indexing `__call__` is made of. + Every row's destination, and whether it has one, is decided in Python + integers from the table read back to the host -- none of the kernel's + arithmetic -- and the surviving rows are written with one indexed copy + per plane. """ expected_index, expected_position = index_cache.clone(), position_cache.clone() table = block_table.tolist() - axes = position_coordinates.shape[-1] + pool = index_cache.shape[0] + kept, pages, slots = [], [], [] rows = zip(request_indices.tolist(), logical_positions.tolist()) for row, (request, position) in enumerate(rows): - page = table[request][position // tokens_per_block] - slot = position % tokens_per_block - expected_index[page, slot] = token_k[row].to(index_cache.dtype) - expected_position[page, slot, :axes] = position_coordinates[row].to( - position_cache.dtype - ) - return expected_index, expected_position - - def is_valid( - self, - token_k: torch.Tensor, - position_coordinates: 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, - ) -> None: - # Each is a shape check over a violation measured to be accepted and - # silently wrong: torch broadcasts instead of raising. - rows = request_indices.shape[0] - assert request_indices.dim() == 1 and logical_positions.dim() == 1, ( - f"request_indices {tuple(request_indices.shape)} and logical_positions " - f"{tuple(logical_positions.shape)} must both be one-dimensional; a " - f"trailing axis broadcasts into a rows x rows index instead of raising" - ) - assert logical_positions.shape[0] == rows, ( - f"logical_positions has {logical_positions.shape[0]} entries against " - f"{rows} request_indices; a length-1 operand broadcasts silently and a " - f"length-1 logical_positions collapses every row onto one slot" - ) - assert token_k.shape[0] == rows, ( - f"token_k has {token_k.shape[0]} rows against {rows} metadata rows; a " - f"single row broadcasts into every destination slot instead of raising" - ) - assert position_coordinates.shape[0] == rows, ( - f"position_coordinates has {position_coordinates.shape[0]} rows against " - f"{rows} metadata rows; a single row broadcasts into every slot" - ) - assert index_cache.shape[1] == tokens_per_block, ( - f"tokens_per_block={tokens_per_block} does not match index_cache's " - f"token axis {index_cache.shape[1]}; a wrong page size writes the " - f"wrong page and the wrong slot" + if not 0 <= request < len(table) or position < 0: + continue + column, slot = divmod(position, tokens_per_block) + if column >= len(table[request]) or not 0 <= table[request][column] < pool: + continue + kept.append(row) + pages.append(table[request][column]) + slots.append(slot) + if not kept: + return expected_index, expected_position + device = index_cache.device + kept_t, page_t, slot_t = ( + torch.tensor(v, dtype=torch.long, device=device) for v in (kept, pages, slots) ) - assert position_cache.shape[1] == tokens_per_block, ( - f"tokens_per_block={tokens_per_block} does not match position_cache's " - f"token axis {position_cache.shape[1]}" - ) - assert index_cache.shape[2] == token_k.shape[1], ( - f"index_cache carries {index_cache.shape[2]} index-KV heads against " - f"token_k's {token_k.shape[1]}; a one-head token_k broadcasts into " - f"every head of a wider cache" + axes = position_coordinates.shape[-1] + expected_index[page_t, slot_t] = token_k[kept_t].to(index_cache.dtype) + expected_position[page_t, slot_t, :axes] = position_coordinates[kept_t].to( + position_cache.dtype ) + return expected_index, expected_position qsa_paged_index_store = QsaPagedIndexStore() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py index 06796397a5f4..be0791c3ccd6 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_kv_store.py @@ -29,13 +29,20 @@ cache_dtype=torch.bfloat16, request_dtype=torch.int32, position_dtype=torch.int64, - k_layout="qkv_slice", - table_layout="dense", hole=None, + # Heads of further columns after v in each row: 27 is the target's fused input + # projection row [q | k | v | gate | index | pad] (55 heads of 256). + row_pad_heads=0, ) +_F8 = torch.float8_e4m3fn + _DECODE = ((0, 1), (31, 1), (32, 1), (33, 1), (100, 1), (255, 1), (7, 1), (64, 1)) +#: 32 requests of one MTP7 verify block each (8 rows), at starts that put some blocks across a +#: page edge and some inside one. +_MTP7_32 = tuple((29 + 61 * i, 8) for i in range(32)) + def _spec(segments, pages, **overrides) -> dict: """`segments` is `(first position, rows)` per request, in row order.""" @@ -61,10 +68,11 @@ class QsaPagedKvStore(OpWrapper): `logical_positions` is absolute, so prefill, decode and MTP rows are the same call. The store is a move, converted once to the pool's dtype (bitwise `Tensor.to`; an fp8-e4m3 pool saturates at 448 and takes no - scale). A row that resolves to no slot -- negative request, request past - the table, negative position, page column past the table, or a table - entry that is negative (`-1`, unallocated) or past the pool -- is skipped - without error. Returns `(k_cache, v_cache)`, both written in place. + scale, so a reader treats it as kv scale 1.0). A row that resolves to no + slot -- negative request, request past the table, negative position, page + column past the table, or a table entry that is negative (`-1`, + unallocated) or past the pool -- is skipped without error. Returns + `(k_cache, v_cache)`, both written in place. Fusion boundary: the scatter only. The K/V projection, norm and RoPE of K, the QSA index plane, selection and attention are the caller's; the @@ -74,89 +82,24 @@ class QsaPagedKvStore(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( + # The decode and prefill targets read k/v off the fused projection row; the MTP + # targets split them off a packed QKV row. Cell( - why="decode: 8 rows at aligned, mid and page-crossing positions", - spec=_spec(_DECODE, pages=64), - ), - Cell( - why="decode MTP: 4 requests x 2 rows, a pair straddling a page", - spec=_spec(((0, 2), (31, 2), (63, 2), (95, 2)), pages=32), - ), - Cell( - why="prefill: 1024 rows, 4 requests x 256 walking 8 pages each", - spec=_spec(((0, 256),) * 4, pages=160), - ), - Cell( - why="mixed: two context chunks and two decode rows, 66 rows", - spec=_spec(((0, 40), (0, 24), (200, 1), (7, 1)), pages=64), - ), - Cell( - why="fp8-e4m3 pool: the row rounded once", - spec=_spec(_DECODE, pages=64, cache_dtype=torch.float8_e4m3fn), - ), - Cell( - why="fp16 rows into a bf16 pool: one conversion", - spec=_spec(((0, 1), (1, 1), (2, 1), (3, 1)), pages=8, row_dtype=torch.float16), - ), - Cell( - why="head-major rows: k/v last-dim stride != 1", - spec=_spec(_DECODE, pages=64, k_layout="head_major"), - ), - Cell( - why="block table as a role slice of a wider buffer", - spec=_spec(_DECODE, pages=64, table_layout="role_slice"), - ), - Cell( - why="block table transposed: column stride != 1", - spec=_spec(_DECODE, pages=64, table_layout="transposed"), - ), - Cell( - why="int32 positions", - spec=_spec(_DECODE, pages=64, position_dtype=torch.int32), - ), - Cell( - why="int64 request indices", - spec=_spec(_DECODE, pages=64, request_dtype=torch.int64), - ), - # ---- head geometry and page size are free axes --------------------- - Cell(why="head_dim 96 (not 2^k), 2 heads", spec=_spec(_DECODE[:6], pages=24, head_dim=96)), - Cell( - why="head_dim 64, 1 head, 64-token pages", - spec=_spec(_DECODE[:6], pages=24, head_dim=64, hkv=1, tokens_per_block=64), - ), - Cell( - why="head_dim 128, 4 heads, 16-token pages", - spec=_spec(_DECODE[:6], pages=48, head_dim=128, hkv=4, tokens_per_block=16), - ), - Cell( - why="head_dim 192 (not 2^k), 2 heads", spec=_spec(_DECODE[:6], pages=24, head_dim=192) - ), - # ---- rows that resolve to no slot are skipped ------------------------ - Cell( - why="skip: table entry -1 (unallocated)", - spec=_spec(_DECODE, pages=64, hole="unallocated_page"), - ), - Cell( - why="skip: table entry past the pool", - spec=_spec(_DECODE, pages=64, hole="page_past_pool"), - ), - Cell( - why="skip: request index past the table", - spec=_spec(_DECODE, pages=64, hole="request_past_table"), - ), - Cell( - why="skip: negative request index", - spec=_spec(_DECODE, pages=64, hole="negative_request"), + why="decode: 8 rows off the fused projection row (row stride 14080), " + "at aligned, mid and page-crossing positions, plus one row on an " + "unallocated (-1) table entry that is skipped", + spec=_spec(_DECODE, pages=64, row_pad_heads=27, hole="unallocated_page"), ), Cell( - why="skip: negative logical position", - spec=_spec(_DECODE, pages=64, hole="negative_position"), + why="prefill 1024 rows off the fused projection row (row stride 14080)", + spec=_spec(((0, 256),) * 4, pages=160, row_pad_heads=27), ), + # MTP7 verify blocks: 8 rows per request, so one block spans up to two pages. + # The non-default pool the forward reaches: fp8-e4m3. Cell( - why="skip: page column past the table", - spec=_spec(_DECODE, pages=64, hole="column_past_table"), + why="fp8 decode MTP7: 32 requests x 8 rows (256), blocks straddling pages", + spec=_spec(_MTP7_32, pages=1200, slots=40, cache_dtype=_F8), ), - Cell(why="zero rows: returns before launching", spec=_spec((), pages=4)), ) note = """ @@ -173,25 +116,32 @@ class QsaPagedKvStore(OpWrapper): - `request_indices` and `logical_positions` are contiguous: the kernel gets no stride for them. A strided view is read through its base pointer and rows land at other rows' slots, silently (`is_valid`). - Every other operand is read through its strides (packed QKV head slices, - head-major rows, pool role views, sliced or transposed tables). The - callable rejects loudly: mismatched or non-3-D/4-D K/V shapes, head - geometry, `k_cache.shape[2] != tokens_per_block`, row dtypes outside - bf16/fp16, pool dtypes outside bf16/fp16/fp8-e4m3, a non-int32 or - non-2-D table, non-integer or wrong-length metadata, CPU operands, an - empty pool or table. `rows == 0` returns before launching. + Every other operand is read through its strides; the cells drive k/v as + head slices of the packed QKV row and of the fused projection row, the + pool's K and V role views, and a contiguous table. The callable rejects + loudly: mismatched or non-3-D/4-D K/V shapes, head geometry, + `k_cache.shape[2] != tokens_per_block`, row dtypes outside bf16/fp16, + pool dtypes outside bf16/fp16/fp8-e4m3, a non-int32 or non-2-D table, + non-integer or wrong-length metadata, CPU operands, an empty pool or + 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; a - replay follows positions rewritten in place. + 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. Hard limits: both metadata values are cast to int32 inside the kernel - (`.to(tl.int32)`); an int64 position at `2^32 + 5` was probed 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 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. Not certified: duplicate destinations, multiple devices, a scaled fp8 - pool (there is no scale argument). + 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). """ 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_nvfp4.py b/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 new file mode 100644 index 000000000000..73e6ec1bebf7 --- /dev/null +++ b/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 @@ -0,0 +1,533 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA sparse path's paged K/V append into an NVFP4 pool: quantize and scatter in one launch.""" + +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 + +#: Values per e4m3 block scale, and e2m1 values per stored byte. +_SF_VEC = 16 +_PER_BYTE = 2 +#: Tokens of a V scale group: V's scales are stored `[tpb / 4, D / 16, 4]` per head page. +_V_SF_QUAD = 4 +#: e2m1's largest magnitude: a block's scale maps its amax onto it. +_E2M1_MAX = 6.0 +#: What `cvt.rn.satfinite.e2m1x2.f32` makes of a NaN of either sign: +6. A NaN reaches +#: it as 0 x inf, from a zero in a block whose nonzero amax rounds its scale to 0. +_E2M1_NAN_CODE = 0x7 +#: The round-to-nearest-even decision points between neighbouring e2m1 magnitudes: +#: `(threshold, inclusive)` steps the code up when `|v| > t`, or `|v| >= t` where the +#: upper neighbour is the even code. +_E2M1_STEPS = ( + (0.25, False), + (0.75, True), + (1.25, False), + (1.75, True), + (2.5, False), + (3.5, True), + (5.0, False), +) + +#: This checkpoint's K/V geometry and the engine's QSA table. `heads_q` places k and v as +#: head slices of the packed `[T, 24 + 2 + 2, 256]` QKV row; `global_k` / `global_v` are +#: `kv_scale_orig_quant[1]` / `[2]`; `slot_pages` role pages of one lifecycle slot hold +#: the call's planes, layer `layer`'s K and V among them (the trunk pool holds 24). +_CKPT = dict( + hkv=2, + head_dim=256, + heads_q=24, + tokens_per_block=32, + slots=32, + width=260, + row_dtype=torch.bfloat16, + request_dtype=torch.int32, + position_dtype=torch.int64, + global_k=1.0, + global_v=1.0, + values="randn", + hole=None, + thop_append=False, + slot_pages=2, + layer=0, +) + +_DECODE = ((0, 1), (31, 1), (32, 1), (33, 1), (100, 1), (255, 1), (7, 1), (64, 1)) + + +def _spec(segments, pages, **overrides) -> dict: + """`segments` is `(first position, rows)` per request, in row order.""" + return {**_CKPT, "segments": segments, "pages": pages, **overrides} + + +@triton.jit +def _rcp_approx(x): + """`rcp.approx.ftz.f32`, the reciprocal the engine's FP4 quantizer uses.""" + return tl.inline_asm_elementwise( + "rcp.approx.ftz.f32 $0, $1;", "=f,f", [x], dtype=tl.float32, is_pure=True, pack=1 + ) + + +@triton.jit +def _e2m1x2(lo, hi): + """One byte of two e2m1 codes, `lo` in the low nibble (`cvt.rn.satfinite`).""" + return tl.inline_asm_elementwise( + "{ .reg .b8 t; .reg .b16 s; cvt.rn.satfinite.e2m1x2.f32 t, $2, $1; " + "mov.b16 s, {t, t}; cvt.u32.u16 $0, s; }", + "=r,f,f", + [lo, hi], + dtype=tl.int32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _quantize_row( + src, + stride_dim, + global_scale, + HEAD_DIM: tl.constexpr, +): + """One head row: (packed e2m1 bytes [D/2], e4m3 scale bytes [D/16]).""" + pairs = tl.arange(0, HEAD_DIM // 2) + lo = tl.load(src + (2 * pairs) * stride_dim).to(tl.float32) + hi = tl.load(src + (2 * pairs + 1) * stride_dim).to(tl.float32) + pair_max = tl.maximum(tl.abs(lo), tl.abs(hi)) + block_max = tl.max(tl.reshape(pair_max, [HEAD_DIM // 16, 8]), axis=1) + sf_value = global_scale * (block_max * _rcp_approx(tl.full([HEAD_DIM // 16], 6.0, tl.float32))) + sf_code = sf_value.to(tl.float8e4nv) + sf_rounded = sf_code.to(tl.float32) + out_scale = tl.where( + block_max != 0.0, + _rcp_approx(sf_rounded * _rcp_approx(tl.zeros_like(sf_rounded) + global_scale)), + 0.0, + ) + lo_scaled = tl.reshape( + tl.reshape(lo, [HEAD_DIM // 16, 8]) * out_scale[:, None], [HEAD_DIM // 2] + ) + hi_scaled = tl.reshape( + tl.reshape(hi, [HEAD_DIM // 16, 8]) * out_scale[:, None], [HEAD_DIM // 2] + ) + packed = (_e2m1x2(lo_scaled, hi_scaled) & 0xFF).to(tl.uint8) + return packed, sf_code.to(tl.uint8, bitcast=True) + + +@triton.jit +def _qsa_paged_kv_store_nvfp4_kernel( + k, + v, + k_cache, + v_cache, + k_scale_cache, + v_scale_cache, + request_indices, + logical_positions, + block_table, + kv_scale_orig_quant, + num_pages, + num_requests, + table_width, + k_stride_row, + k_stride_head, + k_stride_dim, + v_stride_row, + v_stride_head, + v_stride_dim, + kc_stride_page, + kc_stride_head, + kc_stride_token, + vc_stride_page, + vc_stride_head, + vc_stride_token, + ks_stride_page, + ks_stride_head, + vs_stride_page, + vs_stride_head, + bt_stride_request, + bt_stride_page, + HEAD_DIM: tl.constexpr, + TOKENS_PER_BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + head = tl.program_id(1).to(tl.int64) + request = tl.load(request_indices + row).to(tl.int64) + logical = tl.load(logical_positions + row).to(tl.int64) + column = logical // TOKENS_PER_BLOCK + token = logical % TOKENS_PER_BLOCK + lookup = (request >= 0) & (request < num_requests) & (logical >= 0) & (column < table_width) + page = tl.load( + block_table + request * bt_stride_request + column * bt_stride_page, mask=lookup, other=-1 + ).to(tl.int64) + valid = (page >= 0) & (page < num_pages) + page = tl.maximum(page, 0) + + k_data, k_sf = _quantize_row( + k + row * k_stride_row + head * k_stride_head, + k_stride_dim, + tl.load(kv_scale_orig_quant + 1), + HEAD_DIM, + ) + v_data, v_sf = _quantize_row( + v + row * v_stride_row + head * v_stride_head, + v_stride_dim, + tl.load(kv_scale_orig_quant + 2), + HEAD_DIM, + ) + pairs = tl.arange(0, HEAD_DIM // 2) + blocks = tl.arange(0, HEAD_DIM // 16) + tl.store( + k_cache + page * kc_stride_page + head * kc_stride_head + token * kc_stride_token + pairs, + k_data, + mask=valid, + ) + tl.store( + v_cache + page * vc_stride_page + head * vc_stride_head + token * vc_stride_token + pairs, + v_data, + mask=valid, + ) + # K scales token-major within the head page; V scales in four-token groups. + tl.store( + k_scale_cache + + page * ks_stride_page + + head * ks_stride_head + + token * (HEAD_DIM // 16) + + blocks, + k_sf, + mask=valid, + ) + tl.store( + v_scale_cache + + page * vs_stride_page + + head * vs_stride_head + + (token // 4) * (4 * (HEAD_DIM // 16)) + + blocks * 4 + + token % 4, + v_sf, + mask=valid, + ) + + +def _store_nvfp4( + *, + k: torch.Tensor, + v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_scale_cache: torch.Tensor, + v_scale_cache: torch.Tensor, + request_indices: torch.Tensor, + logical_positions: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + kv_scale_orig_quant: torch.Tensor, +) -> None: + """Quantize rows of K and V to NVFP4 and store them into their pages (see the entry).""" + if k.ndim != 3 or v.shape != k.shape: + raise ValueError("NVFP4 K/V store: k and v must be matching [rows, heads, dim]") + rows, heads, head_dim = k.shape + if k.dtype not in (torch.bfloat16, torch.float16) or v.dtype != k.dtype: + raise ValueError("NVFP4 K/V store: k and v must both be BF16 or both FP16") + if head_dim < 32 or head_dim & (head_dim - 1): + raise ValueError("NVFP4 K/V store: head_dim must be a power of two of at least 32") + if tokens_per_block <= 0 or tokens_per_block % _V_SF_QUAD: + raise ValueError("NVFP4 K/V store: tokens_per_block must be a positive multiple of 4") + data_shape = (k_cache.shape[0], heads, tokens_per_block, head_dim // _PER_BYTE) + scale_shape = (k_cache.shape[0], heads, tokens_per_block, head_dim // _SF_VEC) + for name, cache, shape in ( + ("k_cache", k_cache, data_shape), + ("v_cache", v_cache, data_shape), + ("k_scale_cache", k_scale_cache, scale_shape), + ("v_scale_cache", v_scale_cache, scale_shape), + ): + if cache.dtype != torch.uint8 or tuple(cache.shape) != shape: + raise ValueError( + f"NVFP4 K/V store: {name} must be uint8 {shape}, got {cache.dtype} " + f"{tuple(cache.shape)}" + ) + if cache.stride(3) != 1: + raise ValueError(f"NVFP4 K/V store: {name} must be contiguous along its last dim") + for name, cache in (("k_scale_cache", k_scale_cache), ("v_scale_cache", v_scale_cache)): + if cache.stride(2) != head_dim // _SF_VEC: + raise ValueError(f"NVFP4 K/V store: each head page of {name} must be contiguous") + if block_table.ndim != 2 or block_table.dtype != torch.int32: + raise ValueError("NVFP4 K/V store: block_table must be a 2-D int32 tensor") + for name, meta in ( + ("request_indices", request_indices), + ("logical_positions", logical_positions), + ): + if meta.shape != (rows,) or meta.dtype not in (torch.int32, torch.int64): + raise ValueError(f"NVFP4 K/V store: {name} must be one int32/int64 entry per row") + if not meta.is_contiguous(): + raise ValueError(f"NVFP4 K/V store: {name} must be contiguous") + if ( + kv_scale_orig_quant.dtype != torch.float32 + or kv_scale_orig_quant.shape != (3,) + or not kv_scale_orig_quant.is_contiguous() + ): + raise ValueError("NVFP4 K/V store: kv_scale_orig_quant must be a contiguous fp32 [3]") + operands = (k, v, k_cache, v_cache, k_scale_cache, v_scale_cache, request_indices) + operands += (logical_positions, block_table, kv_scale_orig_quant) + if not all(t.is_cuda and t.device == k.device for t in operands): + raise ValueError("NVFP4 K/V store: every operand must be on k's CUDA device") + if rows == 0: + return + _qsa_paged_kv_store_nvfp4_kernel[(rows, heads)]( + k, + v, + k_cache, + v_cache, + k_scale_cache, + v_scale_cache, + request_indices, + logical_positions, + block_table, + kv_scale_orig_quant, + k_cache.shape[0], + block_table.shape[0], + block_table.shape[1], + k.stride(0), + k.stride(1), + k.stride(2), + v.stride(0), + v.stride(1), + v.stride(2), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + k_scale_cache.stride(0), + k_scale_cache.stride(1), + v_scale_cache.stride(0), + v_scale_cache.stride(1), + block_table.stride(0), + block_table.stride(1), + HEAD_DIM=head_dim, + TOKENS_PER_BLOCK=tokens_per_block, + num_warps=4, + ) + + +def _e2m1_codes(values: torch.Tensor) -> torch.Tensor: + """fp32 -> e2m1 codes 0..15, round to nearest even, saturating at 6; a NaN + (0 x inf, in a block whose scale underflows to zero) converts to +6.""" + magnitude = values.abs() + code = torch.zeros(values.shape, dtype=torch.int32, device=values.device) + for threshold, inclusive in _E2M1_STEPS: + code += (magnitude >= threshold if inclusive else magnitude > threshold).to(torch.int32) + code |= torch.signbit(values).to(torch.int32) << 3 + return torch.where(torch.isnan(values), torch.full_like(code, _E2M1_NAN_CODE), code) + + +def _quantize_reference(rows: torch.Tensor, global_scale: torch.Tensor): + """[n, D] rows -> (packed bytes [n, D/2], e4m3 scale bytes [n, D/16]), in plain torch.""" + n, head_dim = rows.shape + x = rows.float().reshape(n, head_dim // _SF_VEC, _SF_VEC) + block_max = x.abs().amax(dim=-1) + sixth = torch.tensor(_E2M1_MAX, dtype=torch.float32, device=rows.device).reciprocal() + scale_code = (global_scale * (block_max * sixth)).to(torch.float8_e4m3fn) + rounded = scale_code.float() + out_scale = torch.where( + block_max != 0, + (rounded * global_scale.reciprocal()).reciprocal(), + torch.zeros_like(rounded), + ) + codes = _e2m1_codes(x * out_scale[..., None]).reshape(n, head_dim // _PER_BYTE, _PER_BYTE) + packed = (codes[..., 0] | (codes[..., 1] << 4)).to(torch.uint8) + return packed, scale_code.view(torch.uint8) + + +class QsaPagedKvStoreNvfp4(OpWrapper): + """Quantize every row's K and V to NVFP4 and store them into the request's page. + + A Triton kernel defined in this file, one launch over `(rows, Hkv)`. For row t + and KV head h, with `page = block_table[request_indices[t], logical_positions[t] + // tpb]` and `s = logical_positions[t] % tpb`, each 16-value block `b` of + `k[t, h]` (and alike for v, with `kv_scale_orig_quant[2]`) is stored as:: + + g = kv_scale_orig_quant[1] + sf = e4m3(g * (amax_b * rcp(6))) # one byte + code_i = e2m1(k[t, h, i] * rcp(float(sf) * rcp(g))) # 0 when amax_b == 0 + k_cache[page, h, s, i // 2] nibble i % 2 (low nibble = even i) + k_scale_cache[page, h, s, b] = sf # token-major + v_scale_cache[page, h, (s // 4) * 4 * (D / 16) + 4 * b + s % 4] # four-token groups + + `rcp` is `rcp.approx.ftz.f32`, `e4m3` and `e2m1` round to nearest even and + saturate. This is the attention op's own NVFP4 append (`cvt_warp_fp16_to_fp4` + behind `quantizeAndWriteFP4KVCache`), the same bits for the same row: a + token's stored value does not depend on which path wrote it. A row that + resolves to no slot -- negative request, request past the table, negative + position, page column past the table, a table entry that is negative or past + the pool -- is skipped. Returns the four caches, written in place. + + Fusion boundary: quantization and the scatter. The projection, norm and RoPE + of K are the caller's; the QSA index plane is not touched. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode: 8 rows at aligned, mid-page and page-crossing positions, with zeros, " + "-0.0 and a 900.0 outlier inside blocks, and a row on an unallocated (-1) page: " + "zero scales, saturation and the skip", + spec=_spec(_DECODE, pages=69, slots=8, values="edge", hole="unallocated_page"), + ), + Cell( + why="prefill: 1024 rows, 4 requests x 256 walking 8 pages each, in the engine's " + "trunk layout: layer 11's planes among the 24 role pages of a slot", + spec=_spec(((0, 256),) * 4, pages=160, slot_pages=24, layer=11), + ), + Cell( + why="the attention op's own append over the same rows (context, chunk, decode " + "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), + ), + ) + + note = """ + Metadata consumed: none -- every input is an argument. The target passes the + step's `qsa_req_idx_per_token` (int32) and `qsa_logical_positions` (int64), + `qsa_block_table` unsliced (raw pool page ids, `-1` unallocated), the layer's + K and V planes of the QSA manager's `get_buffers(i, "HND")` (uint8 `[slots, + Hkv, tpb, D/2]`), the two planes of `get_block_scale_buffers(i, "HND")` + (uint8 `[slots, Hkv, tpb, D/16]`), and its own `kv_scale_orig_quant` -- the + same fp32 `[3]` tensor its dense attention call binds, so both writers apply + one second-level scale. + + Preconditions: + - destination slots are unique across the batch: two rows naming one + (page, slot) both store, unordered. + - the dense path's reader assumes this layout: e2m1 nibbles HND per page, + K scales linear, V scales in four-token groups. The attention op's + append, run over the same rows, leaves bitwise the same pool (a cell), + at the unit second-level scales the target binds. + The callable rejects loudly (ValueError; the test pins the non-contiguous + metadata refusal): non-matching or non-3-D k/v, a row + dtype outside bf16/fp16, a head_dim that is not a power of two >= 32, a + tokens_per_block that is not a multiple of 4, caches that are not uint8 of + the stated shapes or whose last dim or scale head page is strided, a non-2-D + or non-int32 table, metadata that is not one contiguous int32/int64 entry per + row, a `kv_scale_orig_quant` that is not a contiguous fp32 `[3]`, operands on + another device. `rows == 0` returns before launching (not certified). + + Every offset is formed in 64 bits: rows, heads, requests, positions and pages + are widened before any product. CUDA-graph capturable; a replay follows rows, + pages and positions rewritten in place (the test replays 8 MTP verify blocks of + 8 rows bitwise equal to eager). + + Compiled variants: the page count (`num_pages`, the pool's slot count), the + table's height (`num_requests`) and width and every stride are plain-int + arguments, and Triton compiles the kernel once per class of each -- equal to + 1, a multiple of 16, anything else -- besides once per head_dim and + 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. + + 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`). + + Not certified either: second-level scales other than 1.0, int32 positions + or int64 request indices, and of the rows that resolve to no slot any but + an unallocated (`-1`) table entry (the reference skips each). + """ + + def __call__( + self, + k: torch.Tensor, + v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_scale_cache: torch.Tensor, + v_scale_cache: torch.Tensor, + request_indices: torch.Tensor, + logical_positions: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + kv_scale_orig_quant: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + _store_nvfp4( + k=k, + v=v, + k_cache=k_cache, + v_cache=v_cache, + k_scale_cache=k_scale_cache, + v_scale_cache=v_scale_cache, + request_indices=request_indices, + logical_positions=logical_positions, + block_table=block_table, + tokens_per_block=tokens_per_block, + kv_scale_orig_quant=kv_scale_orig_quant, + ) + return k_cache, v_cache, k_scale_cache, v_scale_cache + + def reference( + self, + k: torch.Tensor, + v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_scale_cache: torch.Tensor, + v_scale_cache: torch.Tensor, + request_indices: torch.Tensor, + logical_positions: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + kv_scale_orig_quant: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """The four caches after the store, as fresh tensors; the inputs are untouched. + + Quantization in fp32 torch with IEEE reciprocals where the kernel uses + `rcp.approx.ftz`; the two agree on every operand the cells drive (6.0, the + unit second-level scale, and the e4m3 scale values their rows produce), which + the cells' bitwise gate is what pins. + """ + out = [ + t.clone(memory_format=torch.contiguous_format) + for t in (k_cache, v_cache, k_scale_cache, v_scale_cache) + ] + rows, heads, head_dim = k.shape + requests = request_indices.long() + positions = logical_positions.long() + columns = positions.div(tokens_per_block, rounding_mode="floor") + table_rows, table_width = block_table.shape + ok = (requests >= 0) & (requests < table_rows) & (positions >= 0) & (columns < table_width) + pages = block_table[ + requests.clamp(0, table_rows - 1), columns.clamp(0, table_width - 1) + ].long() + ok &= (pages >= 0) & (pages < k_cache.shape[0]) + page, slot = pages[ok], (positions % tokens_per_block)[ok] + n = int(page.numel()) + if n == 0: + return tuple(out) + groups = head_dim // _SF_VEC + quad_base = (slot // _V_SF_QUAD) * (_V_SF_QUAD * groups) + slot % _V_SF_QUAD + for src, data, scales, global_scale, linear in ( + (k, out[0], out[2], kv_scale_orig_quant[1], True), + (v, out[1], out[3], kv_scale_orig_quant[2], False), + ): + packed, sf = _quantize_reference(src[ok].reshape(n * heads, head_dim), global_scale) + data[page, :, slot, :] = packed.reshape(n, heads, head_dim // _PER_BYTE) + sf = sf.reshape(n, heads, groups) + if linear: + scales[page, :, slot, :] = sf + else: + flat = scales.view(scales.shape[0], heads, tokens_per_block * groups) + offsets = quad_base[:, None] + _V_SF_QUAD * torch.arange(groups, device=sf.device) + for h in range(heads): + flat[page[:, None], h, offsets] = sf[:, h] + return tuple(out) + + +qsa_paged_kv_store_nvfp4 = QsaPagedKvStoreNvfp4() 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 index c1af92df810c..21609fd37846 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.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 @@ -12,6 +12,9 @@ # 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, ) @@ -25,9 +28,8 @@ 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 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. @@ -36,6 +38,7 @@ _SPLIT_K = ("splitk", "merge") _PREFILL = ("prefill",) +_FP8 = dict(cache_dtype=torch.float8_e4m3fn) def _spec(**overrides) -> dict: @@ -53,10 +56,6 @@ def _decode(segments, launches=_SPLIT_K, **overrides) -> dict: return _spec(**{**base, **overrides}) -def _mixed(requests, context_lens) -> dict: - return _spec(mixed=(requests, context_lens), slots=32, prefill_pair=True, launches=_PREFILL) - - class QsaPagedSparseGqa(OpWrapper): """Sparse GQA over the selected tokens of an HND paged K/V pool. @@ -89,9 +88,20 @@ class QsaPagedSparseGqa(OpWrapper): 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. - Fusion boundary: the block-table gather, scale, online softmax, P.V and - the split-K merge are inside. Deciding what to attend to, appending K/V to - the pool, and q's projection, norm, RoPE and output gate are not. + **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}) @@ -99,102 +109,22 @@ class QsaPagedSparseGqa(OpWrapper): CELLS: tuple[Cell, ...] = ( # ---- generation: no prefill pair, split-K --------------------------- Cell( - why="decode: 1 row past the budget, split-K + merge, PDL", - spec=_decode(((4000, 1),)), - ), - Cell( - why="decode: 4 requests x 1 row, split-K + merge, PDL", - spec=_decode(((4000, 1), (4037, 1), (4074, 1), (4111, 1))), - ), - Cell( - why="decode MTP: 4 requests x 2 rows, split-K + merge", - spec=_decode(tuple((4000 + 37 * i, 2) for i in range(4))), - ), - Cell( - why="decode MTP: 8 requests x 2 rows, 8-way split-K", - spec=_decode(tuple((4000 + 37 * i, 2) for i in range(8))), - ), - Cell( - why="decode MTP: all 32 slots x 2 rows, positions 100..8099", - spec=_decode(tuple((100 + (499 * g) % 8000, 2) for g in range(32))), - ), - Cell( - why="300 rows, no prefill pair: split-K, one split, no merge", - spec=_decode(((600, 1),) * 300, launches=("splitk",), slots=384), - ), - Cell( - why="16-wide selection: one tile, the split-K kernel writes out", - spec=_decode(((600, 1),) * 4, launches=("splitk",), width=16), + 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="context 128 rows from 2032: pair passed, still split-K", - spec=_decode(((2032, 128),), prefill_pair=True), - ), - Cell( - why="context 272 rows from 2032: causal-prefix prefill kernel", - spec=_decode(((2032, 272),), launches=_PREFILL, prefill_pair=True), - ), - Cell( - why="context: two requests (150 @1900, 150 @900) in one launch", - spec=_decode(((1900, 150), (900, 150)), launches=_PREFILL, prefill_pair=True), - ), - Cell( - why="context 300 rows from 0: rows see 1..300 tokens", - spec=_decode(((0, 300),), launches=_PREFILL, prefill_pair=True), + why="fp8 context 272 rows from 2032: causal-prefix prefill kernel", + spec=_decode(((2032, 272),), launches=_PREFILL, prefill_pair=True, **_FP8), ), Cell( - why="warm-up capture: 8191 rows of one request from 0", - spec=_decode(((0, 8191),), launches=_PREFILL, prefill_pair=True), - ), - Cell( - why="warm-up capture: 2 requests x 4096 rows from 0", - spec=_decode(((0, 4096), (0, 4096)), launches=_PREFILL, prefill_pair=True), - ), - # ---- mixed steps from the mmmu gate's runtime dump ------------------- - Cell(why="mixed 257 rows: 1 chunk + 31 MTP decodes", spec=_mixed(32, (195,))), - Cell(why="mixed 318 rows: 2 chunks + 30 MTP decodes", spec=_mixed(32, (200, 58))), - Cell(why="mixed 478 rows over 28 slots: 1 chunk", spec=_mixed(28, (424,))), - Cell( - why="mixed 869 rows: 5 chunks + 27 MTP decodes", - spec=_mixed(32, (400, 200, 100, 80, 35)), - ), - Cell( - why="mixed 1334 rows: 4 chunks + 28 MTP decodes", - spec=_mixed(32, (900, 200, 150, 28)), - ), - Cell( - why="mixed 3806 rows: 9 chunks + 23 MTP decodes", - spec=_mixed(32, (1024, 900, 700, 500, 300, 200, 100, 30, 6)), - ), - Cell( - why="mixed 5130 rows: 10 chunks + 22 MTP decodes", - spec=_mixed(32, (2048, 1024, 700, 500, 400, 200, 120, 60, 30, 4)), - ), - Cell( - why="mixed 8169 rows over 31 slots: 10 chunks", - spec=_mixed(31, (4096, 2048, 1024, 500, 200, 120, 80, 40, 15, 4)), - ), - Cell(why="mixed 8190 rows: 4 chunks + 28 MTP decodes", spec=_mixed(32, (8000, 100, 30, 4))), - # ---- formats and layouts -------------------------------------------- - Cell( - why="fp16 pool and fp16 q", - spec=_decode(((900, 1),) * 4, cache_dtype=torch.float16, q_dtype=torch.float16), - ), - Cell( - why="fp8-e4m3 pool under a bf16 q, read unscaled", - spec=_decode(((900, 1),) * 4, cache_dtype=torch.float8_e4m3fn), - ), - Cell( - why="int64 request indices", spec=_decode(((1000, 2),) * 4, request_dtype=torch.int64) - ), - Cell( - why="standalone contiguous K and V buffers, contiguous q", - spec=_decode(((900, 1),) * 4, cache_layout="contiguous", q_layout="contiguous"), - ), - Cell( - why="holes: all-padding row, foreign request, unmapped page", - spec=_decode(((900, 1),) * 4, holes=True), + 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"), ), ) @@ -216,24 +146,37 @@ class docstring). QSA's producer -- `qsa_topk_select` with zero 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 (asserted with a planted index). - - There is no KV scale argument: an fp8 pool is read at face value. + 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 (11 `ValueError`s pinned by the - test, plus a selection width the compress ratio cannot produce, refused - on the prefill kernel only), so `is_valid` adds nothing. `rows == 0` - returns an empty `[0, Hq, D]` tensor. + 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; both launch modes are - inside the cells. Calls are bitwise deterministic. + 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); head geometries other than 24/2 x 256; - `tokens_per_block` other than 32; more than 8192 rows; the split-K kernel - above 300 rows. Hard limits: not probed; the largest certified extent is - 8192 rows x 2051 selected tokens through a [32, 260] table. + 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__( @@ -249,6 +192,7 @@ def __call__( 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, @@ -261,8 +205,19 @@ def __call__( 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, *, @@ -276,8 +231,10 @@ def reference( 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. + """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. @@ -317,7 +274,10 @@ def reference( 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) - return out.to(q.dtype) + 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_decode.py new file mode 100644 index 000000000000..b6180bcb7566 --- /dev/null +++ b/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 @@ -0,0 +1,819 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA sparse GQA for decode rows: a split-K pass that loads its tokens in one round trip, +then a merge spread over the SMs. + +Tier 5 (CuTe DSL written for this target), both launches. + +Split-K pass. One CTA (eight warps) covers `CAP` consecutive selected tokens of one KV +head of one row, so a decode row's 2051 selected tokens take 33 CTAs a KV head: + +1. before its PDL wait the CTA stages the KV head's query heads in shared memory and from + there in every warp's MMA A fragments, stages the request's block-table row (up to + `BT_STAGE` pages) in shared memory, and releases its dependents; +2. after it, its tokens' pages come from shared memory and every token's K and V row + (16 bytes a lane, a warp eight tokens) is requested at once with cp.async, so the CTA + pays one DRAM round trip -- QSA's split-K kernel walks two or three 16-token tiles, + each a chain of dependent loads (selection, page, K/V); +3. scores = q K^T on warp MMAs (bf16 in, fp32 out; a warp eight tokens), scaled into the + log2 domain and masked; one softmax over the CTA's tokens (a warp two heads); P V on + warp MMAs with P in bf16 (a warp 32 head dims); each head's partial times its + reciprocal sum is staged in shared memory and stored in 16-byte pieces with its + log2-sum-exp, as QSA's split-K pass stores them. + +Merge. A CTA covers 128 elements of a head (two CTAs a head) with up to eight warps: warp +`w` reads its splits' partials of four adjacent elements a lane, one 16-byte load a split, +all in flight; every warp forms the split weights itself from two LSEs a lane; the warps' +shares are added in shared memory. + +Both launches release their dependents at entry, before their own waits: the merge and, +through it, the output projection that follows launch while the top-k ahead of the +split-K pass still runs, and the projection streams its weight through that latency-bound +stretch. Before its wait the split-K pass reads q and the merge the output gate: the +caller's contract is that both are complete when the grid ahead of the split-K pass +releases it (the target's top-k does so only after its own wait; a launch without PDL +releases at its end). +""" + +from __future__ import annotations + +import functools +import math +from typing import Optional + +import cuda.bindings.driver as cuda_drv +import cutlass +import cutlass.cute as cute +import torch +from cutlass import const_expr +from cutlass.cute import experimental as cute_ext +from cutlass.cute.nvgpu import cpasync, warp +from cutlass.cute.runtime import from_dlpack, make_ptr + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# 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. +MAX_SPLITS = 64 +#: Rows the target calls this entry with: the launches that run with PDL (QSA enables it up +#: to four rows); larger batches keep the composition. +MAX_ROWS = 4 +#: Selected tokens a split-K CTA covers. +CAP = 64 +#: MMA rows: the query heads of a KV head, padded. +M = 16 +#: bf16 elements of padding after each shared K, V, q and P row (keeps the MMA fragment +#: reads free of bank conflicts). +PAD = 8 +WARPS = 8 +THREADS = WARPS * 32 +#: Block-table entries staged in shared memory before the wait; a wider table is read +#: from global memory after it. +BT_STAGE = 4096 +HEAD_DIM = 256 +LOG2E = 1.4426950408889634 + + +class _SplitKKernel: + def __init__( + self, + *, + group, + width, + splits, + tpb, + scale, + bt_width, + num_pages, + k_strides, + v_strides, + num_requests, + pdl, + ): + assert group <= M and CAP == 2 * 32 and CAP % WARPS == 0 + self.group, self.width, self.splits, self.tpb = group, width, splits, tpb + self.scale = scale * LOG2E + self.bt_width, self.num_pages, self.num_requests = bt_width, num_pages, num_requests + self.stage_bt = bt_width <= BT_STAGE + self.k_strides, self.v_strides = k_strides, v_strides + self.ld = HEAD_DIM + PAD + self.pdl = pdl + + def smem_bytes(self) -> int: + def r(n): + return -(-n // 128) * 128 + + b = ( + r(M * self.ld * 2) + + 2 * r(CAP * self.ld * 2) + + r(M * (CAP + PAD) * 2) + + r(M * (CAP + 4) * 4) + ) + b += r(M * 2 * 4) + 2 * r(CAP * 8) + if self.stage_bt: + b += r(self.bt_width * 4) + return b + 256 + + @cute.jit + def __call__(self, mQ, pK, pV, mBT, mSel, mReq, pPart, mLse, stream: cuda_drv.CUstream): + rows = cute.size(mSel, mode=[0]) + kv_heads = cute.size(mQ, mode=[1]) // self.group + mma = cute.make_tiled_mma( + warp.MmaF16BF16Op(cutlass.BFloat16, cutlass.Float32, (16, 8, 16)), + cute.make_layout((1, WARPS, 1)), + ) + g2s = cute.make_copy_atom(cpasync.CopyG2SOp(), cutlass.BFloat16, num_bits_per_copy=128) + st = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=128 + ) + self.kernel(mQ, pK, pV, mBT, mSel, mReq, pPart, mLse, mma, g2s, st).launch( + grid=[rows * kv_heads * self.splits, 1, 1], + block=[THREADS, 1, 1], + smem=self.smem_bytes(), + stream=stream, + use_pdl=self.pdl, + ) + + @cute.kernel + def kernel( + self, + mQ, + pK, + pV, + mBT, + mSel, + mReq, + pPart, + mLse, + mma: cute.TiledMma, + g2s: cute.CopyAtom, + st: cute.CopyAtom, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + warp_i = cute.arch.warp_idx() + lane = cute.arch.lane_idx() + group: cutlass.Constexpr = self.group + splits: cutlass.Constexpr = self.splits + width: cutlass.Constexpr = self.width + tpb: cutlass.Constexpr = self.tpb + bt_width: cutlass.Constexpr = self.bt_width + ld: cutlass.Constexpr = self.ld + dim: cutlass.Constexpr = HEAD_DIM + heads = cute.size(mQ, mode=[1]) + kv_heads = heads // group + rows = cute.size(mSel, mode=[0]) + row = bidx // (kv_heads * splits) + kvh = (bidx // splits) % kv_heads + split = bidx % splits + c0 = split * CAP + + smem = cutlass.utils.SmemAllocator() + sQ = smem.allocate_tensor( + cutlass.BFloat16, cute.make_layout((M, dim), stride=(ld, 1)), byte_alignment=128 + ) + sK = smem.allocate_tensor( + cutlass.BFloat16, cute.make_layout((CAP, dim), stride=(ld, 1)), byte_alignment=128 + ) + sV = smem.allocate_tensor( + cutlass.BFloat16, cute.make_layout((CAP, dim), stride=(ld, 1)), byte_alignment=128 + ) + sP = smem.allocate_tensor( + cutlass.BFloat16, cute.make_layout((M, CAP), stride=(CAP + PAD, 1)), byte_alignment=128 + ) + sS = smem.allocate_tensor( + cutlass.Float32, cute.make_layout((M, CAP), stride=(CAP + 4, 1)), byte_alignment=128 + ) + sStat = smem.allocate_tensor( + cutlass.Float32, cute.make_layout((M, 2), stride=(2, 1)), byte_alignment=128 + ) + sKo = smem.allocate_tensor( + cutlass.Int64, cute.make_layout((CAP,), stride=(1,)), byte_alignment=128 + ) + sVo = smem.allocate_tensor( + cutlass.Int64, cute.make_layout((CAP,), stride=(1,)), byte_alignment=128 + ) + + # 1. Before the wait: the KV head's query heads (complete before the grid ahead of + # this one released it) into shared memory and the A fragments, the block-table row + # (host-filled), the release. + cpr: cutlass.Constexpr = dim // 8 # 16-byte pieces a row: one a lane + q_vec = cute.logical_divide(mQ, (None, None, 8)) + sq_vec = cute.logical_divide(sQ, (None, 8)) + for i in cutlass.range_constexpr(-(-(group * cpr) // THREADS)): + c = tidx + i * THREADS + cc = cutlass.min(c, group * cpr - 1) + h = cc // cpr + ch = cc % cpr + if c < group * cpr: + cute.copy(g2s, q_vec[row, kvh * group + h, (None, ch)], sq_vec[h, (None, ch)]) + cute.arch.cp_async_commit_group() + for i in cutlass.range_constexpr((M - group) * dim // THREADS): + e = tidx + i * THREADS + sQ[group + e // dim, e % dim] = cutlass.BFloat16(0.0) + req = mReq[row] + req_ok = (req >= 0) & (req < self.num_requests) + safe_req = cutlass.min(cutlass.max(req, 0), self.num_requests - 1) + if const_expr(self.stage_bt): + sBT = smem.allocate_tensor( + cutlass.Int32, cute.make_layout((bt_width,), stride=(1,)), byte_alignment=128 + ) + for i in cutlass.range_constexpr(-(-bt_width // THREADS)): + e = tidx + i * THREADS + ec = cutlass.min(e, bt_width - 1) + v = mBT[safe_req, ec] + if e < bt_width: + sBT[e] = v + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + thr = mma.get_slice(tidx) + tCsQ = thr.partition_A(sQ) + tCrQ = thr.make_fragment_A(tCsQ) + cute.autovec_copy(tCsQ, tCrQ) + if const_expr(self.pdl): + # The merge waits before reading the partials and reads nothing before its wait + # that the grid ahead of this one writes. + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + + # 2. The CTA's tokens: each one's K and V row offsets in elements, or -1. + ks0, ks1, ks2, _ = self.k_strides + vs0, vs1, vs2, _ = self.v_strides + if tidx < CAP: + col = c0 + tidx + colc = cutlass.min(col, width - 1) + tok = mSel[row, colc] + ok = req_ok & (col < width) & (tok >= 0) + safe_tok = cutlass.max(tok, 0) + lp = safe_tok // tpb + ok = ok & (lp < bt_width) + lpc = cutlass.min(lp, bt_width - 1) + page = cutlass.Int32(0) + if const_expr(self.stage_bt): + page = sBT[lpc] + else: + page = mBT[safe_req, lpc] + ok = ok & (page >= 0) & (page < self.num_pages) + pg = cutlass.Int64(cutlass.max(page, 0)) + tip = cutlass.Int64(safe_tok % tpb) + hk = cutlass.Int64(kvh) + sKo[tidx] = cutlass.Int64( + cutlass.select_(ok, pg * ks0 + hk * ks1 + tip * ks2, cutlass.Int64(-1)) + ) + sVo[tidx] = cutlass.Int64( + cutlass.select_(ok, pg * vs0 + hk * vs1 + tip * vs2, cutlass.Int64(-1)) + ) + cute.arch.sync_threads() + + # Every token's K and V row: warp w the tokens 8w .. 8w + 7, a lane 16 bytes of each; + # an invalid token's rows are zeros (its score is masked; V must stay finite). + sk_vec = cute.logical_divide(sK, (None, 8)) + sv_vec = cute.logical_divide(sV, (None, 8)) + k_base = pK.toint() + cutlass.Int64(lane * 16) + v_base = pV.toint() + cutlass.Int64(lane * 16) + tpw: cutlass.Constexpr = CAP // WARPS + for i in cutlass.range_constexpr(tpw): + t = warp_i * tpw + i + ko = sKo[t] + vo = sVo[t] + ksrc = cute.make_tensor( + cute.make_ptr( + cutlass.BFloat16, + k_base + cutlass.max(ko, 0) * 2, + cute.AddressSpace.gmem, + assumed_align=16, + ), + cute.make_layout((8,), stride=(1,)), + ) + vsrc = cute.make_tensor( + cute.make_ptr( + cutlass.BFloat16, + v_base + cutlass.max(vo, 0) * 2, + cute.AddressSpace.gmem, + assumed_align=16, + ), + cute.make_layout((8,), stride=(1,)), + ) + if ko >= 0: + cute.copy(g2s, ksrc, sk_vec[t, (None, lane)]) + cute.copy(g2s, vsrc, sv_vec[t, (None, lane)]) + else: + for e in cutlass.range_constexpr(8): + sK[t, lane * 8 + e] = cutlass.BFloat16(0.0) + sV[t, lane * 8 + e] = cutlass.BFloat16(0.0) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + + # 3. Scores: a warp eight tokens over the whole head dim, masked into shared memory. + tCsK = thr.partition_B(sK) + tCrK = thr.make_fragment_B(tCsK) + cute.autovec_copy(tCsK, tCrK) + scores = cute.make_rmem_tensor(thr.partition_shape_C((M, CAP)), cutlass.Float32) + scores.fill(0.0) + cute.gemm(mma, scores, tCrQ, tCrK, scores) + cs = thr.partition_C(cute.make_identity_tensor((M, CAP))) + neg_inf = cutlass.Float32(-math.inf) + for j in cutlass.range_constexpr(cute.size(scores)): + h, t = cs[j] + v = neg_inf + if sKo[t] >= 0: + v = scores[j] * self.scale + sS[h, t] = v + cute.arch.sync_threads() + # The softmax: warp w the heads w and w + WARPS, a lane two tokens; padding heads + # get P = 0. + for j in cutlass.range_constexpr(M // WARPS): + h = warp_i + WARPS * j + a = sS[h, lane] + b = sS[h, lane + 32] + mx = cute.arch.warp_reduction_max(cute.arch.fmax(a, b)) + safe_mx = cutlass.Float32(cutlass.select_(mx == neg_inf, cutlass.Float32(0.0), mx)) + keep = h < group + pa = cutlass.Float32( + cutlass.select_(keep, cute.math.exp2(a - safe_mx), cutlass.Float32(0.0)) + ) + pb = cutlass.Float32( + cutlass.select_(keep, cute.math.exp2(b - safe_mx), cutlass.Float32(0.0)) + ) + den = cute.arch.warp_reduction_sum(pa + pb) + sP[h, lane] = pa.to(cutlass.BFloat16) + sP[h, lane + 32] = pb.to(cutlass.BFloat16) + if lane == 0: + sStat[h, 0] = mx + sStat[h, 1] = den + cute.arch.sync_threads() + + # P V: a warp 32 of the head dims (four 8-wide tiles). + tCsP = thr.partition_A(sP) + tCrP = thr.make_fragment_A(tCsP) + cute.autovec_copy(tCsP, tCrP) + sVt = cute.make_tensor(sV.iterator, cute.make_layout((dim, CAP), stride=(1, ld))) + tCsV = thr.partition_B(sVt) + tCrV = thr.make_fragment_B(tCsV) + cute.autovec_copy(tCsV, tCrV) + acc = cute.make_rmem_tensor(thr.partition_shape_C((M, dim)), cutlass.Float32) + acc.fill(0.0) + cute.gemm(mma, acc, tCrP, tCrV, acc) + + # Each head's partial times its reciprocal sum, staged over sK (free since the + # softmax barrier) as [group, dim] fp32, then stored in 16-byte pieces: the CTA's + # heads are contiguous in the partial tensor. + sOut = cute.make_tensor( + cute.recast_ptr(sK.iterator, dtype=cutlass.Float32), + cute.make_layout((group, dim), stride=(dim, 1)), + ) + co = thr.partition_C(cute.make_identity_tensor((M, dim))) + h0 = lane // 4 # an m16n8 C fragment's rows: h0 and h0 + 8 + d0 = sStat[h0, 1] + d1 = sStat[h0 + 8, 1] + r0 = cutlass.Float32(0.0) + r1 = cutlass.Float32(0.0) + if d0 > 0.0: + r0 = 1.0 / d0 + if d1 > 0.0: + r1 = 1.0 / d1 + for j in cutlass.range_constexpr(cute.size(acc)): + h, d = co[j] + r = cutlass.Float32(cutlass.select_(h < 8, r0, r1)) + if h < group: + sOut[h, d] = acc[j] * r + cute.arch.sync_threads() + base = ((cutlass.Int64(split) * rows + row) * heads + kvh * group) * dim + out_ptr = pPart.toint() + base * 4 + so_vec = cute.make_tensor( + sOut.iterator, cute.make_layout((group * dim // 4, 4), stride=(4, 1)) + ) + reg = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), cutlass.Float32) + for i in cutlass.range_constexpr(-(-(group * dim // 4) // THREADS)): + c = tidx + i * THREADS + if c < group * dim // 4: + cute.autovec_copy(so_vec[c, None], reg) + dst = cute.make_tensor( + cute.make_ptr( + cutlass.Float32, + out_ptr + cutlass.Int64(c * 16), + cute.AddressSpace.gmem, + assumed_align=16, + ), + cute.make_layout((4,), stride=(1,)), + ) + cute.copy(st, reg, dst) + if tidx < group: + mx = sStat[tidx, 0] + den = sStat[tidx, 1] + lse = neg_inf + if den > 0.0: + lse = mx + cute.math.log2(den) + mLse[split, row, kvh * group + tidx] = lse + + +class _MergeKernel: + def __init__(self, *, heads, dim, splits, gated, pdl): + assert dim % CHUNK == 0 and 1 < splits <= MAX_SPLITS + self.heads, self.dim, self.splits, self.gated, self.pdl = heads, dim, splits, gated, pdl + self.chunks = dim // CHUNK + self.warps = min(8, splits) + self.spw = -(-splits // self.warps) # splits a warp, rounded up + + @cute.jit + def __call__(self, mPart, mLse, mGate, mOut, stream: cuda_drv.CUstream): + rows = cute.size(mOut, mode=[0]) + ld = cute.make_copy_atom(cute.nvgpu.CopyG2ROp(), cutlass.Float32, num_bits_per_copy=128) + self.kernel(mPart, mLse, mGate, mOut, ld).launch( + grid=[rows * self.heads * self.chunks, 1, 1], + block=[self.warps * 32, 1, 1], + smem=self.warps * CHUNK * 4 + 64, + stream=stream, + use_pdl=self.pdl, + ) + + @cute.kernel + def kernel(self, mPart, mLse, mGate, mOut, ld: cute.CopyAtom): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + warp_i = cute.arch.warp_idx() + lane = cute.arch.lane_idx() + per_row: cutlass.Constexpr = self.heads * self.chunks + spw: cutlass.Constexpr = self.spw + row = bidx // per_row + head = (bidx % per_row) // self.chunks + c0 = (bidx % self.chunks) * CHUNK + red = cutlass.utils.SmemAllocator().allocate_tensor( + cutlass.Float32, + cute.make_layout((self.warps, CHUNK), stride=(CHUNK, 1)), + byte_alignment=16, + ) + # The gate, before the wait: see the module notes. + g = cutlass.Float32(0.0) + if const_expr(self.gated): + if tidx < CHUNK: + g = mGate[row, head * self.dim + c0 + tidx].to(cutlass.Float32) + if const_expr(self.pdl): + # The output projection after this reads only its weight before its own wait. + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + + # 1. The split weights, in every warp: exp2(lse - max), and their sum. + neg_inf = cutlass.Float32(-math.inf) + l0 = neg_inf + l1 = neg_inf + if lane < self.splits: + l0 = mLse[lane, row, head] + if lane + 32 < self.splits: + l1 = mLse[lane + 32, row, head] + # 2. Warp w's splits w * spw .. w * spw + spw - 1, four adjacent elements a lane: + # one 16-byte load a split, all in flight. A slot past the split count reads the + # last split again and weighs it 0. + part_vec = cute.logical_divide(mPart, (None, None, None, 4)) + p = cute.make_rmem_tensor(cute.make_layout((spw, 4), stride=(4, 1)), cutlass.Float32) + for j in cutlass.range_constexpr(spw): + sidx = cutlass.min(warp_i * spw + j, self.splits - 1) + cute.copy(ld, part_vec[sidx, row, head, (None, (c0 + 4 * lane) // 4)], p[j, None]) + mx = cute.arch.warp_reduction_max(cute.arch.fmax(l0, l1)) + w0 = cutlass.Float32(0.0) + w1 = cutlass.Float32(0.0) + if mx > neg_inf: + w0 = cute.math.exp2(l0 - mx) + w1 = cute.math.exp2(l1 - mx) + den = cute.arch.warp_reduction_sum(w0 + w1) + acc = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), cutlass.Float32) + acc.fill(0.0) + for j in cutlass.range_constexpr(spw): + sidx = warp_i * spw + j + wa = cute.arch.shuffle_sync(w0, sidx % 32) + wb = cute.arch.shuffle_sync(w1, sidx % 32) + w = cutlass.Float32(cutlass.select_(sidx < 32, wa, wb)) + w = cutlass.Float32(cutlass.select_(sidx < self.splits, w, cutlass.Float32(0.0))) + for e in cutlass.range_constexpr(4): + acc[e] = w * p[j, e] + acc[e] + for e in cutlass.range_constexpr(4): + red[warp_i, 4 * lane + e] = acc[e] + cute.arch.sync_threads() + + # 3. The warps' shares summed in warp order; the output gate on the rounded + # attention value, as the composition applies it. + if tidx < CHUNK: + tot = cutlass.Float32(0.0) + for w_ in cutlass.range_constexpr(self.warps): + tot = tot + red[w_, tidx] + merged = cutlass.Float32(0.0) + if den > 0.0: + merged = tot / den + if const_expr(self.gated): + att = merged.to(cutlass.BFloat16).to(cutlass.Float32) + merged = att * (1.0 / (1.0 + cute.math.exp(-g))) + mOut[row, head, c0 + tidx] = merged.to(mOut.element_type) + + +def _dl(t: torch.Tensor, align: int = 4): + return from_dlpack(t, assumed_align=align) + + +def _ptr(t: torch.Tensor, dtype): + return make_ptr(dtype, t.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + + +def splits(q, k_cache, selected_tokens) -> int: + """The split-K CTAs a row and KV head: one per `CAP` selected tokens.""" + return -(-selected_tokens.shape[1] // CAP) + + +@functools.lru_cache(maxsize=None) +def _compiled_splitk( + rows, + q_shape, + q_stride, + group, + width, + tpb, + scale, + bt_shape, + num_pages, + k_strides, + v_strides, + pdl, +): + heads = q_shape[1] + n_splits = -(-width // CAP) + kern = _SplitKKernel( + group=group, + width=width, + splits=n_splits, + tpb=tpb, + scale=scale, + bt_width=bt_shape[1], + num_pages=num_pages, + k_strides=k_strides, + v_strides=v_strides, + num_requests=bt_shape[0], + pdl=pdl, + ) + dev = "cuda" + q = torch.empty_strided(q_shape, q_stride, dtype=torch.bfloat16, device=dev) + kv = torch.empty(64, dtype=torch.bfloat16, device=dev) + bt = torch.empty(bt_shape, dtype=torch.int32, device=dev) + sel = torch.empty(rows, width, dtype=torch.int32, device=dev) + req = torch.empty(rows, dtype=torch.int32, device=dev) + part = torch.empty(64, dtype=torch.float32, device=dev) + lse = torch.empty(n_splits, rows, heads, dtype=torch.float32, device=dev) + stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) + return cute_ext.compile( + kern, + _dl(q, 16), + _ptr(kv, cutlass.BFloat16), + _ptr(kv, cutlass.BFloat16), + _dl(bt), + _dl(sel), + _dl(req), + _ptr(part, cutlass.Float32), + _dl(lse), + stream, + ) + + +@functools.lru_cache(maxsize=None) +def _compiled_merge(rows, heads, dim, n_splits, gate_stride, dtype, pdl): + kern = _MergeKernel( + heads=heads, dim=dim, splits=n_splits, gated=gate_stride is not None, pdl=pdl + ) + dev = "cuda" + part = torch.empty(n_splits, rows, heads, dim, dtype=torch.float32, device=dev) + lse = torch.empty(n_splits, rows, heads, dtype=torch.float32, device=dev) + gate = torch.empty(rows, gate_stride or heads * dim, dtype=dtype, device=dev)[:, : heads * dim] + out = torch.empty(rows, heads, dim, dtype=dtype, device=dev) + stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) + return cute_ext.compile(kern, _dl(part, 16), _dl(lse), _dl(gate, 2), _dl(out, 2), stream) + + +def launch( + q, + k_cache, + v_cache, + block_table, + selected_tokens, + request_indices, + tokens_per_block, + softmax_scale, + output_gate=None, + pdl=True, +): + """The launch behind `QsaPagedSparseGqaDecode.__call__`; returns `[rows, H, D]` in q's dtype.""" + rows, num_q_heads, head_dim = q.shape + num_kv_heads = k_cache.shape[1] + group_size = num_q_heads // num_kv_heads + width = selected_tokens.shape[1] + n_splits = -(-width // CAP) + output = torch.empty((rows, num_q_heads, head_dim), dtype=q.dtype, device=q.device) + if rows == 0: + return output + dev = q.device + use_pdl = bool(pdl and _qsa_pdl_enabled(rows)) + partial_output = torch.empty( + (n_splits, rows, num_q_heads, head_dim), dtype=torch.float32, device=dev + ) + partial_lse = torch.empty((n_splits, rows, num_q_heads), dtype=torch.float32, device=dev) + fn = _compiled_splitk( + rows, + tuple(q.shape), + tuple(q.stride()), + group_size, + width, + tokens_per_block, + float(softmax_scale), + tuple(block_table.shape), + k_cache.shape[0], + tuple(k_cache.stride()), + tuple(v_cache.stride()), + use_pdl, + ) + stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream) + fn( + _dl(q, 16), + _ptr(k_cache, cutlass.BFloat16), + _ptr(v_cache, cutlass.BFloat16), + _dl(block_table), + _dl(selected_tokens), + _dl(request_indices), + _ptr(partial_output, cutlass.Float32), + _dl(partial_lse), + stream, + ) + gate2d = None if output_gate is None else output_gate.reshape(rows, num_q_heads * head_dim) + mfn = _compiled_merge( + rows, + num_q_heads, + head_dim, + n_splits, + None if gate2d is None else gate2d.stride(0), + q.dtype, + use_pdl, + ) + mfn( + _dl(partial_output, 16), + _dl(partial_lse), + _dl(output.view(rows, num_q_heads * head_dim) if gate2d is None else gate2d, 2), + _dl(output, 2), + stream, + ) + return output + + +def supports(q, k_cache, selected_tokens) -> bool: + """Whether these operands are ones this entry serves (see `is_valid`).""" + return ( + q.dim() == 3 + and 0 < q.shape[0] <= MAX_ROWS + and q.shape[2] == HEAD_DIM + and q.dtype == torch.bfloat16 + and k_cache.dtype == torch.bfloat16 + and q.shape[1] % k_cache.shape[1] == 0 + and q.shape[1] // k_cache.shape[1] <= M + and 1 < splits(q, k_cache, selected_tokens) <= MAX_SPLITS + and all(t.data_ptr() % 16 == 0 for t in (q, k_cache)) + and all(s % 8 == 0 for s in (*q.stride()[:2], *k_cache.stride()[:3])) + ) + + +class QsaPagedSparseGqaDecode(OpWrapper): + """`qsa_paged_sparse_gqa` without a prefill pair, for decode rows: a split-K pass whose + CTAs load their tokens' K/V in one round trip, then the merge -- and, with + `output_gate`, the gate -- as one launch. Returns a new `[rows, H, D]` tensor in q's + dtype. + + Fusion boundary: the split-K attention, the merge, the optional sigmoid output gate. + The selection, the K/V store and the output projection are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="decode, 1 row at position 8700, gated", spec=dict(pos=(8700,), gated=True)), + Cell( + why="decode, 4 rows, one a short sequence, gated", + spec=dict(pos=(8700, 5000, 77, 3000), gated=True), + ), + Cell( + why="decode, 1 row with half the selection -1", + spec=dict(pos=(8700,), gated=True, holes=True), + ), + ) + + note = """ + Metadata consumed: as `qsa_paged_sparse_gqa` (block table, request indices). Preconditions: + as that entry's, plus: no prefill pair; at most `MAX_ROWS` rows; bf16 q and K/V with a head + dim of 256 and at most 16 query heads a KV head; more than one and at most 64 split-K CTAs + a row and KV head (`splits(...)`: 64 selected tokens a CTA); `output_gate` (when given) + `[rows, H * D]` or `[rows, H, D]` with unit inner stride in q's dtype. `supports(...)` is + the predicate a caller can check. + + PDL: the split-K pass reads q and the block table before its wait and the merge reads the + output gate before its own; the selection and the K/V rows are read after the split-K + pass's wait. So q and the gate must be complete when the grid ahead of this entry releases + it: a launch without PDL releases at its end; a PDL producer of the selection (the + target's top-k) must release only after its own wait. + + Numerics: one softmax over a CTA's 64 tokens (QSA's kernel runs an online softmax over + 16-token tiles) and the merge sums the splits in split order (QSA's merge reduces them as + a tree), so the fp32 values can differ in the last bits and the bf16 output by an ulp on + some elements. + + Not certified: a call without `output_gate` (the target always passes one). + """ + + 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, + output_gate: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return launch( + q, + k_cache, + v_cache, + block_table, + selected_tokens, + request_indices, + tokens_per_block, + softmax_scale, + 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, + output_gate: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """The composition entry's reference with the same operands.""" + return _composition.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=None, + compress_ratio=None, + output_gate=output_gate, + ) + + def is_valid( + 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, + output_gate: Optional[torch.Tensor] = None, + ) -> None: + assert q.ndim == 3 and q.shape[2] == HEAD_DIM, "q [rows, H, 256]" + assert 0 < q.shape[0] <= MAX_ROWS, f"1 .. {MAX_ROWS} rows" + assert q.dtype == torch.bfloat16 and k_cache.dtype == v_cache.dtype == torch.bfloat16, ( + "bf16 q and K/V" + ) + assert q.shape[1] % k_cache.shape[1] == 0 and q.shape[1] // k_cache.shape[1] <= M, ( + "<= 16 q heads a KV head" + ) + assert 1 < splits(q, k_cache, selected_tokens) <= MAX_SPLITS, ( + "65 .. 4096 selected tokens a row" + ) + assert q.stride(2) == 1 and k_cache.stride(3) == 1 and v_cache.stride(3) == 1, ( + "unit inner strides" + ) + assert all(t.data_ptr() % 16 == 0 for t in (q, k_cache, v_cache)), ( + "16-byte aligned q and K/V" + ) + assert all( + s % 8 == 0 for s in (*q.stride()[:2], *k_cache.stride()[:3], *v_cache.stride()[:3]) + ), "q and K/V rows 16-byte aligned" + if output_gate is not None: + assert output_gate.dtype == q.dtype and output_gate.stride(-1) == 1, "gate in q's dtype" + assert output_gate.numel() == q.numel(), "gate covers q" + + +qsa_paged_sparse_gqa_decode = QsaPagedSparseGqaDecode() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_sparse_gqa_nvfp4.py new file mode 100644 index 000000000000..6909d8cde45d --- /dev/null +++ b/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 @@ -0,0 +1,1028 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA sparse GQA over an NVFP4 paged pool: attention over explicitly selected tokens.""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# Approved exception (as for `qsa_paged_sparse_gqa`): QSA's sparse path has no torch.ops +# surface; its split-K buckets and row-count PDL policy are reached as Python. +from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import ( + _qsa_pdl_enabled, + _qsa_splitk_launch_config, +) + +#: Base-2 online softmax: scores are scaled by log2(e) once. +_LOG2_E = tl.constexpr(1.4426950408889634) +#: Values per e4m3 block scale, e2m1 values per byte, tokens per V scale group. +_SF_VEC = 16 +_PER_BYTE = 2 +_V_SF_QUAD = 4 +#: e2m1 codes 0..7 -> magnitude; bit 3 is the sign. +_E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + +#: This checkpoint's geometry and the engine's QSA table. `global_k` / `global_v` are +#: `kv_scale_quant_orig[1]` / `[2]`; `slot_pages` role pages of one lifecycle slot hold +#: the call's planes, layer `layer`'s K and V among them (the trunk pool holds 24). +_CHECKPOINT = dict( + q_heads=24, + kv_heads=2, + head_dim=256, + tokens_per_block=32, + compress_ratio=4, + token_topk=2048, + q_dtype=torch.bfloat16, + # 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: 260 pages per slot, -1 past each sequence. + pages=260, + global_k=1.0, + global_v=1.0, + slot_pages=2, + layer=0, +) + +_SPLIT_K = ("splitk", "merge") +_PREFILL = ("prefill",) + + +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; `pool_pages` fixes the pool's page count (the requests then sit + on its last pages); `capture` sweeps the decode site under a CUDA graph. + """ + 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}) + + +@triton.jit +def _e2m1_to_f32(code): + """e2m1 code (0..15, bit 3 the sign) -> its value, exactly.""" + code = code.to(tl.int32) + magnitude_code = code & 7 + exponent = magnitude_code >> 1 + mantissa = (magnitude_code & 1).to(tl.float32) + magnitude = tl.where( + exponent == 0, + 0.5 * mantissa, + (2.0 + mantissa) * (1 << exponent).to(tl.float32) * 0.25, + ) + return tl.where((code & 8) != 0, -magnitude, magnitude) + + +@triton.jit +def _load_keys( + k_cache, + k_scale_cache, + pages, + kv_head, + tokens, + valid, + dims, + k_global, + k_stride_page, + k_stride_head, + k_stride_token, + ks_stride_page, + ks_stride_head, + SCALE_GROUPS: tl.constexpr, +): + """K for the selected tokens as an fp32 [HEAD_DIM, BLOCK_N] tile: e2m1 x scale x global.""" + packed = tl.load( + k_cache + + pages[None, :] * k_stride_page + + kv_head * k_stride_head + + tokens[None, :] * k_stride_token + + dims[:, None] // 2, + mask=valid[None, :], + other=0, + ).to(tl.int32) + codes = (packed >> ((dims[:, None] % 2) * 4)) & 0xF + scales = ( + tl.load( + k_scale_cache + + pages[None, :] * ks_stride_page + + kv_head * ks_stride_head + + tokens[None, :] * SCALE_GROUPS + + dims[:, None] // 16, + mask=valid[None, :], + other=0, + ) + .to(tl.float8e4nv, bitcast=True) + .to(tl.float32) + ) + return _e2m1_to_f32(codes) * scales * k_global + + +@triton.jit +def _load_values( + v_cache, + v_scale_cache, + pages, + kv_head, + tokens, + valid, + dims, + v_global, + v_stride_page, + v_stride_head, + v_stride_token, + vs_stride_page, + vs_stride_head, + SCALE_GROUPS: tl.constexpr, +): + """V for the selected tokens as an fp32 [BLOCK_N, HEAD_DIM] tile; V's scales sit in + four-token groups, `[tpb / 4, D / 16, 4]` per head page.""" + packed = tl.load( + v_cache + + pages[:, None] * v_stride_page + + kv_head * v_stride_head + + tokens[:, None] * v_stride_token + + dims[None, :] // 2, + mask=valid[:, None], + other=0, + ).to(tl.int32) + codes = (packed >> ((dims[None, :] % 2) * 4)) & 0xF + scales = ( + tl.load( + v_scale_cache + + pages[:, None] * vs_stride_page + + kv_head * vs_stride_head + + (tokens[:, None] // 4) * (4 * SCALE_GROUPS) + + (dims[None, :] // 16) * 4 + + tokens[:, None] % 4, + mask=valid[:, None], + other=0, + ) + .to(tl.float8e4nv, bitcast=True) + .to(tl.float32) + ) + return _e2m1_to_f32(codes) * scales * v_global + + +@triton.jit +def _nvfp4_sparse_gqa_prefill_kernel( + q, + k_cache, + v_cache, + k_scale_cache, + v_scale_cache, + kv_scale_quant_orig, + block_table, + selected_tokens, + request_indices, + query_positions, + output, + softmax_scale, + num_cache_pages, + num_requests, + page_table_width, + q_stride_row, + q_stride_head, + k_stride_page, + k_stride_head, + k_stride_token, + v_stride_page, + v_stride_head, + v_stride_token, + ks_stride_page, + ks_stride_head, + vs_stride_page, + vs_stride_head, + block_table_stride_request, + selected_stride_row, + output_stride_row, + output_stride_head, + GROUP_SIZE: tl.constexpr, + TOKENS_PER_BLOCK: tl.constexpr, + TOPK: tl.constexpr, + BLOCK_TOPK: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + HEAD_DIM: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + kv_head = tl.program_id(1) + 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) + k_global = tl.load(kv_scale_quant_orig + 1) + v_global = tl.load(kv_scale_quant_orig + 2) + + head_offsets = tl.arange(0, BLOCK_M) + dim_offsets = tl.arange(0, HEAD_DIM) + query_heads = kv_head * GROUP_SIZE + head_offsets + query_values = tl.load( + q + row * q_stride_row + query_heads[:, None] * q_stride_head + dim_offsets[None, :], + mask=(head_offsets < GROUP_SIZE)[:, None], + other=0.0, + ) + + running_max = tl.full([BLOCK_M], -float("inf"), tl.float32) + running_sum = tl.zeros([BLOCK_M], tl.float32) + accumulator = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32) + token_offsets = tl.arange(0, BLOCK_N) + visible_tokens = tl.load(query_positions + row).to(tl.int64) + 1 + selected_blocks = tl.minimum(visible_tokens // COMPRESS_RATIO, BLOCK_TOPK) + valid_tokens = selected_blocks * COMPRESS_RATIO + visible_tokens % COMPRESS_RATIO + valid_tiles = tl.cdiv(tl.minimum(valid_tokens, TOPK), BLOCK_N) + + for tile in tl.range(0, valid_tiles, num_stages=2): + selected_columns = tile * BLOCK_N + token_offsets + logical_tokens = tl.load( + selected_tokens + row * selected_stride_row + selected_columns, + mask=selected_columns < TOPK, + other=-1, + ) + safe_tokens = tl.maximum(logical_tokens, 0) + logical_pages = safe_tokens // TOKENS_PER_BLOCK + token_in_page = safe_tokens % TOKENS_PER_BLOCK + valid = ( + valid_request + & (selected_columns < TOPK) + & (logical_tokens >= 0) + & (logical_pages < page_table_width) + ) + physical_pages = tl.load( + block_table + + safe_request * block_table_stride_request + + tl.minimum(logical_pages, page_table_width - 1), + mask=valid, + other=-1, + ).to(tl.int64) + valid &= (physical_pages >= 0) & (physical_pages < num_cache_pages) + safe_pages = tl.maximum(physical_pages, 0) + + keys = _load_keys( + k_cache, + k_scale_cache, + safe_pages, + kv_head, + token_in_page, + valid, + dim_offsets, + k_global, + k_stride_page, + k_stride_head, + k_stride_token, + ks_stride_page, + ks_stride_head, + HEAD_DIM // 16, + ).to(query_values.dtype) + values = _load_values( + v_cache, + v_scale_cache, + safe_pages, + kv_head, + token_in_page, + valid, + dim_offsets, + v_global, + v_stride_page, + v_stride_head, + v_stride_token, + vs_stride_page, + vs_stride_head, + HEAD_DIM // 16, + ).to(query_values.dtype) + scores = tl.dot(query_values, keys) * softmax_scale * _LOG2_E + scores = tl.where(valid[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( + valid[None, :], + tl.math.exp2(scores - next_max[:, None]), + 0.0, + ) + accumulator = tl.dot( + probabilities.to(values.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, :], + normalized, + mask=(head_offsets < GROUP_SIZE)[:, None], + ) + + +@triton.jit +def _nvfp4_sparse_gqa_splitk_kernel( + q, + k_cache, + v_cache, + k_scale_cache, + v_scale_cache, + kv_scale_quant_orig, + block_table, + selected_tokens, + request_indices, + partial_output, + partial_lse, + output, + softmax_scale, + num_rows, + num_cache_pages, + num_requests, + page_table_width, + q_stride_row, + q_stride_head, + k_stride_page, + k_stride_head, + k_stride_token, + v_stride_page, + v_stride_head, + v_stride_token, + ks_stride_page, + ks_stride_head, + vs_stride_page, + vs_stride_head, + block_table_stride_request, + selected_stride_row, + output_stride_row, + output_stride_head, + GROUP_SIZE: tl.constexpr, + NUM_QUERY_HEADS: tl.constexpr, + TOKENS_PER_BLOCK: tl.constexpr, + TOPK: tl.constexpr, + NUM_SPLITS: tl.constexpr, + NUM_TILES: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + HEAD_DIM: tl.constexpr, + USE_PDL: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + kv_head = tl.program_id(1) + split = tl.program_id(2) + if USE_PDL: + tl.extra.cuda.gdc_wait() + 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) + k_global = tl.load(kv_scale_quant_orig + 1) + v_global = tl.load(kv_scale_quant_orig + 2) + + head_offsets = tl.arange(0, BLOCK_M) + dim_offsets = tl.arange(0, HEAD_DIM) + token_offsets = tl.arange(0, BLOCK_N) + query_heads = kv_head * GROUP_SIZE + head_offsets + query_values = tl.load( + q + row * q_stride_row + query_heads[:, None] * q_stride_head + dim_offsets[None, :], + mask=(head_offsets < GROUP_SIZE)[:, None], + other=0.0, + ) + + running_max = tl.full([BLOCK_M], -float("inf"), tl.float32) + running_sum = tl.zeros([BLOCK_M], tl.float32) + accumulator = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32) + split_tile_start = split * NUM_TILES // NUM_SPLITS + split_tile_end = (split + 1) * NUM_TILES // NUM_SPLITS + + for tile in range(split_tile_start, split_tile_end): + selected_columns = tile * BLOCK_N + token_offsets + logical_tokens = tl.load( + selected_tokens + row * selected_stride_row + selected_columns, + mask=selected_columns < TOPK, + other=-1, + ) + safe_tokens = tl.maximum(logical_tokens, 0) + logical_pages = safe_tokens // TOKENS_PER_BLOCK + token_in_page = safe_tokens % TOKENS_PER_BLOCK + valid = ( + valid_request + & (selected_columns < TOPK) + & (logical_tokens >= 0) + & (logical_pages < page_table_width) + ) + physical_pages = tl.load( + block_table + + safe_request * block_table_stride_request + + tl.minimum(logical_pages, page_table_width - 1), + mask=valid, + other=-1, + ).to(tl.int64) + valid &= (physical_pages >= 0) & (physical_pages < num_cache_pages) + safe_pages = tl.maximum(physical_pages, 0) + + keys = _load_keys( + k_cache, + k_scale_cache, + safe_pages, + kv_head, + token_in_page, + valid, + dim_offsets, + k_global, + k_stride_page, + k_stride_head, + k_stride_token, + ks_stride_page, + ks_stride_head, + HEAD_DIM // 16, + ).to(query_values.dtype) + values = _load_values( + v_cache, + v_scale_cache, + safe_pages, + kv_head, + token_in_page, + valid, + dim_offsets, + v_global, + v_stride_page, + v_stride_head, + v_stride_token, + vs_stride_page, + vs_stride_head, + HEAD_DIM // 16, + ).to(query_values.dtype) + scores = tl.dot(query_values, keys) + scores *= softmax_scale * _LOG2_E + scores = tl.where(valid[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( + valid[None, :], + tl.math.exp2(scores - next_max[:, None]), + 0.0, + ) + accumulator = tl.dot( + probabilities.to(values.dtype), + values, + accumulator * correction[:, None], + ) + running_sum = running_sum * correction + tl.sum(probabilities, axis=1) + running_max = next_max + + has_values = running_sum > 0 + normalized = tl.where( + has_values[:, None], + accumulator / tl.maximum(running_sum[:, None], 1.0e-20), + 0.0, + ) + output_mask = (head_offsets < GROUP_SIZE)[:, None] + if NUM_SPLITS == 1: + tl.store( + output + + row * output_stride_row + + query_heads[:, None] * output_stride_head + + dim_offsets[None, :], + normalized, + mask=output_mask, + ) + else: + partial_lse_values = tl.where( + has_values, + running_max + tl.math.log2(tl.maximum(running_sum, 1.0e-20)), + -float("inf"), + ) + tl.store( + partial_output + + ((split * num_rows + row) * NUM_QUERY_HEADS + query_heads[:, None]) * HEAD_DIM + + dim_offsets[None, :], + normalized, + mask=output_mask, + ) + tl.store( + partial_lse + (split * num_rows + row) * NUM_QUERY_HEADS + query_heads, + partial_lse_values, + mask=head_offsets < GROUP_SIZE, + ) + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _nvfp4_merge_splitk_kernel( + partial_output, + partial_lse, + output, + num_rows, + output_stride_row, + output_stride_head, + NUM_QUERY_HEADS: tl.constexpr, + NUM_SPLITS: tl.constexpr, + BLOCK_SPLITS: tl.constexpr, + HEAD_DIM: tl.constexpr, + USE_PDL: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + head = tl.program_id(1) + if USE_PDL: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + split_offsets = tl.arange(0, BLOCK_SPLITS) + dim_offsets = tl.arange(0, HEAD_DIM) + split_mask = split_offsets < NUM_SPLITS + lse = tl.load( + partial_lse + (split_offsets * num_rows + row) * NUM_QUERY_HEADS + head, + mask=split_mask, + other=-float("inf"), + ) + max_lse = tl.max(lse, axis=0) + has_values = max_lse > -float("inf") + weights = tl.math.exp2(tl.where(split_mask & has_values, lse - max_lse, -float("inf"))) + denominator = tl.sum(weights, axis=0) + partials = tl.load( + partial_output + + ((split_offsets[:, None] * num_rows + row) * NUM_QUERY_HEADS + head) * HEAD_DIM + + dim_offsets[None, :], + mask=split_mask[:, None], + other=0.0, + ) + merged = tl.sum(partials * weights[:, None], axis=0) + merged = tl.where(denominator > 0, merged / denominator, 0.0) + tl.store(output + row * output_stride_row + head * output_stride_head + dim_offsets, merged) + + +def _sparse_gqa_nvfp4( + *, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_scale_cache: torch.Tensor, + v_scale_cache: torch.Tensor, + kv_scale_quant_orig: 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, +) -> torch.Tensor: + """Sparse GQA over the selected tokens of an NVFP4 HND paged pool (see the entry).""" + if q.ndim != 3: + raise ValueError("NVFP4 sparse GQA: q must be [rows, query_heads, head_dim]") + rows, num_q_heads, head_dim = q.shape + if q.dtype not in (torch.bfloat16, torch.float16) or q.stride(2) != 1: + raise ValueError("NVFP4 sparse GQA: q must be BF16 or FP16 with a unit last-dim stride") + if head_dim < 32 or head_dim & (head_dim - 1): + raise ValueError("NVFP4 sparse GQA: head_dim must be a power of two of at least 32") + if k_cache.ndim != 4 or v_cache.shape != k_cache.shape: + raise ValueError("NVFP4 sparse GQA: k_cache and v_cache must be matching 4-D planes") + num_pages, num_kv_heads = k_cache.shape[0], k_cache.shape[1] + if num_kv_heads <= 0 or num_q_heads % num_kv_heads: + raise ValueError("NVFP4 sparse GQA: query heads must be a multiple of KV heads") + if tokens_per_block <= 0 or tokens_per_block % _V_SF_QUAD: + raise ValueError("NVFP4 sparse GQA: tokens_per_block must be a positive multiple of 4") + data_shape = (num_pages, num_kv_heads, tokens_per_block, head_dim // _PER_BYTE) + scale_shape = (num_pages, num_kv_heads, tokens_per_block, head_dim // _SF_VEC) + for name, plane, shape in ( + ("k_cache", k_cache, data_shape), + ("v_cache", v_cache, data_shape), + ("k_scale_cache", k_scale_cache, scale_shape), + ("v_scale_cache", v_scale_cache, scale_shape), + ): + if plane.dtype != torch.uint8 or tuple(plane.shape) != shape or plane.stride(3) != 1: + raise ValueError( + f"NVFP4 sparse GQA: {name} must be uint8 {shape} with a unit last-dim stride, " + f"got {plane.dtype} {tuple(plane.shape)}" + ) + for name, plane in (("k_scale_cache", k_scale_cache), ("v_scale_cache", v_scale_cache)): + if plane.stride(2) != head_dim // _SF_VEC: + raise ValueError(f"NVFP4 sparse GQA: each head page of {name} must be contiguous") + if ( + kv_scale_quant_orig.dtype != torch.float32 + or kv_scale_quant_orig.shape != (3,) + or not kv_scale_quant_orig.is_contiguous() + ): + raise ValueError("NVFP4 sparse GQA: kv_scale_quant_orig must be a contiguous fp32 [3]") + if block_table.ndim != 2 or block_table.dtype != torch.int32 or block_table.stride(1) != 1: + raise ValueError("NVFP4 sparse GQA: block_table must be a row-contiguous 2-D int32 tensor") + if block_table.shape[0] <= 0 or block_table.shape[1] <= 0 or num_pages <= 0: + raise ValueError("NVFP4 sparse GQA: the pool and the table must be non-empty") + if ( + selected_tokens.ndim != 2 + or selected_tokens.dtype != torch.int32 + or selected_tokens.shape[0] != rows + or selected_tokens.shape[1] <= 0 + or selected_tokens.stride(1) != 1 + ): + raise ValueError( + "NVFP4 sparse GQA: selected_tokens must be a row-contiguous int32 [rows, width]" + ) + if ( + request_indices.shape != (rows,) + or request_indices.dtype not in (torch.int32, torch.int64) + or not request_indices.is_contiguous() + ): + raise ValueError("NVFP4 sparse GQA: request_indices must be one contiguous int per row") + if (query_positions is None) != (compress_ratio is None): + raise ValueError( + "NVFP4 sparse GQA: pass both query_positions and compress_ratio, or neither" + ) + if query_positions is not None and ( + query_positions.shape != (rows,) + or query_positions.dtype not in (torch.int32, torch.int64) + or not query_positions.is_contiguous() + or compress_ratio <= 0 + ): + raise ValueError( + "NVFP4 sparse GQA: query_positions must be one contiguous int per row and " + "compress_ratio positive" + ) + if not (softmax_scale > 0 and softmax_scale < float("inf")): + raise ValueError("NVFP4 sparse GQA: softmax_scale must be finite and positive") + operands = (q, k_cache, v_cache, k_scale_cache, v_scale_cache, kv_scale_quant_orig) + operands += (block_table, selected_tokens, request_indices) + operands += () if query_positions is None else (query_positions,) + if not all(t.is_cuda and t.device == q.device for t in operands): + raise ValueError("NVFP4 sparse GQA: every operand must be on q's CUDA device") + + group_size = num_q_heads // num_kv_heads + block_m = max(16, triton.next_power_of_2(group_size)) + output = torch.empty_like(q, memory_format=torch.contiguous_format) + if rows == 0: + return output + plane_strides = ( + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + k_scale_cache.stride(0), + k_scale_cache.stride(1), + v_scale_cache.stride(0), + v_scale_cache.stride(1), + ) + if query_positions is not None and rows * num_kv_heads > 512: + final_topk = selected_tokens.shape[1] + token_topk = final_topk - compress_ratio + 1 + if token_topk <= 0 or token_topk % compress_ratio: + raise ValueError( + "NVFP4 sparse GQA: the selection width does not come from this compression ratio" + ) + _nvfp4_sparse_gqa_prefill_kernel[(rows, num_kv_heads)]( + q, + k_cache, + v_cache, + k_scale_cache, + v_scale_cache, + kv_scale_quant_orig, + block_table, + selected_tokens, + request_indices, + query_positions, + output, + softmax_scale, + num_pages, + block_table.shape[0], + block_table.shape[1], + q.stride(0), + q.stride(1), + *plane_strides, + block_table.stride(0), + selected_tokens.stride(0), + output.stride(0), + output.stride(1), + GROUP_SIZE=group_size, + TOKENS_PER_BLOCK=tokens_per_block, + TOPK=final_topk, + BLOCK_TOPK=token_topk // compress_ratio, + COMPRESS_RATIO=compress_ratio, + BLOCK_M=block_m, + BLOCK_N=64, + HEAD_DIM=head_dim, + num_warps=2, + num_stages=2, + ) + return output + + block_n, num_splits, num_warps = _qsa_splitk_launch_config( + rows, num_kv_heads, selected_tokens.shape[1] + ) + num_tiles = triton.cdiv(selected_tokens.shape[1], block_n) + use_pdl = num_splits > 1 and _qsa_pdl_enabled(rows) + if num_splits == 1: + partial_output = partial_lse = output + else: + partial_output = torch.empty((num_splits, *q.shape), dtype=torch.float32, device=q.device) + partial_lse = torch.empty( + (num_splits, rows, num_q_heads), dtype=torch.float32, device=q.device + ) + _nvfp4_sparse_gqa_splitk_kernel[(rows, num_kv_heads, num_splits)]( + q, + k_cache, + v_cache, + k_scale_cache, + v_scale_cache, + kv_scale_quant_orig, + block_table, + selected_tokens, + request_indices, + partial_output, + partial_lse, + output, + softmax_scale, + rows, + num_pages, + block_table.shape[0], + block_table.shape[1], + q.stride(0), + q.stride(1), + *plane_strides, + block_table.stride(0), + selected_tokens.stride(0), + output.stride(0), + output.stride(1), + GROUP_SIZE=group_size, + NUM_QUERY_HEADS=num_q_heads, + TOKENS_PER_BLOCK=tokens_per_block, + TOPK=selected_tokens.shape[1], + NUM_SPLITS=num_splits, + NUM_TILES=num_tiles, + BLOCK_M=block_m, + BLOCK_N=block_n, + HEAD_DIM=head_dim, + USE_PDL=use_pdl, + num_warps=num_warps, + num_stages=2, + launch_pdl=use_pdl, + ) + if num_splits == 1: + return output + _nvfp4_merge_splitk_kernel[(rows, num_q_heads)]( + partial_output, + partial_lse, + output, + rows, + output.stride(0), + output.stride(1), + NUM_QUERY_HEADS=num_q_heads, + NUM_SPLITS=num_splits, + BLOCK_SPLITS=triton.next_power_of_2(num_splits), + HEAD_DIM=head_dim, + USE_PDL=use_pdl, + num_warps=2, + num_stages=1, + launch_pdl=use_pdl, + ) + return output + + +def _dequantize_tokens(data, scales, pages, slots, global_scale, grouped: bool) -> torch.Tensor: + """[..., Hkv, D] fp32 values of the tokens at (pages, slots) of one NVFP4 role plane.""" + heads, tokens_per_block, packed = data.shape[1], data.shape[2], data.shape[3] + head_dim = packed * _PER_BYTE + groups = head_dim // _SF_VEC + raw = data[pages, :, slots, :].long() # [..., H, D/2] + lut = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=data.device) + lut = torch.cat([lut, -lut]) + values = torch.stack([lut[raw & 0xF], lut[raw >> 4]], dim=-1).flatten(-2) + if grouped: + flat = scales.reshape(scales.shape[0], heads, tokens_per_block * groups) + base = (slots // _V_SF_QUAD) * (_V_SF_QUAD * groups) + slots % _V_SF_QUAD + columns = base[..., None] + _V_SF_QUAD * torch.arange(groups, device=data.device) + sf = torch.stack([flat[pages[..., None], h, columns] for h in range(heads)], dim=-2) + else: + sf = scales[pages, :, slots, :] + sf = sf.view(torch.float8_e4m3fn).float().repeat_interleave(_SF_VEC, dim=-1) + return values * sf * global_scale + + +class QsaPagedSparseGqaNvfp4(OpWrapper): + """Sparse GQA over the selected tokens of an NVFP4 HND paged pool. + + Triton kernels defined in this file. For row `r`, query head `h`, KV head + `h // (Hq // Hkv)` (contiguous groups), 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] + K[t] = bf16(e2m1(k_cache) * e4m3(k_scale_cache) * kv_scale_quant_orig[1]) + V[t] = bf16(e2m1(v_cache) * e4m3(v_scale_cache) * kv_scale_quant_orig[2]) + s_j = (q[r, h] . K[t]) * softmax_scale + out[r, h] = sum_j softmax(s)_j * V[t] + + The pool layout is the attention op's own NVFP4 append: e2m1 nibbles per + head page `[tpb, D/2]` (low nibble = even element), K's e4m3 scales + token-major `[tpb, D/16]`, V's in four-token groups `[tpb/4, D/16, 4]`. + Each value is dequantized in fp32 and rounded once to `q`'s dtype; at unit + second-level scales every NVFP4 value is exact in bf16, so a call is the + bf16 sparse GQA over the bf16 values the pool holds. After the dequantize the + kernels are QSA's bf16 sparse GQA kernels operation for operation: the same + tiles, online softmax and split-K merge. + + 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 contribute + nothing; a row left with nothing returns zeros. Returns a fresh, contiguous + `[rows, Hq, D]` tensor of `q.dtype`. + + 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. **The prefill pair is a bound, not a mask**, exactly as on the bf16 + entry: the prefill kernel reads only `min(visible // cr, block_topk) * cr + + visible % cr` leading columns of a row. + + Fusion boundary: the block-table gather, the NVFP4 dequantize, scale, + online softmax, P.V and the split-K merge are inside. Selection, appending + K/V to the pool, and q's projection, norm, RoPE and output gate are not. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode: 4 requests x 4 rows, 8-way split-K + merge, in the engine's trunk " + "layout: layer 11's planes among the 24 role pages of a slot", + spec=_decode( + tuple((4000 + 37 * i, 4) for i in range(4)), slot_pages=24, layer=11, pool_pages=600 + ), + ), + Cell( + why="context 4096 rows from 0: rows see 1..4096 tokens", + spec=_decode(((0, 4096),), launches=_PREFILL, prefill_pair=True), + ), + Cell( + 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))), + ), + ) + + note = """ + Metadata consumed: none -- every input is an argument. What the target + passes: `block_table` is `qsa_block_table` (physical page ids, the unscaled + table); `selected_tokens` is `qsa_expand_block_indices`' output, logical + positions, `-1`-padded; `k_cache` / `v_cache` and `k_scale_cache` / + `v_scale_cache` are the layer's planes of the QSA manager's `get_buffers(i, + "HND")` and `get_block_scale_buffers(i, "HND")`; `kv_scale_quant_orig` is the + fp32 [3] tensor its dense attention call binds (index 1 for K, 2 for V). + + 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). + + Launch shapes at Hkv 2 (programs = rows x 2), as on the bf16 entry: split-K + 64-way up to 2 rows, 32-way up to 15 (16-token tiles), 8-way up to 128 and + 4-way up to 256 (64-token tiles), each followed by the merge; past 256 rows + the prefill kernel when the pair is passed, else one split and no merge. + The decode site reaches 256 rows at max_draft_len 7 (32 requests x 8). + + Equivalence: the test also runs QSA's bf16 sparse GQA over a bf16 pool + holding the dequantized values; on every cell but the captured one the two + outputs agree within the test's band. They are not bitwise equal: the + dequantizing load changes the compiled kernel around the same arithmetic. + + CUDA graphs: the decode site captures at G requests x 1..8 rows over the + [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 + 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. + + Preconditions: + - Whenever the prefill kernel can run, each selection row is packed at the + front with exactly the bound's count of valid entries (QSA's producer + does this; see the bf16 entry's note). + - The pool holds what the attention op's NVFP4 append writes; the + `qsa_paged_kv_store_nvfp4` entry writes the same bits. + - `D` a power of two >= 32; `Hq % Hkv == 0`; `tokens_per_block % 4 == 0`. + + The callable rejects loudly (ValueError) every operand whose dtype, shape or + layout it cannot read (the test pins the half-passed prefill pair); `rows == + 0` returns an empty `[0, Hq, D]` tensor (not certified). Calls are bitwise + deterministic. + + 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. + """ + + def __call__( + self, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_scale_cache: torch.Tensor, + v_scale_cache: torch.Tensor, + kv_scale_quant_orig: 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, + ) -> torch.Tensor: + return _sparse_gqa_nvfp4( + q=q, + k_cache=k_cache, + v_cache=v_cache, + k_scale_cache=k_scale_cache, + v_scale_cache=v_scale_cache, + kv_scale_quant_orig=kv_scale_quant_orig, + 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, + k_scale_cache: torch.Tensor, + v_scale_cache: torch.Tensor, + kv_scale_quant_orig: 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, + ) -> torch.Tensor: + """fp32 sparse GQA over every resident selected token, rounded once. + + K and V are dequantized in fp32 and rounded to `q`'s dtype, as the kernel + does. `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 = _dequantize_tokens( + k_cache, k_scale_cache, physical, in_page, kv_scale_quant_orig[1], grouped=False + ) + values = _dequantize_tokens( + v_cache, v_scale_cache, physical, in_page, kv_scale_quant_orig[2], grouped=True + ) + keys = keys.to(q.dtype).float() + values = values.to(q.dtype).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")) + 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) + return out.to(q.dtype) + + +qsa_paged_sparse_gqa_nvfp4 = QsaPagedSparseGqaNvfp4() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py index ced33a120477..415956a05579 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py @@ -15,7 +15,8 @@ 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 for when it flips. +#: (`_QSA_MROPE_SECTION`) and these are certified, on the NVFP4 trunk-plane cell, +#: for when it flips. MROPE_SECTION = (11, 11, 10) #: This checkpoint's index geometry and the engine's QSA metadata: index head @@ -31,12 +32,12 @@ mrope_section=None, logical_dtype=torch.int64, request_dtype=torch.int32, - distinct_axes=True, layout="plain", - duplicate=False, ) -_DECODE = ((4095, 1), (4096, 1), (4097, 1), (4098, 1), (4099, 1), (8191, 1), (3, 1), (2047, 1)) +#: 32 requests of one MTP7 verify block each (8 rows, two completed groups), at starts that put +#: some blocks across a page edge and some inside one. +_MTP7_32 = tuple((29 + 61 * i, 8) for i in range(32)) def _spec(segments, pages, **overrides) -> dict: @@ -44,8 +45,8 @@ def _spec(segments, pages, **overrides) -> dict: return {**_CKPT, "segments": segments, "pages": pages, **overrides} -def _old(segments, pages, **overrides) -> dict: - """The old receipt's scenes: prefill from 0 under mRoPE (11, 11, 10).""" +def _mrope(segments, pages, **overrides) -> dict: + """Prefill from 0 under mRoPE (11, 11, 10).""" return _spec(segments, pages, mrope_section=MROPE_SECTION, **overrides) @@ -73,8 +74,7 @@ class QsaPrefillCompress(OpWrapper): 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). Two rows naming one - anchor in a launch are harmless. + (measured: all 48 of 48 anchors moved, up to 2.94). Fusion boundary: pool, Gemma norm, position gather, rotation, store. The raw-key and position stores (`qsa_paged_index_store`, which must run @@ -84,69 +84,21 @@ class QsaPrefillCompress(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - # First on purpose: the process's first, JIT-compiling call. - Cell( - why="prefill 2 x 128 from 0, mRoPE, distinct t/h/w: cold call", - spec=_old(((0, 128),) * 2, pages=8), - ), Cell( why="prefill 2 x 128 from 0, mrope_section None: the target's call", spec=_spec(((0, 128),) * 2, pages=8), ), Cell( - why="prefill 3 x 64, mRoPE: three requests' pages interleaved", - spec=_old(((0, 64),) * 3, pages=6), - ), - Cell( - why="text rows (t == h == w) under mRoPE", - spec=_old(((0, 128),), pages=4, distinct_axes=False), - ), - Cell( - why="2-token pages: every group straddles two pages", - spec=_old(((0, 64),) * 2, pages=64, tokens_per_block=2), - ), - Cell( - why="decode: 8 rows to 8191, anchors and non-anchors, None", - spec=_spec(_DECODE, pages=1024), - ), - Cell(why="decode: the same rows under mRoPE", spec=_old(_DECODE, pages=1024)), - Cell( - why="decode MTP: 4 requests x 2 rows, one anchor each, None", - spec=_spec(((4094, 2), (4095, 2), (1022, 2), (30, 2)), pages=512), - ), - Cell( - why="mixed: chunk from 0, chunk from 1000, 2 MTP decodes, None", - spec=_spec(((0, 64), (1000, 36), (4094, 2), (255, 2)), pages=256), + why="decode MTP7: 32 requests x 8 rows, two anchors each, None", + spec=_spec(_MTP7_32, pages=1100, slots=40, rounding_flips=1), ), - Cell(why="eps 0.0", spec=_old(((0, 64),), pages=2, eps=0.0)), - Cell(why="eps 1e-5", spec=_old(((0, 64),), pages=2, eps=1e-5)), + # Under an NVFP4 KV pool the position page shares each physical page with the index + # pages, so both caches are strided views of one allocation: index page stride + # 12 x 4096 + 192 = 49344 on the trunk, position strides (24672, 3, 1). Cell( - why="int32 positions, int32 requests", - spec=_old(((0, 64),) * 2, pages=4, logical_dtype=torch.int32), + 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), ), - Cell( - why="int64 positions, int64 requests", - spec=_old(((0, 64),) * 2, pages=4, request_dtype=torch.int64), - ), - Cell( - why="int32 positions, int64 requests", - spec=_old( - ((0, 64),) * 2, pages=4, logical_dtype=torch.int32, request_dtype=torch.int64 - ), - ), - Cell( - why="page-sliced cache, axis-strided plane, column-strided table", - spec=_old(((0, 64),) * 2, pages=4, layout="strided"), - ), - Cell( - why="index cache with a head-dim stride of 2", - spec=_old(((0, 64),), pages=2, layout="head_dim_strided"), - ), - Cell( - why="every row twice: two rows per anchor in one launch", - spec=_old(((0, 64),) * 2, pages=4, duplicate=True), - ), - Cell(why="zero rows: returns before launching", spec=_old((), pages=2)), ) note = """ @@ -169,7 +121,7 @@ class QsaPrefillCompress(OpWrapper): 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. - The callable rejects loudly: metadata shape/dtype mismatches, a non-int32 + 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 `tokens_per_block`, a malformed cos|sin table, `4 R != D`, a non-2^k `D` @@ -179,16 +131,32 @@ class QsaPrefillCompress(OpWrapper): `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. Plain `@triton.jit`, no - autotune; CUDA-graph capturable at a fixed row count. + 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). + + 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. 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 8191 through a - [32, 260] table over 1024 pages. + boundary inside the call. Certified extent: positions to 1928 through a + [40, 260] table over 1100 pages. Not certified: `compress_ratio` other than 4, `D` other than 128, a bf16 - or fp16 `cos_sin`, positions past the table's rows, multiple devices. + 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, + `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`. """ 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_rowtile_attention_cute.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rowtile_attention_cute.py new file mode 100644 index 000000000000..03e2c925518e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_rowtile_attention_cute.py @@ -0,0 +1,995 @@ +# 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. + +One CTA per (group, KV head), 128 MMA rows = `rows_per_tile` query rows x `group_size` +heads, 9 warps: + +* warps 5-8 gather the list's K and V rows (64 a tile) with `cp.async` into one ring of + `kv_slots` 32 KB SW128 shared-memory slots, filled in the MMA warp's order K0, K1, V0, + K2, V1, ...; each slot completes on its mbarrier through `cp.async.mbarrier.arrive`; +* warp 4 issues the tcgen05 MMAs: S[j % 2] = Q K(j)^T (Q from tensor memory), then + O += P(j - 1) V(j - 1) (P from tensor memory, overwriting the S it came from); +* warps 0-3, one thread per MMA row, first copy their Q row into tensor memory, then 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, and at the end + normalize O and store it through the idle ring as whole 512-byte rows. + +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 + + +@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 = 6, + 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 + 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.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, + 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, + 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] + tmem_holding: Int32 + sKV: cute.struct.Align[cute.struct.MemRange[BF16, cute.cosize(k_layout)], 1024] + + self.shared_storage = SharedStorage + + self.kernel( + qk_mma, + qk_mma_acc, + pv_mma, + pv_mma_acc, + q_layout, + k_layout, + v_layout, + p_layout, + q, + k, + v, + slots, + bits, + plan, + out, + rows, + 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_groups, num_kv_heads, 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, + q: cute.Pointer, + k: cute.Pointer, + v: cute.Pointer, + slots: cute.Pointer, + bits: cute.Pointer, + plan: cute.Pointer, + out: cute.Pointer, + rows: 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() + group, kv_head, _ = cute.arch.block_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) + + mPlan = cute.make_tensor(plan, cute.make_layout((2 * 1048576,))) + tiles = mPlan[group * 2] + listed = mPlan[group * 2 + 1] + + if tiles >= 0: + 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() + 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_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) + + if warp_idx >= self.load_warp0: + self.load_role( + tidx - 32 * self.load_warp0, + group, + kv_head, + tiles, + listed, + sK, + k, + v, + slots, + 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, + ) + 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, + tiles, + q_full, + kv_full, + kv_empty, + s_full, + p_full, + o_full, + o_final, + ) + 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, + group, + kv_head, + tiles, + rows, + bits, + bits_stride_row, + out, + out_stride_row, + scale_log2, + q_full, + s_full, + p_full, + o_full, + o_final, + storage.sKV.data_ptr(), + q, + q_stride_row, + ) + 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() + + # ------------------------------------------------------------------ loaders (warps 5-8) + @cute.jit + def load_role( + self, + lt: Int32, + group: Int32, + kv_head: Int32, + tiles: Int32, + listed: Int32, + sK, + k: cute.Pointer, + v: cute.Pointer, + slots: cute.Pointer, + 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, + ): + """Fill the K/V 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.""" + 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), + ), + ) + if tiles > 0: + mSlots = cute.make_tensor( + slots + Int64(group) * slots_stride_group, cute.make_layout((1 << 30,)) + ) + first_token = 16 * (lt // 32) + (lt % 32) // 8 + 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)] + ring = Int32(0) + ring_phase = Int32(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] + # 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) + # Tail: every slot released, so no load is in flight when the CTA retires and the ring + # is free for the epilogue's staging. + 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_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, + tiles: Int32, + q_full, + kv_full, + kv_empty, + s_full, + p_full, + o_full, + o_final, + ): + """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 and is picked by pointer offset; QK(j) overwrites S[j % 2] only after PV(j - 2) + was issued (tcgen05 MMAs of one thread execute in order), and P(j - 2) was its last reader.""" + 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] + if tiles > 0: + cute.arch.mbarrier_wait(q_full, 0) + tcgen05_fence_after() + ring = Int32(0) + ring_phase = Int32(0) + # QK(0) + 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, tStS.layout), + ) + with cute.arch.elect_one(): + tcgen05.commit(s_full) + tcgen05.commit(kv_empty + ring) + ring, ring_phase = self.advance(ring, ring_phase) + for j in cutlass.range(1, tiles + 1, 1, unroll=1): + buf = j % 2 + if j < tiles: + # QK(j) into S[j % 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 = j - 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 prev == 0: + 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) + + @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 softmax_role( + self, + qk_thr, + pv_thr, + tmem_ptr, + st: Int32, + group: Int32, + kv_head: Int32, + tiles: 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, + stage, + q: cute.Pointer, + q_stride_row: Int32, + ): + 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) + + # S: one 64-column load a tile; P: 32 packed-bf16 columns stored over the S it came from. + 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 + + m = st + local_row = m // G + head = m % G + row = group * R + local_row + live = (local_row < R) and (row < rows) + 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,)) + ) + + # Q row -> TMEM columns tmem_q.. (the QK MMA's A operand), 64 elements a store. Dead rows + # store zeros; their S rows are masked out. + if tiles > 0: + q_row = cute.make_tensor( + q + cute.assume(Int64(safe_row) * q_stride_row + (kv_head * G + head) * D, divby=8), + cute.make_layout((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): + cute.autovec_copy( + cute.make_tensor( + q_row.iterator + (c * 2 * p_cols + i * 8), 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) + + row_max = -Float32.inf + row_sum = Float32(0.0) + lo_next = Int32(0) + hi_next = Int32(0) + if live and tiles > 0: + lo_next = mBits[0] + hi_next = mBits[1] + tau = Float32(self.tau) + for j in cutlass.range(0, tiles, 1, unroll=1): + 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 = j % 2 + cute.arch.mbarrier_wait(s_full + buf, (j // 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, (j - 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) + + # Epilogue: O / l in bf16, staged through the (now idle) K/V ring -- a row's 32 16-byte + # chunks XOR-swizzled by the row -- so each warp then stores whole 512-byte rows. + if tiles > 0: + cute.arch.mbarrier_wait(o_final, 0) + tcgen05_fence_after() + inv = Float32(0.0) + if row_sum > Float32(0.0): + inv = Float32(1.0) / row_sum + swz = m % 32 + for c in cutlass.range_constexpr(n_chunks): + rO = cute.make_rmem_tensor(tLcO.shape, Float32) + if tiles > 0: + 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 + else: + rO.fill(0.0) + rB = cute.make_rmem_tensor((self.chunk_cols,), BF16) + rB.store(rO.load().to(BF16)) + for i in cutlass.range_constexpr(self.chunk_cols // 8): + logical = c * (self.chunk_cols // 8) + i + offset = cute.assume(m * D + ((logical ^ swz) * 8), divby=8) + dst = cute.make_tensor(stage + offset, cute.make_layout((8,))) + cute.autovec_copy( + cute.make_tensor(rB.iterator + i * 8, cute.make_layout((8,))), dst + ) + cute.arch.sync_warp() + lane = st % 32 + warp = st // 32 + for r in cutlass.range_constexpr(32): + mr = warp * 32 + r + lr = mr // G + row_r = group * R + lr + if lr < R and row_r < rows: + src_off = cute.assume(mr * D + ((lane ^ r) * 8), divby=8) + rV = cute.make_rmem_tensor((8,), BF16) + cute.autovec_copy(cute.make_tensor(stage + src_off, cute.make_layout((8,))), rV) + dst_off = cute.assume( + Int64(row_r) * out_stride_row + (kv_head * G + mr % G) * D + lane * 8, divby=8 + ) + cute.autovec_copy(rV, cute.make_tensor(out + dst_off, cute.make_layout((8,)))) + + +@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, 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), + 1, + 1, + 1, + 8, + 8, + 8, + 8, + 8, + 8, + 8, + 1, + 1, + 8, + 1.0, + cuda_drv.CUstream(0), + ) + return cute.compile(kernel, *args) + + +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 `(groups, kv_heads)`.""" + + 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, + softmax_scale, + rows, + *, + group_size, + rows_per_tile, + tokens_per_block, + use_pdl, + ): + assert grid == (group_plan.shape[0], k_cache.shape[1]), grid + assert ( + output.is_contiguous() + and packed_bits.dtype == torch.int64 + and slots.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), + rows, + grid[0], + grid[1], + 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__tp1/ops/attention/qsa_sparse_gqa_rowtile.py new file mode 100644 index 000000000000..86a4f4a92b8f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_sparse_gqa_rowtile.py @@ -0,0 +1,796 @@ +# 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. A plan launch forms, +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 runs one program per +(group, KV head) that walks the 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.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 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 the plan pass unions and packs per loop step. +_PLAN_WORDS = 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 +_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 _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, + ROWS_PER_TILE: tl.constexpr, NUM_REQUESTS: tl.constexpr, ROWS_POW2: tl.constexpr, + BLOCK_N: tl.constexpr, SHARED_TILE_COST: tl.constexpr, NUM_CACHE_PAGES: tl.constexpr, + PAGE_TABLE_WIDTH: tl.constexpr, TOKENS_PER_BLOCK: tl.constexpr, PLAN_WORDS: tl.constexpr, +): # fmt: skip + """Per group of rows of one request: list the union of its rows' resident + selected tokens, in logical order, as physical cache slots + (`page * tokens_per_block + offset`) in `slots[group]`, and write each + row's bitmap re-indexed onto that list into `packed_bits[row]` (bit `j` + set when the row selects the list's `j`-th token). `group_plan[group]` is + (tiles of the list for the row-tile pass, or -1; list length). A group the + row-tile pass does not take -- rows of several requests, or a list whose + tiles cost more than the rows' own columns -- appends its rows to the lone + list instead.""" + group = tl.program_id(0).to(tl.int64) + 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) + columns = tl.load(row_info + members * 2 + 1, mask=live, other=0) + own_tiles = tl.sum(tl.cdiv(columns, BLOCK_N), axis=0) + one_request = (lowest == highest_request) & (lowest >= 0) & (lowest < NUM_REQUESTS) + words_needed = tl.where(one_request, (highest_token + 64) // 64, 0) + page_base = ( + block_table + + tl.minimum(tl.maximum(lowest, 0), NUM_REQUESTS - 1) * block_table_stride_request + ) + word_offsets = tl.arange(0, PLAN_WORDS) + lanes = tl.arange(0, 64) + lane_bits = tl.full([64], 1, tl.uint64) << lanes.to(tl.uint64) + listed = tl.sum(tl.zeros([ROWS_POW2], tl.int32), axis=0) + for start in range(0, words_needed, PLAN_WORDS): + words = start + word_offsets + 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) + chosen = ((selected[:, None] >> lanes[None, :].to(tl.uint64)) & 1) != 0 + tokens = words[:, None].to(tl.int32) * 64 + lanes[None, :] + logical_pages = tokens // TOKENS_PER_BLOCK + pages = tl.load( + page_base + tl.minimum(logical_pages, PAGE_TABLE_WIDTH - 1) * block_table_stride_page, + mask=chosen & (logical_pages < PAGE_TABLE_WIDTH), + other=-1, + ) # fmt: skip + kept = chosen & (pages >= 0) & (pages < NUM_CACHE_PAGES) + kept_count = kept.to(tl.int32) + count = tl.sum(kept_count, axis=1) + first = listed + tl.cumsum(count, axis=0) - count + rank = first[:, None] + tl.cumsum(kept_count, axis=1) - kept_count + tl.store( + slots + group * slots_stride_group + rank, + pages * TOKENS_PER_BLOCK + tokens % TOKENS_PER_BLOCK, + mask=kept, + ) # fmt: skip + 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, :] + 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) + count > 64)[None, :], sem="relaxed") # fmt: skip + listed += tl.sum(count, axis=0) + tiles = tl.cdiv(listed, BLOCK_N) + shared = one_request & (tiles * SHARED_TILE_COST <= own_tiles * 4) + tl.store(group_plan + group * 2, tl.where(shared, tiles, -1)) + tl.store(group_plan + group * 2 + 1, listed) + if not shared: + 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_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: tl.constexpr, NUM_REQUESTS: tl.constexpr, 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, LONE_PROGRAMS: 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: tl.constexpr, NUM_REQUESTS: tl.constexpr, 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) + 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) + group_plan = torch.empty((num_groups, 2), dtype=torch.int32, device=q.device) + lone_rows = torch.empty((rows,), dtype=torch.int32, device=q.device) + lone_count = torch.zeros((1,), dtype=torch.int32, device=q.device) + _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), + ROWS_PER_TILE=rows_per_tile, NUM_REQUESTS=block_table.shape[0], ROWS_POW2=rows_pow2, + BLOCK_N=BLOCK_N, SHARED_TILE_COST=_SHARED_TILE_COST, NUM_CACHE_PAGES=k_cache.shape[0], + PAGE_TABLE_WIDTH=table_width, TOKENS_PER_BLOCK=tokens_per_block, + PLAN_WORDS=_PLAN_WORDS, num_warps=_PLAN_WARPS, + ) # 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[(num_groups, num_kv_heads)]( + q, k_cache, v_cache, slots, packed_bits, group_plan, output, 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`. + + 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 + 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 + 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 runs one program per + (group, KV head), loads each listed tile's K/V once for the 120 (row, + head) pairs and masks each row to its own re-indexed bits; 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), 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 two dumped 16384-row context launches of the + target (GB300): 631 and 887 us. + + 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, + 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/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 new file mode 100644 index 000000000000..542d4d9a4e70 --- /dev/null +++ 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 @@ -0,0 +1,628 @@ +# 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 + + +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 + if const_expr(self.pdl): + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + + # 1. The row's metadata and four consecutive keys a thread (a column past the + # window is no candidate: -1); the largest and smallest key in the window. + n = mEnd[row] + seqlen = mLen[row] + visible = cutlass.Int32(mPos[row]) + 1 + 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() + if tidx < at_thr: + kt = cand_key[tidx] + it = cand_idx[tidx] + rank = cutlass.Int32(0) + for j in cutlass.range(0, at_thr, 1, unroll=8): + kj = cand_key[j] + ij = cand_idx[j] + better = (kj > kt) | ((kj == kt) & (ij > it)) + rank = rank + cutlass.Int32(cutlass.select_(better, 1, 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), + ), + ) + + 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/qsa_topk_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py index 1903a9cc6f32..fa7e00d6e456 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py @@ -17,6 +17,11 @@ #: 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 +#: 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. +_CUTE_DSL_MIN_ROWS = 12288 class TopkExpectation(NamedTuple): @@ -41,7 +46,9 @@ class QsaTopkSelect(OpWrapper): """Each row's top-`index_topk` window-local column offsets, written into `indices`. Wraps `torch.ops.trtllm.indexer_topk_prefill` (CUDA, compiled into the - wheel). For row `r` it selects the `index_topk` largest values of + wheel) and, for a contiguous launch of `_CUTE_DSL_MIN_ROWS` rows or more, + its CuTe DSL sibling `torch.ops.trtllm.cute_dsl_indexer_topk_prefill_blackwell` + (the same contract). For row `r` it selects the `index_topk` largest values of `logits[r, row_starts[r]:row_ends[r]]` and writes their offsets *relative to `row_starts[r]`* into `indices[r]`: `min(index_topk, window)` distinct offsets as a prefix, `-1` in every remaining slot. Every slot of every row @@ -68,42 +75,34 @@ class QsaTopkSelect(OpWrapper): # window is honoured, not merely that the tail happened to be small. CELLS: tuple[Cell, ...] = ( Cell( - why="decode: one row at position 7777, 2080 cols, k=512, NaN tail", + why="decode MTP7: 32 requests x 8 rows = 256 rows, some below the budget, NaN tail", spec=dict( - rows=1, + rows=256, cols=_SCORE_COLUMNS, index_topk=_BLOCK_TOPK, - windows=("positions", (7777,)), + windows=( + "positions", + tuple(600 + 233 * g + j for g in range(32) for j in range(8)), + ), logits="relu", tail="nan", ), ), Cell( - why="decode MTP: one request's two rows (7777, 7778), NaN tail", + why="context: 8191 rows at positions 0..8190 (warm-up capture), the CUDA op", spec=dict( - rows=2, - cols=_SCORE_COLUMNS, - index_topk=_BLOCK_TOPK, - windows=("positions", (7777, 7778)), - logits="relu", - tail="nan", - ), - ), - Cell( - why="decode MTP x4: two requests below the budget get -1 padding", - spec=dict( - rows=8, + rows=8191, cols=_SCORE_COLUMNS, index_topk=_BLOCK_TOPK, - windows=("positions", (7777, 7778, 4000, 4001, 1500, 1501, 300, 301)), + windows=("chunk", 0), logits="relu", tail="nan", ), ), Cell( - why="context: 8191 rows at positions 0..8190 (warm-up capture)", + why="context: 16384 rows at positions 0..16383, NaN tail (CuTe DSL op)", spec=dict( - rows=8191, + rows=16384, cols=_SCORE_COLUMNS, index_topk=_BLOCK_TOPK, windows=("chunk", 0), @@ -111,84 +110,6 @@ class QsaTopkSelect(OpWrapper): tail="nan", ), ), - Cell( - why="context: 1024 rows from position 4000 over 2048 cols, relu ties", - spec=dict( - rows=1024, - cols=2048, - index_topk=_BLOCK_TOPK, - windows=("chunk", 4000), - logits="relu", - tail=None, - ), - ), - Cell( - why="nonzero row_starts: offsets come back window-local", - spec=dict( - rows=4, - cols=512, - index_topk=32, - windows=("explicit", (0, 37, 100, 480), (512, 300, 140, 512)), - logits="randn", - tail=None, - ), - ), - Cell( - why="strided logits: column stride 2 takes the scalar load path", - spec=dict( - rows=5, - cols=512, - index_topk=32, - windows=("explicit", (0, 11, 0, 64, 7), (512, 400, 33, 512, 8)), - logits="randn", - tail=None, - strided=True, - ), - ), - Cell( - why="windows of 0, 1, k-1, k, 1 and full width: -1 suffix", - spec=dict( - rows=6, - cols=128, - index_topk=16, - windows=("explicit", (0, 0, 0, 5, 9, 0), (0, 1, 15, 21, 10, 128)), - logits="randn", - tail=None, - ), - ), - Cell( - why="budget wider than every window and than the matrix", - spec=dict( - rows=3, - cols=64, - index_topk=_BLOCK_TOPK, - windows=("explicit", (0, 0, 0), (64, 7, 0)), - logits="randn", - tail=None, - ), - ), - Cell( - why="row tier: 12287 rows, all on the insertion-sort variant", - spec=dict( - rows=6 * 2048 - 1, - cols=512, - index_topk=64, - windows=("random_ends",), - logits="randn", - tail=None, - ), - ), - Cell( - why="row tier: 12416 rows, rows >= 12288 take the radix variant", - spec=dict( - rows=6 * 2048 + 128, - cols=512, - index_topk=64, - windows=("random_ends",), - logits="randn", - tail=None, - ), - ), ) note = """ @@ -200,34 +121,36 @@ class QsaTopkSelect(OpWrapper): 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. Outside the window - anything goes -- the cells poison it with NaN. + 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. - `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. - Loud rejections (measured): bf16/fp16 `logits` ("expected scalar type - Float"), int64 `indices` ("expected scalar type Int"), `row_starts` / - `row_ends` shorter than the row count ("one entry per row"). - - Dispatch is by row index, not column count: rows below 12288 run an - insertion-sort variant, rows at or above a radix variant, in one launch. - This target's `max_num_tokens` is 8192, so serving only reaches the first; - the cells cover both. - - QSA's stock indexer picks the CuTe DSL sibling - `torch.ops.trtllm.cute_dsl_indexer_topk_prefill_blackwell` on this GPU; it - was measured to select identical sets (never identical permutations). This - entry wraps the compiled CUDA op so a first captured decode JIT-compiles - nothing. `torch.ops.trtllm.indexer_topk_decode` is the request-grouped + Loud rejections (measured, not pinned by a test): bf16/fp16 `logits` + ("expected scalar type Float"), int64 `indices` ("expected scalar type + Int"), `row_starts` / `row_ends` shorter than the row count ("one entry per + row"). + + 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 + `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. - CUDA graphs: captured at `[8, 2048]`, `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), + `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). Hard limits: not probed. `indices` is addressed as `row * index_topk + - slot`; the largest certified extent is 12416 x 64. + slot`; the largest certified extent is 16384 x 512. """ def __call__( @@ -238,7 +161,12 @@ def __call__( indices: torch.Tensor, index_topk: int = 2048, ) -> torch.Tensor: - torch.ops.trtllm.indexer_topk_prefill(logits, row_starts, row_ends, indices, index_topk) + if logits.shape[0] >= _CUTE_DSL_MIN_ROWS and logits.is_contiguous(): + torch.ops.trtllm.cute_dsl_indexer_topk_prefill_blackwell( + logits, row_starts, row_ends, indices, index_topk + ) + else: + torch.ops.trtllm.indexer_topk_prefill(logits, row_starts, row_ends, indices, index_topk) return indices def reference( 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qwen4_exp_attn_decode_prologue.py new file mode 100644 index 000000000000..854f7efb1d6c --- /dev/null +++ b/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 @@ -0,0 +1,571 @@ +# 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. + +The composition it replaces ran four launches over two streams: the QK norm + rope of +both planes (`qwen4_exp_qk_norm_rope`), the paged K/V append (`qsa_paged_kv_store`), and on +the side stream the index plane's raw-key and coordinate store (`qsa_paged_index_store`) +and the group compression (`qsa_prefill_compress`), then a cross-stream join. Here the +programs of `qwen4_exp_qk_norm_rope` run as they do there; the ones holding k heads also +write their rotated rows into the K pool; one more program per row stores v, the 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 (the cached copy is the one being written), +normalizes, rotates and stores it at the anchor slot, as the compression kernel does. + +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 the projection row. +""" + +from __future__ import annotations + +from typing import Optional, 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 + +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_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 + + +@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 + """`qwen4_exp_qk_norm_rope`'s tile with its static factors given: rms over each head, + `* (1 + w)`, rotate; fp32 throughout, one rounding at the store. 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_factors(pos_t, pos_h, pos_w, cos_sin, cos_sin_stride, ROTARY_PAIRS: tl.constexpr, + IS_MROPE: tl.constexpr, MROPE_H: tl.constexpr, MROPE_W: tl.constexpr, live): # fmt: skip + """The compressed key's rotation factors (QSA's `_qsa_gemma_norm_rope` reads them at + the group's first-token coordinates).""" + pairs = tl.arange(0, ROTARY_PAIRS) + if IS_MROPE: + height_pair = (pairs % 3 == 1) & (pairs < 3 * MROPE_H) + width_pair = (pairs % 3 == 2) & (pairs < 3 * MROPE_W) + position = tl.where(height_pair, pos_h, tl.where(width_pair, pos_w, pos_t)) + else: + position = pos_t + cosine = tl.load(cos_sin + position * cos_sin_stride + pairs, mask=live, other=0.0) + sine = tl.load(cos_sin + position * cos_sin_stride + ROTARY_PAIRS + pairs, mask=live, other=0.0) + return cosine, sine + + +@triton.jit +def _compress_norm_rope( + x, weight, cosine, sine, eps, HEAD_DIM: tl.constexpr, ROTARY_PAIRS: tl.constexpr +): + """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, + 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, + IS1: tl.constexpr, + IS2: tl.constexpr, + HPB: tl.constexpr, + MAIN_PROGRAMS: tl.constexpr, + TPB: tl.constexpr, + CR: tl.constexpr, + ROTARY_PAIRS: tl.constexpr, + IS_MROPE: tl.constexpr, + MROPE_H: tl.constexpr, + MROPE_W: tl.constexpr, + USE_PDL: tl.constexpr, +): + token = tl.program_id(0).to(tl.int64) + program = tl.program_id(1) + if USE_PDL: + # The index scores read nothing of this launch before their own wait. + tl.extra.cuda.gdc_launch_dependents() + row = proj + token * stride_proj + + # 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() + out = _norm_rope_apply( + row, row, main_heads, main_heads < HQ + HK, w2, wp2, c, s, eps, D, ROT, D + ) + if program * HPB + HPB > HQ: + d = tl.arange(0, D) + is_k = (main_heads >= HQ) & (main_heads < HQ + HK) & live + tl.store( + k_cache + + page * s_kc_page + + (main_heads - HQ).to(tl.int64)[:, None] * s_kc_head + + slot * s_kc_tok + + d[None, :], + out.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, IS1, IS2) + 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 + 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 coordinates, their + # 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) + group = tl.load( + index_cache + + member_pages[:, None] * s_ic_page + + (tl.maximum(members, 0) % TPB)[:, None] * s_ic_slot + + dims[None, :], + mask=group_end & (offsets < CR - 1)[:, 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 + pos_t = tl.load(first_base, mask=group_end, other=0) + pos_h = tl.load(first_base + s_pc_axis, mask=group_end, other=0) + pos_w = tl.load(first_base + 2 * s_pc_axis, mask=group_end, other=0) + cosine, sine = _compress_factors( + pos_t, + pos_h, + pos_w, + cos_sin, + stride_cs, + ROTARY_PAIRS, + IS_MROPE, + MROPE_H, + MROPE_W, + group_end, + ) + key_weight = tl.load(index_k_weight + dims).to(tl.float32) + 1.0 + if USE_PDL: + tl.extra.cuda.gdc_wait() + if live: + # v, as the K/V append stores it. + d = tl.arange(0, D) + for h in tl.static_range(HK): + v = tl.load(row + (HQ + HK + h) * D + d) + tl.store( + v_cache + page * s_vc_page + h * s_vc_head + slot * s_vc_tok + d, + v.to(v_cache.dtype.element_ty), + ) + # 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: + 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) + + +class Qwen4ExpAttnDecodePrologue(OpWrapper): + """`qwen4_exp_qk_norm_rope`, `qsa_paged_kv_store`, `qsa_paged_index_store` and + `qsa_prefill_compress` 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. Writes, as the four + entries do: q and k normalized and rotated in place on `proj`; the index query into + `index_q`; k and v into the paged K/V pool and 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 `(proj, index_q)`. + + Fusion boundary: as the four entries; the projection before and the index scores + after are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="decode, 1 row at position 8699 (a group's last token)", spec=dict(pos=(8699,))), + Cell( + why="decode, 4 rows: three group ends (two on a page's last slot), one with no compression", + spec=dict(pos=(8699, 8700, 4095, 31)), + ), + Cell( + why="decode, 2 rows, one unmapped (request out of range)", + spec=dict(pos=(8699, 1201), unmapped=1), + ), + ) + + 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: those of the four entries; the index plane one KV head of + `index_head_dim` (bf16), the position cache int32 `[pages, tpb, >=3]`; decode rows, + one per request (a group's earlier tokens are already in the index plane). A row + whose page lookup fails writes no cache, as the stores and compression skip it. + + Numerics: the same arithmetic as the four entries; outputs bitwise theirs. + """ + + def __call__( + self, + proj: 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, + index_mrope_section1: int = 0, + index_mrope_section2: int = 0, + compress_mrope_section: Optional[Tuple[int, int, int]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + num_tokens = proj.shape[0] + if num_tokens == 0: + return proj, index_q + hpb, warps = _qk_launch(num_tokens) + main_programs = triton.cdiv(num_heads_q + num_heads_k, hpb) + section = compress_mrope_section or (0, 0, 0) + pdl = get_env_enable_pdl() + _attn_decode_prologue_kernel[(num_tokens, main_programs + 2)]( + proj, + proj.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, + IS1=index_mrope_section1, + IS2=index_mrope_section2, + HPB=hpb, + MAIN_PROGRAMS=main_programs, + TPB=tokens_per_block, + CR=compress_ratio, + ROTARY_PAIRS=cos_sin.shape[1] // 2, + IS_MROPE=compress_mrope_section is not None, + MROPE_H=section[1], + MROPE_W=section[2], + USE_PDL=pdl, + num_warps=warps, + launch_pdl=pdl, + ) + return proj, index_q + + def reference( + self, + proj: 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, + index_mrope_section1: int = 0, + index_mrope_section2: int = 0, + compress_mrope_section: Optional[Tuple[int, int, int]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """The four entries in their composition order, on the given buffers.""" + rows = proj.shape[0] + _qk_norm_rope( + proj, index_q, position_ids, cos_sin, q_weight, k_weight, index_q_weight, + num_heads_q=num_heads_q, num_heads_k=num_heads_k, head_dim=head_dim, + index_heads=index_heads, index_head_dim=index_head_dim, index_q_col=index_q_col, + rotary_dim=rotary_dim, eps=eps, mrope_section1=mrope_section1, + mrope_section2=mrope_section2, index_mrope_section1=index_mrope_section1, + index_mrope_section2=index_mrope_section2, + ) # fmt: skip + heads = proj.view(rows, -1, head_dim) + k = heads[:, num_heads_q : num_heads_q + num_heads_k] + v = heads[:, num_heads_q + num_heads_k : num_heads_q + 2 * num_heads_k] + _kv_store( + k, + v, + k_cache, + v_cache, + request_indices, + logical_positions, + block_table, + tokens_per_block, + ) + token_k = proj[:, index_q_col + index_heads * index_head_dim :][:, :index_head_dim].reshape( + rows, 1, index_head_dim + ) + _index_store( + token_k, coords, index_cache, position_cache, request_indices, logical_positions, + block_table, tokens_per_block, + ) # fmt: skip + _compress( + logical_positions=logical_positions, request_indices=request_indices, + block_table=block_table, index_cache=index_cache, position_cache=position_cache, + k_norm_weight=index_k_weight, cos_sin=cos_sin, eps=eps, + tokens_per_block=tokens_per_block, compress_ratio=compress_ratio, + mrope_section=compress_mrope_section, + ) # fmt: skip + return proj, index_q + + def is_valid(self, proj, 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 position_ids.is_contiguous() and position_ids.dtype == torch.int32, ( + "int32 [3, T] positions" + ) + assert coords.shape == (proj.shape[0], 3), "coords [T, 3]" + assert cos_sin.dtype == torch.float32, ( + "an fp32 cos|sin table (the compressed key's rotation)" + ) + + +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__tp1/ops/attention/qwen4_exp_qk_norm_rope.py b/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 new file mode 100644 index 000000000000..27a17ed46ace --- /dev/null +++ b/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 @@ -0,0 +1,1193 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gemma QK norm + partial NeoX rope for both of an attention layer's planes, one launch. + +The main plane (q and k heads of the fused input projection) is rewritten in +place; the QSA index plane's query heads are read from the same projection row +and written, normalized and rotated, into a separate buffer. The rotation reads +the target's own fp32 `[positions, rotary_dim]` cos|sin table. + +Two kernels, one result: a Triton kernel, and from `_CUTE_MIN_TOKENS` tokens (a mixed +step's prefill chunk) a persistent CuTe DSL kernel that is bitwise the Triton kernel's +eight-head launch and streams the projection at about twice its rate. +""" + +from __future__ import annotations + +import struct +from typing import Dict, Tuple + +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.kernels.blackwell.utils import make_ptr + + +@triton.jit +def _norm_rope_tile( + src, + dst, + heads, + head_mask, + weight, + positions, + num_tokens, + token, + cos_sin, + stride_cs, + eps, + D: tl.constexpr, + ROT: tl.constexpr, + S1: tl.constexpr, + S2: tl.constexpr, + WEIGHT_SPLIT: tl.constexpr, + weight_k, + dst_head_stride, +): + """One `[len(heads), D]` tile: rms over each head, `* (1 + w)`, rotate the + first ROT dims in NeoX halves. fp32 throughout, one rounding at the store.""" + 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 + ) + w = tl.load(weight + d).to(tl.float32) + wp = tl.load(weight + partner).to(tl.float32) + if WEIGHT_SPLIT > 0: + # Heads at or past WEIGHT_SPLIT take the second weight (the k heads). + 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 + rstd = tl.rsqrt(tl.sum(x * x, axis=1) / D + eps) + y = x * rstd[:, None] * w2 + yp = xp * rstd[:, None] * wp2 + # 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. + 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) + out = y * c[None, :] + yp * s[None, :] + tl.store( + dst + heads[:, None] * dst_head_stride + d[None, :], + out.to(dst.dtype.element_ty), + mask=hm, + ) + + +@triton.jit(do_not_specialize=["num_tokens"]) +def _qwen4_exp_qk_norm_rope_kernel( + proj, + stride_proj, + index_q, + stride_iq_token, + stride_iq_head, + positions, + num_tokens, + cos_sin, + stride_cs, + q_weight, + k_weight, + index_q_weight, + eps, + 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, + IS1: tl.constexpr, + IS2: tl.constexpr, + HPB: tl.constexpr, + MAIN_PROGRAMS: tl.constexpr, +): + token = tl.program_id(0).to(tl.int64) + program = tl.program_id(1) + row = proj + token * stride_proj + if program < MAIN_PROGRAMS: + main_heads = program * HPB + tl.arange(0, HPB) + _norm_rope_tile( + row, row, main_heads, main_heads < HQ + HK, q_weight, positions, num_tokens, token, + cos_sin, stride_cs, eps, D, ROT, S1, S2, HQ, k_weight, D, + ) # fmt: skip + else: + index_heads = tl.arange(0, IH) + _norm_rope_tile( + row + IQ_COL, index_q + token * stride_iq_token, index_heads, index_heads < IH, + index_q_weight, positions, num_tokens, token, cos_sin, stride_cs, eps, + ID, ROT, IS1, IS2, 0, index_q_weight, stride_iq_head, + ) # fmt: skip + + +def _launch(num_tokens: int) -> Tuple[int, int]: + """Main-plane heads per program and warps per program, by token count. + + Measured on GB300 (CUDA-graph replay, 12 rotating layers, us per call; + HPB/warps): + + tokens 1/2 2/2 2/4 4/4 8/8 32/8 + 1 1.76 1.84 1.87 1.99 1.92 2.82 + 8 1.83 1.97 1.97 2.02 2.05 2.91 + 32 2.11 2.14 2.18 2.14 2.16 2.99 + 64 2.48 2.39 2.47 2.44 2.48 3.03 + 128 3.24 3.19 3.46 3.12 2.96 3.12 + 2048 29.7 25.4 30.8 24.0 22.6 25.8 + 8192 115 98.5 121 95.7 89.5 122 + """ + if num_tokens <= 16: + return 1, 2 + if num_tokens <= 64: + return 2, 2 + return 8, 8 + + +#: This checkpoint's geometry: 24 q + 2 k heads of 256 at the head of the +#: 14080-wide fused projection row, 4 index q heads of 128 at column 13312, a +#: 64-dim rotary prefix, `mrope_section` [11, 11, 10] on the main plane and +#: coordinate 0 alone on the index plane. +CKPT = dict( + num_heads_q=24, + num_heads_k=2, + head_dim=256, + index_heads=4, + index_head_dim=128, + index_q_col=13312, + proj_width=14080, + index_scratch_heads=6, + rotary_dim=64, + mrope_section1=11, + mrope_section2=10, + index_mrope_section1=0, + index_mrope_section2=0, +) + + +# ---- CuTe DSL path: prefill-size token counts ------------------------------------------- + +#: Token count from which the op runs the CuTe kernel. GB300, graph replay, us per call (Triton +#: -> CuTe): 512 8.4 -> 10.3, 768 12.3 -> 8.2, 1024 14.5 -> 8.3, 4096 45.2 -> 18.5, 16384 +#: 182 -> 82.7. The Triton kernel runs its eight-head launch from 65 tokens on, and the CuTe +#: kernel is bitwise that launch. +_CUTE_MIN_TOKENS = 1024 +#: A token's stage in a warp's ring: the lanes' four chunks (32 x 16 bytes each), the cos and +#: sin of the 32 rotary pairs, and three positions. +_STAGE = 4 * 32 * 16 + 2 * 32 * 4 + 16 +_STAGE_CS = 4 * 32 * 16 +_STAGE_POS = _STAGE_CS + 2 * 32 * 4 +#: Stages in a warp's ring (the next token's loads in flight while one is formed). +_RING = 2 +_CTAS_PER_SM = 3 + + +def _iv(x, loc, ip): + return x.ir_value(loc=loc, ip=ip) + + +def _asm(ret, args, text, constraints, loc, ip, side_effects=False): + return llvm.inline_asm( + ret, + [_iv(a, loc, ip) for a in args], + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _asm4(args, text, constraints, loc, ip, side_effects=False): + ty = T.i32() + res = _asm( + llvm.StructType.get_literal([ty] * 4), args, text, constraints, loc, ip, side_effects + ) + return [cutlass.Int32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(4)] + + +@dsl_user_op +def _ldg_v4(addr, *, loc=None, ip=None): + return _asm4( + [cutlass.Int64(addr)], + "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _ldg_b32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm(T.i32(), [cutlass.Int64(addr)], "ld.global.b32 $0, [$1];", "=r,l", loc, ip, True) + ) + + +@dsl_user_op +def _stg_v4(addr, off: int, a, b, c, d, *, loc=None, ip=None): + """16 bytes at `addr + off`.""" + _asm( + T.i32(), + [cutlass.Int64(addr)] + [cutlass.Int32(v) for v in (a, b, c, d)], + "{ st.global.v4.b32 [$1+%d], {$2, $3, $4, $5}; mov.u32 $0, 0; }" % off, + "=r,l,r,r,r,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _lds_v4(addr, *, loc=None, ip=None): + return _asm4( + [cutlass.Int32(addr)], + "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _lds_b32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm(T.i32(), [cutlass.Int32(addr)], "ld.shared.b32 $0, [$1];", "=r,r", loc, ip, True) + ) + + +@dsl_user_op +def _cp16(dst, src, *, loc=None, ip=None): + _asm( + T.i32(), + [cutlass.Int32(dst), cutlass.Int64(src)], + "{ cp.async.cg.shared.global [$1], [$2], 16; mov.u32 $0, 0; }", + "=r,r,l", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _cp4(dst, src, pred, *, loc=None, ip=None): + """4 bytes where `pred` is not 0.""" + _asm( + T.i32(), + [cutlass.Int32(dst), cutlass.Int64(src), cutlass.Int32(pred)], + "{ .reg .pred p; setp.ne.b32 p, $3, 0; @p cp.async.ca.shared.global [$1], [$2], 4; " + "mov.u32 $0, 0; }", + "=r,r,l,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _cp_commit(*, loc=None, ip=None): + _asm(T.i32(), [], "{ cp.async.commit_group; mov.u32 $0, 0; }", "=r", loc, ip, True) + + +@dsl_user_op +def _cp_wait(n: int, *, loc=None, ip=None): + _asm(T.i32(), [], "{ cp.async.wait_group %d; mov.u32 $0, 0; }" % n, "=r", loc, ip, True) + + +@dsl_user_op +def _bf16_lo(w, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int32(w)], + "{ .reg .b32 t; shl.b32 t, $1, 16; mov.b32 $0, t; }", + "=f,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_hi(w, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int32(w)], + "{ .reg .b32 t; and.b32 t, $1, 0xFFFF0000; mov.b32 $0, t; }", + "=f,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _unpack2(w, *, loc=None, ip=None): + """A bf16x2 word -> its two values as an fp32x2 pair, the low element first.""" + return cutlass.Int64( + _asm( + T.i64(), + [cutlass.Int32(w)], + "{ .reg .b32 a, b; shl.b32 a, $1, 16; and.b32 b, $1, 0xFFFF0000; mov.b64 $0, {a, b}; }", + "=l,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _pair(lo, hi, *, loc=None, ip=None): + """Two fp32 bit patterns -> one fp32x2 pair.""" + return cutlass.Int64( + _asm( + T.i64(), + [cutlass.Int32(lo), cutlass.Int32(hi)], + "mov.b64 $0, {$1, $2};", + "=l,r,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _splat2(x, *, loc=None, ip=None): + return cutlass.Int64( + _asm(T.i64(), [cutlass.Float32(x)], "mov.b64 $0, {$1, $1};", "=l,f", loc, ip) + ) + + +@dsl_user_op +def _mul2(a, b, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + T.i64(), + [cutlass.Int64(a), cutlass.Int64(b)], + "mul.rn.f32x2 $0, $1, $2;", + "=l,l,l", + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma2(a, b, c, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + T.i64(), + [cutlass.Int64(a), cutlass.Int64(b), cutlass.Int64(c)], + "fma.rn.f32x2 $0, $1, $2, $3;", + "=l,l,l,l", + loc, + ip, + ) + ) + + +@dsl_user_op +def _add2(a, b, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + T.i64(), + [cutlass.Int64(a), cutlass.Int64(b)], + "add.rn.f32x2 $0, $1, $2;", + "=l,l,l", + loc, + ip, + ) + ) + + +@dsl_user_op +def _hsum2(a, *, loc=None, ip=None): + """The low plus the high element of a pair.""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int64(a)], + "{ .reg .f32 x, y; mov.b64 {x, y}, $1; add.rn.f32 $0, x, y; }", + "=f,l", + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16x2_of(a, *, loc=None, ip=None): + """A pair rounded (RN) into one bf16x2 word, the low element low.""" + return cutlass.Int32( + _asm( + T.i32(), + [cutlass.Int64(a)], + "{ .reg .f32 x, y; mov.b64 {x, y}, $1; cvt.rn.bf16x2.f32 $0, y, x; }", + "=r,l", + loc, + ip, + ) + ) + + +@dsl_user_op +def _neg_bits(w, *, loc=None, ip=None): + """An fp32 bit pattern with its sign flipped (the Triton kernel's `xor`).""" + return cutlass.Int32( + _asm(T.i32(), [cutlass.Int32(w)], "xor.b32 $0, $1, 0x80000000;", "=r,r", loc, ip) + ) + + +@dsl_user_op +def _fma(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _mul(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _add(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _rstd(ss, eps, width_hex: str, *, loc=None, ip=None): + """`rsqrt(ss / width + eps)` as the Triton kernel's `div.full`, `add`, `rsqrt.approx.ftz`: + the `add` has no rounding mode there, so ptxas fuses it with the division (a multiply by + the reciprocal) into one FFMA -- here the same.""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(ss), cutlass.Float32(eps)], + "{ .reg .f32 t; div.full.f32 t, $1, " + width_hex + "; add.f32 t, t, $2; " + "rsqrt.approx.ftz.f32 $0, t; }", + "=f,f,f", + loc, + ip, + ) + ) + + +def _hex(x: float) -> str: + return "0f%08X" % struct.unpack(" bool: + """The geometry and layouts the CuTe kernel covers; anything else runs the Triton kernel.""" + bf = torch.bfloat16 + return ( + num_heads_q % 4 == 0 + and num_heads_k == 2 + and head_dim == 256 + and index_heads == 4 + and index_head_dim == 128 + and rotary_dim == 64 + and index_q_col % 8 == 0 + and proj.dtype == index_q.dtype == bf + and all(w.dtype == bf and w.is_contiguous() and w.data_ptr() % 16 == 0 for w in weights) + and cos_sin.dtype == torch.float32 + and cos_sin.stride(1) == 1 + and proj.data_ptr() % 16 == 0 + and proj.stride(0) % 8 == 0 + and index_q.data_ptr() % 16 == 0 + and index_q.stride(0) % 8 == 0 + and index_q.stride(1) == index_head_dim + ) + + +def _launch_cute( + proj: torch.Tensor, + index_q: torch.Tensor, + position_ids: torch.Tensor, + cos_sin: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + num_heads_q: int, + index_q_col: int, + eps: float, + sections: Tuple[int, int, int, int], +) -> None: + key = (num_heads_q, index_q_col) + tuple(sections) + fn = _RING_KERNELS.get(key) + if fn is None: + gm = cute.AddressSpace.gmem + bf = cutlass.BFloat16 + fn = cute.compile( + _NormRopeRing(num_heads_q, index_q_col, *sections), + from_dlpack(proj, assumed_align=16).mark_layout_dynamic(leading_dim=1), + make_ptr(bf, index_q.data_ptr(), gm, assumed_align=16), + make_ptr(cutlass.Int32, position_ids.data_ptr(), gm, assumed_align=4), + make_ptr(cutlass.Float32, cos_sin.data_ptr(), gm, assumed_align=4), + make_ptr(bf, q_weight.data_ptr(), gm, assumed_align=16), + make_ptr(bf, k_weight.data_ptr(), gm, assumed_align=16), + make_ptr(bf, index_q_weight.data_ptr(), gm, assumed_align=16), + *(cutlass.Int32(1),) * 4, + cutlass.Float32(1e-6), + cutlass.Int32(1), + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--opt-level 2 --enable-tvm-ffi", + ) + _RING_KERNELS[key] = fn + dev = proj.device.index + if dev not in _SM_COUNT: + _SM_COUNT[dev] = torch.cuda.get_device_properties(proj.device).multi_processor_count + tokens = proj.shape[0] + fn( + proj, + index_q.data_ptr(), + position_ids.data_ptr(), + cos_sin.data_ptr(), + q_weight.data_ptr(), + k_weight.data_ptr(), + index_q_weight.data_ptr(), + tokens, + proj.stride(0), + index_q.stride(0), + cos_sin.stride(0), + eps, + min(tokens, _SM_COUNT[dev] * _CTAS_PER_SM), + ) + + +def _launch_triton( + proj: torch.Tensor, + index_q: torch.Tensor, + position_ids: torch.Tensor, + cos_sin: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + 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, + sections: Tuple[int, int, int, int], +) -> None: + num_tokens = proj.shape[0] + hpb, warps = _launch(num_tokens) + main_programs = triton.cdiv(num_heads_q + num_heads_k, hpb) + _qwen4_exp_qk_norm_rope_kernel[(num_tokens, main_programs + 1)]( + proj, + proj.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, + eps, + 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=sections[0], + S2=sections[1], + IS1=sections[2], + IS2=sections[3], + HPB=hpb, + MAIN_PROGRAMS=main_programs, + num_warps=warps, + ) + + +class Qwen4ExpQkNormRope(OpWrapper): + """Both planes of one attention layer's QK norm + rope, in one launch. + + Main plane, in place on `proj`: heads `[0, Hq)` (weight `q_weight`) and + `[Hq, Hq + Hk)` (weight `k_weight`), each `head_dim` wide at the start of + every row, become `rope(rmsnorm(x) * (1 + w))`. Index plane: the + `index_heads` heads of `index_head_dim` starting at column `index_q_col` of + `proj` are read, normalized with `index_q_weight` and rotated, and written + to `index_q` (`[T, index_heads, index_head_dim]`, any token/head stride, + unit dim stride); `proj`'s index columns are not written. Every other + column of `proj` is left alone. Returns `(proj, index_q)`. + + The rotation is NeoX over the first `rotary_dim` dims of a head: pair `j` + is `(j, j + rotary_dim / 2)`, rotated by `cos_sin[p, j]` and + `cos_sin[p, rotary_dim / 2 + j]` (`cos_sin` is the fp32 `[P, rotary_dim]` + table, cos block then sin block). `position_ids` is the contiguous int32 + `[3, T]`; pair `j` reads coordinate 1 when `j % 3 == 1 and j < 3 * s1`, + coordinate 2 when `j % 3 == 2 and j < 3 * s2`, coordinate 0 otherwise -- + interleaved mRoPE with the plane's own `(s1, s2)`; `(0, 0)` is the plain + rope on coordinate 0. fp32 throughout, one rounding at the store. + + Fusion boundary: norm and rotation of both planes. The projection GEMM + before it, and the KV/index-cache writes and attention after it, are the + caller's. The index plane's raw key is not touched, so the index-cache + store may run before or after this call. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # Four cells, one per compiled launch the forward reaches: the Triton kernel's 1-, 2- + # and 8-head launches (by decode batch size) and the CuTe kernel (prefill). + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode M=1: one token on the 1-head/2-warp launch, three distinct mRoPE rows", + spec=dict(tokens=1, **CKPT), + ), + Cell( + why="decode M=17: first on the 2-head launch (13 main programs)", + spec=dict(tokens=17, **CKPT), + ), + Cell( + why="decode M=128: max concurrency, 8-head launch (last main program half masked), " + "positions up to the table end (262144 rows)", + spec=dict(tokens=128, max_position=262144, **CKPT), + ), + Cell(why="prefill 16384 tokens: max_num_tokens", spec=dict(tokens=16384, **CKPT)), + ) + + note = """ + Metadata consumed: none (positions are an operand). Below `_CUTE_MIN_TOKENS` + tokens one Triton kernel, no autotuner: grid `(T, cdiv(Hq + Hk, HPB) + 1)`; + the last program of a token row runs the whole index plane as one + `[index_heads, index_head_dim]` tile. HPB and warps by token count from + `_launch` (its docstring has the table). Replaces, per attention layer, the + CUDA `fused_qk_norm_rope` main call, the index-scratch copy and the second + `fused_qk_norm_rope` call: 6.67 -> 1.76 us at M=1, 9.15 -> 2.96 at M=128, 185 + -> 89.5 at 8192 tokens (graph replay). + + From `_CUTE_MIN_TOKENS` tokens, where the geometry is this checkpoint's (`_cute_ok`: + head dims 256 / 128, 2 k heads, 4 index heads, a 64-dim rotary prefix, 16-byte + aligned rows, contiguous index heads), the CuTe DSL `_NormRopeRing`: three + persistent CTAs an SM, each forming every G-th token, a warp's next token + prefetched by cp.async (its class docstring has the mapping). 16384 tokens: 182 -> 82.7 us in the + microbenchmark (the Triton kernel was L1-bound; the ring runs at ~65% of DRAM peak, + bounded by load latency at three CTAs an SM). + + Numerics: fp32 throughout, one rounding at the store; the CuTe kernel is bitwise + the Triton kernel's eight-head launch (measured at 1024-16384 tokens, signed zeros + included; not pinned by a test). Against the fp64 + reference over the same table every cell is within one output rounding + step (measured max 1.48 ulp of the row scale, relative RMS <= 0.011 ulp). + Against the CUDA pair it replaces (measured, not pinned by a test), <= 1.41 + ulp of the row scale and 0.09 ulp RMS (that op derives cos/sin with `__sincosf`, whose phase error grows + with position). Three agreeing position rows select the same table row for + every pair, so a text batch is exactly the plain rope on row 0. + + Preconditions in `is_valid`: the main q/k columns end at or before + `index_q_col` (the planes do not overlap, so reading the index plane while + the main plane is rewritten in place is race-free); `index_q` does not + alias `proj`; positions below the table length (an out-of-range position + reads past the table). + + Hard limits: token offsets in int64; positions index the table in int64. + + CUDA graphs: the Triton kernel captured at M=4 replays bitwise an eager call on + inputs rewritten in place. + + Not certified: non-zero index-plane mRoPE sections (the target passes (0, 0)), + capture of the CuTe kernel, `T == 0`, fp16/fp32 activations, `rotary_dim` > either head width, + non-power-of-two head widths or head counts (tl.arange), YaRN, non-Gemma + norm, a k plane on the index side (its scratch slot is not written). + """ + + def __call__( + self, + proj: torch.Tensor, + index_q: torch.Tensor, + position_ids: torch.Tensor, + cos_sin: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + 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, + index_mrope_section1: int = 0, + index_mrope_section2: int = 0, + ) -> Tuple[torch.Tensor, torch.Tensor]: + num_tokens = proj.shape[0] + if num_tokens == 0: + return proj, index_q + sections = (mrope_section1, mrope_section2, index_mrope_section1, index_mrope_section2) + weights = (q_weight, k_weight, index_q_weight) + if num_tokens >= _CUTE_MIN_TOKENS and _cute_ok( + proj, index_q, cos_sin, weights, num_heads_q, num_heads_k, head_dim, index_heads, + index_head_dim, index_q_col, rotary_dim, + ): # fmt: skip + _launch_cute( + proj, index_q, position_ids, cos_sin, *weights, num_heads_q, index_q_col, eps, + sections, + ) # fmt: skip + else: + _launch_triton( + proj, index_q, position_ids, cos_sin, *weights, num_heads_q, num_heads_k, + head_dim, index_heads, index_head_dim, index_q_col, rotary_dim, eps, sections, + ) # fmt: skip + return proj, index_q + + def reference( + self, + proj: torch.Tensor, + index_q: torch.Tensor, + position_ids: torch.Tensor, + cos_sin: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + 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, + index_mrope_section1: int = 0, + index_mrope_section2: int = 0, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """`(proj after, index_q's written values)` in fp64 over the same table, + rounded once; neither input is mutated.""" + num_tokens = proj.shape[0] + half = rotary_dim // 2 + table = cos_sin.double() + pos = position_ids.long() + j = torch.arange(half, device=proj.device) + + def angles(s1: int, s2: int): + coord = torch.zeros(half, dtype=torch.long, device=proj.device) + 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] + rotated = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], -1) + return torch.cat([rotated, x[..., rotary_dim:]], -1) + + out = proj.double().clone() + main_heads = num_heads_q + num_heads_k + main_cols = main_heads * head_dim + w = torch.cat( + [ + q_weight.double().expand(num_heads_q, head_dim), + k_weight.double().expand(num_heads_k, head_dim), + ] + ) + main = out[:, :main_cols].reshape(num_tokens, main_heads, head_dim) + out[:, :main_cols] = plane(main, w, *angles(mrope_section1, mrope_section2)).reshape( + num_tokens, main_cols + ) + iq = proj[:, index_q_col : index_q_col + index_heads * index_head_dim].double() + iq = plane( + iq.reshape(num_tokens, index_heads, index_head_dim), + index_q_weight.double(), + *angles(index_mrope_section1, index_mrope_section2), + ) + return out.to(proj.dtype), iq.to(index_q.dtype) + + def is_valid( + self, + proj: torch.Tensor, + index_q: torch.Tensor, + position_ids: torch.Tensor, + cos_sin: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + 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, + *unused_args: object, + **unused_kwargs: object, + ) -> None: + num_tokens = proj.shape[0] + assert proj.stride(-1) == 1, "proj must have a unit column stride" + assert (num_heads_q + num_heads_k) * head_dim <= index_q_col, ( + "the index q columns overlap the main q/k heads: the in-place main write races " + "the index read" + ) + assert index_q_col + index_heads * index_head_dim <= proj.shape[1], ( + "the index q columns run past the projection row" + ) + assert index_q.shape == (num_tokens, index_heads, index_head_dim), index_q.shape + assert index_q.stride(-1) == 1, "index_q must have a unit dim stride" + lo, hi = index_q.data_ptr(), proj.data_ptr() + assert lo >= hi + proj.numel() * proj.element_size() or hi >= lo + ( + index_q.shape[0] * index_q.stride(0) * index_q.element_size() + ), "index_q aliases proj" + assert position_ids.shape == (3, num_tokens) and position_ids.is_contiguous(), ( + f"position_ids must be a contiguous [3, {num_tokens}]; got {tuple(position_ids.shape)}" + ) + assert position_ids.dtype == torch.int32, position_ids.dtype + assert cos_sin.dtype == torch.float32 and cos_sin.shape[1] == rotary_dim, cos_sin.shape + assert rotary_dim <= min(head_dim, index_head_dim) and rotary_dim % 2 == 0, rotary_dim + if num_tokens: + assert int(position_ids.min()) >= 0, "negative position" + assert int(position_ids.max()) < cos_sin.shape[0], "position past the rope table" + + +qwen4_exp_qk_norm_rope = Qwen4ExpQkNormRope() 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 index 559ba69efc9a..6e5d284f768b 100644 --- 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 @@ -30,24 +30,10 @@ MERGE = 2 THETA = 10000.0 -#: (frames, height, width) in patch units. The `serving_*` rows come 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; 112x112 the square form of the largest segment; 22144 the largest call. -GRIDS = { - "one_block": ((1, 2, 2),), - "tiny_image": ((1, 4, 4),), - "image": ((1, 32, 32),), - "prefill_image": ((1, 64, 64),), - "two_images": ((1, 16, 24), (1, 8, 12)), - "video": ((2, 4, 6),), - "video_and_image": ((3, 8, 8), (1, 4, 4)), - "serving_frame_5580": ((1, 62, 90),), - "serving_frame_12544": ((1, 112, 112),), - "serving_batch_22144": ((1, 112, 112), (1, 96, 100)), -} - -_DTYPES = (torch.bfloat16, torch.float16, torch.float32) +#: (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: @@ -100,60 +86,28 @@ class VisionRope2d(OpWrapper): # "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, ...] = ( - tuple( - Cell(why=f"{name} {str(dtype)[6:]}", spec=dict(grid=grid, dtype=dtype)) - for name, grid in GRIDS.items() - for dtype in _DTYPES - ) - + tuple( - Cell( - why=f"grid side {max(h, w)} ({h}x{w}) {str(dtype)[6:]}: largest angle", - spec=dict(grid=((1, h, w),), dtype=dtype), - ) - for h, w in ((2, 112), (2, 196), (2, 512), (2, 1024), (2, 6272), (6272, 2)) - for dtype in (torch.float32, torch.bfloat16) - ) - + tuple( - Cell( - why=f"head geometry {heads}x{head_dim}", - spec=dict(grid=((1, 8, 12),), dtype=torch.bfloat16, heads=heads, head_dim=head_dim), - ) - for heads, head_dim in ((8, 64), (4, 128), (2, 40)) - ) - + ( - Cell( - why="k carries its own head count (16 q / 2 k)", - spec=dict(grid=GRIDS["two_images"], dtype=torch.bfloat16, heads_k=2), + 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: q/k split views of packed qkv, 62x90", - spec=dict( - grid=GRIDS["serving_frame_5580"], dtype=torch.bfloat16, layout="packed_qkv" - ), + ), + 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, ), - Cell( - why="call site: packed qkv views, video + image", - spec=dict( - grid=((3, 8, 8), (1, 16, 24)), - dtype=torch.bfloat16, - layout="packed_qkv", - grid_dtype=torch.int64, - ), - ), - Cell( - why="grid_thw int64 on the device (host sync)", - spec=dict( - grid=GRIDS["two_images"], - dtype=torch.bfloat16, - grid_dtype=torch.int64, - grid_device="cuda", - ), - ), - Cell( - why="theta 1e6: theta is an argument", - spec=dict(grid=GRIDS["tiny_image"], dtype=torch.float32, theta=1.0e6), - ), - ) + ), ) note = """ @@ -162,22 +116,25 @@ class VisionRope2d(OpWrapper): `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). + mismatches the head and the broadcast raises (loud, so not guarded; not + pinned by a test). - Accuracy vs fp64 angles over the cells: fp32 <= 4.77e-7, fp16 <= 1.95e-3, - bf16 <= 1.56e-2 (one ulp at unit-normal operands), 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`. + 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. - CUDA graphs: capture bakes in the grid; replay under another grid was not - measured. - - Not certified: mixed q/k dtypes; float8; `spatial_merge_size` != 2; grid - sides past 6272 or calls past 22144 tokens. + 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__( 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/ngram_hash.py index 21190a07f730..e3841c25f537 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/ngram_hash.py @@ -67,31 +67,21 @@ class NgramHash(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: batch, length, num_query_tokens; history is uniform random ids + # spec: batch, length, num_query_tokens; history is uniform random int64 ids # with the first `min(CONTEXT_LEN, length)` set to EOS when `lead_eos` # (default True), `eos_positions` forced to EOS, and `eos_density` of the - # remaining ids replaced by EOS; ids_dtype (int64). + # 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. CELLS: tuple[Cell, ...] = ( - _rows(1, 3, 1, "decode: 2 cached ids + 1 new, 1 row"), - _rows(8, 3, 1, "decode: 2 cached ids + 1 new, 8 rows"), _rows(4, 3, 1, "decode row whose cached context is EOS", eos_positions=(0, 1)), - _rows(1, 130, 128, "prefill chunk of 128"), - _rows(2, 130, 128, "prefill chunk of 128, EOS inside", eos_positions=(40, 41, 100)), - _rows(3, 1026, 1024, "prefill chunk of 1024, EOS runs", eos_positions=(7, 8, 9, 512)), - _rows(2, 64, 64, "whole history emitted", eos_positions=(16, 17)), - _rows(1, 2, 0, "empty query: [batch, 0, heads]"), - _rows(4, 6, 4, "4 query tokens per row (draft width)", eos_positions=(2,)), - # The call site: one `[T, ngram_size]` window per packed token from - # `ngram_token_history`, raw ids with EOS anywhere, num_query_tokens 1. - _rows(1, 3, 1, "call site: 1 window, 1 token", lead_eos=False, eos_density=0.2), _rows( 32, 3, 1, "call site: 32 windows, decode / MTP verify", lead_eos=False, eos_density=0.2 ), - _rows(4096, 3, 1, "call site: 4096 windows, prefill", lead_eos=False, eos_density=0.1), _rows( 8192, 3, 1, "call site: 8192 windows, largest prefill", lead_eos=False, eos_density=0.05 ), - _rows(2, 34, 32, "int32 token_history", eos_positions=(10,), ids_dtype=torch.int32), ) note = """ @@ -103,11 +93,12 @@ class NgramHash(OpWrapper): Preconditions: `head_vocab_sizes`, `head_offsets` on `token_history`'s device (raises otherwise; a cpu `layer_multipliers` is accepted and - gives the identical result). `token_history` may be any integer dtype - (int32 and int64 identical). Negative ids are accepted and hashed as - ordinary ids, not padding. `num_query_tokens == 0` returns - `[batch, 0, heads]`. A single multiplier raises `ZeroDivisionError`. The - four silent mistakes are `is_valid`'s. + 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. + A single multiplier raises `ZeroDivisionError`. The four silent mistakes + are `is_valid`'s. A wrong hash is a valid row: an extra shift, no EOS reset, an inclusive segment start, reversed multipliers, no offsets, swapped blocks all stay 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/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__tp1/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__tp1/ops/gemm/bf16_gemv_pdl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py new file mode 100644 index 000000000000..c4b93e549701 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/bf16_gemv_pdl.py @@ -0,0 +1,384 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode-row bf16 projection launched with programmatic dependent launch (PDL), +its weight streaming in while the kernel before it finishes. + +Tier 5 (CuTe DSL written for this target). At decode row counts a projection +reads its weight once and its activation is a few KB, so the GEMV's time is the +weight stream plus its launch and ramp. A normal launch cannot read a byte of +that stream until the producer of the activation has exited. Launched with PDL, +every CTA loads its weight rows -- static, so safe to read early -- into +registers before `griddepcontrol.wait` and triggers its own dependents right +before it; the activation is read, and the output written, only after the wait. + +One CTA computes `OUTS` output columns over the whole K for every row; 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. +Products and sums in fp32: a warp reduction, then a cross-warp one in shared +memory; one rounding at the store. + +A K with a shared-memory tier (`SOUTS`) gives each CTA `SOUTS` more outputs whose +weight rows it stages in shared memory with cp.async before the wait, each thread +exactly the vectors it later reads (so `cp.async.wait_group` alone orders them). The +2560 x 6144 projection then runs as one wave of 640 CTAs whose whole weight is on +chip by the wait, instead of 1280 CTAs of which the last third start loading theirs +only as the first ones exit. +""" + +from __future__ import annotations + +import functools + +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.nvgpu import cpasync +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 + +#: Largest row count one launch serves; every row re-reads the weight from registers. +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)} +#: 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. +SOUTS = {6144: 2} +_COMPILE_OPTIONS = "--ptxas-options -maxrregcount=64" +_COMPILE_OPTIONS_SMEM = "--ptxas-options -maxrregcount=80" +#: CTAs a multiprocessor must hold for a grid with a shared-memory tier: four, the +#: 640-CTA grid's share of 160 SMs. Behind the sparse-attention merge (one row, graph +#: replay) 4 runs the attention tail 13.5 us a layer against 14.6 at 1 and 15.6 at 2-3; +#: the GDN tail is unchanged. +MIN_BLOCKS_SMEM = 4 + + +def tiling(n: int, k: int) -> tuple[int, int, int, int]: + """`(threads, outs, souts, vec)` for a `[*, K] x [K, N]` product: register and shared-memory + outputs a CTA. Without a shared-memory tier that divides N the outputs shrink to divide N.""" + threads, outs, vec = TILING[k] + souts = SOUTS.get(k, 0) + if souts and n % (outs + souts) == 0: + return threads, outs, souts, vec + while n % outs: + outs //= 2 + return threads, outs, 0, vec + + +#: Before its PDL wait a CTA also has the weight rows of CTA `bidx + dist` brought into +#: L2: behind a latency-bound kernel only part of a large grid's CTAs fit, and the later +#: waves' weights would otherwise start streaming only after it. `dist` is about the CTAs +#: resident before the wait: 1216 behind the hyper-connection kernels (measured), at +#: least 912 (a 1280-CTA out projection's first wave) otherwise. Grids smaller than `dist` +#: prefetch nothing. +L2_DIST = 1216 +L2_DIST_MIN = 912 + + +def l2_dist(grid: int) -> int: + return min(L2_DIST, max(grid // 2, L2_DIST_MIN)) + + +@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, + ) + + +class _GemvKernel: + def __init__(self, *, rows, k, threads, outs, vec, pdl, dist=L2_DIST, souts=0): + assert k % (threads * vec) == 0 and threads % 32 == 0 + self.rows, self.threads, self.outs, self.souts, self.vec = rows, threads, outs, souts, vec + self.dist = dist + self.k_tiles = k // (threads * vec) + self.warps = threads // 32 + self.pdl = pdl + + def smem_bytes(self) -> int: + red = self.rows * (self.outs + self.souts) * self.warps * 4 + if not self.souts: + return red + w = self.souts * self.k_tiles * self.threads * self.vec * 2 + return -(-w // 128) * 128 + -(-red // 128) * 128 + 128 + + @cute.jit + def __call__(self, mA, mW, mOut, 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 * self.vec, + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, + ) + load_w = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cutlass.BFloat16, + num_bits_per_copy=16 * self.vec, + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, + ) + g2s = cute.make_copy_atom( + cpasync.CopyG2SOp(), cutlass.BFloat16, num_bits_per_copy=16 * self.vec + ) + self.kernel(mA, mW, mOut, load_a, load_w, g2s).launch( + grid=[cute.ceil_div(n, self.outs + self.souts), 1, 1], + block=[self.threads, 1, 1], + smem=self.smem_bytes(), + stream=stream, + use_pdl=self.pdl, + min_blocks_per_mp=MIN_BLOCKS_SMEM if self.souts else 1, + ) + + @cute.kernel + def kernel( + self, mA, mW, mOut, load_a: cute.CopyAtom, load_w: cute.CopyAtom, g2s: 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 + vec: cutlass.Constexpr = self.vec + k_tiles: cutlass.Constexpr = self.k_tiles + warps: cutlass.Constexpr = self.warps + souts: cutlass.Constexpr = self.souts + tot: cutlass.Constexpr = outs + souts + n0 = bidx * tot + + # (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))) + smem = cutlass.utils.SmemAllocator() + + # The weight is static: it streams before the wait, while the producer runs. The + # shared-memory tier first ([output, k-tile, thread, vec]: a warp's vectors are + # consecutive), so it lands while the register tier streams. + if const_expr(souts > 0): + s_w = smem.allocate_tensor( + cutlass.BFloat16, + cute.make_layout( + (souts, k_tiles, threads, vec), + stride=(k_tiles * threads * vec, threads * vec, vec, 1), + ), + byte_alignment=128, + ) + for j in cutlass.range_constexpr(souts): + s_row = w_vec[n0 + outs + j, (None, (tidx, None))] + for kt in cutlass.range_constexpr(k_tiles): + cute.copy(g2s, s_row[None, kt], s_w[j, kt, tidx, None]) + cute.arch.cp_async_commit_group() + 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]) + if const_expr(self.pdl): + n_rows = cute.size(mW, mode=[0]) + k_len = cute.size(mW, mode=[1]) + if tidx == 0: + if (bidx + self.dist + 1) * tot <= n_rows: + _bulk_prefetch_l2( + mW.iterator + (bidx + self.dist) * tot * k_len, tot * k_len * 2 + ) + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + + acc = cute.make_rmem_tensor(cute.make_layout((rows, tot), stride=(tot, 1)), cutlass.Float32) + acc.fill(0.0) + a_reg = cute.make_rmem_tensor( + cute.make_layout((k_tiles, vec), stride=(vec, 1)), cutlass.BFloat16 + ) + sw_reg = cute.make_rmem_tensor(cute.make_layout((vec,), stride=(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) + if const_expr(souts > 0): + if const_expr(r == 0): + cute.arch.cp_async_wait_group(0) + for j in cutlass.range_constexpr(souts): + for kt in cutlass.range_constexpr(k_tiles): + cute.autovec_copy(s_w[j, kt, tidx, None], sw_reg) + for v in cutlass.range_constexpr(vec): + acc[r, outs + j] = acc[r, outs + j] + a_reg[kt, v].to( + cutlass.Float32 + ) * sw_reg[v].to(cutlass.Float32) + for r in cutlass.range_constexpr(rows): + for j in cutlass.range_constexpr(tot): + acc[r, j] = cute.arch.warp_reduction_sum(acc[r, j]) + part = smem.allocate_tensor( + cutlass.Float32, + cute.make_layout((rows, tot, warps), stride=(tot * warps, warps, 1)), + byte_alignment=16, + ) + with cute.arch.elect_one(): + for r in cutlass.range_constexpr(rows): + for j in cutlass.range_constexpr(tot): + part[r, j, warp] = acc[r, j] + cute.arch.sync_threads() + if tidx < tot: + 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] + mOut[r, n0 + tidx] = total.to(cutlass.BFloat16) + + +@functools.lru_cache(maxsize=None) +def _compiled(rows, n, k, pdl): + threads, outs, souts, vec = tiling(n, k) + kern = _GemvKernel( + rows=rows, + k=k, + threads=threads, + outs=outs, + vec=vec, + pdl=pdl, + dist=l2_dist(-(-n // (outs + souts))), + souts=souts, + ) + a = torch.empty(rows, k, dtype=torch.bfloat16, device="cuda") + w = torch.empty(n, k, dtype=torch.bfloat16, device="cuda") + out = torch.empty(rows, n, dtype=torch.bfloat16, device="cuda") + stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) + return cute_ext.compile( + kern, + from_dlpack(a, assumed_align=32), + from_dlpack(w, assumed_align=32), + from_dlpack(out, assumed_align=32), + stream, + options=_COMPILE_OPTIONS_SMEM if souts else _COMPILE_OPTIONS, + ) + + +def launch(mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool) -> torch.Tensor: + """The launch behind `Bf16GemvPdl.__call__`. `pdl` is honoured unless + `TRTLLM_ENABLE_PDL=0`.""" + m, k = mat_a.shape + n = mat_b.shape[1] + dev = mat_a.device + pdl = pdl and get_env_enable_pdl() + out = torch.empty((m, n), dtype=mat_a.dtype, device=dev) + if m > 0: + fn = _compiled(m, n, k, bool(pdl)) + stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream) + fn( + from_dlpack(mat_a, assumed_align=32), + from_dlpack(mat_b.t(), assumed_align=32), + from_dlpack(out, assumed_align=32), + stream, + ) + return out + + +def _cell(why: str, m: int, n: int, k: int, pdl: bool = True) -> Cell: + return Cell(why=why, spec=dict(m=m, n=n, k=k, pdl=pdl)) + + +class Bf16GemvPdl(OpWrapper): + """`mat_a @ mat_b` in bf16 with fp32 accumulation, one PDL launch per call. + + A drop-in for `ops/gemm/low_m_bf16_mm` at decode row counts: `mat_a` is + dense row-major `[M, K]` with `M <= MAX_ROWS`, `mat_b` the dense column-major + `[K, N]` view of a stored `[N, K]` weight (its `.t()`), the result a fresh + row-major `[M, N]` bf16 tensor. With `pdl` (and PDL available) the kernel is + launched as a programmatic dependent: it may start before the kernel ahead of + it on the stream exits, reads only its weight until `griddepcontrol.wait`, + and lets the kernel after it launch early as well. + """ + + 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. + 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), + ) + + note = """ + Metadata consumed: none. Preconditions: bf16 operands; `mat_a` dense + row-major `[M, K]`, `1 <= M <= MAX_ROWS`; `mat_b` the dense column-major + `[K, N]` view of a `[N, K]` weight (stride `(1, K)`); K a `TILING` key. + + PDL: with `pdl` the launch carries the programmatic-dependent attribute. + Every read of `mat_a` and the store of the output come after + `griddepcontrol.wait`, so a producer launched before this kernel is + complete when they happen, whatever that producer is. Under a CUDA graph a + cross-stream dependency is an ordinary edge and stays a full dependency. + + Compilation: CuTe DSL, JIT on the first call per (M, N, K, pdl); that call + must happen before CUDA-graph capture -- the engine's warmup forward at each + captured batch size does it. + + Numerics: fp32 products and sums, one rounding to bf16; only the summation + order differs from the fp32 reference. + + Hard limits + M rows per launch boundary MAX_ROWS=4 driven at 1 (cells) + K a TILING key 2560, 6144, 640 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. + """ + + def __call__(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = True) -> torch.Tensor: + return launch(mat_a, mat_b, pdl) + + def reference(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = True) -> torch.Tensor: + """fp32-accumulated, TF32 off; `pdl` is a launch attribute and is ignored.""" + with true_fp32_matmul(): + return (mat_a.float() @ mat_b.float()).to(mat_a.dtype) + + def is_valid(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = True) -> None: + assert mat_a.dtype == torch.bfloat16 and mat_b.dtype == torch.bfloat16, "bf16 operands only" + assert mat_a.dim() == 2 and mat_a.is_contiguous(), "mat_a must be dense row-major [M, K]" + assert 1 <= mat_a.shape[0] <= MAX_ROWS, f"M={mat_a.shape[0]} outside 1..{MAX_ROWS}" + assert mat_b.shape[0] == mat_a.shape[1], "inner dimensions differ" + 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_a.shape[1] in TILING, f"K={mat_a.shape[1]} has no tiling" + + +bf16_gemv_pdl = Bf16GemvPdl() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/low_m_bf16_mm.py new file mode 100644 index 000000000000..967a23116a87 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/gemm/low_m_bf16_mm.py @@ -0,0 +1,231 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode-shaped bf16 projection: one kernel per call, its source picked per (M, N, K). + +At decode M a dense bf16 projection streams its weight once, and cuBLAS +reaches only 0.55-0.65 of HBM bandwidth on 0.25-63 MB weights (and ~4.5 us +flat on the small ones). FlashInfer's Blackwell low-M CuTe-DSL kernels -- the +register-prefetch "direct" kernel and the tcgen05 split-K kernel, the two +`flashinfer.mm_bf16(backend="cute-dsl")` dispatches to and +`tensorrt_llm/_torch/modules/low_m_gemm.py` races -- are faster at M <= 32 on +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). + +FlashInfer API: ``flashinfer.gemm.kernels.dense_bf16_gemm_direct.run_direct_dense`` and +``flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk.run_splitk_dense``. +""" + +from __future__ import annotations + +import functools +from typing import NamedTuple + +import torch +from flashinfer.gemm.kernels import dense_bf16_gemm_direct as _fi_direct +from flashinfer.gemm.kernels import dense_bf16_gemm_sm100_splitk as _fi_splitk + +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, +) + + +class _Plan(NamedTuple): + """`kind` is "direct", "splitk" or "cublas"; `tactic` is the kernel's own tactic tuple.""" + + kind: str + tactic: tuple = () + + +def _d(block_size: int, outputs_per_block: int) -> tuple: + """A direct-kernel choice; rows_per_block is the kernel's default for the runtime M.""" + return ("direct", block_size, outputs_per_block) + + +def _s(mma_m: int, mma_n: int, split_k: int, ab_stages: int) -> tuple: + return ("splitk", mma_m, mma_n, split_k, ab_stages) + + +#: (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 +#: 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_z + (6144, 2560): ((1, _d(64, 1)), (2, _d(64, 2))), + # GDN in_a / in_b (48 v heads) + (48, 2560): ((32, _d(64, 1)),), + # GDN out and attention o: [6144 -> 2560] + (2560, 6144): ( + (2, _d(256, 2)), + (8, _s(64, 8, 3, 10)), + (16, _s(64, 16, 3, 8)), + (32, _s(64, 32, 3, 8)), + ), + # attention q (q and its output gate interleaved per head) + (12288, 2560): ((2, _d(64, 2)), (4, _s(128, 8, 1, 6)), (8, _s(128, 16, 1, 6))), + # attention k / v (2 kv heads x 256) and the MoE router (512 experts) + (512, 2560): ((2, _d(64, 2)), (8, _d(64, 1)), (16, _d(32, 1)), (32, _s(64, 8, 4, 10))), + # QSA index projection: 4 index heads x 128 + 1 key x 128 + (640, 2560): ((4, _d(64, 1)), (8, _d(32, 1)), (16, _s(64, 8, 4, 10)), (32, _s(64, 16, 4, 8))), + # shared expert gate_up, and the same with its gate row appended (padded to 1296) + (1280, 2560): ((2, _d(64, 1)), (4, _d(64, 2)), (8, _s(64, 8, 4, 10))), + (1296, 2560): ((2, _d(64, 1)), (4, _d(64, 2)), (8, _s(64, 8, 4, 10))), + # shared expert gate [2560 -> 1] + (1, 2560): ((32, _d(64, 1)),), + # the four GDN input projections concatenated (10240 + 6144 + 48 + 48) + (16480, 2560): ((1, _d(64, 2)),), + # the attention input projections concatenated: q|gate (12288), k, v (512 each), + # index (640), zero-padded to 55 heads of 256 (14080) + (14080, 2560): ((2, _d(64, 2)), (8, _s(128, 8, 1, 6))), + # PLE value projection (decoder layer 1 only); the PLE key projection is [10240, 2560] + (2560, 2560): ((1, _d(64, 1)), (2, _d(64, 2)), (4, _d(32, 2)), (8, _s(64, 8, 2, 10))), +} +# Measured and left on cuBLAS: (2560, 640) shared-expert down (the direct kernel +# cannot tile K=640; split-K and TGV were within 4% of cuBLAS), the GDN concat +# (16480, 2560) at M >= 2 (direct 3.6-5.3% ahead at M=2 across two sweeps, behind +# from M=4), the attention concat (14080, 2560) at M >= 16, and M >= 64 +# everywhere (outside both kernels' M <= 32 policy). Not servable at all: the +# hyper-connection mixer's first projection (fp32 output; both kernels write +# bf16 only) and its up projection [10240, 320] (K=320 tiles neither kernel). + + +@functools.lru_cache(maxsize=None) +def plan(m: int, n: int, k: int) -> _Plan: + """The kernel and tactic `LowMBf16Mm` launches for an [M, K] x [K, N] product.""" + for bound, choice in _TABLE.get((n, k), ()): + if m <= bound: + if choice[0] == "direct": + t = _fi_direct.DirectTactic( + choice[1], choice[2], _fi_direct.default_tactic(m, n, k).rows_per_block + ) + _fi_direct.validate_tactic(t, m, n, k) + return _Plan("direct", tuple(choice[1:]) + (t.rows_per_block,)) + _fi_splitk.validate_tactic(_fi_splitk.SplitKTactic(*choice[1:]), m, n, k) + return _Plan("splitk", tuple(choice[1:])) + return _Plan("cublas") + + +def _cell(why: str, m: int, n: int, k: int, kind: str, pdl: bool = False) -> Cell: + """`kind` is the branch `plan` must pick for this shape; the test holds it to that.""" + return Cell(why=why, spec=dict(m=m, n=n, k=k, kind=kind, pdl=pdl)) + + +class LowMBf16Mm(OpWrapper): + """`mat_a @ mat_b` in bf16 with fp32 accumulation: one kernel launch per call. + + A drop-in for `catalog/gemm/cublas_mm` without bias or `out_dtype`: `mat_a` + is dense row-major `[M, K]`, `mat_b` dense column-major `[K, N]` (a stored + `[N, K]` weight's `.t()`), the result a fresh row-major `[M, N]` bf16 tensor. + `plan(M, N, K)` picks the kernel on the host from three sources: + + - `flashinfer.gemm.kernels.dense_bf16_gemm_direct.run_direct_dense` -- SIMT, + the weight row's whole K prefetched into registers; wins at M <= 2 and on + narrow N; + - `flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk.run_splitk_dense` -- + tcgen05, swap-AB, split-K reduced across a cluster in shared memory (no + extra launch); wins at M 4-32 on the wide shapes; + - `torch.ops.trtllm.cublas_mm` -- everything else, including every M > 32. + + `pdl` launches the FlashInfer kernels with programmatic dependent launch: + they wait on the previous grid (`griddepcontrol.wait`) before reading + `mat_a`, so it is safe after any kernel; it changes nothing on the cuBLAS + branch. + """ + + 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` + # takes for it. The test's host-only plan pins hold the table's other shapes. + CELLS: tuple[Cell, ...] = ( + _cell( + "qwen3.8 decode M=2 GDN in_z: direct, 2-row tile, pdl on", + 2, + 6144, + 2560, + "direct", + pdl=True, + ), + _cell( + "qwen3.8 decode M=32 o_proj: split-K, mma_n 32, pdl on (M boundary)", + 32, + 2560, + 6144, + "splitk", + pdl=True, + ), + _cell("qwen3.8 decode M=128: cuBLAS branch above M=32", 128, 10240, 2560, "cublas"), + ) + + note = """ + Metadata consumed: none. Preconditions: bf16 operands, `mat_a` dense + row-major, `mat_b` dense column-major, both 32-byte aligned (a freshly + allocated activation and a `.t()` weight view are); the FlashInfer kernels + raise ValueError on a layout or alignment they cannot read, the cuBLAS + branch does not (see `is_valid`). + + Compilation: the FlashInfer kernels are CuTe-DSL, JIT-compiled in-process + on the first call per (direct: M, N, K, tactic, pdl; split-K: tactic, pdl) + and cached for the process. That first call must happen before CUDA-graph + capture -- the engine's warmup forward at each captured batch size does it. + + Host cost: the FlashInfer branches spend tens of microseconds of Python per + call (argument validation, DLPack conversion), which a CUDA graph removes + and an eager step does not. Route only graph-captured decode calls here. + + Numerics: every branch accumulates in fp32 and rounds once; the branches + differ from each other only in summation order. + + Hard limits + M FlashInfer branches only at M <= 32 (both kernels' policy) boundary 32 driven at 32 + (cells) and 128 (cuBLAS cell) + N*K largest FlashInfer-branch weight 16480*2560 = 4.2e7 elements boundary 2^31 not probed (48x below it) + """ + + def __call__(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = False) -> torch.Tensor: + m, k = mat_a.shape + n = mat_b.shape[1] + p = plan(m, n, k) + if p.kind == "direct": + out = torch.empty((m, n), dtype=mat_a.dtype, device=mat_a.device) + return _fi_direct.run_direct_dense( + mat_a, mat_b, out, pdl, _fi_direct.DirectTactic(*p.tactic) + ) + if p.kind == "splitk": + out = torch.empty((m, n), dtype=mat_a.dtype, device=mat_a.device) + return _fi_splitk.run_splitk_dense( + mat_a, mat_b, None, out, pdl, _fi_splitk.SplitKTactic(*p.tactic) + ) + return torch.ops.trtllm.cublas_mm(mat_a, mat_b, None, None, 0, None) + + def reference( + self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = False + ) -> torch.Tensor: + """fp32-accumulated, TF32 off; `pdl` is a launch attribute and is ignored.""" + with true_fp32_matmul(): + return (mat_a.float() @ mat_b.float()).to(mat_a.dtype) + + def is_valid(self, mat_a: torch.Tensor, mat_b: torch.Tensor, pdl: bool = False) -> None: + # The cuBLAS branch ignores both leading dimensions: a row-padded + # mat_a or a column-padded mat_b (a slice of a wider tensor) is read as + # if dense -- silently a different product. The FlashInfer branches + # raise on those. A dense row-major mat_b is outside the contract too + # (cuBLAS and the direct kernel raise on it; split-K computes it). + assert mat_a.dtype == torch.bfloat16 and mat_b.dtype == torch.bfloat16, "bf16 operands only" + assert mat_a.is_contiguous(), "mat_a must be dense row-major [M, K]" + 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())" + ) + + +low_m_bf16_mm = LowMBf16Mm() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_gated_residual.py index fc016c9add86..9983a66b1c55 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_gated_residual.py @@ -68,41 +68,22 @@ class Qwen4ExpGatedResidual(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: leading (token dims), hc_count, dtype (bf16), inject (True). Every - # cell uses synthetic weights with this checkpoint's measured element - # statistics, the normed stream reaching ~+-30 as the real one does. + # spec: leading (token dims), inject (True). Every cell is this + # checkpoint's 4x2560 bf16 geometry with synthetic weights of its measured + # element statistics, the normed stream reaching ~+-30 as the real one does. CELLS: tuple[Cell, ...] = ( - tuple( - Cell(why=f"checkpoint 4x2560 bf16, {why}", spec=dict(leading=(t,))) - for t, why in ( - (1, "decode, 1 token"), - (2, "2 tokens"), - (8, "decode batch 8 / MTP draft rows"), - (33, "33 tokens, odd count"), - (512, "prefill 512 tokens"), - (4096, "prefill 4096 tokens"), - (8192, "prefill 8192 tokens, largest measured"), - ) - ) - + ( - Cell( - why="stack-final mixer: block_inject_weight None", - spec=dict(leading=(256,), inject=False), - ), - Cell(why="[3, 17, width] leading dims", spec=dict(leading=(3, 17))), - ) - + tuple( - Cell( - why=f"hc_count {c} at hidden 2560, 64 tokens", spec=dict(leading=(64,), hc_count=c) - ) - for c in (1, 2, 8) - ) - + ( - Cell( - why="float16 activations, 512 tokens", - spec=dict(leading=(512,), dtype=torch.float16), - ), - ) + Cell( + why="checkpoint 4x2560 bf16, decode batch 8 / MTP draft rows", + spec=dict(leading=(8,)), + ), + Cell( + why="checkpoint 4x2560 bf16, prefill 8192 tokens, largest measured", + spec=dict(leading=(8192,)), + ), + Cell( + why="stack-final mixer: block_inject_weight None", + spec=dict(leading=(256,), inject=False), + ), ) note = """ @@ -113,7 +94,7 @@ class Qwen4ExpGatedResidual(OpWrapper): blocks and `shared_head`'s own mixer). Preconditions, all loud (each measured to raise `RuntimeError` before - producing anything): the width divisible by `hc_count` (unflatten), the + 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`). @@ -129,13 +110,16 @@ class Qwen4ExpGatedResidual(OpWrapper): sigmoid saturates (a property of the function, outside the checkpoint's range). - CUDA-graph capture certified, 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. + 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: hidden other than 2560, lowrank other than 320, float32, - weight/activation dtype mixing. Hard limits: not probed (largest operand - 8192 x 10240 = 2^26.3 elements). + 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). """ 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_first_cute.py new file mode 100644 index 000000000000..8e58e9df35cc --- /dev/null +++ b/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 @@ -0,0 +1,605 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode-row hyper-connection site, first half, for one-row batches: the previous +block's injection (with the MoE block's combine and the routed experts' finalize +folded in, when given) and the first projection per stream, one PDL launch. + +Tier 5 (CuTe DSL written for this target). A stream is exactly `THREADS x VEC` +elements, one 16-byte vector per thread (a power-of-two tile would pad a 2560-wide +stream to 8 x 512, three eighths of it masked lanes holding registers), and the grid +is sized so that every CTA is resident at once (one wave), so all of them stream their +weight before the producer finishes. + +Order of work inside a CTA (`OUTS` weight rows x one stream): + +1. its weight rows and its slice of the norm weight into registers -- static, so + before `griddepcontrol.wait`, while the kernel ahead of this one runs; +2. `griddepcontrol.wait`; then the stream's vector of the residual and of the + block output -- with the fold, the slot map and the weights first, then the + slot rows -- and the injection weight; +3. the injected residual rounded once (stored by the stream's first CTA), its sum + of squares and its `(1 + w)`-scaled products with the weight rows, in fp32; +4. a warp then a cross-warp reduction; thread `j < OUTS` stores partial `j`, thread + `OUTS` of the stream's first CTA the sum of squares. +""" + +from __future__ import annotations + +import functools +from typing import Optional, Tuple + +import cuda.bindings.driver as cuda_drv +import cutlass +import cutlass.cute as cute +import torch +from cutlass import const_expr +from cutlass.cute import experimental as cute_ext +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 + +#: One decode row per launch: the forward runs only 1-row decode batches through the +#: site's GEMV pair (`_HC_DECODE_GEMV_MAX_ROWS`). +MAX_ROWS = 1 +VEC = 8 +#: Threads per CTA: one 16-byte vector of a 2560-wide stream each. +THREADS = 320 +#: Streaming multiprocessors the one-wave grid must fit (GB300). +_SMS = 152 +_COMPILE_OPTIONS = "--ptxas-options -maxrregcount=128" + + +def outs_for(n_cols: int, hc_count: int) -> int: + """Weight rows per CTA: the fewest that divide `n_cols` and keep the grid in one wave.""" + for outs in range(1, n_cols + 1): + if n_cols % outs == 0 and (n_cols // outs) * hc_count <= _SMS: + return outs + raise AssertionError(f"no one-wave tiling for {n_cols} columns") + + +class _FirstKernel: + def __init__(self, *, hidden, outs, inject, combine, topk, pdl): + assert hidden == THREADS * VEC, "one vector per thread" + self.hidden, self.outs = hidden, outs + self.inject, self.combine, self.topk, self.pdl = inject, combine, topk, pdl + self.warps = THREADS // 32 + + @cute.jit + def __call__( + self, + mHc, + mBlock, + mInj, + mShared, + mGate, + mTopW, + mTopRows, + mNorm, + mW, + mHcOut, + mPart, + mSs, + stream: cuda_drv.CUstream, + ): + n_tiles = cute.size(mPart, mode=[1]) // self.outs + streams = cute.size(mPart, mode=[0]) + load_static = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cutlass.BFloat16, + num_bits_per_copy=16 * VEC, + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, + ) + load_act = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cutlass.BFloat16, + num_bits_per_copy=16 * VEC, + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, + ) + store = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=16 * VEC + ) + self.kernel( + mHc, + mBlock, + mInj, + mShared, + mGate, + mTopW, + mTopRows, + mNorm, + mW, + mHcOut, + mPart, + mSs, + load_static, + load_act, + store, + ).launch( + grid=[n_tiles, streams, 1], + block=[THREADS, 1, 1], + smem=(self.outs + 1) * self.warps * 4, + stream=stream, + use_pdl=self.pdl, + # At most 96 registers a thread: the mix kernel that follows (512 threads at + # 56) then fits on the SM beside this CTA and streams its weights before + # this kernel ends. + min_blocks_per_mp=2, + ) + + @cute.kernel + def kernel( + self, + mHc, + mBlock, + mInj, + mShared, + mGate, + mTopW, + mTopRows, + mNorm, + mW, + mHcOut, + mPart, + mSs, + load_static: cute.CopyAtom, + load_act: cute.CopyAtom, + store: cute.CopyAtom, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, s, _ = cute.arch.block_idx() + warp = cute.arch.warp_idx() + outs: cutlass.Constexpr = self.outs + warps: cutlass.Constexpr = self.warps + topk: cutlass.Constexpr = self.topk + n0 = bidx * outs + vi = s * THREADS + tidx # this thread's vector of the width; `tidx` of the stream + + w_vec = cute.logical_divide(mW, (None, VEC)) + hc_vec = cute.logical_divide(mHc, (None, VEC)) + nw_vec = cute.logical_divide(mNorm, (None, VEC)) + + # 1. Static operands, before the wait. + w_reg = cute.make_rmem_tensor( + cute.make_layout((outs, VEC), stride=(VEC, 1)), cutlass.BFloat16 + ) + for j in cutlass.range_constexpr(outs): + cute.copy(load_static, w_vec[n0 + j, (None, vi)], w_reg[j, None]) + nw_reg = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16) + cute.copy(load_static, nw_vec[0, (None, vi)], nw_reg) + if const_expr(self.pdl): + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + + # 2. The producers' outputs. + x_reg = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16) + cute.copy(load_act, hc_vec[0, (None, vi)], x_reg) + xr = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.Float32) + for v in cutlass.range_constexpr(VEC): + xr[v] = x_reg[v].to(cutlass.Float32) + if const_expr(self.inject): + b_vec = cute.logical_divide(mBlock, (None, VEC)) + g = mInj[0, s].to(cutlass.Float32) + b = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.Float32) + if const_expr(topk > 0): + # The runner's finalize: fp32 FMAs in slot order from zero, a slot mapped + # to -1 skipped, one rounding. Every slot row is requested before any FMA. + rows = cute.make_rmem_tensor(cute.make_layout((topk,), stride=(1,)), cutlass.Int32) + wts = cute.make_rmem_tensor(cute.make_layout((topk,), stride=(1,)), cutlass.Float32) + for k in cutlass.range_constexpr(topk): + rows[k] = mTopRows[0, k] + wts[k] = mTopW[0, k].to(cutlass.Float32) + fc2 = cute.make_rmem_tensor( + cute.make_layout((topk, VEC), stride=(VEC, 1)), cutlass.BFloat16 + ) + for k in cutlass.range_constexpr(topk): + cute.copy( + load_act, + b_vec[cutlass.max(rows[k], cutlass.Int32(0)), (None, tidx)], + fc2[k, None], + ) + for v in cutlass.range_constexpr(VEC): + acc = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(topk): + if rows[k] >= 0: + acc = cute.math.fma(fc2[k, v].to(cutlass.Float32), wts[k], acc) + b[v] = acc.to(cutlass.BFloat16).to(cutlass.Float32) + else: + blk = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16) + cute.copy(load_act, b_vec[0, (None, tidx)], blk) + for v in cutlass.range_constexpr(VEC): + b[v] = blk[v].to(cutlass.Float32) + if const_expr(self.combine): + # The MoE block output, routed + sigmoid(gate) * shared in fp32, rounded once. + sh_vec = cute.logical_divide(mShared, (None, VEC)) + sh = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16) + cute.copy(load_act, sh_vec[0, (None, tidx)], sh) + z = mGate[0].to(cutlass.Float32) + sg = 1.0 / (1.0 + cute.math.exp(-z)) + for v in cutlass.range_constexpr(VEC): + b[v] = ( + cute.math.fma(sg, sh[v].to(cutlass.Float32), b[v]) + .to(cutlass.BFloat16) + .to(cutlass.Float32) + ) + out = cute.make_rmem_tensor(cute.make_layout((VEC,), stride=(1,)), cutlass.BFloat16) + for v in cutlass.range_constexpr(VEC): + # Rounded once, like `addcmul`: the norm and every later site read this value. + out[v] = cute.math.fma(b[v], g, xr[v]).to(cutlass.BFloat16) + xr[v] = out[v].to(cutlass.Float32) + if bidx == 0: + o_vec = cute.logical_divide(mHcOut, (None, VEC)) + cute.copy(store, out, o_vec[0, (None, vi)]) + + # 3. Sum of squares and the scaled products, fp32. + acc = cute.make_rmem_tensor(cute.make_layout((outs + 1,), stride=(1,)), cutlass.Float32) + acc.fill(0.0) + for v in cutlass.range_constexpr(VEC): + y = xr[v] * (1.0 + nw_reg[v].to(cutlass.Float32)) + for j in cutlass.range_constexpr(outs): + acc[j] = y * w_reg[j, v].to(cutlass.Float32) + acc[j] + acc[outs] = xr[v] * xr[v] + acc[outs] + + # 4. Reductions; thread j < OUTS stores partial j, thread OUTS the sum of squares. + for j in cutlass.range_constexpr(outs + 1): + acc[j] = cute.arch.warp_reduction_sum(acc[j]) + smem = cutlass.utils.SmemAllocator() + red = smem.allocate_tensor( + cutlass.Float32, + cute.make_layout((outs + 1, warps), stride=(warps, 1)), + byte_alignment=16, + ) + with cute.arch.elect_one(): + for j in cutlass.range_constexpr(outs + 1): + red[j, warp] = acc[j] + cute.arch.sync_threads() + if tidx < outs: + total = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(warps): + total = total + red[tidx, w] + mPart[s, n0 + tidx] = total + if tidx == outs: + if bidx == 0: + total = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(warps): + total = total + red[outs, w] + mSs[s] = total + + +def _dl(t: torch.Tensor, align: int = 16): + return from_dlpack(t, assumed_align=align) + + +@functools.lru_cache(maxsize=None) +def _compiled( + hidden, hc_count, n_cols, n_rows, inject, combine, topk, n_perm, pdl, hc_stride, block_stride +): + """One compiled launch per variant and geometry (strides are part of a CuTe layout).""" + outs = outs_for(n_cols, hc_count) + kern = _FirstKernel( + hidden=hidden, outs=outs, inject=inject, combine=combine, topk=topk, pdl=pdl + ) + dev, bf = "cuda", torch.bfloat16 + width = hidden * hc_count + hc = torch.empty(1, hc_stride, dtype=bf, device=dev)[:, :width] + block = torch.empty(max(n_perm, 1), block_stride, dtype=bf, device=dev)[:, :hidden] + inj = torch.empty(1, hc_count, dtype=bf, device=dev) + shared = torch.empty(1, hidden, dtype=bf, device=dev) + gate = torch.empty(1, dtype=bf, device=dev) + topw = torch.empty(1, max(topk, 1), dtype=bf, device=dev) + toprows = torch.empty(1, max(topk, 1), dtype=torch.int32, device=dev) + norm = torch.empty(1, width, dtype=bf, device=dev) + w = torch.empty(n_rows, width, dtype=bf, device=dev) + hc_out = torch.empty(1, width, dtype=bf, device=dev) + part = torch.empty(hc_count, n_cols, dtype=torch.float32, device=dev) + ss = torch.empty(hc_count, dtype=torch.float32, device=dev) + stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) + return cute_ext.compile( + kern, + _dl(hc), + _dl(block), + _dl(inj, 2), + _dl(shared), + _dl(gate, 2), + _dl(topw, 2), + _dl(toprows, 4), + _dl(norm), + _dl(w), + _dl(hc_out), + _dl(part, 4), + _dl(ss, 4), + stream, + options=_COMPILE_OPTIONS, + ) + + +def launch( + hyper_input, + block_out, + inject, + norm_weight, + first, + hc_count, + n_cols, + shared=None, + gate_logits=None, + topk_weights=None, + topk_rows=None, +): + """The launch behind `Qwen4ExpHcDecodeFirstCute.__call__`. Returns `(hc_out, part, ss)`, + `part` fp32 `[C, 1, n_cols]` and `ss` fp32 `[C, 1]`.""" + rows, width = hyper_input.shape + hidden = width // hc_count + dev = hyper_input.device + has_inject = block_out is not None + has_shared = shared is not None + top_k = topk_rows.shape[1] if topk_rows is not None else 0 + hc_out = ( + torch.empty((rows, width), dtype=hyper_input.dtype, device=dev) + if has_inject + else hyper_input + ) + part = torch.empty((hc_count, rows, n_cols), dtype=torch.float32, device=dev) + ss = torch.empty((hc_count, rows), dtype=torch.float32, device=dev) + if rows == 0: + return hc_out, part, ss + pdl = get_env_enable_pdl() + weight = first.t() # the row-major [n_rows, width] weight + gate = gate_logits.reshape(1) if has_shared else hyper_input.new_empty(1) + n_perm = block_out.shape[0] if (has_inject and top_k) else 1 + fn = _compiled( + hidden, + hc_count, + n_cols, + weight.shape[0], + has_inject, + has_shared, + top_k, + n_perm, + bool(pdl), + hyper_input.stride(0), + block_out.stride(0) if (has_inject and top_k) else hidden, + ) + blk = block_out if has_inject else hyper_input.new_empty(1, hidden) + inj = inject if has_inject else hyper_input.new_empty(1, hc_count) + sh = shared if has_shared else hyper_input.new_empty(1, hidden) + topw = topk_weights if top_k else hyper_input.new_empty(1, 1) + toprows = topk_rows if top_k else torch.empty(1, 1, dtype=torch.int32, device=dev) + stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream) + fn( + _dl(hyper_input), + _dl(blk), + _dl(inj, 2), + _dl(sh), + _dl(gate, 2), + _dl(topw, 2), + _dl(toprows, 4), + _dl(norm_weight.reshape(1, -1)), + _dl(weight), + _dl(hc_out), + _dl(part[:, 0, :], 4), + _dl(ss[:, 0], 4), + stream, + ) + return hc_out, part, ss + + +class Qwen4ExpHcDecodeFirstCute(OpWrapper): + """The first half of a one-row decode hyper-connection site: the previous block's + injection and the mixer's first projection (`[down; block_inject]`) per stream. With + `C = hc_count`, `H = width // C`, `N = n_cols`, `first` the `[width, >= N]` view of + the load-time `[down; block_inject; 0]` weight:: + + hc_out[0, s*H+i] = hyper_input[0, s*H+i] + block_out[0, i] * inject[0, s] (rounded once) + part[s, 0, n] = sum_i hc_out[0, s*H+i] * (1 + norm_weight[s*H+i]) * first[s*H+i, n] + ss[s, 0] = sum_i hc_out[0, s*H+i]^2 + + Returns `(hc_out, part, ss)`: `hc_out` fresh in the activation dtype (it *is* + `hyper_input` when `block_out` / `inject` are None -- the norm-only site), `part` + fp32 `[C, 1, N]`, `ss` fp32 `[C, 1]`. The grouped RMS `rstd[s] = (ss[s, 0] / H + + eps)^-1/2` scales stream `s`'s partial: `normed @ first[:, :N] = sum_s rstd[s] * + part[s]`. `qwen4_exp_hc_decode_mix_cute` makes that combination. + + With `shared` and `gate_logits` (the site after a MoE block), `block_out` is the + routed experts' output and the block output injected above is `block_out[0, i] + + sigmoid(gate_logits[0]) * shared[0, i]` in fp32, rounded once to `block_out`'s + dtype -- bitwise what `fused_sigmoid_gate_mul_add` stores. `gate_logits` is one + pre-sigmoid scalar per row, `[T, 1]` or `[T]`. + + With `topk_weights` / `topk_rows` as well (the MoE runner left unfinalized), + `block_out` is its permuted FC2 output `[P, H]` and the routed row is the runner's + finalize, formed here: `routed[0, i] = sum_k topk_weights[0, k] * + block_out[topk_rows[0, k], i]` as fp32 FMAs in slot order from zero (a slot whose + row is -1 is skipped), rounded once to `block_out`'s dtype -- bitwise what the + runner's finalize stores. `topk_rows` is the runner's expanded->permuted index map + `[T, k]` int32, `topk_weights` its `[T, k]` routing weights. The fold needs the + combine (`shared`, `gate_logits`). + + Fusion boundary: the routed experts' finalize (optional), the shared-expert combine + (optional), the previous block's injection, the `(1 + w)` scale and the first + projection per stream, plus the per-stream sum of squares. The routed and shared + experts and the gate projection, the RMS itself, the activations and everything + after are not inside. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: rows (1), inject (True), moe (False: shared + gate_logits [T, 1] too), fold + # (0: top-k of the finalize folded in, block_out then the permuted rows), drop + # (False: some slots -1), n_cols (324). The weight is the target's padded + # [width, 384] (columns past n_cols never read). + # Four cells: the forward reaches four compiled variants (inject at 324 cols, the + # fold at 324, the fold at 320 into the stack-final mixer, norm-only), none of which + # another covers; the -1 slots ride the 324-col fold cell. + CELLS: tuple[Cell, ...] = ( + Cell(why="checkpoint 4x2560, 324 cols of 384, bf16, inject, 1 row", spec=dict(rows=1)), + Cell( + why="finalize fold, top-10, last 3 slots mapped to -1, 1 row", + spec=dict(rows=1, moe=True, fold=10, drop=True), + ), + Cell( + why="finalize fold into the stack-final mixer: 320 cols", + spec=dict(rows=1, moe=True, fold=10, n_cols=320), + ), + Cell(why="norm-only site (stack start / after PLE)", spec=dict(rows=1, inject=False)), + ) + + note = """ + Metadata consumed: none. Call sites: `DecodeTarget`'s trunk, every + hyper-connection site of a one-row decode step (97 per forward), followed by + `qwen4_exp_hc_decode_mix_cute` (or its NVFP4 variant at an MoE site); batches of + more rows take `qwen4_exp_hc_inject_norm` and the mixer. The 47 sites after a MoE + block other than the one before the PLE add pass `shared` and `gate_logits`, and + with them the unfinalized runner's `topk_weights` / `topk_rows`. + + Preconditions: 2-D operands with unit column stride (row strides free); `first` + `[width, >= n_cols]` with unit stride along `width` (the `.t()` view of a row-major + `[N, width]` weight; columns past `n_cols` are never read); one row; `hidden` 2560 + (320 threads x 8); bf16 operands; `n_cols` with a divisor that keeps `n_cols / outs x C <= 152` + CTAs (324 -> 9 rows x 36 tiles x 4 streams = 144; 320 -> 10 rows, 128); the + permuted rows' count is part of the compiled layout (one compile per value). + + Numerics: `hc_out` as `qwen4_exp_hc_inject_norm` (fp32 FMAs, one rounding; the + finalize and the combine likewise); `part` and `ss` are fp32 reductions in a fixed + order, inside the float64 reference's band (the combine and the finalize can each + land one ulp of the block output from it, scaled into `hc_out` by `inject`). + Deterministic (measured). Graph-capturable: compiled on the first call per + variant, which must precede capture (the engine's warmup forward does it). + + PDL: with `TRTLLM_ENABLE_PDL` (default on, sm_90+) the weight rows and the norm + slice stream before `griddepcontrol.wait`, every activation read and every + write after it, and dependents are released before the wait. + + Certified: one bf16 row, the injection at 324 columns, the fold (top-10, three + slots -1) at 324 columns, the fold into the stack-final mixer at 320 columns, + the norm-only site; graph replay under PDL with the fold. + + Not certified: the combine without the fold, a flat `[T]` gate, a row-slice + `hyper_input`, a weight narrower than the padded 384 rows (accepted; no cell + drives them). + """ + + def __call__( + self, + hyper_input: torch.Tensor, + block_out: Optional[torch.Tensor], + inject: Optional[torch.Tensor], + norm_weight: torch.Tensor, + first: torch.Tensor, + hc_count: int, + n_cols: int, + shared: Optional[torch.Tensor] = None, + gate_logits: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + topk_rows: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return launch( + hyper_input, + block_out, + inject, + norm_weight, + first, + hc_count, + n_cols, + shared, + gate_logits, + topk_weights, + topk_rows, + ) + + def reference( + self, + hyper_input: torch.Tensor, + block_out: Optional[torch.Tensor], + inject: Optional[torch.Tensor], + norm_weight: torch.Tensor, + first: torch.Tensor, + hc_count: int, + n_cols: int, + shared: Optional[torch.Tensor] = None, + gate_logits: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + topk_rows: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """float64; the finalize (with `topk_rows`) rounded once to + `block_out`'s dtype, the combine (with `shared`) rounded once to `block_out`'s + dtype, `hc_out` rounded once to the activation dtype, `part` and `ss` + rounded once to fp32.""" + dtype = hyper_input.dtype + rows, width = hyper_input.shape + hidden = width // hc_count + if block_out is None: + hc_out = hyper_input + else: + if topk_rows is not None: + valid = (topk_rows >= 0).unsqueeze(-1) + rows_ = block_out.double()[topk_rows.clamp(min=0).long()] + w = topk_weights.double().unsqueeze(-1) + block_out = torch.where(valid, rows_ * w, 0.0).sum(1).to(block_out.dtype) + if shared is not None: + sg = torch.sigmoid(gate_logits.double().reshape(rows, 1)) + block_out = (block_out.double() + sg * shared.double()).to(block_out.dtype) + streams = hyper_input.double().unflatten(-1, (hc_count, hidden)) + gated = block_out.double().unsqueeze(-2) * inject.double().unsqueeze(-1) + hc_out = (streams + gated).flatten(-2).to(dtype) + x = hc_out.double().unflatten(-1, (hc_count, hidden)) + y = x * (1.0 + norm_weight.double()).unflatten(-1, (hc_count, hidden)) + w = first[:, :n_cols].double().unflatten(0, (hc_count, hidden)) + part = torch.einsum("tsi,sin->stn", y, w).float() + ss = x.pow(2).sum(dim=-1).t().float() + return hc_out, part, ss + + def is_valid( + self, + hyper_input: torch.Tensor, + block_out: Optional[torch.Tensor], + inject: Optional[torch.Tensor], + norm_weight: torch.Tensor, + first: torch.Tensor, + hc_count: int, + n_cols: int, + shared: Optional[torch.Tensor] = None, + gate_logits: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + topk_rows: Optional[torch.Tensor] = None, + ) -> None: + # Raw row strides with unit column stride everywhere; a contiguous + # [width, N] weight would be read with its rows and columns swapped. + assert hyper_input.dim() == 2 and hyper_input.stride(1) == 1, "hyper_input [T, width]" + rows, width = hyper_input.shape + assert width % hc_count == 0 and (width // hc_count) % 16 == 0, "hidden multiple of 16" + assert (block_out is None) == (inject is None), "block_out and inject come together" + assert (shared is None) == (gate_logits is None), "shared and gate_logits come together" + assert shared is None or block_out is not None, "the combine needs block_out" + assert (topk_weights is None) == (topk_rows is None), ( + "topk_weights and topk_rows come together" + ) + assert topk_rows is None or shared is not None, "the finalize fold needs the combine" + if block_out is not None: + if topk_rows is None: + assert block_out.shape == (rows, width // hc_count) and block_out.stride(1) == 1 + else: + # The permuted FC2 rows: any row count, read through topk_rows. + assert block_out.dim() == 2 and block_out.shape[1] == width // hc_count + assert block_out.stride(1) == 1, "permuted rows [P, hidden]" + assert topk_rows.dtype == torch.int32, "topk_rows int32 (the runner's index map)" + assert topk_rows.dim() == 2 and topk_rows.shape[0] == rows, "topk_rows [T, k]" + assert topk_rows.stride(1) == 1 and topk_weights.stride(1) == 1 + assert topk_weights.shape == topk_rows.shape, "topk_weights [T, k]" + assert inject.shape == (rows, hc_count) and inject.stride(1) == 1 + if shared is not None: + assert shared.shape == (rows, width // hc_count), "shared [T, hidden]" + assert shared.stride(1) == 1, "shared [T, hidden]" + assert shared.dtype == block_out.dtype, "shared in block_out's dtype" + assert gate_logits.shape in ((rows,), (rows, 1)), "gate_logits [T, 1] or [T]" + assert rows <= 1 or gate_logits.stride(0) != 0, "one gate logit per row" + assert norm_weight.shape == (width,) and norm_weight.is_contiguous(), "norm_weight [width]" + assert first.shape[0] == width and first.shape[1] >= n_cols, "first [width, >= n_cols]" + assert first.stride(0) == 1, "first must be the .t() view of a row-major [N, width] weight" + assert hyper_input.shape[0] <= MAX_ROWS, "one row" + assert hyper_input.shape[1] // hc_count == THREADS * VEC, "hidden 2560" + assert hyper_input.dtype == torch.bfloat16, "bf16" + outs_for(n_cols, hc_count) + + +qwen4_exp_hc_decode_first_cute = Qwen4ExpHcDecodeFirstCute() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_decode_mix_fp4_cute.py new file mode 100644 index 000000000000..45a3c08057bd --- /dev/null +++ b/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 @@ -0,0 +1,506 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode-row hyper-connection site at an MoE site, second half, for one-row batches: +the grouped RMS, the SiLU bottleneck and the injection weights from the first half's +partials, the up projection, the sigmoid gate, the stream mean and the NVFP4 of the +mean for the routed experts -- one PDL launch. + +Tier 5 (CuTe DSL written for this target). A CTA owns one 16-column NVFP4 block of the +hidden width across the 4 streams (64 rows of the up weight), and 8 threads share a +row, 40 weights each; 160 CTAs of 512 threads, all resident at once. + +Order of work inside a CTA: + +1. its 64 weight rows and its norm-weight elements into registers (static), before + `griddepcontrol.wait`; +2. after the wait: the four sums of squares (every thread), one partial column per + thread for the first 324 threads -> the shared `low` vector and, in CTA 0, the + injection weights; each thread's residual element; +3. each row's dot with `low` over its 8 threads, the sigmoid gate on the normed + element, the mean over streams (16 threads), the store; with NVFP4, warp 0 also + quantizes the block as `torch.ops.trtllm.fp4_quantize` does (linear scales, + `rcp.approx.ftz` reciprocals). +""" + +from __future__ import annotations + +import functools + +import cuda.bindings.driver as cuda_drv +import cutlass +import cutlass.cute as cute +import torch +import torch.nn.functional as F +from cutlass import const_expr +from cutlass.cute import experimental as cute_ext +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.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +MAX_ROWS = 1 +VEC = 8 +#: Hidden columns per CTA: one NVFP4 scale block. +BLOCK_H = 16 +#: Threads sharing one up-weight row. +GROUP = 8 +_COMPILE_OPTIONS = "--ptxas-options -maxrregcount=64" + + +class _MixKernel: + def __init__(self, *, hidden, streams, rank, inject_out, fp4, pdl): + assert rank % (GROUP * VEC) == 0 and hidden % BLOCK_H == 0 + self.hidden, self.streams, self.rank = hidden, streams, rank + self.inject_out, self.fp4, self.pdl = inject_out, fp4, pdl + self.k_tiles = rank // (GROUP * VEC) + self.threads = streams * BLOCK_H * GROUP + assert self.threads >= rank + streams and self.threads % 32 == 0 + + @cute.jit + def __call__( + self, + mPart, + mSs, + mUp, + mHc, + mNorm, + mGs, + mMixed, + mInj, + mXq, + mXsf, + eps: cutlass.Float32, + stream: cuda_drv.CUstream, + ): + load_static = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cutlass.BFloat16, + num_bits_per_copy=16 * VEC, + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, + ) + self.kernel( + mPart, mSs, mUp, mHc, mNorm, mGs, mMixed, mInj, mXq, mXsf, eps, load_static + ).launch( + grid=[self.hidden // BLOCK_H, 1, 1], + block=[self.threads, 1, 1], + smem=((self.rank // VEC) * (VEC + 1) + self.streams * BLOCK_H) * 4 + 32, + stream=stream, + use_pdl=self.pdl, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + mPart, + mSs, + mUp, + mHc, + mNorm, + mGs, + mMixed, + mInj, + mXq, + mXsf, + eps: cutlass.Float32, + load_static: cute.CopyAtom, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + warp = cute.arch.warp_idx() + lane = cute.arch.lane_idx() + hidden: cutlass.Constexpr = self.hidden + streams: cutlass.Constexpr = self.streams + rank: cutlass.Constexpr = self.rank + k_tiles: cutlass.Constexpr = self.k_tiles + i0 = bidx * BLOCK_H + row = tidx // GROUP # (stream, column) row of this CTA: stream-major + q = tidx % GROUP + s = row // BLOCK_H + i = row % BLOCK_H + j = s * hidden + i0 + i # the up-weight row, a column of the residual width + + # 1. Static operands, before the wait. + u_vec = cute.logical_divide(mUp, (None, VEC)) + u_reg = cute.make_rmem_tensor( + cute.make_layout((k_tiles, VEC), stride=(VEC, 1)), cutlass.BFloat16 + ) + for kt in cutlass.range_constexpr(k_tiles): + cute.copy(load_static, u_vec[j, (None, kt * GROUP + q)], u_reg[kt, None]) + gam = 1.0 + mNorm[0, j].to(cutlass.Float32) + gs = cutlass.Float32(0.0) + if const_expr(self.fp4): + gs = mGs[0] + if const_expr(self.pdl): + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + + # 2. Grouped RMS, logits, low (shared), injection weights. Every producer value this + # thread needs is requested before the first is used (a partial column clamped into + # range for the threads past it), so the CTA pays one L2 round trip. + cols: cutlass.Constexpr = rank + (streams if self.inject_out else 0) # partial columns read + pcol = cutlass.min(tidx, cols - 1) + pv = [mPart[t, pcol] for t in range(streams)] + ssv = [mSs[t] for t in range(streams)] + x = mHc[0, j].to(cutlass.Float32) + rstd = cute.make_rmem_tensor(cute.make_layout((streams,), stride=(1,)), cutlass.Float32) + for t in cutlass.range_constexpr(streams): + rstd[t] = cute.math.rsqrt(ssv[t] / hidden + eps) + smem = cutlass.utils.SmemAllocator() + # `low` in rows of VEC with one pad word, so the 8 threads of a row (stride VEC + 1 + # words) read 8 different banks. + low = smem.allocate_tensor( + cutlass.Float32, + cute.make_layout((rank // VEC, VEC), stride=(VEC + 1, 1)), + byte_alignment=16, + ) + vals = smem.allocate_tensor( + cutlass.Float32, + cute.make_layout((streams, BLOCK_H), stride=(BLOCK_H, 1)), + byte_alignment=16, + ) + if tidx < cols: + logit = cutlass.Float32(0.0) + for t in cutlass.range_constexpr(streams): + logit = pv[t] * rstd[t] + logit + z = logit * (1.0 / streams) + sig = 1.0 / (1.0 + cute.math.exp(-z)) + if tidx < rank: + low[tidx // VEC, tidx % VEC] = z * sig + else: + if const_expr(self.inject_out): + if bidx == 0: + mInj[0, tidx - rank] = (2.0 * sig).to(cutlass.BFloat16) + cute.arch.sync_threads() + + # 3. Up projection over the row's 8 threads (two independent sums), gate, normed + # element, stream mean. + sums = [cutlass.Float32(0.0), cutlass.Float32(0.0)] + for kt in cutlass.range_constexpr(k_tiles): + for e in cutlass.range_constexpr(VEC): + sums[e % 2] = ( + u_reg[kt, e].to(cutlass.Float32) * low[kt * GROUP + q, e] + sums[e % 2] + ) + dot = sums[0] + sums[1] + dot = cute.arch.warp_reduction(dot, lambda a, b: a + b, threads_in_group=GROUP) + if q == 0: + gate = 1.0 / (1.0 + cute.math.exp(-dot)) + vals[s, i] = gate * (x * rstd[s] * gam) + cute.arch.sync_threads() + if warp == 0: + m = cutlass.Float32(0.0) + if lane < BLOCK_H: + tot = cutlass.Float32(0.0) + for t in cutlass.range_constexpr(streams): + tot = tot + vals[t, lane] + mb = (tot * (1.0 / streams)).to(cutlass.BFloat16) + mMixed[0, i0 + lane] = mb + m = mb.to(cutlass.Float32) + if const_expr(self.fp4): + # NVFP4 of the 16 rounded values in lanes 0..15 (lanes 16..31 a second, + # discarded group), as `fp4_quantize`: e4m3(gs * amax / 6), then + # x * rcp(sf * rcp(gs)) rounded to e2m1, nearest, ties to even. + amax = cute.arch.warp_reduction_max(cute.math.abs(m), threads_in_group=BLOCK_H) + sf8 = (gs * (amax * cute.arch.rcp_approx(6.0))).to(cutlass.Float8E4M3FN) + osc = cute.arch.rcp_approx(sf8.to(cutlass.Float32) * cute.arch.rcp_approx(gs)) + if amax == 0.0: + osc = cutlass.Float32(0.0) + sv = m * osc + a = cute.math.abs(sv) + code = cutlass.Int32(0) + if a > 0.25: + code = code + 1 + if a >= 0.75: + code = code + 1 + if a > 1.25: + code = code + 1 + if a >= 1.75: + code = code + 1 + if a > 2.5: + code = code + 1 + if a >= 3.5: + code = code + 1 + if a > 5.0: + code = code + 1 + if cute.math.copysign(1.0, sv) < 0.0: # the sign bit, -0.0 included + code = code + 8 + hi = cute.arch.shuffle_sync_down(code, offset=1) + if lane < BLOCK_H: + if lane % 2 == 0: + mXq[0, (i0 + lane) // 2] = (code + hi * 16).to(cutlass.Uint8) + if lane == 0: + mXsf[bidx] = sf8 + + +def _dl(t: torch.Tensor, align: int = 16): + return from_dlpack(t, assumed_align=align) + + +@functools.lru_cache(maxsize=None) +def _compiled( + hidden, streams, rank, part_cols, part_stride, ss_stride, hc_stride, inject_out, fp4, pdl +): + kern = _MixKernel( + hidden=hidden, streams=streams, rank=rank, inject_out=inject_out, fp4=fp4, pdl=pdl + ) + dev, bf = "cuda", torch.bfloat16 + width = hidden * streams + part = torch.empty(streams, part_stride, dtype=torch.float32, device=dev)[:, :part_cols] + ss = torch.empty(streams * ss_stride, dtype=torch.float32, device=dev)[::ss_stride] + up = torch.empty(width, rank, dtype=bf, device=dev) + hc = torch.empty(1, hc_stride, dtype=bf, device=dev)[:, :width] + norm = torch.empty(1, width, dtype=bf, device=dev) + gs = torch.empty(1, dtype=torch.float32, device=dev) + mixed = torch.empty(1, hidden, dtype=bf, device=dev) + inj = torch.empty(1, streams, dtype=bf, device=dev) + xq = torch.empty(1, hidden // 2, dtype=torch.uint8, device=dev) + xsf = torch.empty(hidden // 16, dtype=torch.float8_e4m3fn, device=dev) + stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) + return cute_ext.compile( + kern, + _dl(part, 4), + _dl(ss, 4), + _dl(up), + _dl(hc, 2), + _dl(norm, 2), + _dl(gs, 4), + _dl(mixed, 2), + _dl(inj, 2), + _dl(xq, 1), + _dl(xsf, 1), + cutlass.Float32(1e-6), + stream, + options=_COMPILE_OPTIONS, + ) + + +def launch(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale=None): + """The launch behind `Qwen4ExpHcDecodeMixFp4Cute.__call__`. Returns `(mixed, + inject_out, xq, xsf)`; without `global_scale` (the kernel's plain build, kept for + measurement) `xq` / `xsf` are None.""" + rows, width = hc_out.shape + hidden = width // hc_count + rank = up.shape[0] + dev, dt = hc_out.device, hc_out.dtype + fp4 = global_scale is not None + mixed = torch.empty((rows, hidden), dtype=dt, device=dev) + inject_out = torch.empty((rows, hc_count), dtype=dt, device=dev) if with_inject else None + xq = torch.empty((rows, hidden // 2), dtype=torch.uint8, device=dev) if fp4 else None + xsf = torch.empty((rows * (hidden // 16),), dtype=torch.uint8, device=dev) if fp4 else None + if rows == 0: + return mixed, inject_out, xq, xsf + pdl = get_env_enable_pdl() + part2 = part[:, 0, :] + fn = _compiled( + hidden, + hc_count, + rank, + part2.shape[1], + part2.stride(0), + ss.stride(0), + hc_out.stride(0), + bool(with_inject), + fp4, + bool(pdl), + ) + stream = cuda_drv.CUstream(torch.cuda.current_stream(dev).cuda_stream) + gs = global_scale.reshape(1) if fp4 else mixed.new_empty(1, dtype=torch.float32) + inj = inject_out if with_inject else mixed.new_empty(1, hc_count) + q = xq if fp4 else mixed.new_empty(1, hidden // 2, dtype=torch.uint8) + sf = ( + xsf.view(torch.float8_e4m3fn) + if fp4 + else mixed.new_empty(hidden // 16, dtype=torch.float8_e4m3fn) + ) + fn( + _dl(part2, 4), + _dl(ss[:, 0], 4), + _dl(up.t()), + _dl(hc_out, 2), + _dl(norm_weight.reshape(1, -1), 2), + _dl(gs, 4), + _dl(mixed, 2), + _dl(inj, 2), + _dl(q, 1), + _dl(sf, 1), + cutlass.Float32(eps), + stream, + ) + return mixed, inject_out, xq, xsf + + +def _check(part, ss, up, hc_out, norm_weight, hc_count, with_inject) -> None: + rows, width = hc_out.shape + rank = up.shape[0] + assert hc_out.stride(1) == 1, "hc_out unit column stride" + assert part.dtype == torch.float32 and part.stride(2) == 1, "part fp32, unit last stride" + assert part.shape[:2] == (hc_count, rows), "part [C, T, N]" + assert part.shape[2] >= rank + (hc_count if with_inject else 0), "part too narrow" + assert ss.dtype == torch.float32 and ss.shape == (hc_count, rows), "ss fp32 [C, T]" + assert up.shape == (rank, width), "up [R, width]" + assert norm_weight.shape == (width,) and norm_weight.is_contiguous() + assert rows <= MAX_ROWS, "one row" + assert width // hc_count == 2560 and rank == 320 and hc_count == 4, ( + "the checkpoint's 4 x 2560, rank 320" + ) + assert hc_out.dtype == torch.bfloat16 and up.dtype == torch.bfloat16, "bf16" + assert up.t().is_contiguous(), "up the [R, width] view of a row-major [width, R] weight" + + +def _reference(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject): + """`(mixed, inject_out)` in float64, one rounding per output.""" + dtype = hc_out.dtype + rank = up.shape[0] + hidden = hc_out.shape[1] // hc_count + rstd = torch.rsqrt(ss.double().t() / hidden + eps) + logits = torch.einsum("ts,stn->tn", rstd, part.double()) + low = F.silu(logits[:, :rank] / hc_count) + gate = torch.sigmoid(low @ up.double()) + x = hc_out.double().unflatten(-1, (hc_count, hidden)) * rstd.unsqueeze(-1) + normed = x.flatten(-2) * (1.0 + norm_weight.double()) + mixed = (gate * normed).unflatten(-1, (hc_count, hidden)).mean(dim=-2).to(dtype) + inject_out = None + if with_inject: + inject_out = (2.0 * torch.sigmoid(logits[:, rank : rank + hc_count] / hc_count)).to(dtype) + return mixed, inject_out + + +_INJECT = Cell(why="checkpoint 4x2560 r320 bf16, inject, 1 row (con 1)", spec=dict(rows=1)) +_ROW_SLICE = Cell(why="hc_out a row slice buf[::2], 1 row", spec=dict(rows=1, layout="row_slice")) + + +class Qwen4ExpHcDecodeMixFp4Cute(OpWrapper): + """`qwen4_exp_hc_decode_mix_cute` plus the NVFP4 of its `mixed` for the routed + experts. With `C = hc_count`, `H = width // C`, `R = up.shape[0]`, `part` / `ss` the + per-stream partial logits `[C, T, >= R (+ C)]` and sums of squares `[C, T]` from + `qwen4_exp_hc_decode_first_cute`, `up` the `[R, width]` view of the up weight:: + + rstd[t, s] = 1 / sqrt(ss[s, t] / H + eps) + logits[t, n] = sum_s rstd[t, s] * part[s, t, n] + low = silu(logits[:, :R] / C) + gate = sigmoid(low @ up) # [T, C*H] + normed[t, s*H+i] = hc_out[t, s*H+i] * rstd[t, s] * (1 + norm_weight[s*H+i]) + mixed[t, i] = mean over s of gate[t, s*H+i] * normed[t, s*H+i] + inject_out = 2 * sigmoid(logits[:, R:R+C] / C) # with_inject only + xq, xsf = NVFP4(mixed) # e2m1 pairs, e4m3 scale per 16, linear layout + + fp32 inside, one rounding per output; the quantization reads the rounded `mixed`. + Returns fresh `(mixed, inject_out, xq [T, H/2] uint8, xsf [T*H/16] uint8)`; + `inject_out` is None without `with_inject`. `global_scale` is the 1-element fp32 + NVFP4 global scale. + + Fusion boundary: everything `qwen4_exp_hc_decode_mix_cute` fuses, plus the NVFP4 + quantization of `mixed`. The partials are the producer's. + """ + + ARCHS = frozenset({Arch.SM_103}) + # Its call sites always inject, so no stack-final (no-inject) cell. + CELLS: tuple[Cell, ...] = (_INJECT, _ROW_SLICE) + + note = """ + Metadata consumed: none. Call sites: `DecodeTarget`'s MoE sites at a one-row + decode step, after `qwen4_exp_hc_decode_first_cute`; `xq`/`xsf` feed + `fp4_block_scale_moe_renorm_routed` (linear scales). Batches of more rows take + `qwen4_exp_hc_up_gate_mean_fp4`. + + Preconditions: those of `qwen4_exp_hc_decode_mix_cute`, and `global_scale` fp32 with + one element. + + Numerics: `mixed` and `inject_out` as `qwen4_exp_hc_decode_mix_cute`; the NVFP4 is + `fp4_quantize`'s arithmetic on the rounded `mixed` (`rcp.approx.ftz` reciprocals, + e4m3 scale rounded to nearest, e2m1 nearest-even with saturation at 6). + Deterministic (measured); graph-capturable once compiled (the first call per + variant). PDL: the weight rows and norm elements stream before + `griddepcontrol.wait`. + + Certified: one bf16 row with the injection, `hc_out` compact and a row slice; + graph replay under PDL. Not certified: `with_inject=False`. + """ + + def __call__(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale): + return launch(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale) + + def reference( + self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale + ): + """float64 with one rounding per output, then `fp4_quantize`'s native-torch NVFP4 + of the rounded `mixed` (linear scales).""" + mixed, inject_out = _reference( + part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject + ) + return (mixed, inject_out) + fp4_quantize.reference( + mixed, global_scale, BLOCK_H, False, False + ) + + def is_valid( + self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject, global_scale + ) -> None: + _check(part, ss, up, hc_out, norm_weight, hc_count, with_inject) + # The kernel loads element 0 only; a per-row scale would be applied as row 0's. + assert global_scale.numel() == 1 and global_scale.dtype == torch.float32, ( + "global_scale fp32 [1]" + ) + + +qwen4_exp_hc_decode_mix_fp4_cute = Qwen4ExpHcDecodeMixFp4Cute() + + +class Qwen4ExpHcDecodeMixCute(OpWrapper): + """The second half of a one-row decode hyper-connection site: the grouped RMS, the + SiLU bottleneck and the injection weights from `qwen4_exp_hc_decode_first_cute`'s + per-stream partials, then the up projection, the sigmoid gate on the normed stream + and the mean over streams -- `qwen4_exp_hc_decode_mix_fp4_cute`'s formulas without + the NVFP4 output. Returns fresh `(mixed, inject_out)` in `hc_out`'s dtype; + `inject_out` is None without `with_inject`. + + Fusion boundary: the grouped RMS, the bottleneck SiLU, the injection weights, the + up projection, the sigmoid gate, the normed stream and the stream mean. The + partials are the producer's. + """ + + ARCHS = frozenset({Arch.SM_103}) + CELLS: tuple[Cell, ...] = ( + _INJECT, + Cell(why="stack-final mixer (no inject), 1 row", spec=dict(rows=1, with_inject=False)), + _ROW_SLICE, + ) + + note = """ + Metadata consumed: none. Call sites: `DecodeTarget`'s non-MoE sites at a one-row + decode step, after `qwen4_exp_hc_decode_first_cute`; batches of more rows take + `qwen4_exp_hc_up_gate_mean`. Beside a first kernel capped at 96 registers, its 160 + CTAs (512 threads at 49 registers) are all resident while that kernel runs, so the + up weight is in registers when the wait returns. + + Preconditions: `part` fp32 `[C, T, >= R (+ C)]` and `ss` fp32 `[C, T]`, `part` with + unit last stride; `hc_out` `[T, width]` with unit column stride; `norm_weight` + contiguous `[width]`; one row, the checkpoint's 4 x 2560 width, rank 320, bf16; + `up` the `.t()` view of a contiguous `[width, R]` weight. + + Numerics: fp32 inside (`low`, the gate logits and `normed` are never rounded), one + rounding per output, inside the float64 reference's band. Deterministic + (measured); graph-capturable once compiled (the first call per variant). + + Certified: one bf16 row with and without the injection, `hc_out` compact and a + row slice. Not certified: graph replay of this build (the NVFP4 build's + replay test covers the shared kernel, not this compiled variant). + """ + + def __call__(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject): + mixed, inject_out, _, _ = launch( + part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject + ) + return mixed, inject_out + + def reference(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject): + """float64, one rounding per output.""" + return _reference(part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject) + + def is_valid(self, part, ss, up, hc_out, norm_weight, eps, hc_count, with_inject) -> None: + _check(part, ss, up, hc_out, norm_weight, hc_count, with_inject) + + +qwen4_exp_hc_decode_mix_cute = Qwen4ExpHcDecodeMixCute() 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__tp1/ops/hc/qwen4_exp_hc_decode_tc_first.py new file mode 100644 index 000000000000..df942de81ed0 --- /dev/null +++ b/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 @@ -0,0 +1,660 @@ +# 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..128 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 +`[128 rows, BN columns]` partial over `K / KS` of the 10240-wide stream: a TMA +warp streams the normed rows (rows past `T` zero-filled by TMA, never read) +and the weight rows (only the `n_cols` real ones; the load-time zero padding +to 384 rows is never read), one MMA warp accumulates in tensor memory, four +epilogue warps move the fp32 partial into shared memory. The `KS` CTAs of one +column tile form a thread-block cluster; after a cluster barrier each CTA sums +a slice of the tile's `KS` partials out of its peers' shared memory (DSMEM) in +a fixed order (deterministic) and writes the activations. + +Launched with programmatic dependent launch: the weight stages stream before +the grid-dependency wait (the producer, `qwen4_exp_hc_inject_norm`, releases +dependents at its start), and this kernel releases its own consumer +(`qwen4_exp_hc_up_gate_mean(_fp4)`) at once so that it prefetches its weight. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import torch +import torch.nn.functional as F +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op +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.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 + +#: Largest row count this entry is certified for: one 128-row MMA tile. +MAX_ROWS = 128 + +#: Programmatic dependent launch (sm_90+), off with TRTLLM_ENABLE_PDL=0. Without it +#: the grid-dependency wait and the dependents trigger are no-ops. +_PDL = get_env_enable_pdl() + + +def _asm(ret, args, text, constraints, loc, ip, side_effects=False): + return llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _mapa(addr, rank, *, loc=None, ip=None): + """This CTA's shared-memory address `addr` in cluster CTA `rank`'s window.""" + return cutlass.Int32( + _asm( + T.i32(), + [cutlass.Int32(addr), cutlass.Int32(rank)], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _ld_dsmem_f32(addr, *, loc=None, ip=None): + """`ld.shared::cluster.f32`; volatile so it stays behind the cluster barrier.""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int32(addr)], + "ld.shared::cluster.f32 $0, [$1];", + "=f,r", + loc, + ip, + side_effects=True, + ) + ) + + +class HcDecodeTcFirstKernel: + """`low[t, n] = silu(L[t, n] / C)` for `n < R` and, with `with_inject`, + `inject[t, j] = 2 * sigmoid(L[t, R + j] / C)`, where `L = A @ B^T`, + `A = normed [T, K]`, `B = weight rows [0, n_cols)` of `[>= n_cols, K]`. + + Compile-time: `bn` (MMA N, a multiple of 16), `ks` (K split = cluster + size), `hc`, `rank`, `n_cols`, `k`, `with_inject`. Runtime: the row count + (<= 128) and the pointers. + """ + + def __init__(self, bn, ks, hc, rank, n_cols, k, with_inject, epi_cols=16): + assert bn % epi_cols == 0 and bn % 16 == 0 and 16 <= bn <= 256 + assert k % (64 * ks) == 0, "K must split into 64-wide tiles per CTA" + self.bn, self.ks, self.hc, self.rank, self.n_cols, self.k = bn, ks, hc, rank, n_cols, k + self.with_inject = with_inject + self.n_tiles = (n_cols + bn - 1) // bn + self.acc_dtype = cutlass.Float32 + self.io_dtype = cutlass.BFloat16 + self.mma_tiler = (128, bn, 1) + self.epi_cols = epi_cols + self.cta_group = tcgen05.CtaGroup.ONE + # The MMA itself is single-CTA: no multicast; the launch cluster is along K only. + self.cluster_shape_mn = (1, 1) + self.epilogue_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.threads_per_cta = 32 * 6 + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + self.num_acc_stage = 1 + # Flat elements of one column tile each CTA reduces, at most, per epilogue thread. + per_cta = (128 * bn + ks - 1) // ks + self.red_iters = (per_cta + self.threads_per_cta - 1) // self.threads_per_cta + + def _setup(self, a_major, b_major): + tiled_mma = utils.sm100.make_trivial_tiled_mma( + self.io_dtype, a_major, b_major, self.acc_dtype, self.cta_group, self.mma_tiler[:2] + ) + inst_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + self.mma_tiler = (self.mma_tiler[0], self.mma_tiler[1], inst_k * 4) + self.cta_tile_shape_mnk = self.mma_tiler + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), (tiled_mma.thr_id.shape,) + ) + self.epi_tile = (self.mma_tiler[0], self.epi_cols) + a_one = utils.sm100.make_smem_layout_a(tiled_mma, self.mma_tiler, self.io_dtype, 1) + b_one = utils.sm100.make_smem_layout_b(tiled_mma, self.mma_tiler, self.io_dtype, 1) + ab_bytes = cute.size_in_bytes(self.io_dtype, a_one) + cute.size_in_bytes( + self.io_dtype, b_one + ) + smem_capacity = utils.get_smem_capacity_in_bytes() + self.k_tiles = self.k // self.mma_tiler[2] // self.ks + # The fp32 partial, one padded row per accumulator row (conflict-free row stores). + self.p_stride = self.bn + 1 + p_bytes = 128 * self.p_stride * 4 + self.num_ab_stage = max(2, min(self.k_tiles, (smem_capacity - 2048 - p_bytes) // ab_bytes)) + self.n_pre = min(self.num_ab_stage, self.k_tiles) + self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( + tiled_mma, self.mma_tiler, self.io_dtype, self.num_ab_stage + ) + self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( + tiled_mma, self.mma_tiler, self.io_dtype, self.num_ab_stage + ) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + acc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(acc_fake, arch="sm_100") + return tiled_mma + + @cute.jit + def __call__( + self, + m: cutlass.Int32, + a_ptr: cute.Pointer, + b_ptr: cute.Pointer, + low_t: cute.Tensor, + inj_ptr: cute.Pointer, + stream: cuda.CUstream, + ): + k, ks = self.k, self.ks + a = cute.make_tensor(a_ptr, cute.make_layout((m, k, 1), stride=(k, 1, 0))) + b = cute.make_tensor(b_ptr, cute.make_layout((self.n_cols, k, 1), stride=(k, 1, 0))) + # `low` arrives as a tensor (the TVM-FFI entry takes its stream from a GPU tensor). + low = cute.make_tensor( + low_t.iterator, cute.make_layout((m, self.rank), stride=(self.rank, 1)) + ) + inj = cute.make_tensor(inj_ptr, cute.make_layout((m, self.hc), stride=(self.hc, 1))) + k_major = utils.LayoutEnum.ROW_MAJOR.mma_major_mode() + tiled_mma = self._setup(k_major, k_major) + + a_op = utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + a_smem = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, a, a_smem, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape + ) + b_op = utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + b_smem = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, b, b_smem, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape + ) + self.num_tma_load_bytes = cute.size_in_bytes(self.io_dtype, a_smem) + cute.size_in_bytes( + self.io_dtype, b_smem + ) + self.kernel( + tiled_mma, + tma_atom_a, + tma_a, + tma_atom_b, + tma_b, + low, + inj, + m, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.epi_tile, + ).launch( + grid=[ks, self.n_tiles, 1], + block=[self.threads_per_cta, 1, 1], + cluster=[ks, 1, 1], + stream=stream, + use_pdl=_PDL, + ) + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + mLow: cute.Tensor, + mInj: cute.Tensor, + m: cutlass.Int32, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + epi_tile: cute.Tile, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + # The consumer (the up projection, launched with programmatic dependent launch) + # may start at once: it streams its own static weight, then waits for this grid. + cute.arch.griddepcontrol_launch_dependents() + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + kr = bidx # K slice == rank in the (KS, 1, 1) cluster + nt = bidy # column tile + + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + ab_producer, ab_consumer = PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + ).make_participants() + acc_pipeline = PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len(self.epilogue_warp_id) + ), + cta_layout_vmnk=cluster_layout_vmnk, + ) + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=False, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, + ) + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + sA = smem.allocate_tensor( + element_type=self.io_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + sB = smem.allocate_tensor( + element_type=self.io_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + sP = smem.allocate_tensor( + element_type=self.acc_dtype, + layout=cute.make_layout((128, self.bn), stride=(self.p_stride, 1)), + byte_alignment=16, + ) + + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + k0 = kr * self.k_tiles + + thr_mma = tiled_mma.get_slice(0) + tCgA = thr_mma.partition_A(gA_mkl) + tCgB = thr_mma.partition_B(gB_nkl) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + 0, + cute.make_layout(1), + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + 0, + cute.make_layout(1), + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + if warp_idx == self.tma_warp_id: + tAgA_slice = tAgA[(None, 0, None, 0)] + tBgB_slice = tBgB[(None, nt, None, 0)] + ab_producer.reset() + # The weight does not depend on the previous kernel: the first stages' weight + # tiles stream before the grid-dependency wait (programmatic dependent + # launch), the normed rows -- the previous kernel's output -- after it. + pre = [] + for kk in cutlass.range_constexpr(self.n_pre): + handle = ab_producer.acquire_and_advance() + cute.copy( + tma_atom_b, + tBgB_slice[(None, k0 + kk)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + pre.append(handle) + cute.arch.griddepcontrol_wait() + for kk in cutlass.range_constexpr(self.n_pre): + handle = pre[kk] + cute.copy( + tma_atom_a, + tAgA_slice[(None, k0 + kk)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + for k_tile in cutlass.range(self.n_pre, self.k_tiles, 1, unroll=1): + handle = ab_producer.acquire_and_advance() + cute.copy( + tma_atom_a, + tAgA_slice[(None, k0 + k_tile)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, k0 + k_tile)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + ab_producer.tail() + + if warp_idx == self.mma_warp_id: + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + ab_consumer.reset() + peek = ab_consumer.try_wait() + acc_pipeline.producer_acquire(acc_producer_state) + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for k_tile in cutlass.range(0, self.k_tiles, 1, unroll=1): + handle = ab_consumer.wait_and_advance(peek) + num_kblocks = cute.size(tCrA, mode=[2]) + for kb in cutlass.range(num_kblocks, unroll_full=True): + kcrd = (None, None, kb, handle.index) + cute.gemm(tiled_mma, tCtAcc, tCrA[kcrd], tCrB[kcrd], tCtAcc) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + handle.release() + peek = cutlass.Boolean(1) + if handle.count + 1 < self.k_tiles: + peek = ab_consumer.try_wait() + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + acc_pipeline.producer_tail(acc_producer_state) + + if warp_idx < self.mma_warp_id: + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + copy_atom_t2r = utils.sm100.get_tmem_load_op( + self.cta_tile_shape_mnk, + utils.LayoutEnum.ROW_MAJOR, + self.acc_dtype, + self.acc_dtype, + epi_tile, + False, + ) + tAcc_epi = cute.flat_divide(tCtAcc_base[((None, None), 0, 0, None)], epi_tile) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)]) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_tAcc_base = thr_copy_t2r.partition_S(tAcc_epi) + tTR_sP = thr_copy_t2r.partition_D(cute.flat_divide(sP, epi_tile)) + frag_shape = tTR_sP[(None, None, None, 0, 0)].shape + tTR_rAcc = cute.make_rmem_tensor(frag_shape, self.acc_dtype) + tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_consumer_state.index)] + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + acc_pipeline.consumer_wait(acc_consumer_state) + for sub in cutlass.range_constexpr(self.bn // self.epi_cols): + cute.copy(tiled_copy_t2r, tTR_tAcc[(None, None, None, sub)], tTR_rAcc) + cute.autovec_copy(tTR_rAcc, tTR_sP[(None, None, None, 0, sub)]) + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + cute.arch.barrier( + barrier_id=self.epilog_sync_bar_id, + number_of_threads=32 * len(self.epilogue_warp_id), + ) + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + # Release / acquire at cluster scope, every thread of every CTA: after it, each + # CTA reads all `KS` partials of its column tile out of its peers' shared memory. + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + + # This CTA's share of the tile's `m * BN` outputs, all 6 warps. Every load is + # issued before any sum (indices clamped in bounds, so no branch separates + # them); the `KS` partials are summed in slice order (deterministic). + total = m * self.bn + per = (total + self.ks - 1) // self.ks + base = kr * per + nthr = self.threads_per_cta + vals = cute.make_rmem_tensor((self.red_iters * self.ks,), self.acc_dtype) + p_base = cutlass.Int32(sP.iterator.toint()) + peers = [_mapa(p_base, j) for j in range(self.ks)] + for it in cutlass.range_constexpr(self.red_iters): + fs = cutlass.min(base + tidx + it * nthr, total - 1) + row = fs // self.bn + col = fs - row * self.bn + off = (row * self.p_stride + col) * 4 + for j in cutlass.range_constexpr(self.ks): + vals[it * self.ks + j] = _ld_dsmem_f32(peers[j] + off) + inv_hc = cutlass.Float32(1.0 / self.hc) + for it in cutlass.range_constexpr(self.red_iters): + local = tidx + it * nthr + f = base + local + if local < per and f < total: + row = f // self.bn + col = f - row * self.bn + n = nt * self.bn + col + acc = vals[it * self.ks] + for j in cutlass.range_constexpr(1, self.ks): + acc = acc + vals[it * self.ks + j] + z = acc * inv_hc + sg = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.math.exp(-z)) + if n < self.rank: + mLow[(row, n)] = (z * sg).to(self.io_dtype) + if cutlass.const_expr(self.with_inject): + if n >= self.rank and n < self.rank + self.hc: + mInj[(row, n - self.rank)] = (cutlass.Float32(2.0) * sg).to(self.io_dtype) + + # No CTA may exit (releasing its shared memory) while a peer still reads it. + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + + +_COMPILED: dict = {} + + +def _launcher(bn, ks, hc, rank, n_cols, k, with_inject, ptrs): + """Compile once per static configuration; the row count is a runtime argument + and the TVM-FFI entry launches on the caller's current stream (graph-capturable). + `ptrs` is `(normed, weight, low tensor, inject)`.""" + key = (bn, ks, hc, rank, n_cols, k, with_inject) + fn = _COMPILED.get(key) + if fn is None: + kern = HcDecodeTcFirstKernel(bn, ks, hc, rank, n_cols, k, with_inject) + gmem = cute.AddressSpace.gmem + cptrs = [ + make_ptr(cutlass.BFloat16, ptrs[0], gmem, assumed_align=16), + make_ptr(cutlass.BFloat16, ptrs[1], gmem, assumed_align=16), + from_dlpack(ptrs[2], assumed_align=16).mark_layout_dynamic(leading_dim=1), + make_ptr(cutlass.BFloat16, ptrs[3], gmem, assumed_align=2), + ] + fn = cute.compile( + kern, + MAX_ROWS, + *cptrs, + stream=cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--opt-level 2 --enable-tvm-ffi", + ) + _COMPILED[key] = fn + return fn + + +def hc_decode_tc_first(normed, first, rank, hc_count, with_inject, bn, ks): + """Launch with an explicit tile (`bn` columns, `ks`-way K split). `first` is + the `[K, >= n_cols]` transposed view of the row-major weight, as + `cublas_mm(normed, first)` takes it. Returns `(low, inject)`.""" + rows, k = normed.shape + n_cols = rank + (hc_count if with_inject else 0) + dev, dt = normed.device, normed.dtype + low = torch.empty((rows, rank), dtype=dt, device=dev) + inject = torch.empty((rows, hc_count), dtype=dt, device=dev) if with_inject else None + if rows > 0: + ptrs = ( + normed.data_ptr(), + first.data_ptr(), + low, + (inject if with_inject else low).data_ptr(), + ) + _launcher(bn, ks, hc_count, rank, n_cols, k, with_inject, ptrs)(rows, *ptrs) + return low, inject + + +#: Tile: `_BN` weight rows per CTA (7 tiles cover the 324 real rows of 336) and +#: a `_KS`-way K split, one cluster per column tile. +_BN = 48 +_KS = 16 + + +class Qwen4ExpHcDecodeTcFirst(OpWrapper): + """With `R = rank`, `C = hc_count`, `normed` the `[T, K]` grouped-normed + stream and `first` the `[K, >= R (+ C)]` transposed view of the load-time + `[down; block_inject; 0]` weight, `L = normed @ first` in fp32:: + + low = silu(L[:, :R] / C) # [T, R], normed's dtype + inject = 2 * sigmoid(L[:, R : R + C] / C) # [T, C], only with_inject + + The projection accumulates in fp32 (tensor memory, then a fixed-order sum + of the K-split partials); each activation is fp32 with one rounding. + Returns fresh `(low, inject)`; `inject` is None exactly when `with_inject` + is False. Fusion boundary: what `cublas_mm(normed, first, None, fp32)` + + `qwen4_exp_hc_lowrank_act` compute (cuBLAS's split-K reduce included). + """ + + 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. + CELLS: tuple[Cell, ...] = ( + Cell( + why="checkpoint 4x2560, 324 cols of 384, inject, decode 128 rows (MAX_ROWS)", + spec=dict(rows=128), + ), + Cell( + why="stack-final mixer: 320 cols, no inject, 128 rows", + spec=dict(rows=128, 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` + `[T, K]` contiguous, T <= 128; `first` the `.t()` view of a row-major + `[N, K]` weight with N >= R (+ C); 16-byte aligned; K a multiple of + 64 * KS. 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)) = (16, 7) CTAs of 6 warps, cluster + (KS, 1, 1), one CTA per SM (~227 KB shared memory: every K tile of the + slice in flight plus the padded fp32 partial). 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). + Measured on GB300 (graph replay, 24 cold weights): 6.4 us at 128 rows + against 8.6 us for cuBLAS split-K + reduce + lowrank_act; tile sweep + BN 32 / 48 / 64 / 112 x KS 8 / 10 / 16, 48 x 16 fastest at every row + count. + Certified: bf16, 324 columns with the injection at 2 and 128 rows, 320 + columns without it at 128 rows; graph replay. + 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. + """ + + def __call__( + self, + normed: torch.Tensor, + first: torch.Tensor, + rank: int, + hc_count: int, + with_inject: bool, + ): + return hc_decode_tc_first(normed, first, rank, hc_count, with_inject, _BN, _KS) + + def reference( + self, + normed: torch.Tensor, + first: torch.Tensor, + rank: int, + hc_count: int, + with_inject: bool, + ): + """float64 projection and activations, one rounding per output.""" + n_cols = rank + (hc_count if with_inject else 0) + lg = normed.double() @ first[:, :n_cols].double() + low = F.silu(lg[:, :rank] / hc_count).to(normed.dtype) + inject = None + if with_inject: + inject = (2.0 * torch.sigmoid(lg[:, rank:n_cols] / hc_count)).to(normed.dtype) + return low, inject + + def is_valid( + self, + normed: torch.Tensor, + first: torch.Tensor, + rank: int, + hc_count: int, + with_inject: bool, + ) -> None: + rows, k = normed.shape + n_cols = rank + (hc_count if with_inject else 0) + assert normed.dtype == first.dtype == torch.bfloat16, "bf16 operands" + assert rows <= MAX_ROWS, f"at most {MAX_ROWS} rows (one MMA tile)" + assert normed.stride() == (k, 1), "normed contiguous [T, K]" + assert first.shape[0] == k and first.shape[1] >= n_cols, "first [K, >= n_cols]" + assert first.stride() == (1, k), "first must be the .t() view of a row-major [N, K] weight" + assert k % (64 * _KS) == 0, f"K must be a multiple of {64 * _KS}" + for t in (normed, first): + assert t.data_ptr() % 16 == 0, "16-byte aligned operands" + + +qwen4_exp_hc_decode_tc_first = Qwen4ExpHcDecodeTcFirst() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm.py new file mode 100644 index 000000000000..9ddaa39c7e7f --- /dev/null +++ b/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 @@ -0,0 +1,1058 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hyper-connection post-block injection fused with the next site's grouped +gemma RMS norm: one Triton row kernel writes the updated multi-stream residual +and its normalized copy in the same pass. After a MoE block it also forms the +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. + +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. +""" + +from __future__ import annotations + +import struct +from typing import Optional, Tuple + +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.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr + +#: Trigger dependents at kernel start (`griddepcontrol.launch_dependents`); a no-op +#: unless the next kernel is launched with programmatic dependent launch. +_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. +_FEW_WARPS_MIN_ROWS = 512 + + +@triton.jit +def _hc_inject_norm_kernel( + hc_ptr, + block_ptr, + inject_ptr, + shared_ptr, + gate_ptr, + topk_w_ptr, + topk_row_ptr, + weight_ptr, + hc_out_ptr, + normed_ptr, + stride_hc, + stride_block, + stride_inject, + stride_shared, + stride_gate, + stride_topk_w, + stride_topk_row, + stride_out, + stride_normed, + eps, + HC_COUNT: tl.constexpr, + HIDDEN: tl.constexpr, + BLOCK: tl.constexpr, + HAS_INJECT: tl.constexpr, + HAS_SHARED: tl.constexpr, + TOPK: tl.constexpr, + STREAMS_PER_PROGRAM: tl.constexpr, + PDL_TRIGGER: tl.constexpr, +): + if PDL_TRIGGER: + # Release a programmatic-dependent-launch consumer (the decode mixer's + # first projection) at once: it streams its static weights while this + # kernel runs and waits for the grid to finish before reading `normed`. + tl.extra.cuda.gdc_launch_dependents() + # One program per `STREAMS_PER_PROGRAM` streams of a row; a stream's `HIDDEN` + # columns are one RMS group. A row's programs are adjacent in launch order, so + # the row's block output (and the combine's operands) come from HBM once and + # from L2 for the others; with several streams a program forms it once for them. + pid = tl.program_id(0).to(tl.int64) + row = pid // (HC_COUNT // STREAMS_PER_PROGRAM) + first_stream = ((pid % (HC_COUNT // STREAMS_PER_PROGRAM)) * STREAMS_PER_PROGRAM).to(tl.int32) + cols = tl.arange(0, BLOCK) + mask = cols < HIDDEN + if HAS_INJECT: + if TOPK > 0: + # The routed experts' finalize: this row's TOPK permuted FC2 rows weighted + # by their routing weights, fp32 FMAs in slot order from zero (a slot mapped + # to -1 is skipped), rounded once to the rows' dtype. + b = tl.zeros([BLOCK], dtype=tl.float32) + for k in tl.static_range(TOPK): + prow = tl.load(topk_row_ptr + row * stride_topk_row + k) + wk = tl.load(topk_w_ptr + row * stride_topk_w + k).to(tl.float32) + r = tl.load( + block_ptr + tl.maximum(prow, 0).to(tl.int64) * stride_block + cols, + mask=mask, + other=0.0, + ).to(tl.float32) + b = tl.where(prow >= 0, tl.fma(r, wk, b), b) + b = b.to(block_ptr.dtype.element_ty).to(tl.float32) + else: + 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. + 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) + for j in tl.static_range(STREAMS_PER_PROGRAM): + stream = first_stream + j + offs = stream * HIDDEN + cols + x = tl.load(hc_ptr + row * stride_hc + offs, mask=mask, other=0.0).to(tl.float32) + if HAS_INJECT: + g = tl.load(inject_ptr + row * stride_inject + stream).to(tl.float32) + # Rounded once on the store, like `addcmul`; the norm reads the rounded + # residual, which is the value every later site sees. + xr = (x + b * g).to(hc_out_ptr.dtype.element_ty) + tl.store(hc_out_ptr + row * stride_out + offs, xr, mask=mask) + x = xr.to(tl.float32) + rstd = tl.rsqrt(tl.sum(x * x, axis=0) / HIDDEN + eps) + w = tl.load(weight_ptr + offs, mask=mask, other=0.0).to(tl.float32) + y = x * rstd * (1.0 + w) + tl.store( + normed_ptr + row * stride_normed + offs, y.to(normed_ptr.dtype.element_ty), mask=mask + ) + + +# -------------------------------------------------------------------------------------- +# 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. +# -------------------------------------------------------------------------------------- + +_CUTE_THREADS = 256 +_CUTE_WARPS = _CUTE_THREADS // 32 +_CUTE_VEC = 8 +#: The Triton kernel's block at these hiddens is twice this: a thread's second chunk starts here. +_CUTE_SPLIT = _CUTE_THREADS * _CUTE_VEC +#: Streams a CTA, as the Triton fold's program; a row's `hc_count // 2` CTAs are one cluster. +_CUTE_SPP = 2 + + +def _iv(x, loc, ip): + return x.ir_value(loc=loc, ip=ip) + + +def _asm(ret, args, text, constraints, loc, ip, side_effects=False): + return llvm.inline_asm( + ret, + [_iv(a, loc, ip) for a in args], + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _asm4(args, text, constraints, loc, ip, side_effects=False): + ty = T.i32() + res = _asm( + llvm.StructType.get_literal([ty] * 4), args, text, constraints, loc, ip, side_effects + ) + return [cutlass.Int32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(4)] + + +@dsl_user_op +def _ldg_v4(addr, pred, *, loc=None, ip=None): + """16 bytes, or zeros where `pred` is 0 (nothing read).""" + return _asm4( + [cutlass.Int64(addr), cutlass.Int32(pred)], + "{ .reg .pred p; setp.ne.b32 p, $5, 0; mov.b32 $0, 0; mov.b32 $1, 0; mov.b32 $2, 0; " + "mov.b32 $3, 0; @p ld.global.v4.b32 {$0, $1, $2, $3}, [$4]; }", + "=r,=r,=r,=r,l,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _stg_v4(addr, pred, a, b, c, d, *, loc=None, ip=None): + """16 bytes where `pred` is not 0.""" + _asm( + T.i32(), + [cutlass.Int64(addr), cutlass.Int32(pred)] + [cutlass.Int32(v) for v in (a, b, c, d)], + "{ .reg .pred p; setp.ne.b32 p, $2, 0; @p st.global.v4.b32 [$1], {$3, $4, $5, $6}; " + "mov.u32 $0, 0; }", + "=r,l,r,r,r,r,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _mapa(addr, rank, *, loc=None, ip=None): + """This CTA's shared-memory address `addr` in cluster CTA `rank`'s window.""" + return cutlass.Int32( + _asm( + T.i32(), + [cutlass.Int32(addr), cutlass.Int32(rank)], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_dsmem_v4(addr, pred, a, b, c, d, *, loc=None, ip=None): + """16 bytes into a cluster CTA's shared memory where `pred` is not 0.""" + _asm( + T.i32(), + [cutlass.Int32(addr), cutlass.Int32(pred)] + [cutlass.Int32(v) for v in (a, b, c, d)], + "{ .reg .pred p; setp.ne.b32 p, $2, 0; @p st.shared::cluster.v4.b32 [$1], " + "{$3, $4, $5, $6}; mov.u32 $0, 0; }", + "=r,r,r,r,r,r,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _lds_v4(addr, pred, *, loc=None, ip=None): + """16 bytes of this CTA's shared memory, or zeros where `pred` is 0.""" + return _asm4( + [cutlass.Int32(addr), cutlass.Int32(pred)], + "{ .reg .pred p; setp.ne.b32 p, $5, 0; mov.b32 $0, 0; mov.b32 $1, 0; mov.b32 $2, 0; " + "mov.b32 $3, 0; @p ld.shared.v4.b32 {$0, $1, $2, $3}, [$4]; }", + "=r,=r,=r,=r,r,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _bf16_lo(w, *, loc=None, ip=None): + """The low bf16 of a word as fp32 (exact).""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int32(w)], + "{ .reg .b32 t; shl.b32 t, $1, 16; mov.b32 $0, t; }", + "=f,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_hi(w, *, loc=None, ip=None): + """The high bf16 of a word as fp32 (exact).""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int32(w)], + "{ .reg .b32 t; and.b32 t, $1, 0xFFFF0000; mov.b32 $0, t; }", + "=f,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16x2(lo, hi, *, loc=None, ip=None): + """Two fp32 values rounded (RN) into one bf16x2 word, `lo` low.""" + return cutlass.Int32( + _asm( + T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _round_bf16(x, *, loc=None, ip=None): + """fp32 -> bf16 (RN) -> fp32.""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(x)], + "{ .reg .b16 h; cvt.rn.bf16.f32 h, $1; cvt.f32.bf16 $0, h; }", + "=f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma_if(slot, a, b, c, *, loc=None, ip=None): + """`slot >= 0 ? fma(a, b, c) : c` (the finalize skips a slot mapped to -1).""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int32(slot), cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "{ .reg .pred p; .reg .f32 t; fma.rn.f32 t, $2, $3, $4; setp.ge.s32 p, $1, 0; " + "selp.f32 $0, t, $4, p; }", + "=f,r,f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _mul(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _add(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _sigmoid(x, *, loc=None, ip=None): + """`tl.sigmoid` as the Triton kernel compiles it: `mul` by -log2(e), `ex2.approx`, `add` + 1, `div.full`.""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(x)], + "{ .reg .f32 t; mul.f32 t, $1, 0fBFB8AA3B; ex2.approx.f32 t, t; " + "add.f32 t, t, 0f3F800000; div.full.f32 $0, 0f3F800000, t; }", + "=f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _rstd(ss, eps, hidden_hex: str, *, loc=None, ip=None): + """`rsqrt(ss / hidden + eps)` as the Triton kernel's `div.full`, `add`, `rsqrt.approx.ftz`: + the `add` has no rounding mode there, so ptxas fuses it with the division (a multiply by + the reciprocal) into one FFMA -- here the same.""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(ss), cutlass.Float32(eps)], + "{ .reg .f32 t; div.full.f32 t, $1, " + hidden_hex + "; add.f32 t, t, $2; " + "rsqrt.approx.ftz.f32 $0, t; }", + "=f,f,f", + loc, + ip, + ) + ) + + +def _vals(words): + """Four bf16x2 words -> their eight values as fp32, in column order.""" + out = [] + for w in words: + out += [_bf16_lo(w), _bf16_hi(w)] + return out + + +class _FoldNormKernel: + """CTA (j, row): the row's streams 2j and 2j + 1; the row's CTAs one cluster.""" + + def __init__(self, hc_count: int, hidden: int, top_k: int, pdl_trigger: bool): + self.hc, self.hidden, self.topk, self.pdl = hc_count, hidden, top_k, pdl_trigger + self.ctas = hc_count // _CUTE_SPP + # The block-output columns a CTA forms, one 16-byte vector a thread. + self.per = hidden // self.ctas + self.n_fold = self.per // _CUTE_VEC + # Threads with a second chunk (columns [2048, hidden)). + self.n_b = (hidden - _CUTE_SPLIT) // _CUTE_VEC + self.hidden_hex = "0f%08X" % struct.unpack("> off)) + if tidx % 32 == 0: + red[s, warp] = v + cute.arch.sync_threads() + for s in cutlass.range_constexpr(_CUTE_SPP): + stream = rank * _CUTE_SPP + s + v = red[s, tidx % _CUTE_WARPS] + for off in cutlass.range_constexpr(3): + v = _add(v, cute.arch.shuffle_sync_bfly(v, offset=4 >> off)) + rstd = _rstd(v, eps, self.hidden_hex) + ra, rb = xrs[s] + wa_v, wb_v = _vals(wa[s]), _vals(wb[s]) + ya = [ + _mul(_mul(rstd, ra[e]), _add(wa_v[e], cutlass.Float32(1.0))) + for e in range(_CUTE_VEC) + ] + yb = [ + _mul(_mul(rstd, rb[e]), _add(wb_v[e], cutlass.Float32(1.0))) + for e in range(_CUTE_VEC) + ] + nrow = nrm_base + (row64 * width + stream * self.hidden) * 2 + _stg_v4(nrow + c_a * 2, one, *[_bf16x2(ya[2 * i], ya[2 * i + 1]) for i in range(4)]) + _stg_v4(nrow + c_b * 2, has_b, *[_bf16x2(yb[2 * i], yb[2 * i + 1]) for i in range(4)]) + + +_FOLD_NORM: dict = {} + + +def _fold_cute_ok( + hyper_input, + block_out, + inject, + norm_weight, + hc_count, + shared, + gate_logits, + topk_weights, + topk_rows, +) -> bool: + """The decode MoE site this kernel certifies (the Triton kernel's eight-warp form with + 16-byte vectors); anything else runs the Triton kernel.""" + rows, width = hyper_input.shape + hidden = width // hc_count + bf = torch.bfloat16 + return ( + 0 < rows < _FEW_WARPS_MIN_ROWS + and hc_count == 4 + and _CUTE_SPLIT < hidden <= 2 * _CUTE_SPLIT + and hidden % (_CUTE_SPP * _CUTE_VEC) == 0 + and hyper_input.dtype == block_out.dtype == shared.dtype == bf + and inject.dtype == gate_logits.dtype == topk_weights.dtype == norm_weight.dtype == bf + and topk_rows.dtype == torch.int32 + and all( + t.data_ptr() % 16 == 0 and t.stride(0) % 16 == 0 + for t in (hyper_input, block_out, shared) + ) + and norm_weight.data_ptr() % 16 == 0 + ) + + +def _fold_cute( + hyper_input, + block_out, + inject, + norm_weight, + eps, + hc_count, + shared, + gate_logits, + topk_weights, + topk_rows, + hc_out, + normed, +) -> None: + rows, width = hyper_input.shape + hidden = width // hc_count + top_k = topk_rows.shape[1] + ptrs = ( + hyper_input, + block_out.data_ptr(), + inject.data_ptr(), + shared.data_ptr(), + gate_logits.data_ptr(), + topk_weights.data_ptr(), + topk_rows.data_ptr(), + norm_weight.data_ptr(), + hc_out.data_ptr(), + normed.data_ptr(), + ) + strides = ( + hyper_input.stride(0), + block_out.stride(0), + inject.stride(0), + shared.stride(0), + gate_logits.stride(0), + topk_weights.stride(0), + topk_rows.stride(0), + ) + key = (hc_count, hidden, top_k, _PDL_TRIGGER) + fn = _FOLD_NORM.get(key) + if fn is None: + gm = cute.AddressSpace.gmem + bf = cutlass.BFloat16 + cargs = ( + from_dlpack(hyper_input, assumed_align=16).mark_layout_dynamic(leading_dim=1), + make_ptr(bf, ptrs[1], gm, assumed_align=16), + make_ptr(bf, ptrs[2], gm, assumed_align=2), + make_ptr(bf, ptrs[3], gm, assumed_align=16), + make_ptr(bf, ptrs[4], gm, assumed_align=2), + make_ptr(bf, ptrs[5], gm, assumed_align=2), + make_ptr(cutlass.Int32, ptrs[6], gm, assumed_align=4), + make_ptr(bf, ptrs[7], gm, assumed_align=16), + make_ptr(bf, ptrs[8], gm, assumed_align=16), + make_ptr(bf, ptrs[9], gm, assumed_align=16), + ) + fn = cute.compile( + _FoldNormKernel(hc_count, hidden, top_k, _PDL_TRIGGER), + *cargs, + *(cutlass.Int32(1),) * 8, + cutlass.Float32(1e-6), + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--opt-level 2 --enable-tvm-ffi", + ) + _FOLD_NORM[key] = fn + fn(*ptrs, rows, *strides, eps) + + +class Qwen4ExpHcInjectNorm(OpWrapper): + """With `C = hc_count`, `H = width // C`, every tensor 2-D `[T, ...]`: + + hc_out[t, s*H + i] = hyper_input[t, s*H + i] + block_out[t, i] * inject[t, s] + normed[t, s*H + i] = hc_out[t, s*H + i] * rstd(t, s) * (1 + norm_weight[s*H + i]) + rstd(t, s) = 1 / sqrt(mean_i(hc_out[t, s*H + i]^2) + eps) + + `hc_out` is rounded once to the activation dtype and `normed` is computed + from that rounded value (fp32 inside, one rounding on the store). Returns + `(hc_out, normed)`, both fresh `[T, C*H]` tensors in `hyper_input`'s dtype. + + `block_out=None, inject=None` is the norm alone: `hc_out` *is* + `hyper_input` (not a copy) and only `normed` is written -- the stack's first + site and the site after the PLE add. + + With `shared` and `gate_logits` (after a MoE block), `block_out` is the + routed experts' output and the block output injected above is + + block[t, i] = block_out[t, i] + sigmoid(gate_logits[t]) * shared[t, i] + + in fp32, rounded once to `block_out`'s dtype before the injection reads it: + bitwise what `fused_sigmoid_gate_mul_add` stores. `gate_logits` is one + pre-sigmoid scalar per row, `[T, 1]` or `[T]`; none of the inputs is written. + + With `topk_weights` / `topk_rows` as well (the MoE runner left unfinalized), + `block_out` is its permuted FC2 output `[P, H]` and the routed row is the + runner's finalize, formed here: `routed[t, i] = sum_k topk_weights[t, k] * + block_out[topk_rows[t, k], i]` as fp32 FMAs in slot order from zero (a slot + whose row is -1 is skipped), rounded once to `block_out`'s dtype -- bitwise + what the runner's finalize stores. `topk_rows` is the runner's + expanded->permuted index map `[T, k]` int32, `topk_weights` its `[T, k]` + 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. + """ + + 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 + # 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. + CELLS: tuple[Cell, ...] = ( + Cell( + why="checkpoint 4x2560 bf16, inject, decode batch 128 (eight-warp form)", + spec=dict(rows=128), + ), + Cell( + why="MoE combine, [T, 1] gate, prefill chunk 16384 (the ladder's max_num_tokens)", + spec=dict(rows=16384, moe=True), + ), + Cell( + why="MoE combine + finalize fold (top-10, CuTe kernel), three slots -1, decode 128", + spec=dict(rows=128, moe=True, fold=10, drop=True), + ), + Cell( + why="norm only (no inject: stack start / after PLE), decode 128", + spec=dict(rows=128, 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 + 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 + per forward on the 48-layer trunk; the 48 after a MoE block pass `shared` + and `gate_logits` (in `DecodeTarget` all but the one before the PLE add, + which keeps the unfused combine, and those also `topk_weights` / + `topk_rows`: the runner's finalize folded in). The MTP layer and the PLE + block keep the unfused ops. + + 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`. + + 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 + accumulate is fp32 with one rounding (bitwise what `qwen4_exp_post_inject` + stores for 16-bit operands up to fma contraction); the norm is fp32 with + `rsqrt`, one rounding. Against the float64 reference the combine can land + one ulp of the block output away (its rounding midpoint), which the + accumulate carries scaled by `inject`; the test's band states that. + Graph-capturable by construction (no host sync, outputs allocated per + call); not certified. + + 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 + `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 + `_FEW_WARPS_MIN_ROWS` rows and two from it; the RMS sum's fp32 reduction + order follows the warp count. + + The decode MoE sites -- the fold below `_FEW_WARPS_MIN_ROWS` rows, bf16 + operands, `hc_count` 4, a hidden in (2048, 4096], 16-byte aligned rows with + row strides a multiple of 16 -- run `_FoldNormKernel` (CuTe DSL) instead: + a row's two CTAs one cluster, CTA j forming half the block output (the + 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. + + 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 + drives them). + + Hard limits: row offsets are int64 (`program_id(0)` widened); the grid is + `rows * hc_count` programs along x (2^31 - 1); the column offset + `stream * hidden + i` is int32, bounded by the width. + """ + + def __call__( + self, + hyper_input: torch.Tensor, + block_out: Optional[torch.Tensor], + inject: Optional[torch.Tensor], + norm_weight: torch.Tensor, + eps: float, + hc_count: int, + shared: Optional[torch.Tensor] = None, + gate_logits: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + topk_rows: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + rows, width = hyper_input.shape + hidden = width // hc_count + normed = torch.empty((rows, width), dtype=hyper_input.dtype, device=hyper_input.device) + has_inject = block_out is not None + has_shared = shared is not None + top_k = topk_rows.shape[1] if topk_rows is not None else 0 + hc_out = ( + torch.empty((rows, width), dtype=hyper_input.dtype, device=hyper_input.device) + if has_inject + else hyper_input + ) + if ( + rows > 0 + and top_k + and _fold_cute_ok( + hyper_input, + block_out, + inject, + norm_weight, + hc_count, + shared, + gate_logits, + topk_weights, + topk_rows, + ) + ): + _fold_cute( + hyper_input, + block_out, + inject, + norm_weight, + eps, + hc_count, + shared, + gate_logits, + topk_weights, + topk_rows, + hc_out, + normed, + ) + elif rows > 0: + block = triton.next_power_of_2(hidden) + # The fold gathers TOPK rows to form a row's block output: two streams a + # program form it half as often (when hc_count is even). + spp = 2 if top_k and hc_count % 2 == 0 else 1 + _hc_inject_norm_kernel[(rows * hc_count // spp,)]( + hyper_input, + block_out if has_inject else hyper_input, + inject if has_inject else hyper_input, + shared if has_shared else hyper_input, + gate_logits if has_shared else hyper_input, + topk_weights if top_k else hyper_input, + topk_rows if top_k else hyper_input, + norm_weight, + hc_out, + normed, + hyper_input.stride(0), + block_out.stride(0) if has_inject else 0, + inject.stride(0) if has_inject else 0, + shared.stride(0) if has_shared else 0, + gate_logits.stride(0) if has_shared else 0, + topk_weights.stride(0) if top_k else 0, + topk_rows.stride(0) if top_k else 0, + hc_out.stride(0), + normed.stride(0), + eps, + HC_COUNT=hc_count, + HIDDEN=hidden, + BLOCK=block, + HAS_INJECT=has_inject, + HAS_SHARED=has_shared, + TOPK=top_k, + STREAMS_PER_PROGRAM=spp, + PDL_TRIGGER=_PDL_TRIGGER, + num_warps=(2 if rows >= _FEW_WARPS_MIN_ROWS else 8) if block >= 2048 else 4, + ) + return hc_out, normed + + def reference( + self, + hyper_input: torch.Tensor, + block_out: Optional[torch.Tensor], + inject: Optional[torch.Tensor], + norm_weight: torch.Tensor, + eps: float, + hc_count: int, + shared: Optional[torch.Tensor] = None, + gate_logits: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + 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.""" + dtype = hyper_input.dtype + rows, width = hyper_input.shape + hidden = width // hc_count + if block_out is None: + hc_out = hyper_input + else: + if topk_rows is not None: + valid = (topk_rows >= 0).unsqueeze(-1) + picked = block_out.double()[topk_rows.clamp(min=0).long()] + w = topk_weights.double().unsqueeze(-1) + block_out = torch.where(valid, picked * w, 0.0).sum(1).to(block_out.dtype) + if shared is not None: + s = torch.sigmoid(gate_logits.double().reshape(rows, 1)) + block_out = (block_out.double() + s * shared.double()).to(block_out.dtype) + streams = hyper_input.double().unflatten(-1, (hc_count, hidden)) + gated = block_out.double().unsqueeze(-2) * inject.double().unsqueeze(-1) + hc_out = (streams + gated).flatten(-2).to(dtype) + x = hc_out.double().unflatten(-1, (hc_count, hidden)) + x = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) + normed = (x.flatten(-2) * (1.0 + norm_weight.double())).to(dtype) + return hc_out, normed + + def is_valid( + self, + hyper_input: torch.Tensor, + block_out: Optional[torch.Tensor], + inject: Optional[torch.Tensor], + norm_weight: torch.Tensor, + eps: float, + hc_count: int, + shared: Optional[torch.Tensor] = None, + gate_logits: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + topk_rows: Optional[torch.Tensor] = None, + ) -> None: + # The kernel indexes with raw row strides and unit column stride, and + # reads `inject` / `block_out` / `shared` / the gate per row: a broadcast + # (stride-0) or a wrong-width operand would be read out of bounds or + # silently reused. + assert hyper_input.dim() == 2 and hyper_input.stride(1) == 1, "hyper_input [T, width]" + assert (block_out is None) == (inject is None), "block_out and inject come together" + assert (shared is None) == (gate_logits is None), "shared and gate_logits come together" + assert shared is None or block_out is not None, "the combine needs block_out" + assert (topk_weights is None) == (topk_rows is None), ( + "topk_weights and topk_rows come together" + ) + assert topk_rows is None or shared is not None, "the finalize fold needs the combine" + rows, width = hyper_input.shape + assert width % hc_count == 0, f"width {width} not divisible by hc_count {hc_count}" + assert norm_weight.shape == (width,) and norm_weight.is_contiguous(), "norm_weight [width]" + if block_out is not None: + if topk_rows is None: + assert block_out.shape == (rows, width // hc_count) and block_out.stride(1) == 1 + else: + # The permuted FC2 rows: any row count, read through topk_rows. + assert block_out.dim() == 2 and block_out.shape[1] == width // hc_count + assert block_out.stride(1) == 1, "permuted rows [P, hidden]" + assert topk_rows.dtype == torch.int32, "topk_rows int32 (the runner's index map)" + assert topk_rows.dim() == 2 and topk_rows.shape[0] == rows, "topk_rows [T, k]" + assert topk_rows.stride(1) == 1 and topk_weights.stride(1) == 1 + assert topk_weights.shape == topk_rows.shape, "topk_weights [T, k]" + assert inject.shape == (rows, hc_count) and inject.stride(1) == 1 + if shared is not None: + assert shared.shape == (rows, width // hc_count), "shared [T, hidden]" + assert shared.stride(1) == 1, "shared [T, hidden]" + assert shared.dtype == block_out.dtype, "shared in block_out's dtype" + assert gate_logits.shape in ((rows,), (rows, 1)), "gate_logits [T, 1] or [T]" + assert rows <= 1 or gate_logits.stride(0) != 0, "one gate logit per row" + + +qwen4_exp_hc_inject_norm = Qwen4ExpHcInjectNorm() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_lowrank_act.py new file mode 100644 index 000000000000..493ca58301c4 --- /dev/null +++ b/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 @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hyper-connection mixer activations between its two projections: the SiLU of +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). +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +@triton.jit +def _hc_lowrank_act_kernel( + logits_ptr, + low_ptr, + inject_ptr, + stride_logits, + stride_low, + stride_inject, + RANK: tl.constexpr, + HC: tl.constexpr, + WITH_INJECT: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + k = tl.arange(0, BLOCK) + z = tl.load(logits_ptr + row * stride_logits + k, mask=k < RANK, other=0.0).to(tl.float32) + z = z * (1.0 / HC) + tl.store( + low_ptr + row * stride_low + k, + (z * tl.sigmoid(z)).to(low_ptr.dtype.element_ty), + mask=k < RANK, + ) + if WITH_INJECT: + i = tl.arange(0, HC) + zi = tl.load(logits_ptr + row * stride_logits + RANK + i).to(tl.float32) + tl.store( + inject_ptr + row * stride_inject + i, + (2.0 * tl.sigmoid(zi * (1.0 / HC))).to(inject_ptr.dtype.element_ty), + ) + + +class Qwen4ExpHcLowrankAct(OpWrapper): + """With `R = rank`, `C = hc_count` and `logits` the fp32 `[T, N]` first + projection `normed @ [down; block_inject; 0]^T` (columns `[0, R)` down, + `[R, R + C)` injection when `with_inject`, the rest padding):: + + low = silu(logits[:, :R] / C) # [T, R], `dtype` + inject = 2 * sigmoid(logits[:, R : R + C] / C) # [T, C], `dtype`, only with_inject + + fp32 inside, one rounding per output. Returns fresh `(low, inject)`; + `inject` is `None` exactly when `with_inject` is False. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: rows, inject (True), hc_count (4), rank (320), logits_width (384), dtype (bf16). + 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)), + ) + + 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. + """ + + def __call__( + self, + logits: torch.Tensor, + rank: int, + hc_count: int, + with_inject: bool, + dtype: torch.dtype, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + rows = logits.shape[0] + low = torch.empty((rows, rank), dtype=dtype, device=logits.device) + inject = ( + torch.empty((rows, hc_count), dtype=dtype, device=logits.device) + if with_inject + else None + ) + if rows > 0: + _hc_lowrank_act_kernel[(rows,)]( + logits, + low, + inject if with_inject else low, + logits.stride(0), + low.stride(0), + inject.stride(0) if with_inject else 0, + RANK=rank, + HC=hc_count, + WITH_INJECT=with_inject, + BLOCK=triton.next_power_of_2(rank), + num_warps=4, + ) + return low, inject + + def reference( + self, + logits: torch.Tensor, + rank: int, + hc_count: int, + with_inject: bool, + dtype: torch.dtype, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """float64, one rounding per output.""" + lg = logits.double() + low = F.silu(lg[:, :rank] / hc_count).to(dtype) + inject = None + if with_inject: + inject = (2.0 * torch.sigmoid(lg[:, rank : rank + hc_count] / hc_count)).to(dtype) + return low, inject + + def is_valid( + self, + logits: torch.Tensor, + rank: int, + hc_count: int, + with_inject: bool, + dtype: torch.dtype, + ) -> None: + assert logits.dtype == torch.float32 and logits.stride(1) == 1, "fp32 logits, unit stride" + assert logits.shape[1] >= rank + (hc_count if with_inject else 0), "logits too narrow" + assert hc_count & (hc_count - 1) == 0, "hc_count must be a power of two" + + +qwen4_exp_hc_lowrank_act = Qwen4ExpHcLowrankAct() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_up_gate_mean.py new file mode 100644 index 000000000000..755fc8535e08 --- /dev/null +++ b/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 @@ -0,0 +1,899 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hyper-connection mixer second half as one tcgen05 GEMM with a custom +epilogue: the up projection `low @ up^T` accumulates in tensor memory and the +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. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import torch +import torch.nn.functional as F +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op +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.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 + +#: Programmatic dependent launch (sm_90+), off with TRTLLM_ENABLE_PDL=0. +_PDL = get_env_enable_pdl() + + +@dsl_user_op +def _sigmoid(x, *, loc=None, ip=None): + """sigmoid(x) = 0.5 * tanh(0.5 * x) + 0.5: one MUFU op (tanh.approx.f32, + max relative error 2^-10.99) instead of ex2 + rcp.""" + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], + "{ .reg .f32 t; mul.f32 t, $1, 0f3F000000; tanh.approx.f32 t, t; " + "fma.rn.f32 $0, t, 0f3F000000, 0f3F000000; }", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +# The NVFP4 epilogue (`fp4=True`) is the `torch.ops.trtllm.fp4_quantize` recipe +# (`cvt_warp_fp16_to_fp4`, e4m3 scales) instruction for instruction, every +# multiply an explicit `mul.rn.f32` so that no compiler contraction or +# reassociation can move a rounding. + + +def _asm(ret, args, text, constraints, loc, ip, side_effects=False): + return llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _rcp_approx_ftz(x, *, loc=None, ip=None): + """`rcp.approx.ftz.f32`, the quantizer's `reciprocal_approximate_ftz`.""" + return cutlass.Float32( + _asm(T.f32(), [cutlass.Float32(x)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", loc, ip) + ) + + +@dsl_user_op +def _absmax16(v, *, loc=None, ip=None): + """max |v_i| over 16 fp32 values (exact; NaN-quiet like the bf16 `__hmax2`).""" + body = "abs.f32 m, $1; " + "".join(f"abs.f32 t, ${i}; max.f32 m, m, t; " for i in range(2, 17)) + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(x) for x in v], + "{ .reg .f32 m, t; " + body + "mov.f32 $0, m; }", + "=f" + ",f" * 16, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e4m3_scale(amax, gs, rcp6, *, loc=None, ip=None): + """The block's scale byte: e4m3(gs * (amax * rcp(6))), `cvt.rn.satfinite`.""" + return cutlass.Int32( + _asm( + T.i32(), + [cutlass.Float32(amax), cutlass.Float32(gs), cutlass.Float32(rcp6)], + "{ .reg .f32 t; .reg .b16 h; mul.rn.f32 t, $1, $3; mul.rn.f32 t, $2, t; " + "cvt.rn.satfinite.e4m3x2.f32 h, 0f00000000, t; cvt.u32.u16 $0, h; }", + "=r,f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _e2m1_out_scale(sf, rcp_gs, amax, *, loc=None, ip=None): + """amax != 0 ? rcp(float(e4m3 sf) * rcp(gs)) : 0, the data multiplier.""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int32(sf), cutlass.Float32(rcp_gs), cutlass.Float32(amax)], + "{ .reg .b16 h, lo, hi; .reg .b32 f2; .reg .f32 s, t; .reg .pred p; " + "cvt.u16.u32 h, $1; cvt.rn.f16x2.e4m3x2 f2, h; mov.b32 {lo, hi}, f2; " + "cvt.f32.f16 s, lo; mul.rn.f32 t, s, $2; rcp.approx.ftz.f32 t, t; " + "setp.neu.f32 p, $3, 0f00000000; selp.f32 $0, t, 0f00000000, p; }", + "=f,r,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _e2m1x8(v, scale, *, loc=None, ip=None): + """Eight values times `scale`, each to e2m1 (`cvt.rn.satfinite`), packed + into 4 bytes; element `2k` in byte k's low nibble.""" + muls = "".join(f"mul.rn.f32 a{i}, ${i + 1}, $9; " for i in range(8)) + cvts = "".join(f"cvt.rn.satfinite.e2m1x2.f32 b{k}, a{2 * k + 1}, a{2 * k}; " for k in range(4)) + return cutlass.Int32( + _asm( + T.i32(), + [cutlass.Float32(x) for x in v] + [cutlass.Float32(scale)], + "{ .reg .f32 a<8>; .reg .b8 b<4>; " + muls + cvts + "mov.b32 $0, {b0, b1, b2, b3}; }", + "=r" + ",f" * 9, + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_global_v4(addr, a, b, c, d, *, loc=None, ip=None): + """One 16-byte `st.global.v4.b32` (a row's packed e2m1 for 32 columns): a + row-per-thread epilogue pays per store instruction, not per byte.""" + _asm( + T.i32(), + [cutlass.Int64(addr)] + [cutlass.Int32(x) for x in (a, b, c, d)], + "{ st.global.v4.b32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }", + "=r,l,r,r,r,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _st_global_b16(addr, x, *, loc=None, ip=None): + """One 2-byte store of the low half of `x` (two e4m3 scale bytes).""" + _asm( + T.i32(), + [cutlass.Int64(addr), cutlass.Int32(x)], + "{ .reg .b16 t; cvt.u16.u32 t, $2; st.global.b16 [$1], t; mov.u32 $0, 0; }", + "=r,l,r", + loc, + ip, + side_effects=True, + ) + + +class HcUpGateMeanKernel: + """`mixed[t, h] = 1/C * sum_s sigmoid((low @ up^T)[t, s*H + h]) * normed[t, s*H + h]`. + + Compile-time: `block_h` (hidden columns per tile, 32 or 64), `hc` (C), + `hidden` (H), `rank` (R), `fp4` (also write the NVFP4 of `mixed`: packed + e2m1 `[T, H/2]` and linear e4m3 scales `[T, H/16]` under the global scale + `gs`). Runtime: the row count and the pointers; without `fp4` the `xq`, + `sf` and `gs` pointers are never read. + """ + + def __init__( + self, block_h: int, hc: int, hidden: int, rank: int, num_n_stage: int = 2, fp4: bool = False + ): + assert block_h in (32, 64) and hidden % block_h == 0 and rank % 64 == 0 + self.bh, self.hc, self.hidden, self.rank = block_h, hc, hidden, rank + self.fp4 = fp4 + self.acc_dtype = cutlass.Float32 + self.io_dtype = cutlass.BFloat16 + self.mma_tiler = (128, hc * block_h, 1) + self.epi_cols = 32 + self.cta_group = tcgen05.CtaGroup.ONE + self.cluster_shape_mn = (1, 1) + self.epilogue_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.nrm_warp_id = 6 + self.threads_per_cta = 32 * 7 + self.num_n_stage = num_n_stage + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + self.num_acc_stage = 2 + + def _setup(self, a_major, b_major): + tiled_mma = utils.sm100.make_trivial_tiled_mma( + self.io_dtype, a_major, b_major, self.acc_dtype, self.cta_group, self.mma_tiler[:2] + ) + inst_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + self.mma_tiler = (self.mma_tiler[0], self.mma_tiler[1], inst_k * 4) + self.cta_tile_shape_mnk = self.mma_tiler + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), (tiled_mma.thr_id.shape,) + ) + self.epi_tile = (self.mma_tiler[0], self.epi_cols) + a_one = utils.sm100.make_smem_layout_a(tiled_mma, self.mma_tiler, self.io_dtype, 1) + b_one = utils.sm100.make_smem_layout_b(tiled_mma, self.mma_tiler, self.io_dtype, 1) + ab_bytes = cute.size_in_bytes(self.io_dtype, a_one) + cute.size_in_bytes( + self.io_dtype, b_one + ) + smem_capacity = utils.get_smem_capacity_in_bytes() + # One normed stage is the 32-column slice of every stream for one tile half. + self.n_box = (self.mma_tiler[0], self.epi_cols) + self.num_n_bytes = self.hc * self.n_box[0] * self.n_box[1] * 2 + self.num_ab_stage = min( + 8, (smem_capacity - 2048 - self.num_n_stage * self.num_n_bytes) // ab_bytes + ) + self.n_pre = min(self.num_ab_stage, self.rank // self.mma_tiler[2]) + self.n_smem_layout_staged = utils.sm100.make_smem_layout_epi( + self.io_dtype, utils.LayoutEnum.ROW_MAJOR, self.n_box, self.num_n_stage * self.hc + ) + self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( + tiled_mma, self.mma_tiler, self.io_dtype, self.num_ab_stage + ) + self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( + tiled_mma, self.mma_tiler, self.io_dtype, self.num_ab_stage + ) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + acc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(acc_fake, arch="sm_100") + return tiled_mma + + @cute.jit + def __call__( + self, + m: cutlass.Int32, + low_ptr: cute.Pointer, + up_ptr: cute.Pointer, + normed_ptr: cute.Pointer, + mixed: cute.Tensor, + xq_ptr: cute.Pointer, + sf_ptr: cute.Pointer, + gs_ptr: cute.Pointer, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + ): + bh, hc, hid, rank = self.bh, self.hc, self.hidden, self.rank + nt = hid // bh + width = hc * hid + # low [m, R] (K-major A); the up weight [C*H, R] seen as N = (BH, C, H/BH) so + # that one N tile holds hidden slice [h0, h0+BH) of every stream. + a = cute.make_tensor(low_ptr, cute.make_layout((m, rank, 1), stride=(rank, 1, 0))) + b = cute.make_tensor( + up_ptr, + cute.make_layout(((bh, hc, nt), rank, 1), stride=((rank, hid * rank, bh * rank), 1, 0)), + ) + # normed [m, C*H] is staged through shared memory by TMA in [128, 32] boxes; + # mixed [m, H] is written in the accumulator's N order, its stream mode + # with stride 0 (every stream maps onto the one output column). + nrm = cute.make_tensor(normed_ptr, cute.make_layout((m, width), stride=(width, 1))) + mix = cute.make_tensor( + mixed.iterator, cute.make_layout((m, (bh, hc, nt), 1), stride=(hid, (1, 0, bh), 0)) + ) + # NVFP4 outputs (`fp4` only): 8 e2m1 codes per Int32, one e4m3 byte per 16. + xq = cute.make_tensor(xq_ptr, cute.make_layout((m, hid // 8), stride=(hid // 8, 1))) + xsf = cute.make_tensor(sf_ptr, cute.make_layout((m, hid // 16), stride=(hid // 16, 1))) + gsc = cute.make_tensor(gs_ptr, cute.make_layout(1)) + k_major = utils.LayoutEnum.ROW_MAJOR.mma_major_mode() + tiled_mma = self._setup(k_major, k_major) + + a_op = utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + a_smem = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, a, a_smem, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape + ) + b_op = utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + b_smem = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, b, b_smem, self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape + ) + tma_atom_n, tma_n = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + nrm, + cute.select(self.n_smem_layout_staged, mode=[0, 1]), + self.n_box, + ) + self.num_tma_load_bytes = cute.size_in_bytes(self.io_dtype, a_smem) + cute.size_in_bytes( + self.io_dtype, b_smem + ) + # 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). + sched = utils.PersistentTileSchedulerParams( + ((m + 127) // 128, nt, 1), (*self.cluster_shape_mn, 1), 1, False + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape(sched, max_active_clusters) + self.kernel( + tiled_mma, + tma_atom_a, + tma_a, + tma_atom_b, + tma_b, + tma_atom_n, + tma_n, + mix, + xq, + xsf, + gsc, + m, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.n_smem_layout_staged, + self.epi_tile, + sched, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + use_pdl=_PDL, + ) + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_n: cute.CopyAtom, + mN_mn: cute.Tensor, + mX_mnl: cute.Tensor, + mQ: cute.Tensor, + mSF: cute.Tensor, + mGS: cute.Tensor, + m: cutlass.Int32, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + n_smem_layout_staged: cute.ComposedLayout, + epi_tile: cute.Tile, + sched_params: utils.PersistentTileSchedulerParams, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + if warp_idx == self.nrm_warp_id: + cpasync.prefetch_descriptor(tma_atom_n) + tidx, _, _ = cute.arch.thread_idx() + + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + n_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_n_stage * 2] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + ab_producer, ab_consumer = PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + ).make_participants() + acc_pipeline = PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len(self.epilogue_warp_id) + ), + cta_layout_vmnk=cluster_layout_vmnk, + ) + n_producer, n_consumer = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.n_full_mbar_ptr.data_ptr(), + num_stages=self.num_n_stage, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, len(self.epilogue_warp_id) + ), + tx_count=self.num_n_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + ).make_participants() + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=False, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, + ) + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + sA = smem.allocate_tensor( + element_type=self.io_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + sB = smem.allocate_tensor( + element_type=self.io_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + sN = smem.allocate_tensor( + element_type=self.io_dtype, + layout=n_smem_layout_staged.outer, + byte_alignment=128, + swizzle=n_smem_layout_staged.inner, + ) + + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + tile_mn = cute.slice_(self.mma_tiler, (None, None, 0)) + gX_mnl = cute.local_tile(mX_mnl, tile_mn, (None, None, None)) + cX_mnl = cute.local_tile( + cute.make_identity_tensor(mX_mnl.shape), tile_mn, (None, None, None) + ) + # normed boxes: [128 rows, 32 columns] at (row tile, column block) + gN_box = cute.local_tile(mN_mn, self.n_box, (None, None)) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + thr_mma = tiled_mma.get_slice(0) + tCgA = thr_mma.partition_A(gA_mkl) + tCgB = thr_mma.partition_B(gB_nkl) + tCgX = thr_mma.partition_C(gX_mnl) + tCcX = thr_mma.partition_C(cX_mnl) + + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + 0, + cute.make_layout(1), + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + 0, + cute.make_layout(1), + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + tNsN, tNgN = cpasync.tma_partition( + tma_atom_n, + 0, + cute.make_layout(1), + cute.group_modes(sN, 0, 2), + cute.group_modes(gN_box, 0, 2), + ) + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + tile_sched = utils.StaticPersistentTileScheduler.create( + sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + if warp_idx == self.tma_warp_id: + # Programmatic dependent launch: the first tile's up-weight stages (static) + # stream before the grid-dependency wait, `low` -- the previous kernel's + # output -- after it. Every CTA owns at least one tile (grid <= tiles). + crd = work_tile.tile_idx + tAgA_slice = tAgA[(None, crd[0], None, crd[2])] + tBgB_slice = tBgB[(None, crd[1], None, crd[2])] + ab_producer.reset() + pre = [] + for kk in cutlass.range_constexpr(self.n_pre): + handle = ab_producer.acquire_and_advance() + cute.copy( + tma_atom_b, + tBgB_slice[(None, kk)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + pre.append(handle) + cute.arch.griddepcontrol_wait() + for kk in cutlass.range_constexpr(self.n_pre): + handle = pre[kk] + cute.copy( + tma_atom_a, + tAgA_slice[(None, kk)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + for k_tile in cutlass.range(self.n_pre, k_tile_cnt, 1, unroll=1): + handle = ab_producer.acquire_and_advance() + cute.copy( + tma_atom_a, + tAgA_slice[(None, k_tile)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, k_tile)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + while work_tile.is_valid_tile: + crd = work_tile.tile_idx + tAgA_slice = tAgA[(None, crd[0], None, crd[2])] + tBgB_slice = tBgB[(None, crd[1], None, crd[2])] + ab_producer.reset() + peek = ab_producer.try_acquire() + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + handle = ab_producer.acquire_and_advance(peek) + cute.copy( + tma_atom_a, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + ) + peek = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek = ab_producer.try_acquire() + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + ab_producer.tail() + + if warp_idx == self.nrm_warp_id: + cute.arch.griddepcontrol_wait() + halves = self.bh // self.epi_cols + col_blocks = self.hidden // self.epi_cols + while work_tile.is_valid_tile: + crd = work_tile.tile_idx + for half in cutlass.range_constexpr(halves): + handle = n_producer.acquire_and_advance() + for s in cutlass.range_constexpr(self.hc): + cute.copy( + tma_atom_n, + tNgN[(None, crd[0], s * col_blocks + crd[1] * halves + half)], + tNsN[(None, handle.index * self.hc + s)], + tma_bar_ptr=handle.barrier, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + n_producer.tail() + + if warp_idx == self.mma_warp_id: + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + while work_tile.is_valid_tile: + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + ab_consumer.reset() + peek = ab_consumer.try_wait() + acc_pipeline.producer_acquire(acc_producer_state) + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + handle = ab_consumer.wait_and_advance(peek) + num_kblocks = cute.size(tCrA, mode=[2]) + for kb in cutlass.range(num_kblocks, unroll_full=True): + kcrd = (None, None, kb, handle.index) + cute.gemm(tiled_mma, tCtAcc, tCrA[kcrd], tCrB[kcrd], tCtAcc) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + handle.release() + peek = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek = ab_consumer.try_wait() + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + acc_pipeline.producer_tail(acc_producer_state) + + if warp_idx < self.mma_warp_id: + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + copy_atom_t2r = utils.sm100.get_tmem_load_op( + self.cta_tile_shape_mnk, + utils.LayoutEnum.ROW_MAJOR, + self.io_dtype, + self.acc_dtype, + epi_tile, + False, + ) + tAcc_epi = cute.flat_divide(tCtAcc_base[((None, None), 0, 0, None)], epi_tile) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)]) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_tAcc_base = thr_copy_t2r.partition_S(tAcc_epi) + tTR_sN = thr_copy_t2r.partition_D(sN) + tTR_gX = thr_copy_t2r.partition_D( + cute.flat_divide(tCgX[((None, None), 0, 0, None, None, None)], epi_tile) + ) + tTR_cX = thr_copy_t2r.partition_D( + cute.flat_divide(tCcX[((None, None), 0, 0, None, None, None)], epi_tile) + ) + frag_shape = tTR_gX[(None, None, None, 0, 0, 0, 0, 0)].shape + tTR_rAcc = cute.make_rmem_tensor(frag_shape, self.acc_dtype) + tTR_rX = cute.make_rmem_tensor(frag_shape, self.io_dtype) + tTR_rSum = cute.make_rmem_tensor(frag_shape, self.acc_dtype) + nfrag = cute.size(tTR_rAcc) + halves = self.bh // self.epi_cols + rN = cute.make_rmem_tensor(frag_shape, self.io_dtype) + inv_hc = cutlass.Float32(1.0 / self.hc) + if cutlass.const_expr(self.fp4): + # A thread owns one row's `epi_cols` consecutive columns (the + # 32dp32b TMEM load), i.e. whole 16-element scale blocks. + gs = mGS[0] + rcp_gs = _rcp_approx_ftz(gs) + rcp6 = _rcp_approx_ftz(cutlass.Float32(6.0)) + rV = cute.make_rmem_tensor((nfrag,), self.acc_dtype) + assert nfrag == 32, "the NVFP4 epilogue stores two scale blocks per row" + q_base = mQ.iterator.toint() + sf_base = mSF.iterator.toint() + + while work_tile.is_valid_tile: + crd = work_tile.tile_idx + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_consumer_state.index)] + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + gX_t = tTR_gX[(None, None, None, None, None, *crd)] + gX_t = cute.group_modes(gX_t, 3, cute.rank(gX_t)) + cX_t = tTR_cX[(None, None, None, None, None, *crd)] + cX_t = cute.group_modes(cX_t, 3, cute.rank(cX_t)) + row_ok = cX_t[(0, 0, 0, 0)][0] < m + + acc_pipeline.consumer_wait(acc_consumer_state) + for half in cutlass.range_constexpr(halves): + n_handle = n_consumer.wait_and_advance() + for s in cutlass.range_constexpr(self.hc): + cute.autovec_copy( + tTR_sN[(None, None, None, n_handle.index * self.hc + s)], rN + ) + cute.copy( + tiled_copy_t2r, + tTR_tAcc[(None, None, None, s * halves + half)], + tTR_rAcc, + ) + for i in cutlass.range_constexpr(nfrag): + gate = _sigmoid(tTR_rAcc[i]) + if cutlass.const_expr(s == 0): + tTR_rSum[i] = gate * rN[i].to(self.acc_dtype) + else: + tTR_rSum[i] = tTR_rSum[i] + gate * rN[i].to(self.acc_dtype) + cute.arch.sync_warp() + n_handle.release() + for i in cutlass.range_constexpr(nfrag): + tTR_rX[i] = (tTR_rSum[i] * inv_hc).to(self.io_dtype) + if row_ok: + cute.autovec_copy(tTR_rX, gX_t[(None, None, None, half)]) + if cutlass.const_expr(self.fp4): + row = cX_t[(0, 0, 0, 0)][0] + blk0 = (crd[1] * self.bh + half * self.epi_cols) // 16 + for i in cutlass.range_constexpr(nfrag): + rV[i] = tTR_rX[i].to(self.acc_dtype) + words, sfs = [], [] + for b in cutlass.range_constexpr(2): + v = [rV[16 * b + j] for j in range(16)] + amax = _absmax16(v) + sf = _e4m3_scale(amax, gs, rcp6) + osc = _e2m1_out_scale(sf, rcp_gs, amax) + words += [_e2m1x8(v[:8], osc), _e2m1x8(v[8:], osc)] + sfs.append(sf) + row64 = cutlass.Int64(row) + _st_global_v4( + q_base + row64 * (self.hidden // 2) + cutlass.Int64(blk0) * 8, + *words, + ) + _st_global_b16( + sf_base + row64 * (self.hidden // 16) + cutlass.Int64(blk0), + sfs[0] | (sfs[1] << 8), + ) + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + cute.arch.barrier( + barrier_id=self.epilog_sync_bar_id, + number_of_threads=32 * len(self.epilogue_warp_id), + ) + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + +_COMPILED: dict = {} + + +def _launcher(block_h: int, hc: int, hidden: int, rank: int, ptrs, fp4: bool = False): + """Compile once per static configuration; later calls reuse the TVM-FFI entry, + which launches on the caller's current stream (graph-capturable). `ptrs` is + `(low, up, normed, mixed tensor, xq, sf, gs)`.""" + key = (block_h, hc, hidden, rank, fp4) + fn = _COMPILED.get(key) + if fn is None: + kern = HcUpGateMeanKernel(block_h, hc, hidden, rank, fp4=fp4) + cptrs = [ + make_ptr(cutlass.BFloat16, p, cute.AddressSpace.gmem, assumed_align=16) + for p in ptrs[:3] + ] + out = from_dlpack(ptrs[3], assumed_align=16).mark_layout_dynamic(leading_dim=1) + qptrs = [ + make_ptr(cutlass.Int32, ptrs[4], cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.Uint8, ptrs[5], cute.AddressSpace.gmem, assumed_align=2), + make_ptr(cutlass.Float32, ptrs[6], cute.AddressSpace.gmem, assumed_align=4), + ] + fn = cute.compile( + kern, + 128, + *cptrs, + out, + *qptrs, + max_active_clusters=utils.HardwareInfo().get_max_active_clusters(1), + stream=cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--opt-level 2 --enable-tvm-ffi", + ) + _COMPILED[key] = fn + return fn + + +def hc_up_gate_mean( + low: torch.Tensor, up: torch.Tensor, normed: torch.Tensor, hc_count: int, block_h: int +) -> torch.Tensor: + """Launch the fused kernel; `up` is the `[R, C*H]` transposed view of the + `[C*H, R]` up weight, as `cublas_mm(low, up)` takes it.""" + rows, rank = low.shape + hidden = normed.shape[1] // hc_count + mixed = torch.empty((rows, hidden), dtype=normed.dtype, device=normed.device) + if rows > 0: + ptrs = (low.data_ptr(), up.data_ptr(), normed.data_ptr()) + # The NVFP4 pointers are never read without `fp4`. + unused = (mixed.data_ptr(),) * 3 + _launcher(block_h, hc_count, hidden, rank, (*ptrs, mixed, *unused))( + rows, *ptrs, mixed, *unused + ) + return mixed + + +def hc_up_gate_mean_fp4( + low: torch.Tensor, + up: torch.Tensor, + normed: torch.Tensor, + hc_count: int, + global_scale: torch.Tensor, + block_h: int, +): + """`hc_up_gate_mean` plus the NVFP4 of `mixed` in the same epilogue. + Returns `(mixed [T, H], xq [T, H/2] uint8, xsf [T*H/16] uint8)`.""" + rows, rank = low.shape + hidden = normed.shape[1] // hc_count + dev = normed.device + mixed = torch.empty((rows, hidden), dtype=normed.dtype, device=dev) + xq = torch.empty((rows, hidden // 2), dtype=torch.uint8, device=dev) + xsf = torch.empty((rows * (hidden // 16),), dtype=torch.uint8, device=dev) + if rows > 0: + ptrs = (low.data_ptr(), up.data_ptr(), normed.data_ptr()) + qptrs = (xq.data_ptr(), xsf.data_ptr(), global_scale.data_ptr()) + _launcher(block_h, hc_count, hidden, rank, (*ptrs, mixed, *qptrs), fp4=True)( + rows, *ptrs, mixed, *qptrs + ) + 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. +_BLOCK_H = 32 + + +class Qwen4ExpHcUpGateMean(OpWrapper): + """With `C = hc_count`, `H = width // C`, `low` the `[T, R]` SiLU bottleneck + and `up` the `[R, C*H]` transposed view of the `[C*H, R]` up weight:: + + mixed[t, i] = mean over s of sigmoid((low @ up)[t, s*H + i]) * normed[t, s*H + i] + + The GEMM accumulates in fp32 in tensor memory and is never rounded or + written; sigmoid, product and stream mean are fp32 with one rounding to + `normed`'s dtype. Returns a fresh `[T, H]` tensor. Fusion boundary: the up + projection, the sigmoid gate and the stream mean (two launches as `cublas_mm(low, + up)` followed by a separate gate-and-mean pass). + """ + + 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)"), + ) + ) + (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). + """ + + def __call__( + self, low: torch.Tensor, up: torch.Tensor, normed: torch.Tensor, hc_count: int + ) -> torch.Tensor: + return hc_up_gate_mean(low, up, normed, hc_count, _BLOCK_H) + + def reference( + self, low: torch.Tensor, up: torch.Tensor, normed: torch.Tensor, hc_count: int + ) -> torch.Tensor: + """float64, one rounding.""" + hidden = normed.shape[1] // hc_count + gate = torch.sigmoid(F.linear(low.double(), up.double().t())).unflatten( + -1, (hc_count, hidden) + ) + x = normed.double().unflatten(-1, (hc_count, hidden)) + return (gate * x).mean(dim=-2).to(normed.dtype) + + def is_valid( + self, low: torch.Tensor, up: torch.Tensor, normed: torch.Tensor, hc_count: int + ) -> None: + rows, rank = low.shape + width = normed.shape[1] + assert low.dtype == up.dtype == normed.dtype == torch.bfloat16, "bf16 operands" + assert normed.shape[0] == rows and up.shape == (rank, width), "shapes" + assert low.stride() == (rank, 1) and normed.stride() == (width, 1), "contiguous rows" + assert up.stride() == (1, rank), "up must be the transposed view of a [C*H, R] weight" + assert width % hc_count == 0 and (width // hc_count) % 32 == 0 and hc_count * 32 <= 256 + assert rank % 64 == 0, "rank must be a multiple of 64" + for t in (low, up, normed): + assert t.data_ptr() % 16 == 0, "16-byte aligned operands" + + +qwen4_exp_hc_up_gate_mean = Qwen4ExpHcUpGateMean() 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__tp1/ops/hc/qwen4_exp_hc_up_gate_mean_fp4.py new file mode 100644 index 000000000000..6d4cfb844b64 --- /dev/null +++ b/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 @@ -0,0 +1,127 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hyper-connection mixer second half at an MoE site: `qwen4_exp_hc_up_gate_mean` +(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. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) + +from .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 + + +class Qwen4ExpHcUpGateMeanFp4(OpWrapper): + """With `C = hc_count`, `H = width // C`, `low` the `[T, R]` SiLU bottleneck, + `up` the `[R, C*H]` transposed view of the `[C*H, R]` up weight:: + + mixed[t, i] = mean over s of sigmoid((low @ up)[t, s*H + i]) * normed[t, s*H + i] + xq, xsf = NVFP4(mixed) # e2m1 pairs, e4m3 scale per 16, linear layout + + `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. + """ + + 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 (max_con)"), + (16384, "prefill chunk 16384 (max_num_tokens)"), + ) + ) + (Cell(why="ragged tail, 16320 + 70 rows (mixed step)", spec=dict(rows=16390)),) + + 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). + + 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. + + 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. + + 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). + """ + + def __call__( + self, + low: torch.Tensor, + up: torch.Tensor, + normed: torch.Tensor, + hc_count: int, + global_scale: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return hc_up_gate_mean_fp4(low, up, normed, hc_count, global_scale, _BLOCK_H) + + def reference( + self, + low: torch.Tensor, + up: torch.Tensor, + normed: torch.Tensor, + hc_count: int, + global_scale: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """`mixed` in float64 with one rounding, then the NVFP4 of that rounded + value in native torch.""" + mixed = qwen4_exp_hc_up_gate_mean.reference(low, up, normed, hc_count) + return (mixed,) + fp4_quantize.reference(mixed, global_scale, SF_VEC, False, False) + + def is_valid( + self, + low: torch.Tensor, + up: torch.Tensor, + normed: torch.Tensor, + hc_count: int, + global_scale: torch.Tensor, + ) -> None: + qwen4_exp_hc_up_gate_mean.is_valid(low, up, normed, hc_count) + assert (normed.shape[1] // hc_count) % SF_VEC == 0, "hidden must be whole 16-blocks" + # The kernel loads element 0 only; a per-row scale would be applied as row 0's. + assert global_scale.numel() == 1 and global_scale.dtype == torch.float32 + assert global_scale.is_cuda + + +qwen4_exp_hc_up_gate_mean_fp4 = Qwen4ExpHcUpGateMeanFp4() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_post_inject.py index 679d4ce50521..d51ddde1e610 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_post_inject.py @@ -43,55 +43,21 @@ class Qwen4ExpPostInject(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # spec: leading (token dims), hc_count, hidden, dtype; layout - # "contiguous" (default) | "row_slice" (`buf[::2]`) | "col_slice" - # (`buf[:, :width]`), both last-dim stride 1. - CELLS: tuple[Cell, ...] = ( - tuple( - Cell( - why=f"checkpoint 4x2560 bf16, {why}", - spec=dict(leading=(t,), hc_count=CKPT_HC_COUNT, hidden=CKPT_HIDDEN), - ) - for t, why in ( - (1, "decode, 1 token"), - (2, "2 tokens"), - (7, "7 tokens, odd count"), - (64, "decode batch / MTP rows, 64 tokens"), - (2048, "prefill 2048 tokens"), - (8192, "prefill 8192 tokens, largest measured"), - ) + # "contiguous" (default) | "col_slice" (`buf[:, :width]`, last-dim stride 1). + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"checkpoint 4x2560 bf16, {why}", + spec=dict(leading=(t,), hc_count=CKPT_HC_COUNT, hidden=CKPT_HIDDEN), ) - + ( - Cell( - why="[2, 5, width] leading dims", - spec=dict(leading=(2, 5), hc_count=CKPT_HC_COUNT, hidden=CKPT_HIDDEN), - ), - Cell( - why="hyper_input a row slice buf[::2] of a wider buffer", - spec=dict( - leading=(8,), hc_count=CKPT_HC_COUNT, hidden=CKPT_HIDDEN, layout="row_slice" - ), - ), - Cell( - why="hyper_input a column slice buf[:, :width]", - spec=dict( - leading=(8,), hc_count=CKPT_HC_COUNT, hidden=CKPT_HIDDEN, layout="col_slice" - ), - ), - ) - + tuple( - Cell( - why=f"hc_count {c} read from inject, hidden 128", - spec=dict(leading=(48,), hc_count=c, hidden=128), - ) - for c in (1, 2, 3, 8) - ) - + tuple( - Cell( - why=f"{str(d).removeprefix('torch.')} operands, 64 tokens", - spec=dict(leading=(64,), hc_count=CKPT_HC_COUNT, hidden=CKPT_HIDDEN, dtype=d), - ) - for d in (torch.float16, torch.float32) + for t, why in ( + (64, "decode batch / MTP rows, 64 tokens"), + (8192, "prefill 8192 tokens, largest measured"), ) + ) + ( + Cell( + why="hyper_input a column slice buf[:, :width]", + spec=dict(leading=(8,), hc_count=CKPT_HC_COUNT, hidden=CKPT_HIDDEN, layout="col_slice"), + ), ) note = """ @@ -122,8 +88,13 @@ class Qwen4ExpPostInject(OpWrapper): broadcast operands take a non-vectorized path. An in-place form was not measured. - Not certified: float8 operands. Hard limits: not probed (the largest - measured operand is 8192 x 10240 = 2^26.3 elements). + 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. + + 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). """ def __call__( 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py new file mode 100644 index 000000000000..9ba3b8122735 --- /dev/null +++ b/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 @@ -0,0 +1,484 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CuteDSL NVFP4 FC1 of the routed experts: gather, grouped GEMM, SwiGLU and NVFP4 +requantization of the intermediate, in one trtllm call.""" + +from __future__ import annotations + +from typing import NamedTuple, 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 +from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import SWIGLU_LIMIT_SCALAR_DISABLED +from tensorrt_llm._torch.utils import ActivationType + +_SWIGLU = int(ActivationType.Swiglu) + +#: e2m1 magnitudes by 3-bit code, and the midpoints between neighbours. A value +#: exactly on a midpoint rounds to the even code (`_TIE_UP[i]`: the even side is up). +E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +E2M1_MAX = 6.0 +E4M3_MAX = 448.0 +_E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0) +_TIE_UP = (False, True, False, True, False, True, False) + +#: The pairing the fused SwiGLU epilogue reads: the weight's rows `[128 g, 128 g + 64)` +#: are up rows, the next 64 their gate rows, and output column `64 g + r` pairs row +#: `128 g + r` with row `128 g + 64 + r`. Fixed by the epilogue's 64-column subtile, +#: whatever MMA N tile the autotuner picks. +GATE_UP_GROUP = 64 + + +def _ckpt( + tokens: int, + tile_size: int, + mma_n: int = 128, + dist: str = "router", + gsf: str = "calibrated", + zero_tokens: int = 0, +) -> dict: + """The checkpoint's routed experts: 512 experts, top-10, hidden 2560, I 640, all + local (tp1), SwiGLU, no clamp, 16-element e4m3 block scales. + + `mma_n` is the autotuner tactic's MMA N tile (128 is also the fallback when the + tuner has no entry). `dist` names how the test draws the expert ids: `router` = + each row's top-10 of uniform scores. `gsf` sets `global_sf` against the + activations' magnitude: `calibrated` = the activation maximum the test + calibrates to maps to the e4m3 maximum, as the checkpoint's per-layer `1 / + down_proj.input_scale` does; `subnormal` = 2^-12 of it, so block scales fall + into e4m3 subnormals and zero. `zero_tokens` makes that many leading tokens + all-zero rows (zero block scales). + """ + return dict( + tokens=tokens, + tile_size=tile_size, + mma_n=mma_n, + dist=dist, + gsf=gsf, + zero_tokens=zero_tokens, + num_experts=512, + top_k=10, + hidden=2560, + intermediate=640, + ) + + +class Fc1Expected(NamedTuple): + """What `reference` returns: the op's two outputs, plus what the test's compare reads. + + `c` / `c_sf` are laid out exactly as the op returns them; rows the op does not + write hold zeros. `rows` marks the permuted rows the op writes and the compare + gates. `scaled` is `y * global_sf` on those rows, before rounding -- the value + the block scale and the codes are derived from, which rounding loses. + """ + + c: torch.Tensor + c_sf: torch.Tensor + rows: torch.Tensor + scaled: torch.Tensor + + +class CuteDslNvfp4GatherGroupedGemmActFusion(OpWrapper): + """`torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell`: + FC1 of the routed experts on the CuteDSL MoE path, one kernel launch. + + For every permuted row `p` the sort filled -- tile `i = p // tile_size` below + `num_non_exiting_tiles` and `p < tile_idx_to_mn_limit[i]`: + + t = permuted_idx_to_expanded_idx[p] // top_k the token + e = tile_idx_to_group_idx[i] the local expert + acc = alpha[e] * deq(input[t]) @ deq(weight[e]).T [2 I], fp32 + y[p, j] = up * silu(gate), up = acc[128 g + r], gate = acc[128 g + 64 + r], + j = 64 g + r + out[p], out_sf[p] = NVFP4(y[p] * global_sf) 16-element blocks + + where `deq` is e2m1 times its e4m3 block scale, and NVFP4 takes per block + `sf = e4m3(max|v| / 6)` (saturating at 448) and codes `e2m1(v / sf)` (saturating + at 6). So `out * out_sf / global_sf` approximates `y`. + + Arguments, as the CuteDSL MoE passes them (`run_moe_nvfp4_impl`, + `_torch/moe/fused_moe/fused_moe_cute_dsl.py`): + + - `input` `[T, H/2]` float4_e2m1fn_x2 and `input_scale` `[T, H/16]` uint8: the + un-permuted activation as `fp4_quantize(x, 1 / input_scale_1, 16, False, + False)` returns it -- LINEAR (row-major) e4m3 scales, viewed 2-D. Gathered + inside; the caller never permutes it. + - `weight` `[E, 2I, H/2]` float4_e2m1fn_x2: per expert `cat([up, gate])` + interleaved in `GATE_UP_GROUP`-row groups (up rows 0:64, gate rows 0:64, up + rows 64:128, ...). No 32-row shuffle (that is trtllm-gen's layout). + - `weight_scale` `[E, 2I, H/16]` uint8: the e4m3 block scales with the same row + interleave, then 128x4-swizzled per expert (`block_scale_interleave`). + - `alpha` `[E]` fp32: per expert, `input_scale_1 * weight_scale_2[e]` -- one + factor for both halves, applied to the accumulator before SwiGLU. + - the four sort tensors, exactly as `moe_sort` returns them for the same + `top_k` and `tile_tokens_dim == tile_size`. + - `global_sf` fp32, one element: the FC2 input's global scale, `1 / + max(down_proj.input_scale)` of the layer. + - `num_experts`, `num_local_experts`, `local_expert_offset`: read only by the + autotuner's shape buckets and cache key; the kernel takes the expert count + from `weight.size(0)` and the expert from the tile. + + Returns `(c, c_sf)`: `c` `[M, I/2]` float4_e2m1fn_x2 and `c_sf` `[M * I/16]` + uint8, `M = permuted_idx_to_expanded_idx.numel()`, one row per PERMUTED row and + the block scales 128x4-swizzled over `[M, I/16]` -- exactly the `input` / + `input_scale` the plain NVFP4 grouped GEMM (FC2) takes, with the same four sort + tensors. Rows the sort did not fill are not part of the result (see `note`). + + Fusion boundary: the gather of the activation rows, the grouped GEMM, the + per-expert alpha, SwiGLU and the requantization happen inside. The caller owns + the activation quantization, the sort, FC2 and the routing-weighted combine; + routing weights are not read here. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="tile 128: decode, 1 request without MTP", + spec=_ckpt(1, 128), + ), + Cell( + why="tile 128: 8192 rows, max_num_tokens prefill", + spec=_ckpt(8192, 128), + ), + Cell( + why="tile 128: tiny global_sf and 7 all-zero token rows; subnormal and zero " + "block scales", + spec=_ckpt(2048, 128, gsf="subnormal", zero_tokens=7), + ), + ) + + note = """ + Path certified: the CuTe DSL kernel `BlockScaledContiguousGatherGroupedGemmKernel` + (`_torch/moe/kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py`), + built by `Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner` + (`_torch/custom_ops/cute_dsl_custom_ops.py`), which JIT-compiles it with + `cute.compile` on the first call per (tile, tactic, top_k, activation) and + keeps it for the process. The op is registered only when the CuTe DSL imports + (`IS_CUTLASS_DSL_AVAILABLE`), and the runner accepts SM 100 and 103 only. + + Metadata consumed: none beyond the arguments. Stateless: allocates `c` and + `c_sf` on every call, on the current stream, with no host synchronization, so + it captures into a CUDA graph together with moe_sort. + + Call sites: the engine's CuteDSL MoE (`run_moe_nvfp4_impl`) and the target's + routed-expert step, eager for prefill and mixed batches (T up to + max_num_tokens 8192) and captured for decode (T = requests x MTP verify width, + up to 32 x 4 = 128). Tile 128 is the target's and the cells' tile; 256 is + certified through the per-pair invariance below. + + Tactics. The op asks the AutoTuner for an MMA tile of (tile_size, 128) or + (tile_size, 256) and with no tuning entry falls back to (tile_size, 128); + serving warmup tunes per row bucket and may pick either. Every pair's codes + and block scales are bitwise identical across tile 128 / 256 x MMA N 128 / 256 + x the fallback (the test checks it at 8192 rows), so cells driven at one + tactic certify the others. + + Determinism. A (token, expert) pair's output row depends on that pair alone: + it is bitwise unchanged by reordering an expert's rows (moe_sort orders them by + atomics from T = 5), by repeating the call, and by cutting a batch to a prefix. + + What is written: the rows the sort filled hold the result; the padding rows of + a written tile -- from its tile_idx_to_mn_limit to its end -- get zero block + scales and zero codes, so a plain grouped GEMM over whole tiles (FC2) reads + zeros there; the tiles at and past num_non_exiting_tiles are not written at + all, and under a graph keep whatever the buffer last held. FC2 must be tiled by + the same tile arrays, and the combine reads only + expanded_idx_to_permuted_idx's rows. + + Scale range: with a tiny global_sf the block scales fall to e4m3 subnormals + and to zero (a cell). An all-zero token gives scale 0 and codes of magnitude + 0 (the same cell). With global_sf 25x the calibrated value the block scales + clip at 448 and the codes at 6 (no NaN) -- measured, not certified. A block + whose scale is 0 dequantizes to zero whatever its codes hold, and the test's + compare does not read them. + + Preconditions the caller owns (a wrong preparation of each operand -- the + weight's rows, its scales, the activation's scales, alpha -- runs without + error, with finite scales, and puts most block scales and codes outside the + test's gate): + - `weight` / `weight_scale` in the layout the class docstring states: not + trtllm-gen's row interleave, up rows first, scales swizzled after the same + row interleave as the weight, alpha not shifted by an expert. + - One `alpha` per expert scales BOTH halves, so up_proj and gate_proj must + share `weight_scale_2`, and `input_scale` must be one value per layer; this + checkpoint satisfies both. `global_sf` inverts the largest + `down_proj.input_scale`; the spread of the rest is FC2's `alpha`. + - `input_scale` LINEAR. Swizzled scales have the same size when T % 128 == 0 + and are read as garbage without error. + - The sort from moe_sort with the same `top_k` and tile_tokens_dim == + `tile_size`, over the tokens `input` holds. Nothing checks `top_k` against it. + + Refused by `is_valid`, each silently wrong unguarded: a non-contiguous input, + input_scale, weight, weight_scale or alpha (data_ptr() read as dense); an + alpha shorter than weight.size(0); a weight_scale with fewer bytes than the + weight needs (read past its end: zeros of the same allocation come back as zero + experts, a fresh allocation faults). + + Loud (AssertionError from the runner before the launch): tile_size 256 with a + tile-128 sort and 128 with a tile-256 sort; input_scale not H/16 wide; + global_sf of 2 elements; alpha not fp32; weight_scale typed float8_e4m3fn; + input typed uint8. tile_size 64: ValueError. + + Measured at bring-up and not re-driven by the test: the CUDA-graph capture + with moe_sort, the determinism and write statements, the wrong preparations' + distance from the gate, the short-alpha and short-weight_scale refusals and + the loud rejections. The test re-drives the cells, the tile / tactic + invariance and the `input` contiguity refusal. + + Uncertified: Relu2 and SiTU; a SwiGLU clamp (swiglu_limit_scalar >= 0); + expert_counts (DeepEP count-native metadata); expert parallelism (the kernel + reads neither num_experts nor local_expert_offset, so EP changes only the sort + and the weight slice); other hidden / intermediate sizes and top_k; + scaling_vector_size other than 16; other architectures. + + Hard limits (sm_103) + written rows x I/16 B of c_sf boundary 2^31 B: an illegal memory access + once rows written reach 53687092 at I = 640 (T ~ 5.37M at + top-10; T = 8192 writes at most 212480 rows) + input, input_scale, weight past 2^31 elements / bytes: exact (probed + up to 1536 experts) + permuted index int32, within moe_sort's 2^31-pair limit; no index passes + through fp32 + """ + + def __call__( + self, + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + scaling_vector_size: int = 16, + activation_type: int = _SWIGLU, + swiglu_limit_scalar: float = SWIGLU_LIMIT_SCALAR_DISABLED, + expert_counts: Optional[torch.Tensor] = None, + expert_capacity: int = 0, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell( + input=input, + weight=weight, + input_scale=input_scale, + weight_scale=weight_scale, + alpha=alpha, + tile_idx_to_group_idx=tile_idx_to_group_idx, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + global_sf=global_sf, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=local_expert_offset, + tile_size=tile_size, + scaling_vector_size=scaling_vector_size, + activation_type=activation_type, + swiglu_limit_scalar=swiglu_limit_scalar, + expert_counts=expert_counts, + expert_capacity=expert_capacity, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + ) + + def reference( + self, + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + scaling_vector_size: int = 16, + activation_type: int = _SWIGLU, + swiglu_limit_scalar: float = SWIGLU_LIMIT_SCALAR_DISABLED, + expert_counts: Optional[torch.Tensor] = None, + expert_capacity: int = 0, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, + ) -> Fc1Expected: + """Gather, dequantize, GEMM in float64, SwiGLU, exact NVFP4 rounding. + + Products of e2m1 and e4m3 values are exact, and float64 sums of 2560 of + them are exact to ~1e-16 relative, so this side carries no accumulation + error of its own; the block scale and codes are the exact round-to-nearest- + even of the exact `y * global_sf`. Where a block scale rounds to zero the + codes here are zero: they dequantize to zero whatever they hold, and the + test's compare does not read them. + + Certified surface only: SwiGLU without a clamp, the sort-metadata path + (`expert_counts` None), no SiTU. Anything else raises NotImplementedError. + `num_experts`, `num_local_experts`, `local_expert_offset` and + `expert_capacity` do not enter the computation. + """ + if ( + int(activation_type) != _SWIGLU + or swiglu_limit_scalar >= 0 + or expert_counts is not None + or situ_beta is not None + or situ_linear_beta is not None + ): + raise NotImplementedError( + "reference covers SwiGLU without a clamp on sort metadata only" + ) + vec = scaling_vector_size + dev = input.device + m = permuted_idx_to_expanded_idx.numel() + n2 = weight.shape[1] + inter = n2 // 2 + hidden = input.shape[1] * 2 + n_tiles = int(num_non_exiting_tiles.reshape(-1)[0]) + span = n_tiles * tile_size + + rows = torch.zeros(m, dtype=torch.bool, device=dev) + limit = tile_idx_to_mn_limit[:n_tiles].long().repeat_interleave(tile_size) + rows[:span] = torch.arange(span, device=dev) < limit + real = torch.nonzero(rows).squeeze(1) + token = permuted_idx_to_expanded_idx.long()[real] // top_k + expert = tile_idx_to_group_idx[:n_tiles].long().repeat_interleave(tile_size)[real] + + a = _dequant(input.view(torch.uint8)[token], input_scale.view(torch.uint8)[token], vec) + y = torch.empty(real.numel(), inter, dtype=torch.float64, device=dev) + w_sf = weight_scale.view(torch.uint8).reshape(weight.shape[0], -1) + group = GATE_UP_GROUP + for e in torch.unique(expert).tolist(): + sel = torch.nonzero(expert == e).squeeze(1) + w = _dequant(weight[e].view(torch.uint8), unswizzle(w_sf[e], n2, hidden // vec), vec) + acc = (a[sel] @ w.T) * float(alpha[e]) + acc = acc.view(-1, n2 // (2 * group), 2, group) + up = acc[:, :, 0].reshape(-1, inter) + gate = acc[:, :, 1].reshape(-1, inter) + y[sel] = up * gate * torch.sigmoid(gate) + scaled = y * float(global_sf.reshape(-1)[0]) + + blocks = scaled.view(-1, inter // vec, vec) + pre = blocks.abs().amax(-1) / E2M1_MAX + sf = pre.clamp(max=E4M3_MAX).float().to(torch.float8_e4m3fn) + sf_val = sf.double().unsqueeze(-1) + q = torch.where(sf_val > 0, blocks / sf_val, torch.zeros_like(blocks)) + codes = e2m1_codes(q.reshape(-1, inter)) + c = torch.zeros(m, inter // 2, dtype=torch.uint8, device=dev) + c[real] = codes[:, 0::2] | (codes[:, 1::2] << 4) + sf_rows = torch.zeros(m, inter // vec, dtype=torch.uint8, device=dev) + sf_rows[real] = sf.view(torch.uint8) + return Fc1Expected( + c=c.view(torch.float4_e2m1fn_x2), c_sf=swizzle(sf_rows), rows=rows, scaled=scaled + ) + + def is_valid( + self, + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + scaling_vector_size: int = 16, + activation_type: int = _SWIGLU, + swiglu_limit_scalar: float = SWIGLU_LIMIT_SCALAR_DISABLED, + expert_counts: Optional[torch.Tensor] = None, + expert_capacity: int = 0, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, + ) -> None: + # The runner hands data_ptr() to the kernel with dense row-major layouts + # assumed; it checks dims and dtypes, never strides. + for name, t in ( + ("input", input), + ("input_scale", input_scale), + ("weight", weight), + ("weight_scale", weight_scale), + ("alpha", alpha), + ): + assert t.is_contiguous(), ( + f"{name} must be contiguous; the kernel reads it from data_ptr() as a " + "dense row-major buffer and silently reads the wrong elements" + ) + # One alpha per expert of `weight`; the runner checks only alpha.dim() == 1. + assert alpha.numel() == weight.size(0), ( + f"alpha has {alpha.numel()} elements for {weight.size(0)} experts; the " + "kernel reads alpha[expert] past the end without raising" + ) + # The runner checks weight_scale.dim() == 3 only, never its size. + expected = weight.numel() * 2 // scaling_vector_size + assert weight_scale.numel() == expected, ( + f"weight_scale holds {weight_scale.numel()} bytes, the weight needs " + f"{expected}; the kernel reads past its end -- silently within the " + "allocation, with an illegal memory access beyond it" + ) + + +# ── native-torch NVFP4 helpers for reference() and the test's compare ──────── + + +def unswizzle(buf: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """128x4-swizzled scale bytes -> `[rows, cols]` (rows % 128 == 0, cols % 4 == 0). + + Byte `(r, c)` sits at `(r // 128) * 128 * cols + (c // 4) * 512 + (r % 32) * 16 + + ((r % 128) // 32) * 4 + c % 4`. + """ + return buf.reshape(rows // 128, cols // 4, 32, 4, 4).permute(0, 3, 2, 1, 4).reshape(rows, cols) + + +def swizzle(sf: torch.Tensor) -> torch.Tensor: + """`[rows, cols]` scale bytes -> the flat 128x4-swizzled buffer; inverse of `unswizzle`.""" + rows, cols = sf.shape + return sf.reshape(rows // 128, 4, 32, cols // 4, 4).permute(0, 3, 2, 1, 4).reshape(-1) + + +def e4m3_value(byte: torch.Tensor) -> torch.Tensor: + return byte.to(torch.uint8).view(torch.float8_e4m3fn).double() + + +def _dequant(packed: torch.Tensor, sf: torch.Tensor, vec: int) -> torch.Tensor: + """`[R, K/2]` packed e2m1 (element 2j in the low nibble) and `[R, K/vec]` e4m3 + bytes -> float64 `[R, K]`. Exact.""" + lut = torch.tensor(E2M1, dtype=torch.float64, device=packed.device) + lo, hi = (packed & 0xF).long(), (packed >> 4).long() + codes = torch.stack([lo, hi], dim=-1).reshape(packed.shape[0], -1) + vals = lut[codes & 7] * torch.where(codes >= 8, -1.0, 1.0).double() + return vals * e4m3_value(sf).repeat_interleave(vec, dim=1) + + +def e2m1_codes(v: torch.Tensor) -> torch.Tensor: + """float -> 4-bit e2m1 codes (bit 3 = sign), round-to-nearest-even, saturating at 6.""" + mid = torch.tensor(_E2M1_MIDPOINTS, device=v.device, dtype=v.dtype) + tie_up = torch.tensor(_TIE_UP, device=v.device) + a = v.abs().unsqueeze(-1) + code = ((a > mid).sum(-1) + ((a == mid) & tie_up).any(-1)).to(torch.uint8) + negative = (v < 0) | ((v == 0) & torch.signbit(v)) + return code | (negative.to(torch.uint8) << 3) + + +cute_dsl_nvfp4_gather_grouped_gemm_act_fusion = CuteDslNvfp4GatherGroupedGemmActFusion() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/cute_dsl_nvfp4_grouped_gemm.py new file mode 100644 index 000000000000..3d7512323e9c --- /dev/null +++ b/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 @@ -0,0 +1,367 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CuteDSL NVFP4 FC2 of the routed experts: the plain grouped GEMM over the permuted +FC1 output, one bf16 row per permuted row, in one trtllm call.""" + +from __future__ import annotations + +from typing import NamedTuple + +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 + +from .cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import _dequant, unswizzle + + +def _ckpt( + tokens: int, + tile_size: int, + mma_n: int = 128, + cluster_n: int = 1, + dist: str = "router", + gsf: str = "calibrated", + zero_tokens: int = 0, +) -> dict: + """The checkpoint's routed experts: 512 experts, top-10, hidden 2560, I 640, all + local (tp1), bf16 out, 16-element e4m3 block scales. + + `mma_n` / `cluster_n` select the autotuner tactic: MMA tile `(tile_size, + mma_n)`, cluster `(tile_size // 128, cluster_n)`; (128, 1) is also the + fallback when the tuner has no entry. `dist` names how the test draws the + expert ids: `router` = each row's top-10 of uniform scores. `gsf` is the FC1 + `global_sf` that quantized this op's input, against the SwiGLU output's + magnitude: `calibrated` = the output maximum the test calibrates to maps to + the e4m3 maximum, as the checkpoint's per-layer `1 / down_proj.input_scale` + does; `subnormal` = 2^-12 of it, so input block scales fall into e4m3 + subnormals and zero. `zero_tokens` makes that many leading tokens' rows all + zero, as FC1 returns them for an all-zero token. + """ + return dict( + tokens=tokens, + tile_size=tile_size, + mma_n=mma_n, + cluster_n=cluster_n, + dist=dist, + gsf=gsf, + zero_tokens=zero_tokens, + num_experts=512, + top_k=10, + hidden=2560, + intermediate=640, + ) + + +class Fc2Expected(NamedTuple): + """What `reference` returns: the op's output, plus what the test's compare reads. + + `out` is laid out exactly as the op returns it; rows the op does not write + hold zeros. `rows` marks the rows the op writes -- always the prefix of + `num_non_exiting_tiles * tile_size` rows. `exact` is `alpha[e] * sum_k a_k + w_k` on those rows, before any rounding; `magnitude` is `|alpha[e]| * sum_k + |a_k w_k|`, the scale every fp32 rounding of the accumulation is relative to. + """ + + out: torch.Tensor + rows: torch.Tensor + exact: torch.Tensor + magnitude: torch.Tensor + + +class CuteDslNvfp4GroupedGemm(OpWrapper): + """`torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_blackwell`: FC2 of the routed + experts on the CuteDSL MoE path, one kernel launch. + + For every row `p` of the first `num_non_exiting_tiles` tiles (tile `i = p // + tile_size`): + + e = tile_idx_to_group_idx[i] the local expert + acc[p] = deq(input[p]) @ deq(weight[e]).T [H], fp32 + out[p] = bf16(alpha[e] * acc[p]) fp32 product + + where `deq` is e2m1 times its e4m3 block scale. Every row of a live tile is + computed, padding rows included; no row is gathered, weighted or summed. + + Arguments, as the CuteDSL MoE passes them (`run_moe_nvfp4_impl`, + `_torch/moe/fused_moe/fused_moe_cute_dsl.py`, its non-fused-finalize branch): + + - `input` `[M, I/2]` float4_e2m1fn_x2 and `input_scale` `[M * I/16]` uint8: + FC1's output exactly as the CuteDSL FC1 + (`cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell`) returns it -- + one row per PERMUTED row, the block scales 128x4-swizzled over `[M, + I/16]`. `M = permuted_idx_to_expanded_idx.numel()` of the sort. + - `weight` `[E, H, I/2]` float4_e2m1fn_x2: `down_proj` per expert in the + checkpoint's row order. No 32-row shuffle (that is trtllm-gen's layout). + - `weight_scale` `[E, H, I/16]` uint8: its e4m3 block scales, 128x4-swizzled + per expert (`block_scale_interleave`). + - `alpha` `[E]` fp32: per expert, `max(down_proj.input_scale) * weight_scale_2[e]` + -- the inverse of FC1's `global_sf` times the expert's weight scale. + - `tile_idx_to_group_idx` `[M / tile_size]` int32 and `num_non_exiting_tiles` + `[1]` int32: moe_sort's `tile_idx_to_expert_idx` and + `num_non_exiting_tiles` from the sort FC1 ran on, with tile_tokens_dim == + `tile_size`. + - `num_experts`, `top_k`, `num_local_experts`, `local_expert_offset`: read only + by the autotuner's shape buckets and cache key; the kernel takes the + expert count from `weight.size(0)` and the expert from the tile. + - `output_dtype`: bf16, the only one the runner builds. + + Returns `[M, H]` bfloat16, one row per permuted row. The pair `(t, k)` is row + `expanded_idx_to_permuted_idx[t, k]`; rows outside the live tiles are not + written (see `note`). + + Fusion boundary: the grouped GEMM, the per-expert alpha and the bf16 + conversion happen inside. The rows arrive already permuted (FC1 gathered + them); the routing-weighted top-k combine is the caller's + (`moe_unpermute_inplace`). Routing weights are not read here. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="tile 128: decode, 1 request without MTP; the fallback tactic", + spec=_ckpt(1, 128), + ), + Cell( + why="tile 128: MMA N 256 and cluster N 2 tactic, 8192-row max_num_tokens prefill", + spec=_ckpt(8192, 128, mma_n=256, cluster_n=2), + ), + Cell( + why="tile 128: tiny FC1 global_sf and 7 all-zero token rows; subnormal and zero " + "input scales", + spec=_ckpt(2048, 128, gsf="subnormal", zero_tokens=7), + ), + ) + + note = """ + Path certified: the CuTe DSL kernel `Sm100BlockScaledContiguousGroupedGemmKernel` + (`_torch/moe/kernels/blackwell/blockscaled_contiguous_grouped_gemm.py`), + built by `Sm100BlockScaledContiguousGroupedGemmRunner` + (`_torch/custom_ops/cute_dsl_custom_ops.py`), which JIT-compiles it with + `cute.compile` on the first call per (tile, tactic) and keeps it for the + process. The op is registered only when the CuTe DSL imports + (`IS_CUTLASS_DSL_AVAILABLE`), and the runner accepts SM 100 and 103 only. + + Metadata consumed: none beyond the arguments. Stateless: allocates its output + with torch.empty on every call, on the current stream, with no host + synchronization; `num_non_exiting_tiles` is read on the device, so it + captures into a CUDA graph together with moe_sort and FC1. + + Call sites: the engine's CuteDSL MoE (`run_moe_nvfp4_impl`, when the fused + finalize is off) and the target's routed-expert step, both followed by + `moe_unpermute_inplace`; eager for prefill and mixed batches (T up to + max_num_tokens 8192) and captured for decode (T = requests x MTP verify + width, up to 32 x 4 = 128). Tile 128 is the target's and the cells' tile; + 256 is certified through the per-pair invariance below. + + Numerics: each written element equals the reference's `out` bit for bit on + every cell (and, at bring-up, on the real FC1 output at 8192 tokens and on the + checkpoint's own down_proj experts). Operands whose block scales span a wider range can round a + partial sum; the test's compare admits that at bf16 ties only. + + Tactics. The op asks the AutoTuner for an MMA tile of (tile_size, 128) or + (tile_size, 256) and a cluster of (tile_size // 128, 1) or (tile_size // + 128, 2), and with no tuning entry falls back to ((tile_size, 128), + (tile_size // 128, 1)); serving warmup tunes per row bucket and may pick any + of the four. Every pair's row is bitwise identical across tile 128 / 256 x + the four tactics x the fallback (the test checks it at 8192 rows), so cells + driven at one tactic certify the others. + + Determinism. A (token, expert) pair's output row depends on that pair's input + row and its expert alone: it is bitwise unchanged by moving it to another row + of its expert's run (moe_sort orders them by atomics from T = 5), by repeating + the call, and by cutting a batch to a prefix. This op does not make the MoE + run-to-run nondeterministic; the fused-finalize variant + (`cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell`) reduces across a + token's experts in launch order and does. + + What is read and written: every row of the first num_non_exiting_tiles tiles + is written; the pair rows hold their results and the padding rows of a live + tile are computed from whatever the input holds there -- zero for FC1's + output, which zeroes them. Rows at and past num_non_exiting_tiles * tile_size + are never written: eager they hold whatever torch.empty returned, under a + graph whatever the pool last held. Input rows past that point reach no written + row. The combine must read only expanded_idx_to_permuted_idx's rows. + + Preconditions the caller owns: + - `weight` / `weight_scale` in the layout the class docstring states, and + `input_scale` swizzled. A wrong preparation of each runs without error, + with finite output, and puts most of the pair rows' elements outside the + test's gate: trtllm-gen's 32-row shuffle of the weight, input scales + linear, one alpha for the whole layer. + - `alpha` per expert, `max(down_proj.input_scale) * weight_scale_2[e]`, with + the maximum the value FC1's `global_sf` inverts. + - `tile_idx_to_group_idx` / `num_non_exiting_tiles` from the same sort as + FC1's input, tile_tokens_dim == `tile_size`, and every expert index below + `weight.size(0)`: nothing checks the values (a tile naming expert 512 of + 512 ran without error and came back all zeros). + - Every tensor on the device. Nothing checks it: a host-resident `alpha` or + tile table was read correctly on a coherent-memory Grace host, unmeasured + where the GPU cannot dereference a host pointer. + + Refused by `is_valid`, each silently wrong unguarded: a non-contiguous input, + input_scale, weight, weight_scale, alpha or tile_idx_to_group_idx (data_ptr() + read as dense; finite, wrong rows). + + Loud (AssertionError from the runner before the launch): a tile_size other + than the sort's, either way round; output_dtype float16; alpha bf16 or short + of weight.size(0); weight_scale with fewer experts or typed float8_e4m3fn; + input typed uint8; input_scale 2-D or short; scaling_vector_size 32; + num_non_exiting_tiles of 2 elements; tile_idx_to_group_idx int64 or one tile + short. tile_size 64: ValueError. + + Measured at bring-up and not re-driven by the test: the CUDA-graph capture + with moe_sort and FC1, the determinism and read / write statements, the wrong + preparations' distance from the gate, and the loud rejections. The test + re-drives the cells, the tile / tactic invariance and the `input` contiguity + refusal. + + Uncertified: output dtypes other than bf16 (the runner builds none); expert + parallelism (the kernel reads neither local_expert_offset nor num_experts, so + EP changes only the sort and the weight slice); other hidden / intermediate + sizes; scaling_vector_size other than 16; other architectures. + + Hard limits (sm_103) + input, input_scale, output, weight, weight_scale past 2^31 elements / + bytes: exact (probed up to M = 53687296 rows and 20973 + experts) + tile index int32 tables, within moe_sort's 2^31-pair limit; no index + passes through fp32 + """ + + def __call__( + self, + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + scaling_vector_size: int = 16, + ) -> torch.Tensor: + return torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_blackwell( + input=input, + weight=weight, + input_scale=input_scale, + weight_scale=weight_scale, + alpha=alpha, + tile_idx_to_group_idx=tile_idx_to_group_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=local_expert_offset, + tile_size=tile_size, + output_dtype=output_dtype, + scaling_vector_size=scaling_vector_size, + ) + + def reference( + self, + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + scaling_vector_size: int = 16, + ) -> Fc2Expected: + """Dequantize, GEMM in float64, scale by alpha in fp32, round once to bf16. + + Products of e2m1 and e4m3 values are exact, and float64 sums of 640 of + them are exact to ~1e-16 relative, so `exact` carries no accumulation + error of its own; `out` is that sum rounded to fp32, multiplied by the + fp32 alpha, rounded to bf16 -- what an fp32-accumulating GEMM with this + epilogue returns whenever its fp32 partial sums are exact. Rows whose + input dequantizes to all zeros are zero without a GEMM. + + bf16 output only (anything else raises NotImplementedError). + `num_experts`, `top_k`, `num_local_experts` and `local_expert_offset` do + not enter the computation. + """ + if output_dtype != torch.bfloat16: + raise NotImplementedError("reference covers bf16 output only") + vec = scaling_vector_size + dev = input.device + m = input.shape[0] + hidden = weight.shape[1] + k_blocks = input.shape[1] * 2 // vec + n_tiles = int(num_non_exiting_tiles.reshape(-1)[0]) + span = n_tiles * tile_size + + rows = torch.zeros(m, dtype=torch.bool, device=dev) + rows[:span] = True + a = _dequant( + input.view(torch.uint8)[:span], + unswizzle(input_scale.view(torch.uint8), m, k_blocks)[:span], + vec, + ) + live = a.abs().amax(dim=1) > 0 + expert = tile_idx_to_group_idx[:n_tiles].long().repeat_interleave(tile_size) + acc = torch.zeros(span, hidden, dtype=torch.float64, device=dev) + mag = torch.zeros(span, hidden, dtype=torch.float64, device=dev) + w_sf = weight_scale.view(torch.uint8).reshape(weight.shape[0], -1) + for e in torch.unique(expert[live]).tolist(): + sel = torch.nonzero(live & (expert == e)).squeeze(1) + w = _dequant(weight[e].view(torch.uint8), unswizzle(w_sf[e], hidden, k_blocks), vec) + acc[sel] = a[sel] @ w.T + mag[sel] = a[sel].abs() @ w.abs().T + scale = alpha.float()[expert].unsqueeze(1) + out = torch.zeros(m, hidden, dtype=torch.bfloat16, device=dev) + out[:span] = (acc.float() * scale).to(torch.bfloat16) + return Fc2Expected( + out=out, + rows=rows, + exact=acc * scale.double(), + magnitude=mag * scale.double().abs(), + ) + + def is_valid( + self, + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + scaling_vector_size: int = 16, + ) -> None: + # The runner hands data_ptr() to the kernel with dense row-major layouts + # assumed; it checks dims, sizes and dtypes, never strides. + for name, t in ( + ("input", input), + ("input_scale", input_scale), + ("weight", weight), + ("weight_scale", weight_scale), + ("alpha", alpha), + ("tile_idx_to_group_idx", tile_idx_to_group_idx), + ): + assert t.is_contiguous(), ( + f"{name} must be contiguous; the kernel reads it from data_ptr() as a " + "dense row-major buffer and silently reads the wrong elements" + ) + + +cute_dsl_nvfp4_grouped_gemm = CuteDslNvfp4GroupedGemm() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_permuted.py new file mode 100644 index 000000000000..f6efc5755b5a --- /dev/null +++ b/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 @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""NVFP4 MoE layer entered with router logits, left unfinalized: the trtllm-gen +runner's Renormalize routing + grouped FC1 GEMM + SwiGLU + NVFP4 +requantization + grouped FC2 GEMM, in one call, returning the permuted FC2 rows +with the routing weights and the expanded->permuted index map instead of the +combined `[T, H]` rows. The consumer forms the top-k weighted sum itself.""" + +from __future__ import annotations + +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 + +from .fp4_block_scale_moe_renorm_routed import ACT_SWIGLU, ROUTING_RENORMALIZE, _ckpt +from .fp4_block_scale_moe_renorm_routed import fp4_block_scale_moe_renorm_routed as _finalized + + +def finalize( + rows: torch.Tensor, topk_weights: torch.Tensor, topk_rows: torch.Tensor +) -> torch.Tensor: + """The runner's finalize, emulated: per token `sum_k w[t, k] * rows[topk_rows[t, k]]` + as fp32 FMAs in slot order from zero (a -1 slot skipped), rounded once to the + rows' dtype. Each fp32 FMA is formed exactly in float64 (a product of two + fp32 values is exact there) and rounded to fp32. Returns `[T, H]`.""" + acc = torch.zeros(topk_rows.shape[0], rows.shape[1], dtype=torch.float32, device=rows.device) + for k in range(topk_rows.shape[1]): + idx = topk_rows[:, k].long() + r = rows[idx.clamp(min=0)].float().double() + w = topk_weights[:, k].float().double().unsqueeze(-1) + step = (r * w + acc.double()).float() + acc = torch.where((idx >= 0).unsqueeze(-1), step, acc) + return acc.to(rows.dtype) + + +class Fp4BlockScaleMoeRenormPermuted(OpWrapper): + """`fp4_block_scale_moe_renorm_routed` with `do_finalize=False`. + + Same operands, same routing, same GEMMs; the runner skips its finalize + launch and returns three tensors: + + - `rows` `[P, H]` bf16: the FC2 output in the runner's permuted, padded row + order (`P` is the runner's padded maximum; rows no slot maps to hold + garbage); + - `topk_weights` `[T, k]` bf16: the routing weights, slot order; + - `topk_rows` `[T, k]` int32: slot `(t, k)`'s row in `rows`, or -1. + + `finalize(rows, topk_weights, topk_rows)` (this module) is bitwise the + finalized entry's output on the same operands. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="_moe_block decode, finalize folded into the next site: one token; the block " + "routing kernel", + spec=_ckpt(1), + ), + Cell( + why="128 rows: the con=128 decode batch; the cluster routing kernel", + spec=_ckpt(128), + ), + Cell(why="coarse 3-level logits: ties decide most slots", spec=_ckpt(64, dist="coarse")), + ) + + note = """ + Measured against the finalized entry on identical operands: the emulated + finalize of the three outputs is bitwise its output at every cell, so the + routing semantics, the GEMMs and the weight layout are the finalized + entry's (see its note). Which finalize kernel the runner would have + launched (scalar below ~1184 blocks, vectorized above) does not change the + arithmetic: both accumulate fp32 FMAs in slot order and round once. + + `rows` is the runner's own `[P, H]` buffer, allocated per call at the + padded maximum the runner computes from `T`, `k` and the expert count; + only rows some slot maps to hold data. `topk_weights` carries the output + dtype (bf16), as the finalize reads it. + + Silent, refused by `is_valid`: the finalized entry's refusals (a + non-contiguous operand, an uncertified (num_experts, top_k) pair, hidden + not a multiple of 256). + + Hard limits + T certified at 1, 64 and 128 (decode); the decode forward folds at + every site that takes the combine + num_experts 512, top_k 10 + """ + + def __call__( + self, + routing_logits: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + output1_scale_scalar: torch.Tensor, + output1_scale_gate_scalar: torch.Tensor, + output2_scale_scalar: torch.Tensor, + num_experts: int, + top_k: int, + intermediate_size: int, + ) -> tuple: + rows, topk_weights, topk_rows = torch.ops.trtllm.fp4_block_scale_moe_runner( + routing_logits, + None, + hidden_states, + hidden_states_scale, + gemm1_weights, + gemm1_weights_scale, + None, + None, + None, + None, + gemm2_weights, + gemm2_weights_scale, + None, + output1_scale_scalar, + output1_scale_gate_scalar, + output2_scale_scalar, + num_experts, + top_k, + None, + None, + intermediate_size, + 0, + num_experts, + None, + ROUTING_RENORMALIZE, + False, + ACT_SWIGLU, + ) + return rows, topk_weights, topk_rows + + def reference(self, *args, **kwargs) -> torch.Tensor: + """The finalized entry's fp32 reference `[T, H]`: what + `finalize(*self(...))` must answer (same arguments as that entry's + `reference`).""" + return _finalized.reference(*args, **kwargs) + + def is_valid(self, *args, **kwargs) -> None: + _finalized.is_valid(*args, **kwargs) + + +fp4_block_scale_moe_renorm_permuted = Fp4BlockScaleMoeRenormPermuted() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_block_scale_moe_renorm_routed.py new file mode 100644 index 000000000000..d4a76c8155c4 --- /dev/null +++ b/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 @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""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.""" + +from __future__ import annotations + +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 +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 .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. +ACT_SWIGLU = 0 + +#: The (num_experts, top_k) pairs this entry certifies. The Renormalize routing +#: policy dispatches on a compiled tier list, and a pair outside every tier is a +#: logged error with no exception and **no routing launch** -- the GEMMs then +#: read stale routing buffers. So the guard admits what was measured, not what +#: the tier list might also cover. +CERTIFIED_PAIRS = frozenset({(512, 10)}) + + +def _ckpt(tokens: int, dist: str = "randn") -> dict: + """qwen4_exp's routed geometry: 512 experts, top-10, H 2560 / I 640.""" + return dict( + tokens=tokens, + hidden=2560, + intermediate=640, + num_experts=512, + top_k=10, + logits_dtype=torch.bfloat16, + dist=dist, + ) + + +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. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # One cell per routing kernel the call picks by row count: block, single-cluster + # and the multi-kernel histogram path. + CELLS: tuple[Cell, ...] = ( + Cell(why="_moe_block decode: one token; the block routing kernel", spec=_ckpt(1)), + Cell(why="128 rows: the con=128 decode batch; the cluster routing kernel", spec=_ckpt(128)), + Cell( + why="16384 rows: max_num_tokens of the tuned config; past the 8192 tuning bucket", + spec=_ckpt(16384), + ), + ) + + note = """ + Routing semantics, measured against the composition it replaces + (`renorm_moe_routing_op` -> pre-routed runner): bitwise-equal finalized + output at every cell (the test re-drives this), and at bring-up on randn / + coarse / large-magnitude / all-equal logits, bf16 and fp32 logits, + 1..16384 rows. Both kernels rank by a packed + (twiddled value, 65535 - index) key (ties to the lower expert index), take + an fp32 `expf` softmax over the k with the same 32-lane warp reduction, and + round the weight to bf16 once. + + Kernel choice inside the call is by row count (block kernel up to a small + count, single-cluster up to the cluster cap, then the multi-kernel + histogram path); every one of them shares the top-k and softmax code. + `routing_bias` is None (the thop defaults the bias dtype to bf16 and never + reads it on this policy). PDL (`TRTLLM_ENABLE_PDL`) is read by the routing + launch. + + Silent, refused by `is_valid`: a non-contiguous operand (raw pointers, + dense layout assumed) and a (num_experts, top_k) pair outside + `CERTIFIED_PAIRS` (a tier-list miss logs and launches nothing; the test + holds this refusal). + + Not certified: fp32 logits (the runner's other accepted logits dtype). + + Hard limits + 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 + pass: a call past 8192 rows runs the 8192 bucket's tactic + num_experts 512, top_k 10 (CERTIFIED_PAIRS) + """ + + def __call__( + self, + routing_logits: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + output1_scale_scalar: torch.Tensor, + output1_scale_gate_scalar: torch.Tensor, + output2_scale_scalar: torch.Tensor, + num_experts: int, + top_k: int, + intermediate_size: int, + ) -> torch.Tensor: + return torch.ops.trtllm.fp4_block_scale_moe_runner( + routing_logits, + None, + hidden_states, + hidden_states_scale, + gemm1_weights, + gemm1_weights_scale, + None, + None, + None, + None, + gemm2_weights, + gemm2_weights_scale, + None, + output1_scale_scalar, + output1_scale_gate_scalar, + output2_scale_scalar, + num_experts, + top_k, + None, + None, + intermediate_size, + 0, + num_experts, + None, + ROUTING_RENORMALIZE, + True, + ACT_SWIGLU, + )[0] + + def reference( + self, + routing_logits: torch.Tensor, + hidden_states: torch.Tensor, + top_k: int, + up: tuple, + gate: tuple, + down: tuple, + fc2_global_scale: float, + ) -> torch.Tensor: + """The routing, then the 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 + runner's routing writes). The layer: the pre-routed runner entry's fp32 + reference over the dequantized `(codes, scales)` per-expert operands, as + that entry documents (`hidden_states` is the exact fp32 value of the + 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) + + def is_valid( + self, + routing_logits: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + output1_scale_scalar: torch.Tensor, + output1_scale_gate_scalar: torch.Tensor, + output2_scale_scalar: torch.Tensor, + num_experts: int, + top_k: int, + intermediate_size: int, + ) -> 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 + # realistic way to get here. + for name, tensor in ( + ("routing_logits", routing_logits), + ("hidden_states", hidden_states), + ("hidden_states_scale", hidden_states_scale), + ("gemm1_weights", gemm1_weights), + ("gemm1_weights_scale", gemm1_weights_scale), + ("gemm2_weights", gemm2_weights), + ("gemm2_weights_scale", gemm2_weights_scale), + ("output1_scale_scalar", output1_scale_scalar), + ("output1_scale_gate_scalar", output1_scale_gate_scalar), + ("output2_scale_scalar", output2_scale_scalar), + ): + assert tensor.is_contiguous(), ( + f"{name} must be contiguous; a strided view is read as if dense " + "and silently produces wrong results" + ) + # A tier-list miss in the Renormalize routing policy logs an error and + # launches nothing; the GEMMs then consume stale routing buffers. + assert (num_experts, top_k) in CERTIFIED_PAIRS, ( + f"(num_experts, top_k) = ({num_experts}, {top_k}) is not a certified " + f"Renormalize routing pair {sorted(CERTIFIED_PAIRS)}" + ) + # Same layout rule as the pre-routed entry: the prepared block-scale + # layout is silently misread below a hidden multiple of 256. + assert gemm1_weights.shape[-1] * 2 % 256 == 0, ( + "hidden size (gemm1_weights.shape[-1] * 2) must be a multiple of 256" + ) + + +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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp4_moe_renorm_cute_dsl.py new file mode 100644 index 000000000000..85cbcb72ae9f --- /dev/null +++ b/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 @@ -0,0 +1,225 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""NVFP4 MoE layer entered with router logits, on TRT-LLM's CuTe DSL Blackwell +MoE kernels: Renormalize routing (`renorm_moe_routing_op`), the token sort, +a gathered grouped FC1 GEMM with SwiGLU and NVFP4 requantization in its +epilogue, and a grouped FC2 GEMM that scatter-adds the routing-weighted rows +into the output -- the layer `fp4_block_scale_moe_renorm_routed` computes, for +prefill-sized row counts, from a second copy of the experts in this kernel +family's layout.""" + +from __future__ import annotations + +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 + +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 +from .renorm_moe_routing_op import renorm_moe_routing_op + +#: Token tile of the grouped GEMMs (the kernels' M tile). +TILE_TOKENS = 128 +_FP4 = torch.float4_e2m1fn_x2 + + +class Fp4MoeRenormCuteDsl(OpWrapper): + """One NVFP4 MoE layer whose routing happens inside the call, on the CuTe + DSL MoE kernels. Same function as `fp4_block_scale_moe_renorm_routed`: + per token the `top_k` largest logits (ties to the LOWER expert index), + weighted by an fp32 softmax over just those `top_k` rounded to bf16; the + experts' SwiGLU FC1, the NVFP4 requantization of its output, FC2, and the + weighted sum. Returns the combined `[T, H]` bf16 rows. + + Operands, in this kernel family's layout (see `weights.py`): `fc1_w` + `[E, 2I, H/2]` uint8, `[up ; gate]` rows interleaved in 64-row groups; + `fc1_s` `[E, 2I, H/16]` e4m3 bytes in that row order, 128x4-swizzled; + `fc2_w` `[E, H, I/2]` as stored; `fc2_s` `[E, H, I/16]` swizzled. + `hidden_states` / `hidden_states_scale` are the NVFP4 activation with + **linear** block scales, quantized with the layer's FC1 activation global + scale. Scalars: `fc1_alpha` `[E]` dequantizes the FC1 accumulator (the + routed entry's `output1_scale_gate_scalar`); `fc2_input_scale` `[1]` is the + FC2 activation's quantization multiplier (the reciprocal of its dequant + scale), one number for every expert; `fc2_alpha` `[E]` dequantizes the FC2 + accumulator (the routed entry's `output2_scale_scalar`). + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="4096 rows: where the target switches to this entry; the sort's cluster kernel", + spec=_ckpt(4096), + ), + Cell( + why="16377 rows: a con=128 mixed step (16256 context + 121 generation)", + spec=_ckpt(16377), + ), + Cell( + why="16384 rows: max_num_tokens of the tuned config; the sort's cooperative kernel", + spec=_ckpt(16384), + ), + ) + + note = """ + Semantics against `fp4_block_scale_moe_renorm_routed` (trtllm-gen) on the + same experts: the same routing (both are the `renorm_moe_routing_op` + function, weights rounded to bf16 once) and the same NVFP4 arithmetic per + expert (e2m1 x e4m3-block x fp32-alpha products accumulated in fp32, SwiGLU + on the FC1 accumulator, the FC2 input requantized to NVFP4 with 16-element + blocks), but not bitwise: the GEMMs' accumulation order differs, and the + weighted sum over a token's experts is a scatter-add with fp32 atomics, so + its order -- and the output's last bit -- varies between launches. Both sit + inside the routed entry's band against its fp32 reference. + + One FC2 activation scale for all experts: the routed entry takes it per + expert (folded into `output1_scale_scalar`); a checkpoint whose experts + disagree on `down_proj.input_scale` cannot use this entry as is. + + Kernels (all TRT-LLM's): `trtllm::moe_sort` (tile metadata), + `trtllm::moe_output_memset_inplace` (zeroes the rows the scatter-add + touches), `trtllm::cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell`, + `trtllm::cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell`; the two + GEMMs pick their tactics through the AutoTuner. Not CUDA-graph certified + (prefill steps run eagerly). + + Hard limits + T certified 4096..16384 (the target calls it from + `_CUTE_DSL_MOE_MIN_ROWS` = 4096 rows) + num_experts 512, top_k 10 (the routed entry's CERTIFIED_PAIRS) + """ + + def __call__( + self, + routing_logits: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + fc1_w: torch.Tensor, + fc1_s: torch.Tensor, + fc1_alpha: torch.Tensor, + fc2_w: torch.Tensor, + fc2_s: torch.Tensor, + fc2_alpha: torch.Tensor, + fc2_input_scale: torch.Tensor, + num_experts: int, + top_k: int, + ) -> torch.Tensor: + tokens = hidden_states.shape[0] + hidden = hidden_states.shape[1] * 2 + ids, weights = renorm_moe_routing_op(routing_logits, top_k, torch.bfloat16) + scales = weights.float() + t2e, t2lim, e2p, p2e, _, live = torch.ops.trtllm.moe_sort( + token_selected_experts=ids, + token_final_scales=scales, + num_experts=num_experts, + top_k=top_k, + local_expert_offset=0, + local_num_experts=num_experts, + tile_tokens_dim=TILE_TOKENS, + ) + out = torch.empty((tokens, hidden), dtype=torch.bfloat16, device=hidden_states.device) + torch.ops.trtllm.moe_output_memset_inplace( + input=out, + tile_idx_to_mn_limit=t2lim, + expanded_idx_to_permuted_idx=e2p, + permuted_idx_to_expanded_idx=p2e, + num_non_exiting_tiles=live, + tile_tokens_dim=TILE_TOKENS, + top_k=top_k, + ep_size=1, + enable_alltoall=False, + ) + h, h_sf = torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell( + input=hidden_states.view(_FP4), + weight=fc1_w.view(_FP4), + input_scale=hidden_states_scale.view(torch.uint8).view(tokens, hidden // 16), + weight_scale=fc1_s.view(torch.uint8), + alpha=fc1_alpha, + tile_idx_to_group_idx=t2e, + tile_idx_to_mn_limit=t2lim, + permuted_idx_to_expanded_idx=p2e, + num_non_exiting_tiles=live, + global_sf=fc2_input_scale, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_experts, + local_expert_offset=0, + tile_size=TILE_TOKENS, + ) + torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell( + input=h.view(_FP4), + weight=fc2_w.view(_FP4), + input_scale=h_sf.view(torch.uint8), + weight_scale=fc2_s.view(torch.uint8), + alpha=fc2_alpha, + output=out, + tile_idx_to_group_idx=t2e, + tile_idx_to_mn_limit=t2lim, + permuted_idx_to_expanded_idx=p2e, + num_non_exiting_tiles=live, + token_final_scales=scales, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_experts, + local_expert_offset=0, + tile_size=TILE_TOKENS, + output_dtype=torch.bfloat16, + ) + return out + + def reference( + self, + routing_logits: torch.Tensor, + hidden_states: torch.Tensor, + top_k: int, + up: tuple, + gate: tuple, + down: tuple, + fc2_global_scale: float, + ) -> torch.Tensor: + """The routed entry's fp32 reference (routing, then the layer over the + dequantized `(codes, scales)` per-expert operands); fp32 `[T, H]`.""" + return _routed.reference( + routing_logits, hidden_states, top_k, up, gate, down, fc2_global_scale + ) + + def is_valid( + self, + routing_logits: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + fc1_w: torch.Tensor, + fc1_s: torch.Tensor, + fc1_alpha: torch.Tensor, + fc2_w: torch.Tensor, + fc2_s: torch.Tensor, + fc2_alpha: torch.Tensor, + fc2_input_scale: torch.Tensor, + num_experts: int, + top_k: int, + ) -> None: + for name, tensor in ( + ("routing_logits", routing_logits), + ("hidden_states", hidden_states), + ("hidden_states_scale", hidden_states_scale), + ("fc1_w", fc1_w), + ("fc1_s", fc1_s), + ("fc2_w", fc2_w), + ("fc2_s", fc2_s), + ): + assert tensor.is_contiguous(), f"{name} must be contiguous" + assert (num_experts, top_k) in CERTIFIED_PAIRS, ( + f"(num_experts, top_k) = ({num_experts}, {top_k}) is not certified" + ) + tokens, half_hidden = hidden_states.shape + assert hidden_states_scale.numel() == tokens * half_hidden * 2 // 16, ( + "hidden_states_scale must hold one linear e4m3 scale a 16-element block" + ) + assert fc1_alpha.shape == (num_experts,) and fc2_alpha.shape == (num_experts,), ( + "per-expert alphas" + ) + assert fc2_input_scale.numel() == 1, "one FC2 activation scale for every expert" + + +fp4_moe_renorm_cute_dsl = Fp4MoeRenormCuteDsl() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp8_block_scale_moe_runner.py index fcb1f0eac64a..2bb80ff63ecb 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fp8_block_scale_moe_runner.py @@ -101,85 +101,27 @@ class Fp8BlockScaleMoeRunner(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - # The token column at the checkpoint geometry. The cold-cache fallback - # tactic is keyed on the token count, and rows per active expert - # (T * 10 / 512) is what moves the kernel's tile choice. + # The token column at the checkpoint geometry, the only one the draft layer + # calls. The cold-cache fallback tactic is keyed on the token count, and rows + # per active expert (T * 10 / 512) is what moves the kernel's tile choice. Cell( why="MTPLayer._moe decode: one token, 10 of 512 experts", spec=dict(tokens=1, **_CKPT) ), - Cell(why="2 tokens: up to 20 experts, one row each", spec=dict(tokens=2, **_CKPT)), - Cell(why="8 tokens: ~0.16 rows per expert", spec=dict(tokens=8, **_CKPT)), - Cell(why="16 tokens: ~0.3 rows per expert", spec=dict(tokens=16, **_CKPT)), Cell( why="20 tokens == H/128: the activation scale buffer is square", spec=dict(tokens=20, **_CKPT), ), Cell( - why="64 tokens: ~1.25 rows per expert, most experts active", - spec=dict(tokens=64, **_CKPT), - ), - Cell(why="256 tokens: ~5 rows per expert", spec=dict(tokens=256, **_CKPT)), - Cell(why="1024 tokens: ~20 rows per expert", spec=dict(tokens=1024, **_CKPT)), - Cell(why="4096 tokens: ~80 rows per expert", spec=dict(tokens=4096, **_CKPT)), - Cell( - why="8192 tokens: max_num_tokens, the top of the certified column", - spec=dict(tokens=8192, **_CKPT), - ), - Cell( - why="gemm1_clamp_limit=1.0: live, ~63 ulp RMS from the unclamped result", - spec=dict(tokens=16, gemm1_clamp_limit=1.0, **_CKPT), - ), - Cell( - why="EP window (384, 128) of 512: rank 3 of a 4-way split", - spec=dict(tokens=64, **_geom(2560, 640, 512, 10, offset=384, local=128)), - ), - # Geometries away from the checkpoint's, each a multiple of 128 on both - # block axes. - Cell( - why="H 256 / I 128, 8 experts, top-1: one slot per token", - spec=dict(tokens=16, **_geom(256, 128, 8, 1)), - ), - Cell(why="H 256 / I 128, 8 experts, top-2", spec=dict(tokens=16, **_geom(256, 128, 8, 2))), - Cell( - why="H 512 / I 512, 16 experts, top-8: the worst RMS measured", - spec=dict(tokens=16, **_geom(512, 512, 16, 8)), - ), - Cell( - why="H 1024 / I 256, 128 experts, top-6", - spec=dict(tokens=16, **_geom(1024, 256, 128, 6)), - ), - Cell(why="checkpoint widths at top-2", spec=dict(tokens=16, **_geom(2560, 640, 512, 2))), - Cell( - why="checkpoint widths at top-16, past the checkpoint's 10", - spec=dict(tokens=16, **_geom(2560, 640, 512, 16)), - ), - Cell( - why="H 2560 / I 1280: twice the checkpoint's intermediate", - spec=dict(tokens=16, **_geom(2560, 1280, 512, 10)), - ), - Cell( - why="H 7168 / I 2048, 64 experts, top-8: the widest K", - spec=dict(tokens=16, **_geom(7168, 2048, 64, 8)), - ), - # The expert window is a range test plus an index shift. - Cell( - why="window (4, 4) of 8: ids 0..3 dropped, 4..7 shifted to slots 0..3", - spec=dict(tokens=64, **_geom(256, 128, 8, 2, offset=4, local=4)), - ), - Cell( - why="window (4, 6) of 8 overruns num_experts: accepted, surplus slots idle", - spec=dict(tokens=64, **_geom(256, 128, 8, 2, offset=4, local=6)), - ), - Cell( - why="ids -1, num_experts and 40 in the routing: silently dropped", - spec=dict(tokens=16, foreign_ids=True, **_geom(256, 128, 8, 2)), + why="16384 rows: max_num_tokens of the tuned config; past the last tuning bucket", + spec=dict(tokens=16384, **_CKPT), ), ) 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 was measured bitwise equal to the cold fallback. + tuned tactic is bitwise equal to the cold fallback (the test holds this at + 256 tokens). Preconditions. `hidden_states` is `[T, H]` e4m3 and `hidden_states_scale` `[H/128, T]` fp32 (see the class docstring); weights are the checkpoint's @@ -194,9 +136,10 @@ class Fp8BlockScaleMoeRunner(OpWrapper): 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`; ids outside it (negative, - `>= num_experts`) are dropped, a token with no slot in it comes back - bitwise zero, and a window starting at `num_experts` returns all zeros. + 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). @@ -206,8 +149,9 @@ class Fp8BlockScaleMoeRunner(OpWrapper): `use_dp`. `gemm1_clamp_limit` is live and one scalar for every expert, applied after the FC1 requantization; None is bitwise equal to omitting it. - Silent on the raw op, refused by `is_valid`: any non-contiguous tensor (read - as dense, 55-115 ulp RMS off); `act_type != 0` (the Python runner forces + 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). `num_fused_shared_experts > 0` is rejected loudly on the pre-routed path, @@ -217,12 +161,21 @@ class Fp8BlockScaleMoeRunner(OpWrapper): states): an upstream defect, eager is unaffected. Hard limits - T certified 1..8192 (max_num_tokens); chunk above it, not probed - top_k <= 32 (raises past it); certified 1, 2, 6, 8, 10, 16 - E, H, I certified up to 512, 7168, 2048 - - Not certified: the in-kernel routing path (`routing_logits` given), - multi-rank execution (the windows ran in one process), sm_100. + T certified at 1, 20 and 16384 (max_num_tokens of the tuned config); + chunk above it, 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. """ def __call__( 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fused_sigmoid_gate_mul_add.py index f05a5c6e24f1..367a6139f8a8 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fused_sigmoid_gate_mul_add.py @@ -59,30 +59,12 @@ class FusedSigmoidGateMulAdd(OpWrapper): why="_moe_block / MTPLayer._moe decode: one token, in place; 8 warps", spec=_ckpt(1), ), - Cell(why="2 tokens: the smallest multi-row grid", spec=_ckpt(2)), - Cell(why="17 tokens: an odd count", spec=_ckpt(17)), Cell(why="2048 tokens: prefill-shaped", spec=_ckpt(2048)), - Cell(why="fp16 activations", spec=_ckpt(33, dtype=torch.float16, gate_dtype=torch.float16)), - Cell(why="fp32 activations", spec=_ckpt(33, dtype=torch.float32, gate_dtype=torch.float32)), - Cell(why="fp32 gate against bf16 activations", spec=_ckpt(9, gate_dtype=torch.float32)), - Cell(why="gate as a flat [T] vector", spec=_ckpt(33, gate_layout="flat")), - Cell(why="gate as a [1, T] row", spec=_ckpt(33, gate_layout="row")), Cell( - why="gate as a strided column of a wider tensor", spec=_ckpt(33, gate_layout="strided") + why="gate as a strided column of a wider tensor, as the trunk's fused shared-expert " + "GEMM leaves it", + spec=_ckpt(33, gate_layout="strided"), ), - Cell(why="into an output buffer: both inputs untouched", spec=_ckpt(17, dest="output")), - Cell( - why="row-padded windows: every tensor's row stride exceeds hidden", - spec=_ckpt(33, dest="padded"), - ), - # Launch regimes: the column block is next_pow2(hidden) capped at 8192, - # and the warp count steps with it. - Cell(why="hidden 1: the 16-wide minimum block, 2 warps", spec=_ckpt(8, hidden=1)), - Cell(why="hidden 7: an odd width inside the minimum block", spec=_ckpt(8, hidden=7)), - Cell(why="hidden 640: a 1024 block, 4 warps", spec=_ckpt(8, hidden=640)), - Cell(why="hidden 6145: one 8192 block, mostly masked", spec=_ckpt(8, hidden=6145)), - Cell(why="hidden 12288: two column blocks", spec=_ckpt(8, hidden=12288)), - Cell(why="hidden 20480: three column blocks, masked tail", spec=_ckpt(8, hidden=20480)), ) note = """ @@ -90,15 +72,16 @@ class FusedSigmoidGateMulAdd(OpWrapper): The Triton kernel JIT-compiles on first call. Preconditions, all enforced loudly by the callable (AssertionError before - any launch): `final_hidden_states` 2-D; `shared_expert_output` and `output` - match it in shape and dtype; `gate_logits.numel() == num_tokens`; the - destination has `stride(-1) == 1` (its row stride is free). A - column-strided `final_hidden_states` or `shared_expert_output` is densified - inside the call (an extra copy kernel), as is a non-contiguous gate. Host - operands raise ValueError from Triton. `num_tokens == 0` returns the - destination without launching. `output` may alias `shared_expert_output`. - - Numerics: the default `assert_close` band holds for bf16, fp16 and fp32. + any launch; the test holds the wrong-length gate): `final_hidden_states` + 2-D; `shared_expert_output` and `output` match it in shape and dtype; + `gate_logits.numel() == num_tokens`; the destination has `stride(-1) == 1` + (its row stride is free). A column-strided `final_hidden_states` or + `shared_expert_output` is densified inside the call (an extra copy kernel), + as is a non-contiguous gate. `num_tokens == 0` returns the destination + without launching (certified). + + Numerics: the default `assert_close` band holds for bf16, the dtype the target + 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, @@ -108,8 +91,11 @@ class FusedSigmoidGateMulAdd(OpWrapper): Hard limits num_tokens row offset widened to int64 in the kernel; largest run 8192 - Not certified: CUDA-graph capture, torch.compile, operands on two devices, - float8 activations. + 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. """ 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_sort.py new file mode 100644 index 000000000000..a22c0614c1d0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_sort.py @@ -0,0 +1,272 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MoE sort: bucket the routed (token, expert) pairs into per-expert row tiles +for the CuteDSL grouped GEMMs, in one trtllm call.""" + +from __future__ import annotations + +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 + + +def _ckpt(tokens: int, tile_tokens_dim: int, dist: str = "router") -> dict: + """The checkpoint's trunk MoE: 512 experts, top-10, all local (tp1), fp32 scales. + + `dist` names how the test draws the expert ids: `router` = each row's top-10 of + uniform scores over the 512 experts; `tile_edges` = three experts hold exactly + tile-1, tile and tile+1 rows, the rest `router`. Both put 10 distinct experts on + every row, as the router does. + """ + return dict( + tokens=tokens, + dist=dist, + num_experts=512, + top_k=10, + local_expert_offset=0, + local_num_experts=512, + tile_tokens_dim=tile_tokens_dim, + scales_dtype=torch.float32, + ) + + +class MoeSort(OpWrapper): + """`torch.ops.trtllm.moe_sort`: routed (token, slot) pairs -> per-expert row tiles. + + Input: `token_selected_experts` `[T, top_k]` int32 global expert ids (the + router's output), `token_final_scales` `[T, top_k]` (its weights). Call it + with the op's own keyword names, as the CuteDSL MoE does. + + Numbering: an *expanded* index is `t * top_k + k` (row-major over + `token_selected_experts`); a *permuted* index is a row of the grouped GEMM's + M dimension. Each local expert `e = id - local_expert_offset` owns a + contiguous run of `count_e` permuted rows starting on a tile boundary, + experts in increasing order, the run padded up to whole tiles of + `tile_tokens_dim` rows. Returns a list of six int32 tensors, in this order: + + - `tile_idx_to_expert_idx` `[max_tiles]`: the local expert of tile `i`; + - `tile_idx_to_mn_limit` `[max_tiles]`: one past the last real row of tile + `i`, as an absolute permuted index: `min((i + 1) * tile, start_e + count_e)`; + - `expanded_idx_to_permuted_idx` `[T, top_k]`: the permuted row of each + pair, `-1` for a pair whose expert is not local; + - `permuted_idx_to_expanded_idx` `[max_tiles * tile]`: the inverse on real rows; + - `total_num_padded_tokens` `[1]` = `num_non_exiting_tiles * tile`; + - `num_non_exiting_tiles` `[1]`: tiles actually used. + + `max_tiles = m + (T * top_k - m) // tile` with `m = min(local_num_experts, + T * top_k)` is computed on the host from the shapes and scalars alone, so + every output shape is fixed by `T`, which is what lets a CUDA graph capture + it. Only the first `num_non_exiting_tiles` entries of the tile arrays and + the rows below each tile's `mn_limit` are written; everything else is left + as allocated (see `note`). + + Fusion boundary: sorting only. Routing (the top-k) is the caller's, and so + are the permutation of the activations (the FC1 gather does it) and the + routing-weighted unpermute; this op does not read `token_final_scales` into + any output. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # One cell per kernel path the target reaches at its tile, 128: block (T <= 4), + # cluster (T 5..8192) and cooperative (T >= 8193). + CELLS: tuple[Cell, ...] = ( + Cell( + why="tile 128: decode, 1 request without MTP; block kernel", + spec=_ckpt(1, 128), + ), + Cell( + why="tile 128: expert counts tile-1, tile, tile+1; the mn_limit clamp; cluster kernel", + spec=_ckpt(300, 128, "tile_edges"), + ), + Cell( + why="tile 128: 8193 rows; first cooperative-kernel T (prefill)", + spec=_ckpt(8193, 128), + ), + ) + + note = """ + Metadata consumed: none. Stateless; allocates its six outputs plus scratch on + every call and runs on the current stream, with no host synchronization, so + it captures into a CUDA graph. + + Call it as the engine's CuteDSL MoE does (`run_moe_nvfp4_impl`, + `_torch/moe/fused_moe/fused_moe_cute_dsl.py`): `num_experts` = the slot + count, `local_expert_offset` / `local_num_experts` = this rank's slots, and + `tile_tokens_dim` = the grouped GEMMs' tile, 128 or 256 on sm_103. The + downstream FC1 gather GEMM, output memset and FC2 take these six tensors as + they are. + + Kernel path, by `T`: + T 1..4 routingCustom::routingIndicesBlockKernel + T 5..8192 routingCustom::routingIndicesClusterKernel + T 8193..471859 routingInitExpertCounts + routingIndicesCoopKernel + T >= 471860 routingInitExpertCounts + routingIndicesHistogramKernel + + routingIndicesOffsetsKernel + The cooperative ceiling is (SMs - 8) * 512 * 64 / top_k rows with 512 + experts: 471859 on a 152-SM GB300, a different T on another SM count. + + Order within an expert's run is unspecified. The block kernel (T <= 4) lays + each expert's rows out in token order and is bitwise repeatable; every other + path assigns them by atomics, so `expanded_idx_to_permuted_idx` differs + between identical calls. A test therefore compares each expert's rows as a + set and the two mappings as exact inverses; the tile arrays and counts are + deterministic. + + Not written: the tile arrays past `num_non_exiting_tiles`, and the rows of + `permuted_idx_to_expanded_idx` from each tile's `mn_limit` to the tile's + end. Eager, those hold whatever the allocator handed back; under a CUDA + graph they hold the previous replay's indices, which look valid. Every + consumer must bound its reads by `num_non_exiting_tiles` and + `tile_idx_to_mn_limit`. + + `token_final_scales`: shape-checked and never written; its values never + reach an output. Its pointer must be non-null, which is why T = 0 raises. + + Data hazards `is_valid` cannot see (values): a row naming the same expert + twice keeps both pairs on the cluster kernel but loses one on the block + kernel (that pair gets -1). An id outside [local_expert_offset, + local_expert_offset + local_num_experts) gets -1, except that the cluster + kernel reads ids through int16, so an id of 65536 + e sorts as expert e. The + router produces neither. + + Loud: int64 ids, 1-D ids, a `[T, top_k]` size mismatch on either input, + CPU tensors, top_k above 32, num_experts above 2048, T = 0. + + Silent, refused by `is_valid`: a non-contiguous `token_selected_experts` + (read from `data_ptr()` as dense `[T, top_k]`); 2^31 or more elements (the + host counts pairs in int32: it raises a negative-dimension error from 2^31 + and silently sorts only the wrapped remainder from 2^32); `tile_tokens_dim` + <= 0 (the host divides by it). + + Certified only at the checkpoint's values: num_experts 512, top_k 10, all + experts local (offset 0), tile 128 (the target's), fp32 scales, and T 1, + 300 and 8193 (block, cluster and the first cooperative T); CUDA-graph replay + on new ids at T = 128; the contiguity guard. The order, unwritten-region, + data-hazard and loud-rejection statements above were measured at bring-up + and are not re-driven by the test. Expert parallelism, tile 256 and other + tiles, other expert counts and the histogram path are uncertified. + + Hard limits (sm_103) + T * top_k boundary 2^31 pairs (int32 product) + max_tiles * tile boundary 2^31 padded rows (raises past it) + expert id int16 in the cluster kernel; ids < 2048 never reach it + """ + + def __call__( + self, + token_selected_experts: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + local_expert_offset: int, + local_num_experts: int, + tile_tokens_dim: int, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.moe_sort( + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + num_experts=num_experts, + top_k=top_k, + local_expert_offset=local_expert_offset, + local_num_experts=local_num_experts, + tile_tokens_dim=tile_tokens_dim, + ) + + def reference( + self, + token_selected_experts: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + local_expert_offset: int, + local_num_experts: int, + tile_tokens_dim: int, + ) -> list[torch.Tensor]: + """Counting sort by local expert, stable in the expanded index. + + Within an expert, rows are in increasing expanded index; the test's compare + accepts any order there. What the op leaves unwritten is `-1` here. + `token_final_scales`, `num_experts` and `top_k` are not read: the scales + never reach an output, and the ids' shape carries the rest. + """ + ids = token_selected_experts + dev = ids.device + tile = tile_tokens_dim + num_tokens, width = ids.shape + num_expanded = num_tokens * width + local = ids.reshape(-1).long() - local_expert_offset + is_local = (local >= 0) & (local < local_num_experts) + counts = torch.bincount(local[is_local], minlength=local_num_experts) + tiles_per_expert = (counts + tile - 1) // tile + num_tiles = int(tiles_per_expert.sum()) + first = min(local_num_experts, num_expanded) + max_tiles = first + (num_expanded - first) // tile + + run_start = (torch.cumsum(tiles_per_expert, 0) - tiles_per_expert) * tile + tile_expert = torch.repeat_interleave( + torch.arange(local_num_experts, device=dev), tiles_per_expert + ) + tile_idx = torch.arange(num_tiles, device=dev) + mn_limit = torch.minimum( + (tile_idx + 1) * tile, run_start[tile_expert] + counts[tile_expert] + ) + + expanded = torch.arange(num_expanded, device=dev)[is_local] + expert = local[is_local] + order = torch.sort(expert, stable=True).indices + expanded, expert = expanded[order], expert[order] + rank = torch.arange(expert.numel(), device=dev) - (torch.cumsum(counts, 0) - counts)[expert] + permuted = run_start[expert] + rank + + def unwritten(n: int) -> torch.Tensor: + return torch.full((n,), -1, dtype=torch.long, device=dev) + + tile_idx_to_expert_idx = unwritten(max_tiles) + tile_idx_to_expert_idx[:num_tiles] = tile_expert + tile_idx_to_mn_limit = unwritten(max_tiles) + tile_idx_to_mn_limit[:num_tiles] = mn_limit + permuted_idx_to_expanded_idx = unwritten(max_tiles * tile) + permuted_idx_to_expanded_idx[permuted] = expanded + expanded_idx_to_permuted_idx = unwritten(num_expanded) + expanded_idx_to_permuted_idx[expanded] = permuted + outs = ( + tile_idx_to_expert_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx.view(num_tokens, width), + permuted_idx_to_expanded_idx, + torch.tensor([num_tiles * tile], device=dev), + torch.tensor([num_tiles], device=dev), + ) + return [t.to(torch.int32) for t in outs] + + def is_valid( + self, + token_selected_experts: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + local_expert_offset: int, + local_num_experts: int, + tile_tokens_dim: int, + ) -> None: + # The host computes T * top_k in int32: from 2^31 the op raises a + # negative-dimension error, from 2^32 it sorts only the wrapped remainder. + numel = token_selected_experts.numel() + assert numel < 2**31, ( + f"token_selected_experts has {numel} >= 2^31 elements; moe_sort counts " + "pairs in int32 and silently sorts only the wrapped remainder past 2^32" + ) + # The kernels read the ids from data_ptr() as a dense [T, top_k] buffer. + assert token_selected_experts.is_contiguous(), ( + "token_selected_experts must be contiguous; a strided view is read as a " + "dense [T, top_k] buffer and silently sorts the wrong ids" + ) + assert tile_tokens_dim > 0, ( + f"tile_tokens_dim must be positive, got {tile_tokens_dim}; 0 returns a " + "negative num_non_exiting_tiles without raising" + ) + + +moe_sort = MoeSort() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_unpermute_inplace.py new file mode 100644 index 000000000000..e82eaa5dfdc4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/moe_unpermute_inplace.py @@ -0,0 +1,246 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MoE unpermute: each token's routing-weighted sum of its top-k expert rows, taken +from the grouped GEMM's permuted output in a fixed slot order, in one trtllm call. +""" + +from __future__ import annotations + +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 + + +def _ckpt(tokens: int, values: str = "normal") -> dict: + """The checkpoint's routed-expert combine: top-10 of 512 experts, hidden 2560, + bf16 FC2 rows, fp32 routing weights. + + The test lays `permuted_input` out exactly as the CuteDSL FC2 receives it: it + draws each row's top-10 of uniform scores over the 512 experts, sorts the ids + with `moe_sort` at `tile_tokens_dim` and writes each (token, slot) row at the + permuted index the sort gave it. `values` is how the FC2 rows are drawn: + `normal` = standard normal; `wide` = standard normal times 2^u, u uniform in + [-30, 30] per element, so one sum adds terms up to 2^60 apart in magnitude. + The scales are a softmax over each row's ten slots, as the renormalizing + router emits them. + """ + return dict( + tokens=tokens, + tile_tokens_dim=128, + values=values, + num_experts=512, + top_k=10, + hidden_size=2560, + dtype=torch.bfloat16, + scales_dtype=torch.float32, + ) + + +def _fma_f32(acc: torch.Tensor, prod: torch.Tensor) -> torch.Tensor: + """`acc + prod` rounded once to fp32, as a fused multiply-add rounds it. + + `acc` is fp32; `prod` is the fp64 product of a bf16 and an fp32 value, which is + exact (at most 8 + 24 significant bits). The fp64 sum is rounded to odd -- its + exact rounding error from TwoSum, then the last bit forced to 1 if the sum was + inexact -- so that the conversion to fp32 is one correct rounding of the exact + sum rather than two roundings in a row (53 >= 24 + 2 bits). + """ + a = acc.double() + s = a + prod + bb = s - a + err = (a - (s - bb)) + (prod - bb) + inexact = (err != 0) & torch.isfinite(s) + # A sum rounded away from zero truncates to its neighbour one step toward zero. + toward_zero = inexact & ((err < 0) != (s < 0)) + bits = s.view(torch.int64) + bits = torch.where(inexact, (bits - toward_zero.long()) | 1, bits) + return bits.view(torch.float64).float() + + +class MoeUnpermuteInplace(OpWrapper): + """`torch.ops.trtllm.moe_unpermute_inplace`: the routing-weighted top-k combine. + + For each token `t`, writes row `t` of `output`: + + output[t] = bf16( sum over k = 0 .. top_k - 1, in that order, of + permuted_input[expanded_idx_to_permuted_idx[t, k]] + * topk_scales[t, k] ) + + accumulated in fp32 with one rounding per slot (each product fused into the + running sum) and rounded once to the output dtype. A slot whose permuted index + is negative is skipped; a token with every slot skipped gets a zero row. + + Inputs, as the CuteDSL MoE hands them over: `permuted_input` `[M, H]` bf16 is + the FC2 grouped GEMM's output, one row per (token, expert) pair in permuted + order, `M >= T * top_k` (the sort's padded row count); `output` `[T, H]` bf16 + is written in full; `expanded_idx_to_permuted_idx` `[T, top_k]` int32 is the + third output of `moe_sort`; `topk_scales` `[T, top_k]` fp32 are the routing + weights. Call it with these keyword names, the op's own. Returns `output`. + + Fusion boundary: the gather, the weighting and the sum. The FC2 GEMM and the + sort are the caller's, and so is anything added to the routed output (the + shared expert). The op is deterministic: one thread block per token sums its + slots in slot order, so the result does not depend on where the sort placed + each row, on which other tokens share the batch, or on the call. That is what + separates it from the FC2 GEMM's fused finalize, whose reduction order varies + between launches. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode, 1 request without MTP; rows from the sort's block kernel", + spec=_ckpt(1), + ), + Cell(why="8192 rows, max_num_tokens prefill", spec=_ckpt(8192)), + Cell( + why="FC2 rows spanning 2^-30..2^30; terms 2^60 apart in one sum", + spec=_ckpt(8192, values="wide"), + ), + ) + + note = """ + Metadata consumed: none. Stateless; allocates nothing, writes `output` in place + on the current stream with one kernel, `moeUnpermuteKernel<__nv_bfloat16, + float, 256>` (`cpp/tensorrt_llm/kernels/moe/cuteDsl/moeUtils.cu`, registered + in `cpp/tensorrt_llm/thop/moe/cuteDsl/cuteDslMoeUtilsOp.cpp`), one thread block + per token. No host synchronization: it captures into a CUDA graph, and replay + on new rows, scales and ids (moe_sort captured with it) is bitwise equal to + eager and to `reference`. + + Call sites: the engine's CuteDSL MoE calls it when the FC2 finalize is not + fused (`run_moe_nvfp4_impl`, `_torch/moe/fused_moe/fused_moe_cute_dsl.py`), and + this target's `_moe_block` calls it in the prefill and the decode forward with + the same four tensors: the FC2 output, the forward's dense `[num_tokens, + hidden]` routed-output buffer, the sort's third output and the router's fp32 + weights. Prefill runs eagerly at up to max_num_tokens (8192) rows; decode runs + captured at requests x MTP verify width, up to 128. + + What it reads: only the `T * top_k` rows `expanded_idx_to_permuted_idx` names, + and the scales. The padding rows of each tile and the tiles past the sort's + `num_non_exiting_tiles` are never read (every cell fills every other row of + `permuted_input` with NaN, and under a graph those rows hold the previous + replay's rows), so the FC2 output's unwritten regions need no memset before + this op. What it writes: every element of `output`, whatever it held, and none + of its inputs. + + Determinism (bitwise, measured at bring-up; the test re-drives only graph + replay against eager): repeated calls; a second moe_sort of the same ids, which + places an expert's rows in another order; a random permutation of all permuted + rows; a prefix of a larger batch, with or without a fresh sort of the prefix + alone; a shuffle of the tokens. The FC2 GEMM's fused finalize + (`cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell`) differs from its + own first call in about half the output elements at decode, while the plain FC2 + GEMM followed by this op does not, and is bitwise equal to the slot-order fp32 + sum of the real FC2 rows. + + Accuracy against the exact sum: the fp32 sum's error is bounded by its terms' + magnitudes, so an element whose terms nearly cancel can miss the float64 sum + rounded once to bf16 by many ulps of the element itself (827 of 20971520 + elements at T = 8192, standard-normal rows and softmax weights, up to 137.5). + + Silent, refused by `is_valid`: a non-contiguous `permuted_input`, `output`, + `expanded_idx_to_permuted_idx` or `topk_scales` (all four are read or written + from `data_ptr()` as dense row-major); an `output` whose dtype is not + `permuted_input`'s (an fp32 output receives the bf16 bit patterns in the first + half of its bytes); 2^31 or more tokens (nothing is written and nothing + raises). The test holds the `permuted_input` refusal. + + Data hazards `is_valid` cannot see (values): a permuted index at or past + `permuted_input.size(0)` is read without a bounds check; moe_sort never emits + one. + + Loud (measured, not re-driven): int64 indices; hidden size not a multiple of 8; fp32 rows, or fp16 scales + with bf16 rows; fewer than `T * top_k` permuted rows; a scales or index shape + that disagrees with `output`. + + Accepted, not certified: T = 0 returns and writes nothing; bf16 scales; fp16 + rows; a CPU `output`. + + Certified only at the checkpoint's values: top_k 10, hidden 2560, bf16 rows and + output, fp32 scales, the sort's layout at tile 128, every slot local; T 1 and + 8192 eager, 128 under CUDA-graph replay. Slots marked -1 (another rank's + expert) are not certified. + + Hard limits (sm_103) + tokens one thread block per token, the count narrowed to int32: + boundary 2^31, `is_valid` refuses it + output elements T * H, offset in int64: exact past 2^31 and 2^32 + permuted rows int32 index, offset in int64; no index passes through fp32 + """ + + def __call__( + self, + permuted_input: torch.Tensor, + output: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + topk_scales: torch.Tensor, + ) -> torch.Tensor: + torch.ops.trtllm.moe_unpermute_inplace( + permuted_input=permuted_input, + output=output, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + topk_scales=topk_scales, + ) + return output + + def reference( + self, + permuted_input: torch.Tensor, + output: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + topk_scales: torch.Tensor, + ) -> torch.Tensor: + """The slot-order fp32 sum, one rounding per slot, rounded once at the end. + + Each product of a row and its weight is exact in float64; `_fma_f32` adds + it to the fp32 running sum with one rounding, as a fused multiply-add + does. A slot with a negative index leaves the sum untouched. `output` is + ignored: the result has `permuted_input`'s dtype, which is what the op + writes. + """ + idx = expanded_idx_to_permuted_idx.long() + acc = torch.zeros( + idx.shape[0], permuted_input.shape[1], dtype=torch.float32, device=idx.device + ) + for k in range(idx.shape[1]): + slot = idx[:, k] + rows = permuted_input.index_select(0, slot.clamp_min(0)).double() + prod = rows * topk_scales[:, k : k + 1].double() + acc = torch.where((slot >= 0).unsqueeze(1), _fma_f32(acc, prod), acc) + return acc.to(permuted_input.dtype) + + def is_valid( + self, + permuted_input: torch.Tensor, + output: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + topk_scales: torch.Tensor, + ) -> None: + # One thread block per token, the count narrowed to int32: from 2^31 + # tokens nothing is written and no error surfaces. + assert output.size(0) < 2**31, ( + f"output has {output.size(0)} >= 2^31 rows; moe_unpermute_inplace launches one " + "block per token from an int32 count and silently writes nothing" + ) + # The kernel reads and writes all four through data_ptr() as dense + # row-major buffers, and writes `output` with permuted_input's dtype. + for name, t in ( + ("permuted_input", permuted_input), + ("output", output), + ("expanded_idx_to_permuted_idx", expanded_idx_to_permuted_idx), + ("topk_scales", topk_scales), + ): + assert t.is_contiguous(), ( + f"{name} must be contiguous; moe_unpermute_inplace addresses it as a dense " + "row-major buffer and silently combines or writes the wrong elements" + ) + assert output.dtype == permuted_input.dtype, ( + f"output dtype {output.dtype} must be permuted_input's ({permuted_input.dtype}); " + "the op writes permuted_input's element type into it without checking" + ) + + +moe_unpermute_inplace = MoeUnpermuteInplace() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py index 3e100830050c..c97fd51ca4e9 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/renorm_moe_routing_op.py @@ -48,47 +48,16 @@ class RenormMoeRoutingOp(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - Cell(why="_moe_block / MTPLayer._moe decode: one token, bf16 weights", spec=_ckpt(1)), - Cell(why="7 tokens: an odd count", spec=_ckpt(7)), - Cell(why="2048 tokens: prefill-shaped", spec=_ckpt(2048)), + Cell(why="MTPLayer._moe decode: one token, bf16 weights", spec=_ckpt(1)), Cell( - why="4096 tokens: thousands of rows carry an exact bf16 tie in the top-10", - spec=_ckpt(4096), - ), - Cell(why="8192 tokens: max_num_tokens", spec=_ckpt(8192)), - Cell( - why="the default fp32 output dtype at the checkpoint router", - spec=_ckpt(512, torch.float32), - ), - Cell( - why="fp32 logits, bf16 out: the cast must follow the fp32 softmax", - spec=dict(_ckpt(1024), logits_dtype=torch.float32, magnitude=3.0), - ), - Cell( - why="fp16 logits: the third accepted input dtype", - spec=dict(_ckpt(256), logits_dtype=torch.float16, magnitude=2.0), + 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", + spec=_ckpt(16384, torch.float32), ), Cell( why="coarse logits (3 values): ties decide most of the ranking", spec=dict(_ckpt(64), output_dtype=torch.float32, dist="coarse"), ), - Cell( - why="topk 16: the op's ceiling", - spec=dict(_ckpt(256), topk=16), - ), - Cell(why="topk 1: a single slot at weight exactly 1", spec=dict(_ckpt(256), topk=1)), - Cell( - why="33 experts: one past a warp's 32 lanes", - spec=dict(_ckpt(64), num_experts=33, output_dtype=torch.float32), - ), - Cell( - why="129 experts: past 128, where the block grows to 256 threads", - spec=dict(_ckpt(64), num_experts=129, output_dtype=torch.float32), - ), - Cell( - why="topk == num_experts (8 of 8): the full ranking", - spec=dict(_ckpt(16), num_experts=8, topk=8, output_dtype=torch.float32), - ), ) note = """ @@ -101,15 +70,12 @@ class RenormMoeRoutingOp(OpWrapper): 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 - that ceiling. `topk == 0` and `num_tokens == 0` return empties. + that ceiling. Silent, refused by `is_valid`: a non-contiguous `router_logits` (read as a - dense `[T, E]` buffer from `data_ptr()`) and `topk > num_experts` (the - surplus slots come back as expert 0 at weight 0 -- duplicated ids). - - Non-finite logits poison the whole row to NaN (`+inf`, `NaN`, an all - `-inf` row); a `-inf` logit among finite ones is a soft mask, at weight 0 - once `topk` exceeds the finite experts. + dense `[T, E]` buffer from `data_ptr()`; the test holds this refusal) and + `topk > num_experts` (the surplus slots come back as expert 0 at weight 0 + -- duplicated ids). `torch.ops.trtllm.default_moe_routing_op` has the identical signature and softmaxes before the top-k: same ids, un-renormalized weights (rows sum to @@ -117,14 +83,21 @@ 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 (in dtype x out - dtype x num_experts {1..512} x topk {0..16} x tokens {0..8192}); the cells - and the test's sweep re-drive the boundaries of that space. + 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, + 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 0..8192 (16384 measured once) + T certified 1..16384 (max_num_tokens of the tuned config) + + 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). """ def __call__( 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/flashinfer_gemma_rmsnorm.py index 57e25f06609e..4ef212a17f9b 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/flashinfer_gemma_rmsnorm.py @@ -32,10 +32,6 @@ class FlashinferGemmaRmsnorm(OpWrapper): # 10240-wide residual, each a column slice at row stride 10240 -- that, not # a contiguous row, is what the strided-kernel cells below certify. CELLS: tuple[Cell, ...] = ( - Cell( - why="decode: 1 token, stream 3 of the 4x2560 residual (column 7680)", - spec=dict(shape=(1, 2560), col_offset=7680, **_STREAM), - ), Cell( why="decode batch 33: stream 1, row stride 10240 > width 2560", spec=dict(shape=(33, 2560), col_offset=2560, **_STREAM), @@ -48,60 +44,21 @@ class FlashinferGemmaRmsnorm(OpWrapper): why="MTP pre_e_norm: the contiguous [T, 2560] embedding", spec=dict(shape=(4, 2560), dtype=torch.bfloat16, eps=1e-6), ), - Cell( - why="hidden 5120, 1 token, contiguous", - spec=dict(shape=(1, 5120), dtype=torch.bfloat16, eps=1e-6), - ), - Cell( - why="hidden 5120, 33 tokens, contiguous", - spec=dict(shape=(33, 5120), dtype=torch.bfloat16, eps=1e-6), - ), - Cell( - why="hidden 5120, 8192 tokens: the largest row count measured", - spec=dict(shape=(8192, 5120), dtype=torch.bfloat16, eps=1e-6), - ), - Cell( - why="trained-like weight ~0.02: the +1 dominates the scale", - spec=dict(shape=(4, 5120), dtype=torch.bfloat16, eps=1e-6, weight_scale=0.02), - ), - Cell(why="fp16", spec=dict(shape=(4, 5120), dtype=torch.float16, eps=1e-6)), - Cell(why="fp32", spec=dict(shape=(4, 5120), dtype=torch.float32, eps=1e-6)), - Cell( - why="3-D [16, 32, 128]: the per-head kernel", - spec=dict(shape=(16, 32, 128), dtype=torch.bfloat16, eps=1e-5), - ), - Cell( - why="column slice buf[:, :5120] of a [8, 10240] buffer", - spec=dict(shape=(8, 5120), col_offset=0, **_STREAM), - ), - Cell( - why="row slice rows[::2]: row stride 2 x width", - spec=dict(shape=(4, 5120), row_step=2, dtype=torch.bfloat16, eps=1e-6), - ), - Cell( - why="hidden 111: odd, not a multiple of the 8-element bf16 vector", - spec=dict(shape=(16, 111), dtype=torch.bfloat16, eps=1e-6), - ), - Cell( - why="hidden 1152: not a power of two", - spec=dict(shape=(16, 1152), dtype=torch.bfloat16, eps=1e-6), - ), - Cell( - why="hidden 6144", - spec=dict(shape=(16, 6144), dtype=torch.bfloat16, eps=1e-6), - ), ) note = """ Metadata consumed: none. Registered only when flashinfer is importable. - Certified path: flashinfer's CuTe DSL kernels (`rmsnorm_cute` for 2-D, - `qk_rmsnorm_cute` for 3-D, `weight_bias=1.0`); a strided row selects a - strided kernel variant. The CUDA-JIT fallback (`FLASHINFER_USE_CUDA_NORM=1`, + 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`, 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. - Loud, so not guarded: 1-D and 4-D `x`, `x.stride(-1) != 1`, a weight of + 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; float64/int8/uint8 raise KeyError from the DSL dtype table. float8_e4m3fn and float8_e5m2 are in that table and return a plausible fp8 tensor -- @@ -111,9 +68,10 @@ class FlashinferGemmaRmsnorm(OpWrapper): width (8 elements for bf16/fp16, 4 for fp32) and assumes 16-byte aligned pointers, so odd row strides and misaligned slices were not measured. - Numerics: default tolerance holds for all dtypes, nothing loosened; the - largest deviation 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. + 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. Hard limits T * H contiguous inputs above 2^31 - 1 elements route to the strided 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py index 886a2e830c40..35201417be8a 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py @@ -41,44 +41,19 @@ class LayerNorm(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: shape, dtype; `param_dtype` (default dtype) for weight and bias; - # `bias` (default True); `layout` "contiguous" (default) or "column_slice" - # (the second half of a [rows, 2N] per-token buffer). Weight and bias are - # drawn over the checkpoint's measured ranges. - CELLS: tuple[Cell, ...] = ( - tuple( - Cell( - why=f"vision norm, {rows} patch rows x 1152 bf16", - spec=dict(shape=(rows, WIDTH), dtype=torch.bfloat16), - ) - for rows in (1, 7, 256, 1024, 4096) - ) - + tuple( - Cell( - why=f"{str(dtype)[6:]} activations, {rows} rows", - spec=dict(shape=(rows, WIDTH), dtype=dtype), - ) - for dtype in (torch.float16, torch.float32) - for rows in (7, 1024) - ) - + ( - Cell( - why="fp32 weight/bias against bf16 x: widened on load", - spec=dict(shape=(512, WIDTH), dtype=torch.bfloat16, param_dtype=torch.float32), - ), - Cell( - why="bias=None drops the additive term", - spec=dict(shape=(256, WIDTH), dtype=torch.bfloat16, bias=False), - ), - Cell( - why="3-D input keeps its shape", - spec=dict(shape=(4, 64, WIDTH), dtype=torch.bfloat16), - ), - Cell( - why="column slice of a [rows, 2N] buffer, read in place", - spec=dict(shape=(512, WIDTH), dtype=torch.bfloat16, layout="column_slice"), - ), + # spec: shape, dtype (weight and bias share it). Weight and bias are drawn + # over the checkpoint's measured ranges. + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"vision norm, {rows} patch rows x 1152 bf16", + spec=dict(shape=(rows, WIDTH), dtype=torch.bfloat16), ) + for rows in (1, 1024) + ) + ( + Cell( + why="vision norm at the tower call ceiling, 32768 patch rows x 1152 bf16", + spec=dict(shape=(32768, WIDTH), dtype=torch.bfloat16), + ), ) note = """ @@ -87,23 +62,29 @@ class LayerNorm(OpWrapper): Layout: `x.stride(-1) == 1` is read in place through its row stride; any other `x`, and non-contiguous `weight`/`bias`, are silently `.contiguous()`d - first (one extra copy, correct result). `weight`/`bias` dtype is free. + first (one extra copy). `weight`/`bias` dtype is free. + + Certified: contiguous bf16 `[rows, 1152]` `x` at 1, 1024 and 32768 rows + with a bf16 weight and bias, as the vision tower calls it. Not certified: + strided or 3-D `x`, `bias=None`, fp16/fp32 activations, weight/bias in + another dtype than `x`. Loud failures (left to the callable): weight or bias length != N raises a bare AssertionError with no message; `next_pow2(N) * x.element_size() > 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: 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. + 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. 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. Runs unchanged under - `inference_mode`; CUDA-graph capturable after a warm-up call compiles the - Triton kernel. + 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. """ 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py index dca5f65af9aa..fb79cf780dfd 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major.py @@ -8,17 +8,50 @@ from typing import Optional import torch +import triton # Registration is lazy: the op is defined beside the mamba/GDN mixer, and # `import tensorrt_llm._torch.custom_ops` alone does not make # torch.ops.trtllm.rms_norm_gated_token_major resolve. -import tensorrt_llm._torch.modules.mamba.layernorm_gated # noqa: F401 +import tensorrt_llm._torch.modules.mamba.layernorm_gated as _layernorm_gated from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper #: 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") -#: The same geometry under the op's default silu gate. -_SILU = dict(_CKPT, gate_activation="silu") +#: 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. +_PREFILL_MIN_ROWS = 32768 # capacity: fixed: a dispatch threshold, not a buffer size +_PREFILL_ROWS = 8 + + +def _prefill_multirow(x: torch.Tensor, z: torch.Tensor, gate_activation: str) -> bool: + # A `z` that does not cover `x` goes to the registered op, which rejects it. + n = x.shape[-1] + return ( + x.shape[0] >= _PREFILL_MIN_ROWS + and x.shape[0] == z.shape[0] * z.shape[1] + and z.shape[2] == n + and x.stride(-1) == 1 + and z.stride(2) == 1 + and z.stride(1) == n + and n <= 256 + and n & (n - 1) == 0 + and gate_activation in ("silu", "sigmoid") + ) + + +def _prefill_launch(x, z, weight, eps, gate_activation): + """The multi-row kernel at `_PREFILL_ROWS` rows a program (no PDL: the grid covers every SM).""" + rows, n = x.shape + out = torch.empty_like(x) + _layernorm_gated._rms_norm_gated_fwd_multirow_kernel[(triton.cdiv(rows, _PREFILL_ROWS),)]( + x, out, weight, z, None, None, x.stride(0), out.stride(0), z.stride(0), rows, eps, + N=n, ROWS=_PREFILL_ROWS, HEADS_PER_TOK=z.shape[1], GATE_SIGMOID=gate_activation == "sigmoid", + WEIGHT_IS_DELTA=False, LAUNCH_WITH_PDL=False, num_warps=4, + ) # fmt: skip + return out class RmsNormGatedTokenMajor(OpWrapper): @@ -35,6 +68,9 @@ 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]`. + Fusion boundary: norm, weight, gate and the optional quantize. The caller owns both projections producing `x` and `z`, and every reshape. """ @@ -42,61 +78,31 @@ class RmsNormGatedTokenMajor(OpWrapper): ARCHS = frozenset({Arch.SM_103}) CELLS: tuple[Cell, ...] = ( - Cell(why="decode: 1 token, checkpoint GDN epilogue, sigmoid", spec=dict(tokens=1, **_CKPT)), - Cell(why="17 tokens, checkpoint geometry, sigmoid", spec=dict(tokens=17, **_CKPT)), - Cell( - why="prefill 2048 tokens, checkpoint geometry, sigmoid", spec=dict(tokens=2048, **_CKPT) - ), - Cell(why="silu gate (the op's default), 2 tokens", spec=dict(tokens=2, **_SILU)), - Cell(why="silu gate, 512 tokens", spec=dict(tokens=512, **_SILU)), - Cell( - why="z a column slice of a 16384-wide projection, 1 token", - spec=dict(tokens=1, z="column_slice", **_SILU), - ), Cell( - why="z a column slice: token stride 16384, 512 tokens", - spec=dict(tokens=512, z="column_slice", **_CKPT), + why="decode: z the in_proj column view (token stride 16480), 128 tokens", + spec=dict(tokens=128, z="in_proj", **_CKPT), ), Cell( - why="x row stride 2N: only stride(-1) must be 1", - spec=dict(tokens=17, x="strided_rows", **_SILU), + why="prefill: z the in_proj column view, 16384 tokens at 8 rows a program", + spec=dict(tokens=16384, z="in_proj", **_CKPT), ), Cell( - why="7 rows: not a multiple of the kernel's 4 rows per program", - spec=dict(_SILU, tokens=7, heads=1), + why="683 tokens: the first launch at 8 rows a program", + spec=dict(tokens=683, z="in_proj", **_CKPT), ), - Cell(why="15 rows over 3 heads: row mask", spec=dict(_SILU, tokens=5, heads=3)), - Cell(why="6 rows, 1 token x 6 heads: row mask", spec=dict(_SILU, tokens=1, heads=6)), - Cell( - why="N 192: not a power of two, generic fallback kernel", - spec=dict(_SILU, tokens=9, heads=4, width=192), - ), - Cell(why="N 512: wider than 256, fallback", spec=dict(_SILU, tokens=9, heads=4, width=512)), - Cell(why="N 1024: fallback", spec=dict(_SILU, tokens=9, heads=4, width=1024)), - Cell( - why="z head stride 2N: not per-token contiguous, fallback", - spec=dict(_SILU, tokens=9, heads=4, z="head_strided"), - ), - Cell(why="fp16", spec=dict(_SILU, tokens=17, dtype=torch.float16)), - Cell(why="fp32", spec=dict(_SILU, tokens=17, dtype=torch.float32)), - Cell( - why="eps 1.0 on rows scaled by 0.01: eps dominates the mean square", - spec=dict(_SILU, tokens=4, heads=8, eps=1.0, x_scale=0.01), - ), - Cell(why="fp8 out, scale 1.0", spec=dict(tokens=17, fp8_scale=1.0, **_SILU)), - Cell(why="fp8 out, scale 0.5", spec=dict(tokens=17, fp8_scale=0.5, **_SILU)), - Cell(why="fp8 out, scale 0.0234375", spec=dict(tokens=17, fp8_scale=0.0234375, **_SILU)), ) note = """ Metadata consumed: none. Two Triton kernels: a multi-row kernel (4 rows per - program) when `x.stride(-1) == 1`, `z.stride(2) == 1`, `z.stride(1) == N` - and N is a power of two <= 256 -- the checkpoint's N = 128 -- else the - generic single-row kernel on a packed copy of `z`. PDL on the multi-row - path only when the grid is smaller than the SM count and - `TRTLLM_ENABLE_PDL != 0`; results are bitwise identical with it off. - - Loud, so not guarded: 3-D `x` (ValueError), a `z` not covering `x` + program; 8 from 32768 rows, without fp8 output) when `x.stride(-1) == 1`, + `z.stride(2) == 1`, `z.stride(1) == N` and N is a power of two <= 256 -- + the checkpoint's N = 128 -- else the generic single-row kernel on a packed + copy of `z`. PDL on the 4-row path only when the grid is smaller than the + SM count and `TRTLLM_ENABLE_PDL != 0`; results are bitwise identical with + it off. The 8-row programs reduce a row's mean square in another order + than the 4-row ones (one-ulp differences on some outputs). + + Loud (measured), so not guarded: 3-D `x` (ValueError), a `z` not covering `x` (AssertionError "does not match x shape"; 2-D `z` ValueError), an unknown `gate_activation` (ValueError before launch), and on the fallback a non-unit `x.stride(-1)` (bare AssertionError). `is_valid` covers what the @@ -105,12 +111,17 @@ class RmsNormGatedTokenMajor(OpWrapper): extra elements (only `[0]` read) or not fp32 (widened silently). `weight` and `z` are widened to fp32 and need not match `x.dtype`. The - output is freshly allocated and contiguous even for a strided `x`; no input - is mutated. + 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`. - Numerics: bf16 bitwise identical to the fp32 reference up to 816 rows; at - 24576 and 196608 rows one-ulp differences on ~1e-5 of elements, inside the - default tolerance. fp8 output byte-identical. Nothing loosened. + 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. Hard limits N * itemsize fallback holds a row in 64 KB boundary 65536 B @@ -126,6 +137,8 @@ def __call__( fp8_scale: Optional[torch.Tensor] = None, gate_activation: str = "silu", ) -> torch.Tensor: + if fp8_scale is None and z.dim() == 3 and _prefill_multirow(x, z, gate_activation): + return _prefill_launch(x, z, weight, eps, gate_activation) return torch.ops.trtllm.rms_norm_gated_token_major( x, z, weight, eps, fp8_scale, gate_activation ) 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 index 7cb589e7c5a1..af6161baa6ee 100644 --- 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 @@ -68,59 +68,27 @@ class GatedNgramInjection(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # spec: kind "ckpt" (default) -- operands with this checkpoint's norm-gamma - # and value statistics: leading, hc_count, dtype, value_dtype (= dtype), - # 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. + # 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, ...] = ( - tuple( - Cell(why=f"checkpoint 4x2560 bf16, {why}", spec=dict(leading=(t,))) - for t, why in ( - (1, "decode, 1 token"), - (2, "2 tokens"), - (8, "decode batch 8"), - (33, "33 tokens / MTP verify 8 x 4 + 1"), - (512, "prefill 512 tokens"), - (4096, "prefill 4096 tokens"), - (8192, "prefill 8192 tokens, largest measured"), - ) - ) - + (Cell(why="[3, 17, width] leading dims", spec=dict(leading=(3, 17))),) - + tuple( - Cell(why=f"hc_count {c} at D 2560, 64 tokens", spec=dict(leading=(64,), hc_count=c)) - for c in (1, 2, 8) - ) - + tuple( - Cell( - why=f"{str(d).removeprefix('torch.')} operands, 512 tokens", - spec=dict(leading=(512,), dtype=d), - ) - for d in (torch.float16, torch.float32) - ) - + ( - Cell( - why="bf16 key/query with an fp32 value returns fp32", - spec=dict(leading=(128,), value_dtype=torch.float32), + 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, ), - 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, - ), - ), - Cell( - why="same dots with clamp_min 0.0, fp32 value", - spec=dict( - kind="single_channel", - dots=(0.0, 5e-8, -5e-8, 0.05), - value_dtype=torch.float32, - clamp_min=0.0, - ), - ), - ) + ), ) note = """ 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py index 9ad4c4d8230a..a6f8c97dbd5e 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py @@ -25,9 +25,10 @@ CKPT_NGRAM_SIZE = 3 CKPT_EOS_TOKEN_ID = 248044 CKPT_VOCAB_SIZE = 248320 -#: Largest `ple_ngram` pool the cache manager was seen to allocate for this -#: checkpoint (10, 35, 366, 368, 396, 399, 432, 461 over twelve runtime-value -#: dumps, batch never above 32). The target hands the whole pool over unsliced. +#: 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. ENGINE_SLOTS = 461 @@ -87,160 +88,37 @@ class NgramTokenHistory(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # spec: lengths, slots, carries, pool_slots (the cache manager's width, not - # the batch's), ngram_size (3), eos_density (0.0, of the packed ids), - # ids_dtype (int64), index_dtype (int32, for state_indices), qsl_dtype - # (int32), qsl_from "literal" | "seq_lens" (zeros + int64 cumsum of an - # int32 seq_lens prefix -- `cu_seqlens_from_seq_lens`' form), layout - # "contiguous" | "strided_ids" | "pool_view" (`pool[::2]`). + # the batch's), eos_density (0.0, of the packed ids), ids_dtype (int64), + # index_dtype (int32, for state_indices), qsl_dtype (int32), qsl_from + # "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, ...] = ( - tuple( - _cell( - f"decode {n} x 1 token, pool sized to the batch", - (1,) * n, - range(n), - (True,) * n, - n + 2, - ) - for n in (1, 8, 32) - ) - + tuple( - _cell( - f"decode {n} x 1 token at the 461-slot engine pool", - (1,) * n, - _scattered(n), - (True,) * n, - ENGINE_SLOTS, - **_CALL_SITE, - ) - for n in (1, 8, 32) - ) - + tuple( - _cell( - f"prefill {'+'.join(map(str, ls))}, fresh, EOS in stream", - ls, - range(len(ls)), - (False,) * len(ls), - len(ls) + 2, - eos_density=0.05, - ) - for ls in ((37,), (1024,), (37, 1, 128), (7, 7, 7, 7)) - ) - + ( - _cell( - "mixed 64+13 contexts + 6 decodes, 10-slot pool", - (64, 13) + (1,) * 6, - (4, 0, 7, 2, 5, 1, 6, 3), - (False,) + (True,) * 7, - 10, - eos_density=0.03, - ), - _cell( - "mixed 64+13 contexts + 6 decodes at 461 slots", - (64, 13) + (1,) * 6, - (0, ENGINE_SLOTS - 1, 4, 17, 230, 231, 459, 1), - (False,) + (True,) * 7, - ENGINE_SLOTS, - eos_density=0.03, - **_CALL_SITE, - ), - _cell( - "MTP verify p=2: 8 decodes x 2 rows, qsl rebuilt", - (2,) * 8, - _scattered(8), - (True,) * 8, - ENGINE_SLOTS, - qsl_from="seq_lens", - **_CALL_SITE, - ), - _cell( - "MTP verify p=4: 32 decodes x 4 rows, qsl rebuilt", - (4,) * 32, - _scattered(32), - (True,) * 32, - ENGINE_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( - "fresh sequences seed EOS, never the stale slot", (3, 5), (2, 0), (False, False), 4 - ), - _cell("short sequence publishes part carry, part new", (1,), (2,), (True,), 5), - _cell( - "zero-token rows: carry kept, fresh reseeded", - (2, 0, 0), - (0, 1, 2), - (True, True, False), - 3, - ), - _cell( - "unsorted slots, carry flags mixed, 8-slot pool", - (4, 6, 2), - (5, 0, 3), - (True, False, True), - 8, - ), - _cell("unnamed slots untouched, 8-slot pool", (5, 3), (1, 4), (True, True), 8), - _cell( - "in-stream EOS passed through raw", - (17, 1, 5, 1), - (3, 0, 2, 1), - (False, True, True, True), - 6, - eos_density=0.3, - ), - ) - + tuple( - _cell( - f"ids {str(i)[6:]}, slots/qsl {str(x)[6:]}", - (9, 1, 4), - (2, 0, 1), - (True, False, True), - 5, - ids_dtype=i, - index_dtype=x, - qsl_dtype=x, - ) - for i in (torch.int32, torch.int64) - for x in (torch.int32, torch.int64) - ) - + ( - _cell( - "ngram_size 4: a 3-id carry row", - (5, 1, 12), - (0, 2, 1), - (True, True, False), - 5, - ngram_size=4, - eos_density=0.05, - ), - _cell("ngram_size 2: a 1-id carry row", (6, 1), (1, 0), (False, True), 4, ngram_size=2), - _cell( - "input_ids a stride-2 view", - (9, 1, 4), - (2, 0, 1), - (True, False, True), - 5, - layout="strided_ids", - ), - _cell( - "context_pool a pool[::2] slot view, written through", - (9, 1, 4), - (2, 0, 1), - (True, False, True), - 5, - layout="pool_view", - ), - ) + _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( + "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( + "zero-token rows: carry kept, fresh reseeded", + (2, 0, 0), + (0, 1, 2), + (True, True, False), + 3, + ), ) note = """ @@ -277,13 +155,19 @@ class NgramTokenHistory(OpWrapper): pool. At 461 slots the result is bitwise identical to an 11-slot pool with the same slot content. + Certified at the 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. + 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. - Hard limits: not probed (`num_tokens` measured to 1024 here, 4096 in - timing). + Hard limits: not probed (`num_tokens` measured to 4096 in timing). """ 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_gate_conv_inject_norm.py new file mode 100644 index 000000000000..6afb43f1c077 --- /dev/null +++ b/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 @@ -0,0 +1,657 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""PLE block tail for decode-shaped batches: the key and query grouped norms, +the per-stream sigmoid gate, the conv-input grouped norm, the dilated +depthwise short conv against the paged conv state (published back), the +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``). + +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 +state publish are per channel, and each sequence's conv state is 184 KB of +`[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 + +import math +from typing import Optional, Tuple + +import torch +import triton +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 + + +@triton.jit +def _rms(x, w, HIDDEN: tl.constexpr, eps): + """Gemma RMS norm in fp32: `x * rsqrt(mean(x^2) + eps) * (1 + w)`.""" + return x * tl.rsqrt(tl.sum(x * x, axis=0) / HIDDEN + eps) * (1.0 + w) + + +@triton.jit +def _ple_gate_kernel( + hc_ptr, + key_ptr, + value_ptr, + nk_ptr, + nq_ptr, + nc_ptr, + gated_ptr, + xc_ptr, + stride_hc, + stride_key, + stride_value, + stride_out, + eps, + clamp_min, + sqrt_d, + HIDDEN: tl.constexpr, + BLOCK: tl.constexpr, +): + """Per (row, stream): the key and query grouped norms, the gate, the gated + value and the conv-input grouped norm. Writes `gated` and `xc`.""" + row = tl.program_id(0).to(tl.int64) + stream = tl.program_id(1) + cols = tl.arange(0, BLOCK) + m = cols < HIDDEN + ch = stream * HIDDEN + cols + x = tl.load(hc_ptr + row * stride_hc + ch, mask=m, other=0.0).to(tl.float32) + kp = tl.load(key_ptr + row * stride_key + ch, mask=m, other=0.0).to(tl.float32) + v = tl.load(value_ptr + row * stride_value + cols, mask=m, other=0.0).to(tl.float32) + wk = tl.load(nk_ptr + ch, mask=m, other=0.0).to(tl.float32) + wq = tl.load(nq_ptr + ch, mask=m, other=0.0).to(tl.float32) + wc = tl.load(nc_ptr + ch, mask=m, other=0.0).to(tl.float32) + kn = _rms(kp, wk, HIDDEN, eps).to(tl.bfloat16).to(tl.float32) + qn = _rms(x, wq, HIDDEN, eps).to(tl.bfloat16).to(tl.float32) + g = tl.sum(kn * qn, axis=0) / sqrt_d + sg = tl.sqrt(tl.maximum(tl.abs(g), clamp_min)) + sg = tl.where(g > 0, sg, tl.where(g < 0, -sg, 0.0)) + gated = ((1.0 / (1.0 + tl.exp(-sg))) * v).to(tl.bfloat16) + xc = _rms(gated.to(tl.float32), wc, HIDDEN, eps).to(tl.bfloat16) + tl.store(gated_ptr + row * stride_out + ch, gated, mask=m) + tl.store(xc_ptr + row * stride_out + ch, xc, mask=m) + + +@triton.jit +def _ple_conv_kernel( + xc_ptr, + gated_ptr, + ple_ptr, + cw_ptr, + state_ptr, + slots_ptr, + carry_ptr, + qsl_ptr, + stride_out, + stride_state, + pad_slot, + KSIZE: tl.constexpr, + DIL: tl.constexpr, + STATE: tl.constexpr, + HEAD: tl.constexpr, + MAX_P: tl.constexpr, + ROWS: tl.constexpr, + HAS_CARRY: tl.constexpr, +): + """Per (sequence, block of ROWS channels): the dilated depthwise conv of + each of the sequence's rows against `[state | conv inputs]`, its SiLU added + to the gated value (`ple`), and the state advanced in place. + + A channel's state is `STATE` contiguous columns. The publish moves it as a + `[ROWS, HEAD]` tile of the first HEAD (a power of two) columns plus, when + `STATE == HEAD + 1`, one `[ROWS]` vector for the last, so no lane is spent + on padding and every index is a plain 2-D offset. Every load is issued + before the first store.""" + s = tl.program_id(0) + ch0 = tl.program_id(1) * ROWS + begin = tl.load(qsl_ptr + s).to(tl.int64) + length = (tl.load(qsl_ptr + s + 1).to(tl.int64) - begin).to(tl.int32) + slot = tl.load(slots_ptr + s).to(tl.int64) + live = slot != pad_slot + # A sequence without carry reads a zero state (it still publishes). + if HAS_CARRY: + seeded = live & (tl.load(carry_ptr + s) != 0) + else: + seeded = live + slot = tl.where(live, slot, 0) + block = state_ptr + slot * stride_state + ch0 * STATE + row0 = begin * stride_out + ch0 + xrows = xc_ptr + row0 + r = tl.arange(0, ROWS) + i = tl.arange(0, HEAD) + + # The published state, `window[length : length + STATE]` per channel: the + # surviving state columns, then this sequence's newest conv inputs. + src = length + i + keep = tl.load( + block + r[:, None] * STATE + src[None, :], mask=seeded & (src < STATE)[None, :], other=0.0 + ) + fresh = tl.load( + xrows + (src - STATE)[None, :] * stride_out + r[:, None], + mask=live & (src >= STATE)[None, :], + other=0.0, + ) + head = tl.where((src < STATE)[None, :], keep, fresh) + if STATE > HEAD: + src_t = length + HEAD + keep_t = tl.load(block + r * STATE + src_t, mask=seeded & (src_t < STATE), other=0.0) + fresh_t = tl.load( + xrows + (src_t - STATE) * stride_out + r, mask=live & (src_t >= STATE), other=0.0 + ) + tail = tl.where(src_t < STATE, keep_t, fresh_t) + + ch = ch0 + r + for j in tl.static_range(MAX_P): + valid = j < length + gated = tl.load(gated_ptr + row0 + j * stride_out + r, mask=valid, other=0.0) + acc = tl.zeros([ROWS], dtype=tl.float32) + # Tap k of row j reads window column j + k * DIL of `[state | xc rows]`. + for k in tl.static_range(KSIZE): + q = j + k * DIL + wtap = tl.load(cw_ptr + ch * KSIZE + k).to(tl.float32) + if q < STATE: + tap = tl.load(block + r * STATE + q, mask=seeded, other=0.0) + else: + tap = tl.load(xrows + (q - STATE) * stride_out + r, mask=valid, other=0.0) + acc += wtap * tap.to(tl.float32) + conv = acc.to(tl.bfloat16).to(tl.float32) + y = (conv / (1.0 + tl.exp(-conv))).to(tl.bfloat16).to(tl.float32) + out = (gated.to(tl.float32) + y).to(tl.bfloat16) + tl.store(ple_ptr + row0 + j * stride_out + r, out, mask=valid) + + # Every read of these state rows -- `keep` and the taps -- lands before + # the first store over them. + tl.debug_barrier() + tl.store(block + r[:, None] * STATE + i[None, :], head, mask=live) + if STATE > HEAD: + tl.store(block + r * STATE + HEAD, tail, mask=live) + + +@triton.jit +def _ple_conv_rows_kernel( + xc_ptr, + gated_ptr, + ple_ptr, + cw_ptr, + state_ptr, + slots_ptr, + carry_ptr, + qsl_ptr, + num_rows, + num_seqs, + stride_out, + stride_state, + pad_slot, + KSIZE: tl.constexpr, + DIL: tl.constexpr, + STATE: tl.constexpr, + BLOCK_T: tl.constexpr, + ROWS: tl.constexpr, + LOG_S: tl.constexpr, + HAS_CARRY: tl.constexpr, +): + """Context-shaped batches: per (BLOCK_T packed rows, ROWS channels), each + row's dilated conv over `[state | conv inputs]` of its own sequence, its + SiLU added to the gated value. Reads the state only; the publish is the + per-sequence kernel with no rows, launched after.""" + t = tl.program_id(0).to(tl.int64) * BLOCK_T + tl.arange(0, BLOCK_T) + tmask = t < num_rows + ch = tl.program_id(1) * ROWS + tl.arange(0, ROWS) + s = _seq_of_rows(qsl_ptr, t, num_seqs, LOG_S) + begin = tl.load(qsl_ptr + s, mask=tmask, other=0).to(tl.int64) + tpos = t - begin + slot = tl.load(slots_ptr + s, mask=tmask, other=pad_slot).to(tl.int64) + seeded = slot != pad_slot + if HAS_CARRY: + seeded = seeded & (tl.load(carry_ptr + s, mask=tmask, other=0) != 0) + slot = tl.where(seeded, slot, 0) + m2 = tmask[:, None] + offs = t[:, None] * stride_out + ch[None, :] + gated = tl.load(gated_ptr + offs, mask=m2, other=0.0) + acc = tl.zeros([BLOCK_T, ROWS], dtype=tl.float32) + for k in tl.static_range(KSIZE): + wtap = tl.load(cw_ptr + ch * KSIZE + k).to(tl.float32) + q = tpos + k * DIL + from_state = q < STATE + tap_s = tl.load( + state_ptr + slot[:, None] * stride_state + ch[None, :] * STATE + q[:, None], + mask=m2 & (from_state & seeded)[:, None], + other=0.0, + ) + tap_x = tl.load( + xc_ptr + offs + (k * DIL - STATE) * stride_out, + mask=m2 & (~from_state)[:, None], + other=0.0, + ) + tap = tl.where(from_state[:, None], tap_s, tap_x) + acc += wtap[None, :] * tap.to(tl.float32) + conv = acc.to(tl.bfloat16).to(tl.float32) + y = (conv / (1.0 + tl.exp(-conv))).to(tl.bfloat16).to(tl.float32) + out = (gated.to(tl.float32) + y).to(tl.bfloat16) + tl.store(ple_ptr + offs, out, mask=m2) + + +@triton.jit +def _ple_residual_norm_kernel( + hc_ptr, + ple_ptr, + nx_ptr, + hc_out_ptr, + normed_ptr, + stride_hc, + stride_ple, + stride_hc_out, + stride_normed, + eps, + HIDDEN: tl.constexpr, + BLOCK: tl.constexpr, +): + """Per (row, stream): the PLE output added into the stream, then the next + site's grouped norm of the result.""" + row = tl.program_id(0).to(tl.int64) + stream = tl.program_id(1) + cols = tl.arange(0, BLOCK) + m = cols < HIDDEN + ch = stream * HIDDEN + cols + x = tl.load(hc_ptr + row * stride_hc + ch, mask=m, other=0.0).to(tl.float32) + ple = tl.load(ple_ptr + row * stride_ple + ch, mask=m, other=0.0).to(tl.float32) + wx = tl.load(nx_ptr + ch, mask=m, other=0.0).to(tl.float32) + hn = (x + ple).to(tl.bfloat16) + tl.store(hc_out_ptr + row * stride_hc_out + ch, hn, mask=m) + nrm = _rms(hn.to(tl.float32), wx, HIDDEN, eps) + tl.store(normed_ptr + row * stride_normed + ch, nrm.to(tl.bfloat16), mask=m) + + +def _state_head(state_len: int) -> int: + """The largest power of two not above `state_len`: the conv kernel's tile + width; the one column past it, if any, moves as a vector.""" + return 1 << (state_len.bit_length() - 1) + + +#: Channels per conv program. +_CONV_ROWS = 128 +_CONV_WARPS = 4 +#: Context-shaped conv tiles: packed rows x channels per program. +_ROWS_BLOCK_T = 32 +_ROWS_CHANNELS = 128 +_ROWS_WARPS = 4 +_MIN_LOG_S = 8 + +#: This checkpoint's PLE tail geometry: 4 streams of 2560, conv kernel 4 at +#: dilation 3 (a 9-column state), the gate floor `_PLE_CLAMP_MIN`. +CKPT_HC_COUNT = 4 +CKPT_HIDDEN = 2560 +CKPT_DILATION = 3 +CKPT_CLAMP_MIN = 1e-6 +#: Largest PLE-conv pool seen for this checkpoint (see `dilated_short_conv_paged`). +ENGINE_SLOTS = 461 + + +def _cell(why: str, lengths, **spec) -> Cell: + return Cell(why=why, spec=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`:: + + kn = round(n(key[t, c], key_norm[c])) + qn = round(n(hc[t, c], query_norm[c])) + g = sum(kn * qn) / sqrt(D); sg = sign(g) * sqrt(max(|g|, clamp_min)) + gt = round(sigmoid(sg) * value[t]) # gated_ngram_injection + xc = round(n(gt, conv_norm[c])) + conv = round(sum_k conv_weight[c, k] * window[c, t' + k * dilation]) + ple = round(gt + round(silu(conv))) # dilated_short_conv_paged + hc_out[t, c] = round(hc[t, c] + ple) + normed[t, c] = round(n(hc_out[t, c], next_norm[c])) # hc_inject_norm, norm only + + 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 + carries its state (a decode batch); a `pad_slot_id` sequence reads a zero + state and publishes nothing. + + `max_rows > PER_SEQ_MAX_ROWS` (context-shaped steps) takes a fourth + launch: the conv over packed (row tile, channel block) programs, the + per-sequence kernel then only publishing the state. + + Returns `(hc_out, normed)`, fresh `[T, C*D]` in `hc`'s dtype. + `conv_states` is advanced in place. + + Fusion boundary: everything from the key/query norms to the next site's + norm. The n-gram embedding and the key / value projections are the + caller's. + """ + + 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 + # 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), + # 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, + ), + ) + + note = """ + Metadata consumed: none. Call site: the decode target's PLE block, once + per forward: `hc` the `[T, 10240]` stream, `key` the key projection + `[T, 10240]` and `value` the value projection `[T, 2560]` (both cuBLAS, + bf16), the three PLE norm weights and the next site's attention-side norm + weight as contiguous `[10240]`, the `[10240, 1, 4]` conv weight, the whole + `[slots, 10240, 9]` conv pool, int32 slots, int64 `query_start_loc_long` + (the rebuild under MTP verify), `max_rows = gen_p`. + + Preconditions (data; not guarded): at most `max_rows` rows per sequence; + distinct live slots; `query_start_loc[-1] == T`. Operands 2-D with unit + 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, an 8192-row fresh prefill and a mixed 2-context + + 64-decode step, all on the 461-slot pool. Not certified: more than one + row per sequence on the per-sequence path (`1 < max_rows <= + PER_SEQ_MAX_ROWS`) or sequences shorter than `max_rows` -- a verify step + takes the unfused ops, so the target's per-sequence path runs one row. + + Numerics: every rounding point of the composition is kept, so the only + differences are fp32 reduction order (the three norms' sums of squares + and the gate's dot) and fp32 `exp` -- each can move a value across a + rounding midpoint by one ulp, and that ulp propagates down the chain. + + CUDA-graph capturable: no host sync, outputs and scratch (`gated` / `ple`, + `xc`) allocated per call. Three launches: gate (rows x streams), conv + (sequences x channel blocks of `_CONV_ROWS`), residual + norm (rows x + streams). + """ + + def __call__( + self, + hc: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + key_norm: torch.Tensor, + query_norm: torch.Tensor, + conv_norm: torch.Tensor, + conv_weight: torch.Tensor, + conv_states: torch.Tensor, + state_indices: torch.Tensor, + query_start_loc: torch.Tensor, + next_norm: torch.Tensor, + eps: float, + hc_count: int, + clamp_min: float, + dilation: int, + pad_slot_id: int, + max_rows: int, + has_initial_states: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + rows, width = hc.shape + hidden = width // hc_count + ksize = conv_weight.shape[-1] + state_len = conv_states.shape[-1] + hc_out = torch.empty((rows, width), dtype=hc.dtype, device=hc.device) + normed = torch.empty((rows, width), dtype=hc.dtype, device=hc.device) + num_seqs = state_indices.shape[0] + if num_seqs == 0 or rows == 0: + return hc_out, normed + gated = torch.empty((rows, width), dtype=hc.dtype, device=hc.device) + xc = torch.empty((rows, width), dtype=hc.dtype, device=hc.device) + block = triton.next_power_of_2(hidden) + warps = 8 if block >= 2048 else 4 + _ple_gate_kernel[(rows, hc_count)]( + hc, + key, + value, + key_norm, + query_norm, + conv_norm, + gated, + xc, + hc.stride(0), + key.stride(0), + value.stride(0), + gated.stride(0), + eps, + clamp_min, + math.sqrt(hidden), + HIDDEN=hidden, + BLOCK=block, + num_warps=warps, + ) + # `ple` reuses `gated`'s storage row for row: each element is read by + # the program that then writes it. + conv_rows = min(_CONV_ROWS, width) + has_carry = has_initial_states is not None + carry = has_initial_states if has_carry else state_indices + per_seq = max_rows <= PER_SEQ_MAX_ROWS + if not per_seq: + tile_ch = min(_ROWS_CHANNELS, width) + _ple_conv_rows_kernel[(triton.cdiv(rows, _ROWS_BLOCK_T), width // tile_ch)]( + xc, + gated, + gated, + conv_weight, + conv_states, + state_indices, + carry, + query_start_loc, + rows, + num_seqs, + gated.stride(0), + conv_states.stride(0), + pad_slot_id, + KSIZE=ksize, + DIL=dilation, + STATE=state_len, + BLOCK_T=_ROWS_BLOCK_T, + ROWS=tile_ch, + LOG_S=max(_MIN_LOG_S, (num_seqs - 1).bit_length()), + HAS_CARRY=has_carry, + num_warps=_ROWS_WARPS, + ) + # Per sequence: the rows (decode) and the state publish. On the + # context-shaped path it carries no rows and only publishes, after + # every program above has read the state. + _ple_conv_kernel[(num_seqs, width // conv_rows)]( + xc, + gated, + gated, + conv_weight, + conv_states, + state_indices, + carry, + query_start_loc, + gated.stride(0), + conv_states.stride(0), + pad_slot_id, + KSIZE=ksize, + DIL=dilation, + STATE=state_len, + HEAD=_state_head(state_len), + MAX_P=max_rows if per_seq else 0, + ROWS=conv_rows, + HAS_CARRY=has_carry, + num_warps=_CONV_WARPS, + ) + _ple_residual_norm_kernel[(rows, hc_count)]( + hc, + gated, + next_norm, + hc_out, + normed, + hc.stride(0), + gated.stride(0), + hc_out.stride(0), + normed.stride(0), + eps, + HIDDEN=hidden, + BLOCK=block, + num_warps=warps, + ) + return hc_out, normed + + def reference( + self, + hc: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + key_norm: torch.Tensor, + query_norm: torch.Tensor, + conv_norm: torch.Tensor, + conv_weight: torch.Tensor, + conv_states: torch.Tensor, + state_indices: torch.Tensor, + query_start_loc: torch.Tensor, + next_norm: torch.Tensor, + eps: float, + hc_count: int, + clamp_min: float, + dilation: int, + pad_slot_id: int, + max_rows: int, + has_initial_states: Optional[torch.Tensor] = None, + ) -> 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 + an explicit tap sum. Mutates nothing.""" + dtype = hc.dtype + rows, width = hc.shape + hidden = width // hc_count + + def rnd(x): + return x.to(dtype).double() + + def norm(x, w): + w = w.double().view(hc_count, hidden) + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + w) + + streams = hc.double().view(rows, hc_count, hidden) + kn = rnd(norm(key.double().view(rows, hc_count, hidden), key_norm)) + qn = rnd(norm(streams, query_norm)) + g = (kn * qn).sum(-1, keepdim=True) / math.sqrt(hidden) + sg = g.abs().clamp_min(clamp_min).sqrt() * g.sign() + gated = rnd(torch.sigmoid(sg) * value.double().unsqueeze(-2)) + xc = rnd(norm(gated, conv_norm)).reshape(rows, width) + + ksize = conv_weight.shape[-1] + state_len = conv_states.shape[-1] + w = conv_weight.double().reshape(width, ksize) + conv = torch.zeros(rows, width, dtype=torch.float64, device=hc.device) + pool = conv_states.clone() + starts = query_start_loc.tolist() + carries = ( + [True] * state_indices.shape[0] + if has_initial_states is None + else [bool(c) for c in has_initial_states.tolist()] + ) + for s, slot in enumerate(state_indices.tolist()): + begin, end = starts[s], starts[s + 1] + live = slot != pad_slot_id + prefix = ( + conv_states[slot].double() + if live and carries[s] + else torch.zeros(width, state_len, dtype=torch.float64, device=hc.device) + ) + window = torch.cat([prefix, xc[begin:end].t()], dim=1) + n = end - begin + acc = torch.zeros(width, n, dtype=torch.float64, device=hc.device) + for k in range(ksize): + acc += w[:, k : k + 1] * window[:, k * dilation : k * dilation + n] + conv[begin:end] = rnd(acc.t()) + if live: + pool[slot] = window[:, end - begin : end - begin + state_len].to(dtype) + ple = rnd(gated.reshape(rows, width) + rnd(torch.nn.functional.silu(conv))) + hc_out = (hc.double() + ple).to(dtype) + normed = norm(hc_out.double().view(rows, hc_count, hidden), next_norm) + return hc_out, normed.reshape(rows, width).to(dtype), pool + + def is_valid( + self, + hc: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + key_norm: torch.Tensor, + query_norm: torch.Tensor, + conv_norm: torch.Tensor, + conv_weight: torch.Tensor, + conv_states: torch.Tensor, + state_indices: torch.Tensor, + query_start_loc: torch.Tensor, + next_norm: torch.Tensor, + eps: float, + hc_count: int, + clamp_min: float, + dilation: int, + pad_slot_id: int, + max_rows: int, + has_initial_states: Optional[torch.Tensor] = None, + ) -> None: + rows, width = hc.shape + assert width % hc_count == 0, f"width {width} not divisible by hc_count {hc_count}" + hidden = width // hc_count + for name, t in (("hc", hc), ("key", key), ("value", value)): + assert t.dim() == 2 and t.stride(1) == 1, f"{name} must be 2-D, unit column stride" + assert key.shape == hc.shape, f"key {tuple(key.shape)} must match hc {tuple(hc.shape)}" + assert value.shape == (rows, hidden), f"value must be [{rows}, {hidden}]" + for name, t in ( + ("key_norm", key_norm), + ("query_norm", query_norm), + ("conv_norm", conv_norm), + ("next_norm", next_norm), + ): + assert t.shape == (width,) and t.is_contiguous(), f"{name} must be contiguous [{width}]" + assert conv_weight.dim() == 3 and conv_weight.shape[:2] == (width, 1), ( + "conv_weight [W, 1, K]" + ) + assert conv_weight.is_contiguous(), "conv_weight contiguous" + ksize = conv_weight.shape[-1] + assert conv_states.dim() == 3 and conv_states.shape[1:] == ( + width, + (ksize - 1) * dilation, + ), f"conv_states must be [slots, {width}, {(ksize - 1) * dilation}]" + assert conv_states.stride(2) == 1 and conv_states.stride(1) == conv_states.shape[2], ( + "conv_states rows must be dense" + ) + assert conv_states.dtype == hc.dtype, "conv state dtype must match the activation" + assert query_start_loc.shape[0] == state_indices.shape[0] + 1, ( + "query_start_loc must be [num_seqs + 1]" + ) + state_len = conv_states.shape[-1] + assert state_len - _state_head(state_len) <= 1, ( + f"a {state_len}-column conv state is not a power of two plus at most one" + ) + assert width % min(_CONV_ROWS, width) == 0, ( + f"width {width} must be a multiple of the conv block ({_CONV_ROWS} channels)" + ) + assert max_rows >= 1 + assert has_initial_states is None or has_initial_states.shape == state_indices.shape, ( + "has_initial_states must be one flag per sequence" + ) + + +ple_gate_conv_inject_norm = PleGateConvInjectNorm() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py new file mode 100644 index 000000000000..dea3bb7c4136 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ple_ngram_embed.py @@ -0,0 +1,555 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""PLE n-gram embedding for decode-shaped batches: id history, n-gram hash, +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). + +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 +observe a half-published row. + +Context-shaped batches (more than ``PER_SEQ_MAX_ROWS`` rows per sequence) take +a second form: one program per 32 packed rows, each row binary-searching +``query_start_loc`` for its sequence, then a one-program-per-sequence publish +launched after it -- two launches, bitwise equal to the first form. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +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 + +#: This checkpoint's PLE geometry: 3-gram, 16 heads of 160 fp8 columns. +CKPT_NGRAM_SIZE = 3 +CKPT_HEAD_DIM = 160 +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`). +ENGINE_SLOTS = 461 +#: 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 +_ROWS_BLOCK_T = 32 +#: The row-to-sequence search covers at least 2^8 sequences, so one compiled +#: variant serves every batch up to that size. +_MIN_LOG_S = 8 + + +@triton.jit +def _ple_ngram_embed_kernel( + ids_ptr, + pool_ptr, + slots_ptr, + carry_ptr, + qsl_ptr, + table_ptr, + scale_ptr, + mult_ptr, + vocab_ptr, + offs_ptr, + out_ptr, + stride_pool, + stride_table, + stride_out, + eos, + NGRAM: tl.constexpr, + HEADS_PER: tl.constexpr, + HEADS: tl.constexpr, + BLOCK_H: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_P: tl.constexpr, + BLOCK_C: tl.constexpr, + HAS_CARRY: tl.constexpr, +): + CTX: tl.constexpr = NGRAM - 1 + s = tl.program_id(0) + begin = tl.load(qsl_ptr + s).to(tl.int64) + length = tl.load(qsl_ptr + s + 1).to(tl.int64) - begin + slot = tl.load(slots_ptr + s).to(tl.int64) + pool_row = pool_ptr + slot * stride_pool + eos64 = tl.full([BLOCK_P], eos, tl.int64) + if HAS_CARRY: + fresh = tl.load(carry_ptr + s) == 0 + + t = tl.arange(0, BLOCK_P) + tmask = t < length + heads = tl.arange(0, BLOCK_H) + hmask = heads < HEADS + order = heads // HEADS_PER + 2 + vocab = tl.load(vocab_ptr + heads, mask=hmask, other=1) + offs = tl.load(offs_ptr + heads, mask=hmask, other=0) + + # Shift 0 is the token itself; shift k reads `k` back in the sequence's own + # stream, i.e. the packed ids or, before the slice starts, the carry row. + # Once a shift lands on an EOS, every deeper shift reads EOS: the hash's + # segment starts one past the most recent EOS strictly before the token. + mix = tl.load(ids_ptr + begin + t, mask=tmask, other=0).to(tl.int64) * tl.load(mult_ptr) + mixed = tl.zeros([BLOCK_P, BLOCK_H], dtype=tl.int64) + hit = t < 0 + for k in tl.static_range(1, NGRAM): + pos = t - k + from_ids = tl.load(ids_ptr + begin + pos, mask=tmask & (pos >= 0), other=0).to(tl.int64) + from_pool = tl.load(pool_row + (CTX + pos), mask=tmask & (pos < 0), other=0) + if HAS_CARRY: + from_pool = tl.where(fresh, eos64, from_pool) + ck = tl.where(pos >= 0, from_ids, from_pool) + sk = tl.where(hit, eos64, ck) + hit = hit | (ck == eos64) + mix = mix ^ (sk * tl.load(mult_ptr + k)) + mixed = tl.where((order == k + 1)[None, :], mix[:, None], mixed) + + rem = mixed % vocab[None, :] + rows = tl.where(rem < 0, rem + vocab[None, :], rem) + offs[None, :] + + d = tl.arange(0, BLOCK_D) + m3 = tmask[:, None, None] & hmask[None, :, None] & (d < HEAD_DIM)[None, None, :] + q = tl.load(table_ptr + rows[:, :, None] * stride_table + d[None, None, :], mask=m3, other=0.0) + val = q.to(tl.float32) * tl.load(scale_ptr).to(tl.float32) + cols = heads[None, :, None] * HEAD_DIM + d[None, None, :] + tl.store( + out_ptr + (begin + t)[:, None, None] * stride_out + cols, + val.to(out_ptr.dtype.element_ty), + mask=m3, + ) + + # Publish the last CTX ids of `carry + slice` -- after every thread's carry + # read, since the row being overwritten is the one just read. + c = tl.arange(0, BLOCK_C) + cmask = c < CTX + src = length + c + from_slice = tl.load(ids_ptr + begin + src - CTX, mask=cmask & (src >= CTX), other=0).to( + tl.int64 + ) + old = tl.load(pool_row + src, mask=cmask & (src < CTX), other=0) + if HAS_CARRY: + old = tl.where(fresh, tl.full([BLOCK_C], eos, tl.int64), old) + published = tl.where(src >= CTX, from_slice, old) + tl.debug_barrier() + tl.store(pool_row + c, published, mask=cmask) + + +@triton.jit +def _seq_of_rows(qsl_ptr, t, num_seqs, LOG_S: tl.constexpr): + """Per packed row `t`: the largest `s < num_seqs` with `qsl[s] <= t`, i.e. + the sequence that owns the row (zero-row sequences are skipped over).""" + s = tl.zeros_like(t) + for i in tl.static_range(LOG_S): + cand = s + (1 << (LOG_S - 1 - i)) + ok = cand < num_seqs + start = tl.load(qsl_ptr + cand, mask=ok, other=0) + s = tl.where(ok & (start <= t), cand, s) + return s + + +@triton.jit +def _ple_ngram_embed_rows_kernel( + ids_ptr, + pool_ptr, + slots_ptr, + carry_ptr, + qsl_ptr, + table_ptr, + scale_ptr, + mult_ptr, + vocab_ptr, + offs_ptr, + out_ptr, + num_rows, + num_seqs, + stride_pool, + stride_table, + stride_out, + eos, + NGRAM: tl.constexpr, + HEADS_PER: tl.constexpr, + HEADS: tl.constexpr, + BLOCK_H: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, + LOG_S: tl.constexpr, + HAS_CARRY: tl.constexpr, +): + """Context-shaped batches: one program per BLOCK_T packed rows, each row + finding its sequence in `query_start_loc`. Window, hash, gather and + dequantize only; the publish is `_ple_ngram_publish_kernel`, launched + after, because the rows that read a slot's carry may span two programs.""" + CTX: tl.constexpr = NGRAM - 1 + t = tl.program_id(0).to(tl.int64) * BLOCK_T + tl.arange(0, BLOCK_T) + tmask = t < num_rows + s = _seq_of_rows(qsl_ptr, t, num_seqs, LOG_S) + begin = tl.load(qsl_ptr + s, mask=tmask, other=0).to(tl.int64) + tpos = t - begin + slot = tl.load(slots_ptr + s, mask=tmask, other=0).to(tl.int64) + eos64 = tl.full([BLOCK_T], eos, tl.int64) + if HAS_CARRY: + fresh = tl.load(carry_ptr + s, mask=tmask, other=1) == 0 + + heads = tl.arange(0, BLOCK_H) + hmask = heads < HEADS + order = heads // HEADS_PER + 2 + vocab = tl.load(vocab_ptr + heads, mask=hmask, other=1) + offs = tl.load(offs_ptr + heads, mask=hmask, other=0) + + mix = tl.load(ids_ptr + t, mask=tmask, other=0).to(tl.int64) * tl.load(mult_ptr) + mixed = tl.zeros([BLOCK_T, BLOCK_H], dtype=tl.int64) + hit = tmask & (t < 0) + for k in tl.static_range(1, NGRAM): + pos = tpos - k + from_ids = tl.load(ids_ptr + t - k, mask=tmask & (pos >= 0), other=0).to(tl.int64) + from_pool = tl.load( + pool_ptr + slot * stride_pool + (CTX + pos), mask=tmask & (pos < 0), other=0 + ) + if HAS_CARRY: + from_pool = tl.where(fresh, eos64, from_pool) + ck = tl.where(pos >= 0, from_ids, from_pool) + sk = tl.where(hit, eos64, ck) + hit = hit | (ck == eos64) + mix = mix ^ (sk * tl.load(mult_ptr + k)) + mixed = tl.where((order == k + 1)[None, :], mix[:, None], mixed) + + rem = mixed % vocab[None, :] + rows = tl.where(rem < 0, rem + vocab[None, :], rem) + offs[None, :] + + d = tl.arange(0, BLOCK_D) + dmask = d < HEAD_DIM + m2 = tmask[:, None] & dmask[None, :] + scale = tl.load(scale_ptr).to(tl.float32) + for h in tl.static_range(HEADS): + row_h = tl.sum(tl.where((heads == h)[None, :], rows, 0), axis=1) + q = tl.load(table_ptr + row_h[:, None] * stride_table + d[None, :], mask=m2, other=0.0) + val = q.to(tl.float32) * scale + tl.store( + out_ptr + t[:, None] * stride_out + (h * HEAD_DIM + d)[None, :], + val.to(out_ptr.dtype.element_ty), + mask=m2, + ) + + +@triton.jit +def _ple_ngram_publish_kernel( + ids_ptr, + pool_ptr, + slots_ptr, + carry_ptr, + qsl_ptr, + stride_pool, + eos, + NGRAM: tl.constexpr, + BLOCK_C: tl.constexpr, + HAS_CARRY: tl.constexpr, +): + """One program per sequence: the last `NGRAM - 1` ids of `carry + slice` + published to its slot, after every row's carry read (stream order).""" + CTX: tl.constexpr = NGRAM - 1 + s = tl.program_id(0) + begin = tl.load(qsl_ptr + s).to(tl.int64) + length = tl.load(qsl_ptr + s + 1).to(tl.int64) - begin + slot = tl.load(slots_ptr + s).to(tl.int64) + pool_row = pool_ptr + slot * stride_pool + c = tl.arange(0, BLOCK_C) + cmask = c < CTX + src = length + c + from_slice = tl.load(ids_ptr + begin + src - CTX, mask=cmask & (src >= CTX), other=0).to( + tl.int64 + ) + old = tl.load(pool_row + src, mask=cmask & (src < CTX), other=0) + if HAS_CARRY: + fresh = tl.load(carry_ptr + s) == 0 + old = tl.where(fresh, tl.full([BLOCK_C], eos, tl.int64), old) + published = tl.where(src >= CTX, from_slice, old) + tl.debug_barrier() + 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)) + + +class PleNgramEmbed(OpWrapper): + """For sequence `s` over packed rows `[b, e)` with `slot = state_indices[s]`, + its id stream is `carry + input_ids[b:e]`, where `carry` is + `context_pool[slot]` (oldest first) when `has_initial_states[s]` -- or + always, when `has_initial_states` is None -- and `ngram_size - 1` copies of + `eos_token_id` when not. Row `t` of that slice gets the window + `stream[t : t + ngram_size]` (newest last), exactly `ngram_token_history`'s; + each window hashes to one global row per head exactly as `ngram_hash` with + `num_query_tokens=1`; and:: + + out[t, h * D : (h + 1) * D] = to(scale.dtype)(float(table[row[t, h]]) * float(scale)) + + with `D = table.shape[1]` -- `dequantize_e4m3_per_tensor`'s one rounding. + Then, **in place**, `context_pool[slot]` becomes the last `ngram_size - 1` + ids of the stream (a zero-row sequence republishes its carry). Unnamed slots + are untouched. + + Returns a fresh contiguous `[num_tokens, heads * D]` in `scale`'s dtype. + `context_pool` is the only argument modified. + + Fusion boundary: carry read and seed, window, hash, gather, dequantize, + publish. The key/value projections that consume the result are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: lengths (rows per sequence), pool_slots (default num_seqs + 2), + # carries (None -> the argument is None), max_rows (default max(lengths)), + # eos_density (0.0), table "small" | "host" | "ckpt" (the full + # 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( + "decode 128 x 1 row, the full checkpoint-height table", + (1,) * 128, + pool_slots=ENGINE_SLOTS, + table="ckpt", + 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, + ), + ) + + note = """ + Metadata consumed: none -- batch layout, pool, slot map are explicit. + Call site: the decode target's PLE block, once per forward: + `input_ids` int32 `[T]`, the whole `[slots, 2]` int64 n-gram pool, + int32 `state_indices[:batch]`, `has_initial_states=None` (every decode + sequence carries), int64 `query_start_loc_long` (or the + `cu_seqlens_from_seq_lens` rebuild under MTP verify), the + `[320001536, 160]` fp8 table (host-resident by default), the `[1, 1, 1]` + bf16 scale, `max_rows = gen_p`. + + Preconditions (data; not guarded -- each would need a sync): + - every sequence has at most `max_rows` rows; rows past it are neither + written nor published; + - slots real (`0 <= slot < slots`) and distinct; -1 is not supported; + - `query_start_loc[-1] == num_tokens`; + - token id * max multiplier < 2^63 (int64 wraps, as in `ngram_hash`). + + Certified at a 128 x 1-row decode on a host-resident table and on the + full checkpoint-height device table, and a mixed step of 2 contexts (one + fresh, one carrying) + 64 decode rows on the packed-row form. Not + certified: more than one row per sequence on the per-sequence path + (`1 < max_rows <= PER_SEQ_MAX_ROWS`) or zero-row sequences -- a verify + step takes the unfused ops, so the target's per-sequence path runs one + row. + + 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. + Reads the table through its device address, so a host-mapped table needs + no copy. + """ + + def __call__( + self, + input_ids: torch.Tensor, + context_pool: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: Optional[torch.Tensor], + query_start_loc: torch.Tensor, + table: torch.Tensor, + scale: torch.Tensor, + layer_multipliers: torch.Tensor, + head_vocab_sizes: torch.Tensor, + head_offsets: torch.Tensor, + eos_token_id: int, + max_rows: int, + ) -> torch.Tensor: + num_tokens = input_ids.shape[0] + num_seqs = state_indices.shape[0] + ngram_size = layer_multipliers.numel() + heads = head_vocab_sizes.numel() + head_dim = table.shape[1] + out = torch.empty( + (num_tokens, heads * head_dim), dtype=scale.dtype, device=input_ids.device + ) + if num_seqs > 0 and max_rows > PER_SEQ_MAX_ROWS: + carry = has_initial_states if has_initial_states is not None else state_indices + if num_tokens > 0: + _ple_ngram_embed_rows_kernel[(triton.cdiv(num_tokens, _ROWS_BLOCK_T),)]( + input_ids, + context_pool, + state_indices, + carry, + query_start_loc, + table, + scale, + layer_multipliers, + head_vocab_sizes, + head_offsets, + out, + num_tokens, + num_seqs, + context_pool.stride(0), + table.stride(0), + out.stride(0), + eos_token_id, + NGRAM=ngram_size, + HEADS_PER=heads // (ngram_size - 1), + HEADS=heads, + BLOCK_H=triton.next_power_of_2(heads), + HEAD_DIM=head_dim, + BLOCK_D=triton.next_power_of_2(head_dim), + BLOCK_T=_ROWS_BLOCK_T, + LOG_S=max(_MIN_LOG_S, (num_seqs - 1).bit_length()), + HAS_CARRY=has_initial_states is not None, + num_warps=4, + ) + _ple_ngram_publish_kernel[(num_seqs,)]( + input_ids, + context_pool, + state_indices, + carry, + query_start_loc, + context_pool.stride(0), + eos_token_id, + NGRAM=ngram_size, + BLOCK_C=triton.next_power_of_2(ngram_size - 1), + HAS_CARRY=has_initial_states is not None, + num_warps=1, + ) + elif num_seqs > 0: + _ple_ngram_embed_kernel[(num_seqs,)]( + input_ids, + context_pool, + state_indices, + has_initial_states if has_initial_states is not None else state_indices, + query_start_loc, + table, + scale, + layer_multipliers, + head_vocab_sizes, + head_offsets, + out, + context_pool.stride(0), + table.stride(0), + out.stride(0), + eos_token_id, + NGRAM=ngram_size, + HEADS_PER=heads // (ngram_size - 1), + HEADS=heads, + BLOCK_H=triton.next_power_of_2(heads), + HEAD_DIM=head_dim, + BLOCK_D=triton.next_power_of_2(head_dim), + BLOCK_P=triton.next_power_of_2(max(max_rows, 1)), + BLOCK_C=triton.next_power_of_2(ngram_size - 1), + HAS_CARRY=has_initial_states is not None, + num_warps=4, + ) + return out + + def reference( + self, + input_ids: torch.Tensor, + context_pool: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: Optional[torch.Tensor], + query_start_loc: torch.Tensor, + table: torch.Tensor, + scale: torch.Tensor, + layer_multipliers: torch.Tensor, + head_vocab_sizes: torch.Tensor, + head_offsets: torch.Tensor, + eos_token_id: int, + 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` + (window and published pool) and `ngram_hash` (rows) -- and a float64 + gather-and-scale rounded once. Mutates nothing.""" + carries = ( + torch.ones_like(state_indices, dtype=torch.bool) + if has_initial_states is None + else has_initial_states + ) + ngram_size = layer_multipliers.numel() + window, pool = ngram_token_history.reference( + input_ids, + context_pool, + state_indices, + carries, + query_start_loc, + eos_token_id, + ngram_size, + ) + rows = ngram_hash.reference( + window, layer_multipliers, head_vocab_sizes, head_offsets, eos_token_id, 1 + ).reshape(window.shape[0], -1) + gathered = table[rows.reshape(-1)].double().reshape(rows.shape[0], -1) + out = (gathered * scale.double().reshape(())).to(scale.dtype) + return out, pool + + def is_valid( + self, + input_ids: torch.Tensor, + context_pool: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: Optional[torch.Tensor], + query_start_loc: torch.Tensor, + table: torch.Tensor, + scale: torch.Tensor, + layer_multipliers: torch.Tensor, + head_vocab_sizes: torch.Tensor, + head_offsets: torch.Tensor, + eos_token_id: int, + max_rows: int, + ) -> None: + num_seqs = state_indices.shape[0] + ngram_size = layer_multipliers.numel() + assert input_ids.dim() == 1 and input_ids.is_contiguous(), "input_ids [T], contiguous" + assert ngram_size >= 2 and layer_multipliers.dtype == torch.int64, ( + "layer_multipliers int64, at least two orders" + ) + assert context_pool.dtype == torch.int64 and context_pool.stride(1) == 1, ( + "context_pool int64 with unit column stride" + ) + assert context_pool.shape[1] == ngram_size - 1, ( + f"context_pool must be ngram_size - 1 = {ngram_size - 1} wide" + ) + assert query_start_loc.shape[0] == num_seqs + 1, ( + "query_start_loc must be [num_seqs + 1]; a short one misattributes rows" + ) + assert has_initial_states is None or has_initial_states.shape[0] == num_seqs, ( + "has_initial_states must be one flag per sequence" + ) + assert head_vocab_sizes.numel() % (ngram_size - 1) == 0, ( + "heads must split evenly over the n-gram orders" + ) + assert head_offsets.numel() == head_vocab_sizes.numel(), "one offset per head" + assert head_vocab_sizes.dtype == head_offsets.dtype == torch.int64, "int64 vocab/offsets" + assert table.dim() == 2 and table.stride(1) == 1, "table [rows, D], unit column stride" + assert scale.numel() == 1, "one per-tensor scale" + assert max_rows >= 1, "max_rows >= 1" + + +ple_ngram_embed = PleNgramEmbed() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/bf16_to_fp32.py new file mode 100644 index 000000000000..61edc7eb2a1d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/bf16_to_fp32.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Widening bf16 -> fp32 of a contiguous tensor in one vectorized kernel: the logits +processor's cast of the step's `[rows, vocab]` logits. + +Tier 5 (Triton written for this target). torch's `.float()` runs this as an unrolled, +unvectorized `direct_copy` at under half the bandwidth; the values are the same bit for +bit. +""" + +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 vocabulary: the width of the logits rows. +VOCAB = 248320 +#: Elements a program (8 warps: 16 a thread, two 16-byte loads). +_BLOCK = 4096 +_WARPS = 8 + + +@triton.jit +def _bf16_to_fp32_kernel(x_ptr, y_ptr, n, BLOCK: tl.constexpr): + offs = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + tl.store(y_ptr + offs, tl.load(x_ptr + offs, mask=mask).to(tl.float32), mask=mask) + + +class Bf16ToFp32(OpWrapper): + """`x.float()` for a contiguous bf16 `x`: a fresh contiguous fp32 tensor of `x`'s + shape. Widening is exact, so every element is `x`'s value; NaN payloads widen as + torch widens them (the 16 bits shifted up). + + Fusion boundary: the cast only; the head's GEMM before it and the sampler after it + are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="decode logits: 128 rows of the 248320 vocabulary", spec=dict(shape=(128, VOCAB))), + Cell(why="one row; a ragged last program", spec=dict(shape=(1, VOCAB))), + Cell( + why="every bf16 bit pattern: NaNs, infinities, subnormals, signed zeros", + spec=dict(shape=(1, 65536), source="all_codes"), + ), + ) + + note = """ + Metadata consumed: none; stateless. One Triton kernel, `cdiv(numel, 4096)` programs + of 8 warps over the flat tensor. Call site: the target's logits processor + (`_WideningLogitsProcessor`), on the head's bf16 logits; the shell's `.float()` + after it is then a no-op. + + GB300, graph replay: 128 x 248320 31.3 us against 82.1 for `.float()`, 124 rows + 30.7 against 79.9, 33 rows 8.3 against 16.5, one row 4.1 against 4.2. + + Preconditions: bf16, contiguous (the kernel reads the flat storage). + """ + + def __call__(self, x: torch.Tensor) -> torch.Tensor: + y = torch.empty(x.shape, dtype=torch.float32, device=x.device) + n = x.numel() + if n: + _bf16_to_fp32_kernel[(triton.cdiv(n, _BLOCK),)](x, y, n, BLOCK=_BLOCK, num_warps=_WARPS) + return y + + def reference(self, x: torch.Tensor) -> torch.Tensor: + return x.float() + + def is_valid(self, x: torch.Tensor) -> None: + assert x.dtype == torch.bfloat16, "bf16 input" + assert x.is_contiguous(), "contiguous input: the kernel reads the flat storage" + + +bf16_to_fp32 = Bf16ToFp32() 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 index 68888e530047..7f3922dea06a 100644 --- 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 @@ -42,67 +42,15 @@ class DequantizeE4m3PerTensor(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell(why="_ple_block decode: [1, 16, 160] gathered rows, bf16 scale", spec=_ple(1)), - Cell(why="_ple_block at 8 tokens", spec=_ple(8)), - Cell(why="_ple_block prefill: [4096, 16, 160]", spec=_ple(4096)), 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" - ), - ), - Cell( - why="every e4m3 byte into fp16", - spec=dict( - shape=(1, 256), dtype=torch.float16, scales=_DOMAIN_SCALES, source="all_codes" - ), - ), - Cell( - why="every e4m3 byte into fp32", - spec=dict( - shape=(1, 256), dtype=torch.float32, scales=_DOMAIN_SCALES, source="all_codes" - ), - ), - Cell( - why="rank 2: the scale's rank tracks the input's", - spec=dict(shape=(4, 8), dtype=torch.bfloat16, scales=(PLE_SCALE,), source="randn"), - ), - Cell( - why="rank 5: the deepest rank run", - spec=dict( - shape=(2, 2, 2, 4, 8), dtype=torch.bfloat16, scales=(PLE_SCALE,), source="randn" - ), - ), - Cell( - why="1 x 1: far below the fixed 1024-block grid", - spec=dict(shape=(1, 1), dtype=torch.bfloat16, scales=(PLE_SCALE,), source="randn"), - ), - Cell( - why="3 x 7: odd shape, no width or alignment rule", - spec=dict(shape=(3, 7), dtype=torch.bfloat16, scales=(PLE_SCALE,), source="randn"), - ), - Cell( - why="[2^20, 160]: 167,772,160 elements through the grid-stride loop", - spec=dict( - shape=(1 << 20, 160), dtype=torch.bfloat16, scales=(PLE_SCALE,), source="randn" - ), - ), - Cell( - why="scale as a one-element slice at a nonzero storage offset", - spec=dict( - shape=(8, 160), - dtype=torch.bfloat16, - scales=(PLE_SCALE,), - source="randn", - scale_slice=True, - ), + why="_ple_block prefill at max_num_tokens: [16384, 16, 160] through the grid-stride " + "loop", + spec=_ple(16384), ), Cell( - why="input at a non-16-byte-aligned address: same kernel", + why="every e4m3 byte into bf16, five scale values, NaN codes included", spec=dict( - shape=(64, 160), - dtype=torch.bfloat16, - scales=(PLE_SCALE, 1 / 3, 7.3e-3), - source="unaligned", + shape=(1, 256), dtype=torch.bfloat16, scales=_DOMAIN_SCALES, source="all_codes" ), ), ) @@ -113,26 +61,26 @@ class DequantizeE4m3PerTensor(OpWrapper): grid with a grid-stride loop; there is no vectorized variant, so alignment is not an axis. - Preconditions, all enforced loudly by the op (`symmetric_dequantize_per_ - tensor` in thop/fp8Op.cpp): `input` float8_e4m3fn (e5m2, int8, bf16 raise), - CUDA, contiguous, rank >= 2, non-empty (zero rows or columns raise -- an - empty batch must branch first); `scales` CUDA, bf16 / fp16 / fp32 (fp64 - raises), of `input`'s rank with every dim 1. + 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. - Degenerate scale values are accepted and computed exactly: 0 zeroes every - lane (signs kept, no NaN), a negative scale flips signs, +inf sends finite - lanes to +-inf and zero lanes to NaN, NaN makes everything NaN. Overflow is - inf, not a clamp (fp16 scale 1000 sends 448 to inf). NaN codes stay NaN. - `448 * PLE_SCALE` reads back as bf16 0.08935546875. + 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 probed to 167,772,160; the int32 boundary 2^31 not probed + numel certified to 41,943,040 (16384 tokens); probed to 167,772,160; the + int32 boundary 2^31 not probed - Not certified: `TRTLLM_ENABLE_PDL=1`, `scales` on a different device than - `input`. + 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: 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py index b006f4e73adc..51e97e8e5444 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py @@ -71,47 +71,13 @@ class Fp8Quantize1x128(OpWrapper): CELLS: tuple[Cell, ...] = ( Cell(why="MTPLayer._moe decode: one token at hidden 2560", spec=dict(tokens=1, **_CKPT)), - Cell( - why="3 tokens: inside the data buffer's 4-row allocation pad", - spec=dict(tokens=3, **_CKPT), - ), - Cell(why="4 tokens: exactly one 4-row pad unit", spec=dict(tokens=4, **_CKPT)), - Cell( - why="5 tokens: the first count spilling into a second pad unit", - spec=dict(tokens=5, **_CKPT), - ), Cell( why="20 tokens == H/128: the scale buffer is square, its transpose undetectable", spec=dict(tokens=20, **_CKPT), ), - Cell(why="64 tokens: a multiple of every pad", spec=dict(tokens=64, **_CKPT)), - Cell(why="1024 tokens: prefill-shaped", spec=dict(tokens=1024, **_CKPT)), Cell( - why="8192 tokens: max_num_tokens, the top of the certified column", - spec=dict(tokens=8192, **_CKPT), - ), - Cell( - why="H 128: one whole block per row", spec=dict(tokens=7, hidden=128, use_ue8m0=False) - ), - Cell( - why="H 7168: the widest width measured", - spec=dict(tokens=64, hidden=7168, use_ue8m0=False), - ), - Cell( - why="H 16: the narrowest width the op takes, one short block", - spec=dict(tokens=3, hidden=16, use_ue8m0=False), - ), - Cell( - why="H 144: a 16-wide tail block takes the amax of its own elements", - spec=dict(tokens=64, hidden=144, use_ue8m0=False), - ), - Cell( - why="H 2576: the checkpoint width plus a 16-wide tail", - spec=dict(tokens=64, hidden=2576, use_ue8m0=False), - ), - Cell( - why="use_ue8m0=True: the round-up power-of-two scale, bit-exact", - spec=dict(tokens=17, hidden=2560, use_ue8m0=True), + why="16384 rows: max_num_tokens of the tuned config", + spec=dict(tokens=16384, **_CKPT), ), ) @@ -121,25 +87,26 @@ class Fp8Quantize1x128(OpWrapper): sm_100 family the op returns a 1-D padded buffer instead. That branch was read, not run. - Preconditions, all enforced loudly by the op: `input` bf16 (fp16, fp32 and - fp8 raise), 2-D (flatten `[B, S, H]` yourself), contiguous, CUDA, and - `H % 16 == 0` -- the width, not the block count, so H need not be a - multiple of 128. `T == 0` is accepted and returns `[0, H]` / `[NB, 0]`. - A consumer needing `H % 128 == 0` (the block-scale MoE runner) owns that - constraint. - - Edge behaviour, measured: an all-zero block takes scale - `bf16(1e-10) / 448` and all-zero codes; `inf` poisons its whole block - (scale inf, its own code NaN, every other code 0); NaN stays in its own - element (the running max ignores it); the block max maps to exactly +-448 - and nothing overflows. `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 capturable; stream- - independent; the registered fake matches eager shapes. + Preconditions, enforced loudly by the op: `input` bf16, 2-D (flatten + `[B, S, H]` yourself), contiguous (the test holds the column-slice + rejection), CUDA, and `H % 16 == 0` -- the width, not the block count, so + H need not be a multiple of 128. A consumer needing `H % 128 == 0` (the + block-scale MoE runner) owns that constraint. + + `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. Hard limits - T certified 0..8192 at H 2560; above that not probed - H multiple of 16 (raises otherwise); certified 16..7168 + T certified at 1, 20 and 16384 at H 2560 (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. """ def __call__( 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 index d1932d41571f..af83f235c2f7 100644 --- 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 @@ -23,13 +23,6 @@ ENGINE_SLOTS = 461 -def _engine_batch_slots() -> tuple[int, ...]: - """32 distinct slots over the whole engine pool, slot 0 and the last included.""" - slots = (0, ENGINE_SLOTS - 1) + tuple((i * 13 + 5) % (ENGINE_SLOTS - 2) + 1 for i in range(30)) - assert len(set(slots)) == 32 - return slots - - class CausalConv1dFwd(OpWrapper): """Depthwise causal conv1d (+ optional SiLU) of every sequence in a batch, against its carried window, publishing the new window -- one launch. @@ -58,177 +51,34 @@ class CausalConv1dFwd(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec keys: dim, width, dtype, lengths (varlen) or batch3d=(batch, seqlen), - # pool_slots (None -> no conv_states), slots (None -> cache_indices None), - # initial (None -> has_initial_state None), silu, bias, pad_slot_id, - # layout: "in_place" | "out" | "channel_last". + # 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 37+1+128 tokens in place, 8-slot pool", - spec=dict( - dim=CKPT_DIM, - lengths=(37, 1, 128), - pool_slots=8, - slots=(3, 0, 5), - initial=(True, False, True), - ), - ), - 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="32 sequences (the batch ceiling) scattered over the 461-slot pool", - spec=dict( - dim=CKPT_DIM, - lengths=tuple(1 + (i * 7) % 96 for i in range(32)), - pool_slots=ENGINE_SLOTS, - slots=_engine_batch_slots(), - initial=tuple(i % 3 != 0 for i in range(32)), - ), - ), - Cell( - why="explicit channel-major out: result lands there, x is read-only", - spec=dict( - dim=CKPT_DIM, - lengths=(64, 9), - pool_slots=4, - slots=(1, 2), - initial=(True, True), - layout="out", - ), - ), - Cell( - why="channel-last x (a bare .t() of token-major), into its own out", - spec=dict( - dim=CKPT_DIM, - lengths=(37, 1, 128), - pool_slots=8, - slots=(3, 0, 5), - initial=(True, False, True), - layout="channel_last", - ), - ), - 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="pad_slot_id 99: the sentinel is compared by value", - spec=dict( - dim=1024, - lengths=(16, 8, 16), - pool_slots=4, - slots=(0, 99, 2), - initial=(True, True, True), - pad_slot_id=99, - ), - ), - Cell( - why="sequences shorter than the window: state part carry, part tokens", - spec=dict( - dim=512, - lengths=(2, 1, 2), - pool_slots=4, - slots=(0, 1, 2), - initial=(True, True, False), - ), - ), - Cell( - why="a zero-token sequence republishes its carry unchanged", - spec=dict( - dim=512, lengths=(0, 40), pool_slots=4, slots=(2, 3), initial=(True, False) - ), - ), - ) - + tuple( - Cell( - why=f"{str(dtype).removeprefix('torch.')} with bias, silu={silu}", - spec=dict( - dim=256, - lengths=(30, 34), - pool_slots=2, - slots=(0, 1), - initial=(True, True), - dtype=dtype, - bias=True, - silu=silu, - ), - ) - for dtype in (torch.bfloat16, torch.float16, torch.float32) - for silu in (True, False) - ) - + ( - Cell( - why="conv_states None: every sequence starts from zeros", - spec=dict(dim=512, lengths=(12, 20), pool_slots=None, slots=None, initial=None), - ), - Cell( - why="cache_indices/has_initial_state None: slot s, no carry", - spec=dict(dim=512, lengths=(24, 40), pool_slots=2, slots=None, initial=None), - ), - Cell( - why="dense [batch, dim, seqlen] layout, query_start_loc None", - spec=dict( - dim=256, batch3d=(3, 40), pool_slots=3, slots=None, initial=(True, False, True) - ), - ), - Cell( - why="one 2048-token sequence: the vectorised channel-major load", - spec=dict(dim=CKPT_DIM, lengths=(2048,), pool_slots=1, slots=(0,), initial=(True,)), - ), - Cell( - why="one 2047-token sequence: the scalar channel-major load", - spec=dict(dim=CKPT_DIM, lengths=(2047,), pool_slots=1, slots=(0,), initial=(True,)), - ), - Cell( - why="several long sequences: the 128-thread block", - spec=dict( - dim=1024, - lengths=(700, 701, 648), - pool_slots=3, - slots=(0, 1, 2), - initial=(True, True, True), - ), - ), - Cell( - why="every sequence a single token", - spec=dict( - dim=1024, - lengths=(1, 1, 1, 1), - pool_slots=4, - slots=(0, 1, 2, 3), - initial=(True, True, True, True), - ), - ), - ) - + tuple( - Cell( - why=f"conv width {width}", - spec=dict( - dim=512, - width=width, - lengths=(33, 7), - pool_slots=2, - slots=(0, 1), - initial=(True, True), - ), - ) - for width in (2, 3) - ) + 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 = """ @@ -239,10 +89,13 @@ class CausalConv1dFwd(OpWrapper): 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 first three cells are that call; 461 slots tops the certified pool - interval (8 and 461 bracket every width the cache manager was seen to - allocate). At 461 the activation and state are bitwise identical to the - 8-slot result for the same slot content. + 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`, @@ -252,11 +105,15 @@ class CausalConv1dFwd(OpWrapper): 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 of length 0 - and shorter than `W - 1` are legal. Unnamed slots are bitwise untouched. - - `TRTLLM_CAUSAL_CONV1D_SILU_MODE` selects faster SiLU forms on the - channel-last kernel; only the default (unset) is certified. + `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 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py index bbeb9ca56599..cbcd77719624 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_triton_update.py @@ -24,9 +24,10 @@ #: 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). The verify path hands -#: the whole pool over unsliced. +#: A 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 (a bf16 recurrent pool holds more slots in the +#: same memory, and the conv pool shares its slot space) takes the same path. The verify +#: path hands the whole pool over unsliced. ENGINE_SLOTS = 461 #: The engine's `max_num_sequences`: the snapshot buffer's row count. ENGINE_SNAPSHOT_ROWS = 32 @@ -34,8 +35,6 @@ def _pinned(batch: int) -> tuple[int, ...]: """`batch` distinct slots of the engine pool, slot 0 and the last pinned in.""" - if batch == 1: - return (0,) return (0,) + tuple(i * 14 + 3 for i in range(1, batch - 1)) + (ENGINE_SLOTS - 1,) @@ -69,170 +68,67 @@ class CausalConv1dTritonUpdate(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec keys: dim, width (4), state_len (pool columns, 3), seqlen (2), batch, - # pool_slots, slots ("random" | tuple | None), dtype, bias, activation - # ("silu"), pad_slot_id, x_form ("token_major" | "contiguous" | "2d"), - # snapshots (rows of the snapshot buffer, or 0 for none), accepted, - # launch_dependent_kernels. - CELLS: tuple[Cell, ...] = ( - tuple( - Cell( - why=f"verify at the 461-slot pool, batch {batch}, {seqlen} tokens/row", - spec=dict( - dim=CKPT_DIM, - batch=batch, - seqlen=seqlen, - pool_slots=ENGINE_SLOTS, - slots=_pinned(batch), - snapshots=ENGINE_SNAPSHOT_ROWS, - ), - ) - for batch, seqlen in ((1, 2), (8, 3), (32, 2), (17, 4)) - ) - + tuple( - Cell( - why=f"verify, 128-slot pool, batch {batch}, {seqlen} tokens/row", - spec=dict( - dim=CKPT_DIM, - batch=batch, - seqlen=seqlen, - pool_slots=128, - slots="random", - snapshots=batch, - ), - ) - for batch, seqlen in ((1, 2), (4, 2), (8, 3), (17, 4), (64, 2)) - ) - + ( - Cell( - why="pad_slot_id -1 rows: pool slot and snapshot rows untouched", - spec=dict( - dim=512, batch=6, pool_slots=16, slots=(3, -1, 7, 11, -1, 2), snapshots=6 - ), - ), - Cell( - why="pad_slot_id 9: the sentinel is compared by value", - spec=dict(dim=512, batch=4, pool_slots=16, slots=(5, 9, 0, 9), pad_slot_id=9), - ), - Cell( - why="conv_state_indices None: row b reads and writes slot b", - spec=dict(dim=512, batch=4, pool_slots=4, slots=None), - ), - ) - + tuple( - Cell( - why=f"activation={activation!r}", - spec=dict(dim=512, batch=4, pool_slots=16, slots="random", activation=activation), - ) - for activation in (None, False, True, "swish") - ) - + tuple( - Cell( - why=f"bias, activation={activation!r}: added in the fp32 accumulator", - spec=dict( - dim=512, - batch=4, - pool_slots=16, - slots="random", - bias=True, - activation=activation, - ), - ) - for activation in ("silu", None) - ) - + ( - Cell( - why="2-D [batch, dim] one-token step returns 2-D", - spec=dict(dim=512, batch=12, seqlen=1, pool_slots=16, slots="random", x_form="2d"), + # 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, + # SiLU, pad_slot_id -1, no accepted-token roll. + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"verify, x over the [n*p, 16480] in_proj, batch 32, {seqlen} tok/row", + spec=dict( + dim=CKPT_DIM, + batch=32, + seqlen=seqlen, + pool_slots=ENGINE_SLOTS, + slots=_pinned(32), + snapshots=ENGINE_SNAPSHOT_ROWS, ), ) - + tuple( - Cell( - why=f"pool {state_len} columns wide: columns past W-1 are inert", - spec=dict(dim=512, batch=3, pool_slots=16, slots="random", state_len=state_len), - ) - for state_len in (4, 6) - ) - + ( - Cell( - why="num_accepted_tokens 1/2/3/1 rolls the window to column 0", - spec=dict( - dim=512, - batch=4, - seqlen=3, - state_len=5, # W-1 + (seqlen-1): the pool this form needs - pool_slots=16, - slots="random", - accepted=(1, 2, 3, 1), - ), - ), - ) - + tuple( - Cell( - why=f"conv width {width}", - spec=dict( - dim=512, - batch=3, - width=width, - state_len=width - 1, - pool_slots=16, - slots="random", - ), - ) - for width in (2, 3) - ) - + tuple( - Cell( - why=f"dim {dim}, off the 128-channel block", - spec=dict(dim=dim, batch=2, pool_slots=8, slots="random"), - ) - for dim in (1, 127, 128, 129, 513) - ) - + tuple( - Cell( - why=f"{str(dtype).removeprefix('torch.')} with bias", - spec=dict(dim=512, batch=4, pool_slots=16, slots="random", dtype=dtype, bias=True), - ) - for dtype in (torch.bfloat16, torch.float16, torch.float32) - ) - + ( - Cell( - why="x contiguous [batch, dim, seqlen] rather than token-major", - spec=dict( - dim=512, batch=3, seqlen=4, pool_slots=16, slots=(2, 5, 1), x_form="contiguous" - ), - ), - Cell( - why="launch_dependent_kernels=True (PDL trigger)", - spec=dict( - dim=512, batch=4, pool_slots=16, slots="random", launch_dependent_kernels=True - ), + # seqlen 8: max_draft_len 7, the widest verify step. + for seqlen in (3, 8) + ) + ( + Cell( + why="seqlen 5 into an 8-step snapshot buffer: steps 5..7 untouched", + spec=dict( + dim=CKPT_DIM, + batch=8, + seqlen=5, + pool_slots=ENGINE_SLOTS, + slots=_pinned(8), + snapshots=ENGINE_SNAPSHOT_ROWS, + snapshot_steps=8, ), - ) + ), ) note = """ Metadata consumed: none. Call site: the GDN speculative verify step -- `x` the `[n, 10240, p]` - transposed view of `[n, p, 10240]` (p = runtime_draft_len + 1), the layer's + transposed view of `[n, p, 10240]` (p = runtime_draft_len + 1, 2..8 here), the layer's whole conv pool `[slots, 10240, 3]`, W 4, no bias, `"silu"`, the verify rows' slots, `pad_slot_id` -1, snapshots into `[max_num_sequences, p, 10240, 3]` with `intermediate_state_indices = arange(n)`, `num_accepted_tokens` None. The op advances the pool by all `p` tokens whatever is accepted later; the - snapshots are what a rollback has to read. Certified at 128 and 461 pool - slots (bitwise identical for the same slot content), snapshot buffers of - `batch` and of 32 rows. + snapshots are what a rollback has to read. Certified at the 461-slot pool + with a 32-row snapshot buffer and `x` as the target's q|k|v column view + 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). Raise loudly, so not guarded: width 1 (triton CompilationError); a missing `intermediate_state_indices` with snapshots and slot indices given. Not guarded (content): duplicate slots or duplicate snapshot rows race; - `num_accepted_tokens[b]` must lie in `[1, seqlen]` (certified 1..3 at - seqlen 3). Out-of-range slot values: illegal access, not probed. + `num_accepted_tokens[b]` must lie in `[1, seqlen]` (the accepted-token + roll is uncertified; the target passes None). Out-of-range slot values: + illegal access, not probed. - Accepted and inert: `launch_dependent_kernels` (bitwise identical either - way). Accepted but not exposed: `cache_seqlens` (ignored by the kernel -- + Accepted, uncertified: `launch_dependent_kernels` (the target passes + False). 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. 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py index d397c03bb956..14413d38be2b 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py @@ -1,15 +1,27 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Single-step causal depthwise conv1d against a slot-indexed conv-state pool.""" +"""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`. +""" from __future__ import annotations from typing import Optional +import cutlass +import cutlass.cute as cute import torch +from cutlass._mlir.dialects import llvm +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op 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 +from tensorrt_llm._torch.kernels.blackwell.utils import make_ptr #: This checkpoint's GDN `[q|k|v]` width; the conv is width 4 over a 3-column pool. CKPT_DIM = 10240 @@ -19,10 +31,285 @@ ENGINE_SLOTS = 461 +# -------------------------------------------------------------------------------------- +# 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 +# (`__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). +_CUTE_MIN_ROWS = 64 +_CH = 8 +_THREADS = 128 + + +def _iv(x, loc, ip): + return x.ir_value(loc=loc, ip=ip) + + +def _asm(ret, args, text, constraints, loc, ip, side_effects=False): + return llvm.inline_asm( + ret, + [_iv(a, loc, ip) for a in args], + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _asm4(args, text, constraints, loc, ip, side_effects=False): + ty = T.i32() + res = _asm( + llvm.StructType.get_literal([ty] * 4), args, text, constraints, loc, ip, side_effects + ) + return [cutlass.Int32(llvm.extractvalue(ty, res, [i], loc=loc, ip=ip)) for i in range(4)] + + +@dsl_user_op +def _ldg_v4(addr, *, loc=None, ip=None): + """16 bytes; volatile so that no load moves across the pool's in-place update.""" + return _asm4( + [cutlass.Int64(addr)], + "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _ldg_nc_v4(addr, *, loc=None, ip=None): + """16 bytes of a read-only operand (the weights), non-coherent path.""" + return _asm4( + [cutlass.Int64(addr)], + "ld.global.nc.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + loc, + ip, + ) + + +@dsl_user_op +def _stg_v4(addr, a, b, c, d, *, loc=None, ip=None): + _asm( + T.i32(), + [cutlass.Int64(addr)] + [cutlass.Int32(v) for v in (a, b, c, d)], + "{ st.global.v4.b32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }", + "=r,l,r,r,r,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _bf16_lo(w, *, loc=None, ip=None): + """The low bf16 of a word as fp32 (exact).""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int32(w)], + "{ .reg .b32 t; shl.b32 t, $1, 16; mov.b32 $0, t; }", + "=f,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_hi(w, *, loc=None, ip=None): + """The high bf16 of a word as fp32 (exact).""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int32(w)], + "{ .reg .b32 t; and.b32 t, $1, 0xFFFF0000; mov.b32 $0, t; }", + "=f,r", + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16x2(lo, hi, *, loc=None, ip=None): + """`cvt.rn.bf16x2.f32`: `lo` in the low half (exact for values already in bf16).""" + return cutlass.Int32( + _asm( + T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _conv_silu(w0, w1, w2, w3, s0, s1, s2, x, *, loc=None, ip=None): + """`__fmaf_rn(w[i], s[i], v)` over the four taps from 0, then `v * __frcp_rn(1 + + __expf(-v))` (`__expf` is `ex2.approx.ftz` of the `mul.ftz` by log2e).""" + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(v) for v in (w0, w1, w2, w3, s0, s1, s2, x)], + "{ .reg .f32 v, t, e; fma.rn.f32 v, $1, $5, 0f00000000; fma.rn.f32 v, $2, $6, v; " + "fma.rn.f32 v, $3, $7, v; fma.rn.f32 v, $4, $8, v; neg.f32 t, v; " + "mul.ftz.f32 t, t, 0f3FB8AA3B; ex2.approx.ftz.f32 e, t; add.f32 e, e, 0f3F800000; " + "rcp.rn.f32 e, e; mul.rn.f32 $0, v, e; }", + "=f,f,f,f,f,f,f,f,f", + loc, + ip, + ) + ) + + +class _DecodeConvKernel: + """One token a row: `x[row, c] = silu(conv)` in place and the row's slot shifted.""" + + def __init__(self, dim: int, x_stride: int): + assert dim % (_CH * _THREADS) == 0 + self.dim, self.sx = dim, x_stride + + @cute.jit + def __call__( + self, + x_t: cute.Tensor, + w_ptr: cute.Pointer, + st_ptr: cute.Pointer, + slot_ptr: cute.Pointer, + rows: cutlass.Int32, + slot_stride: cutlass.Int32, + pad_slot_id: cutlass.Int32, + stream, + ): + slots = cute.make_tensor(slot_ptr, cute.make_layout((rows,), stride=(1,))) + # `x` arrives as a tensor: the TVM-FFI entry takes its stream from a GPU tensor. + self.kernel( + x_t.iterator.toint(), w_ptr.toint(), st_ptr.toint(), slots, slot_stride, pad_slot_id + ).launch( + grid=[rows, self.dim // (_CH * _THREADS), 1], block=[_THREADS, 1, 1], stream=stream + ) + + @cute.kernel + def kernel( + self, + x_base: cutlass.Int64, + w_base: cutlass.Int64, + st_base: cutlass.Int64, + mSlot: cute.Tensor, + slot_stride: cutlass.Int32, + pad_slot_id: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + row, cb, _ = cute.arch.block_idx() + c0 = (cb * _THREADS + tidx) * _CH + # The weights: 8 channels x 4 taps, 32 contiguous bf16. + wv = [] + for q in cutlass.range_constexpr(4): + wv += _ldg_nc_v4(w_base + cutlass.Int64(c0 * 4 + q * 8) * 2) + slot = mSlot[row] + if slot != pad_slot_id: + sb = st_base + (cutlass.Int64(slot) * slot_stride + c0 * 3) * 2 + sv = [] + for q in cutlass.range_constexpr(3): + sv += _ldg_v4(sb + q * 16) + xa = x_base + (cutlass.Int64(row) * self.sx + c0) * 2 + xv = _ldg_v4(xa) + + def elem(words, e): + return _bf16_lo(words[e // 2]) if e % 2 == 0 else _bf16_hi(words[e // 2]) + + outs, shifted = [], [] + for ch in cutlass.range_constexpr(_CH): + w = [elem(wv, ch * 4 + k) for k in range(4)] + st = [elem(sv, ch * 3 + k) for k in range(3)] + xc = elem(xv, ch) + outs.append(_conv_silu(*w, *st, xc)) + shifted += [st[1], st[2], xc] + _stg_v4(xa, *[_bf16x2(outs[2 * i], outs[2 * i + 1]) for i in range(4)]) + for q in cutlass.range_constexpr(3): + _stg_v4( + sb + q * 16, + *[ + _bf16x2(shifted[8 * q + 2 * i], shifted[8 * q + 2 * i + 1]) + for i in range(4) + ], + ) + + +_DECODE_CONV: dict = {} + + +def _decode_conv_ok(x, conv_state, weight, bias, silu_activation, cache_seqlens, indices) -> bool: + """The decode shape this kernel certifies; anything else runs the TRT-LLM op.""" + return ( + x.dim() == 2 + and x.shape[0] >= _CUTE_MIN_ROWS + and bias is None + and silu_activation + and cache_seqlens is None + and indices is not None + and indices.dtype == torch.int32 + and indices.is_contiguous() + and indices.shape == (x.shape[0],) + and x.dtype == conv_state.dtype == weight.dtype == torch.bfloat16 + and weight.shape == (x.shape[1], 4) + and weight.is_contiguous() + and conv_state.dim() == 3 + and conv_state.shape[1:] == (x.shape[1], 3) + and conv_state.stride()[1:] == (3, 1) + and x.shape[1] % (_CH * _THREADS) == 0 + and x.stride(1) == 1 + and (x.stride(0) * 2) % 16 == 0 + and (conv_state.stride(0) * 2) % 16 == 0 + and x.data_ptr() % 16 == 0 + and conv_state.data_ptr() % 16 == 0 + and weight.data_ptr() % 16 == 0 + ) + + +def _decode_conv(x, conv_state, weight, indices, pad_slot_id) -> None: + dim = x.shape[1] + key = (dim, x.stride(0)) + args = (x, weight.data_ptr(), conv_state.data_ptr(), indices.data_ptr()) + fn = _DECODE_CONV.get(key) + if fn is None: + gm = cute.AddressSpace.gmem + cargs = ( + from_dlpack(x, assumed_align=16).mark_layout_dynamic(leading_dim=1), + make_ptr(cutlass.BFloat16, args[1], gm, assumed_align=16), + make_ptr(cutlass.BFloat16, args[2], gm, assumed_align=16), + make_ptr(cutlass.Int32, args[3], gm, assumed_align=4), + ) + fn = cute.compile( + _DecodeConvKernel(dim, x.stride(0)), + *cargs, + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Int32(-1), + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--opt-level 2 --enable-tvm-ffi", + ) + _DECODE_CONV[key] = fn + fn(*args, x.shape[0], conv_state.stride(0), pad_slot_id) + + def _pinned(batch: int) -> tuple[int, ...]: """`batch` distinct slots of the engine pool, slot 0 and the last pinned in.""" - if batch == 1: - return (0,) return (0,) + tuple(i * 14 + 3 for i in range(1, batch - 1)) + (ENGINE_SLOTS - 1,) @@ -49,147 +336,56 @@ class CausalConv1dUpdate(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec keys: dim, width (4), state_len (3), seqlen (1), batch, pool_slots, - # slots ("random" | tuple | None), dtype, bias, silu, cache_seqlens, - # pad_slot_id, x_form ("3d" | "2d" | "2d_row_slice" | "token_major"), - # strided_state. + # 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. CELLS: tuple[Cell, ...] = ( - tuple( - Cell( - why=f"decode batch {batch}, call-site 2-D x, 461-slot pool", - spec=dict( - dim=CKPT_DIM, - batch=batch, - pool_slots=ENGINE_SLOTS, - slots=_pinned(batch), - x_form="2d", - ), - ) - for batch in (1, 2, 17, 32) - ) - + ( - Cell( - why="mixed batch: 2-D x is the decode-row slice of the projection", - spec=dict( - dim=CKPT_DIM, - batch=8, - pool_slots=ENGINE_SLOTS, - slots=_pinned(8), - x_form="2d_row_slice", - ), - ), - ) - + tuple( - Cell( - why=f"decode batch {batch}, [batch, dim, 1] x, 128-slot pool", - spec=dict(dim=CKPT_DIM, batch=batch, pool_slots=128, slots="random"), - ) - for batch in (1, 2, 17, 64) - ) - + ( - Cell( - why="pad_slot_id -1 rows: x rows and the pool untouched", - spec=dict(dim=CKPT_DIM, batch=6, pool_slots=128, slots=(3, -1, 7, 11, -1, 2)), - ), - Cell( - why="pad_slot_id 9: the sentinel is compared by value", - spec=dict(dim=CKPT_DIM, batch=4, pool_slots=128, slots=(5, 9, 0, 9), pad_slot_id=9), - ), - ) - + tuple( - Cell( - why=f"bias={bias}, silu={silu} (general kernel unless silu)", - spec=dict( - dim=CKPT_DIM, batch=4, pool_slots=128, slots="random", bias=bias, silu=silu - ), - ) - for bias in (False, True) - for silu in (True, False) - ) - + tuple( - Cell( - why=f"state_len {state_len} > W-1: the whole state shifts left", - spec=dict(dim=512, batch=3, pool_slots=16, slots="random", state_len=state_len), - ) - for state_len in (4, 5, 8) - ) - + tuple( - Cell( - why=f"multi-token step: seqlen {seqlen}, state_len {state_len}", - spec=dict( - dim=512, - batch=3, - pool_slots=16, - slots="random", - seqlen=seqlen, - state_len=state_len, - ), - ) - for seqlen in (2, 3, 5) - for state_len in (3, 5) - ) - + tuple( - Cell( - why=f"circular ring, state_len {state_len}, seqlen {seqlen}, a pad row", - spec=dict( - dim=512, - batch=4, - pool_slots=16, - slots=(3, 7, -1, 11), - seqlen=seqlen, - state_len=state_len, - cache_seqlens=(0, 1, state_len - 1, 37), - ), - ) - for state_len in (4, 8) - for seqlen in (1, 3) - ) - + tuple( - Cell( - why=f"{str(dtype).removeprefix('torch.')} with bias", - spec=dict(dim=512, batch=4, pool_slots=16, slots="random", dtype=dtype, bias=True), - ) - for dtype in (torch.bfloat16, torch.float16, torch.float32) - ) - + ( - Cell( - why="conv_state_indices None: row b uses slot b, batch-sized pool", - spec=dict(dim=512, batch=4, pool_slots=4, slots=None), - ), - Cell( - why="x a [batch, tokens, dim].transpose(1, 2) view", - spec=dict( - dim=512, batch=3, seqlen=4, pool_slots=16, slots=(2, 5, 1), x_form="token_major" - ), + Cell( + why="x the q|k|v column view of the [rows, 16480] in_proj, batch 32", + spec=dict(dim=CKPT_DIM, batch=32, pool_slots=ENGINE_SLOTS, slots=_pinned(32)), + ), + Cell( + why="x the q|k|v column view of the [rows, 16480] in_proj, batch 128", + spec=dict(dim=CKPT_DIM, batch=128, pool_slots=256, slots="random"), + ), + Cell( + why="64 rows (the CuTe kernel's first) over a layer view, pad -1 rows untouched", + spec=dict( + dim=CKPT_DIM, + batch=64, + pool_slots=ENGINE_SLOTS, + slots=tuple(-1 if i % 21 == 5 else i * 7 for i in range(64)), + pool_layers=2, ), - Cell( - why="conv_state a [:, :, ::2] view: interleaved columns untouched", - spec=dict(dim=512, batch=3, pool_slots=16, slots=(2, 5, 1), strided_state=True), - ), - ) - + tuple( - Cell( - why=f"dim {dim}, off the 128-channel block", - spec=dict(dim=dim, batch=2, pool_slots=8, slots="random"), - ) - for dim in (1, 127, 128, 129, 513) - ) + ), ) note = """ Metadata consumed: none. - Call site: the GDN decode step -- 2-D `[num_decodes, 10240]` x (the decode - rows of the projection), the layer's whole conv pool `[slots, 10240, 3]`, - W 4, no bias, SiLU on, `cache_seqlens` None, int32 slots, `pad_slot_id` -1. - That exact shape (silu, no ring, `state_len == W - 1`) runs a dedicated - seqlen-1 kernel; every other cell runs the general one. Certified at 128 - and 461 pool slots; at 461 the result is bitwise identical to 128 for the - same slot content. - - Everything metadata the op checks itself, loudly: width 2..4, `state_len - >= W - 1`, one dtype across x/conv_state/weight (bias matches weight), - weight shape, int32 contiguous length-batch `conv_state_indices`, int32 + Call site: the GDN decode step -- 2-D `[num_decodes, 10240]` x, the q|k|v + column view (row stride 16480) of the decode rows of the layer's one + `[T, 16480]` input-projection output (the op writes those columns only); + the layer's whole conv pool `[slots, 10240, 3]`, W 4, no bias, SiLU on, + `cache_seqlens` None, int32 slots, `pad_slot_id` -1. + That exact shape (silu, no ring, `state_len == W - 1`) runs TRT-LLM's + dedicated seqlen-1 kernel; any other shape runs its general one, which the + target never reaches and which is not certified. From `_CUTE_MIN_ROWS` + (64) rows, with bf16 operands, a contiguous weight, a pool whose slot is + `[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. + + 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, + int32 contiguous length-batch `conv_state_indices`, int32 `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 @@ -197,9 +393,9 @@ class CausalConv1dUpdate(OpWrapper): 10 runs while both outputs matched. 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` is skipped: its x rows and the pool - stay bitwise unchanged (linear and ring mode). Any value works as the - sentinel, including one that collides with a live slot. + 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. @@ -216,6 +412,11 @@ def __call__( conv_state_indices: Optional[torch.Tensor], pad_slot_id: int, ) -> torch.Tensor: + if _decode_conv_ok( + x, conv_state, weight, bias, silu_activation, cache_seqlens, conv_state_indices + ): + _decode_conv(x, conv_state, weight, conv_state_indices, pad_slot_id) + return x torch.ops.trtllm.causal_conv1d_update( x.unsqueeze(-1) if x.dim() == 2 else x, conv_state, 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 index 703afea315b0..701fa42c82ed 100644 --- 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 @@ -27,10 +27,9 @@ #: This checkpoint's Gated DeltaNet geometry at tp1. NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM = 16, 48, 128, 128 -#: Largest slot count the cache manager was seen to allocate the recurrent pool -#: at for this checkpoint (10, 35, 366, 368, 396, 399, 432, 461 over twelve -#: runtime-value dumps; the batch never above 32). The target passes the -#: layer's whole pool unsliced. +#: 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 @@ -74,75 +73,24 @@ class ChunkGatedDeltaRule(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: seqlens (packed with cu_seqlens) or batch/num_tokens (no - # cu_seqlens); initial: "indexed" (a `slots`-row pool, `indices` or arange) - # | "positional" | None; output: None | "buffer" | "window" (wide[:, 50:146] - # of a 200-token buffer); output_final_state; scale; normalize_qk (False - # with l2norm feeds raw q/k). + # 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, ...] = ( - tuple( - Cell( - why=f"target prefill call, 461-slot pool, seqlens {list(s)}", - spec=dict( - seqlens=s, initial="indexed", slots=ENGINE_SLOTS, indices=_engine_slots(len(s)) - ), - ) - for s in ((1,), (64,), (200, 37, 91), (1024,)) - ) - + ( - Cell( - why="461-slot pool, output= buffer (the stock mixer's call)", - spec=dict( - seqlens=(200, 37, 91), - initial="indexed", - slots=ENGINE_SLOTS, - indices=_engine_slots(3), - output="buffer", - ), - ), - ) - + tuple( - Cell( - why=f"indexed, seqlens {list(s)}, {len(s) + 3}-slot pool, output= buffer", - spec=dict(seqlens=s, initial="indexed", slots=len(s) + 3, output="buffer"), - ) - for s in ((1,), (8,), (64,), (137,), (128, 64), (200, 37, 91), (8, 8, 8), (1024,)) - ) - + tuple( - Cell( - why=f"positional initial state, final state out, {list(s)}", - spec=dict(seqlens=s, initial="positional", output_final_state=True), - ) - for s in ((96, 48), (333,)) - ) - + ( - Cell( - why="initial_state None: zero start", - spec=dict(seqlens=(77,), initial=None, output_final_state=True), - ), - Cell( - why="use_qk_l2norm_in_kernel on raw q/k", - spec=dict( - seqlens=(100,), - initial=None, - output_final_state=True, - l2norm=True, - normalize_qk=False, - ), + 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 ), - Cell( - why="no cu_seqlens: two equal-length rows", - spec=dict(batch=2, num_tokens=96, initial="positional", output_final_state=True), - ), - Cell( - why="explicit scale 0.05 replaces the default", - spec=dict(seqlens=(80,), initial=None, scale=0.05), - ), - Cell( - why="output= a token-window slice of a wider buffer", - spec=dict(seqlens=(96,), initial="positional", output="window"), - ), - ) + ), ) note = """ @@ -150,17 +98,23 @@ class ChunkGatedDeltaRule(OpWrapper): 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 fp32 pool `[slots, 48, 128, 128]`, - int32 slots, int64 `cu_seqlens`; the MTP-verify context half passes the - `si_p` / `qsl_long_p` prefixes. Certified at small pools and at 461 slots; - at 461 output and state are bitwise identical to a 6-slot pool with the - same slot content. + 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`). Certified - in a fresh interpreter at this geometry and at 5 / 20 heads. The positional path - never aliased its state. + 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 @@ -170,12 +124,16 @@ class ChunkGatedDeltaRule(OpWrapper): 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: 1..1024, off and on the 64 boundary, - mixed ragged, all correct. Memory: a `[B, NT, H, K, V]` per-chunk state - scratch in `q.dtype`, 24 MiB per 1024 tokens here. Deterministic. + 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`/ - pool, `K = 256` (accepted, unmeasured), empty `T` or zero-length spans. + 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) 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/cu_seqlens_from_seq_lens.py index 6a10c7d409cd..904246678bc0 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/cu_seqlens_from_seq_lens.py @@ -20,26 +20,16 @@ #: bool and floating dtypes and answers wrongly with them. _INTEGER_DTYPES = (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64) -#: Batch shapes this checkpoint's serving path produces. ``p`` is the trunk's -#: query rows per generation request: 1 unspeculated, ``1 + runtime_draft_len`` -#: under MTP (2 at this checkpoint's one MTP layer). +#: 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. BATCHES = { - "single_prefill": (4096,), - "decode_p1": (1,) * 32, - "decode_p2": (2,) * 32, - "mixed_p1": (1024, 37, 1, 1, 1, 1), "mixed_p2": (1024, 37, 2, 2, 2, 2), - "chunked_prefill_p2": (512, 512, 96, 2, 2, 2, 2, 2, 2, 2), "padded_decode_p2": (2, 2, 2, 0, 0, 0), - "one_sequence": (7,), - "max_batch": (2,) * 256, } -def _name(dtype: torch.dtype) -> str: - return str(dtype).removeprefix("torch.") - - class CuSeqlensFromSeqLens(OpWrapper): """`out[0] = 0`, `out[s] = seq_lens[0] + ... + seq_lens[s - 1]`: a fresh contiguous `[num_seqs + 1]` tensor of `dtype` on `seq_lens.device`. @@ -52,54 +42,31 @@ class CuSeqlensFromSeqLens(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: lengths, in_dtype, out_dtype; `layout` "contiguous" (default), + # spec: lengths, in_dtype, out_dtype; `layout` "contiguous" (default) or # "prefix" (the first len(lengths) of a longer buffer -- the call site's - # `split(seq_lens_cuda, [batch, ...])[0]`), or "strided" (a [::2] view). + # `seq_lens_cuda[:batch]`). CELLS: tuple[Cell, ...] = ( - tuple( - Cell( - why=f"{name} {_name(i)}->{_name(o)}", - spec=dict(lengths=lengths, in_dtype=i, out_dtype=o), - ) - for name, lengths in BATCHES.items() - for i in (torch.int32, torch.int64) - for o in (torch.int32, torch.int64) - ) - + ( - Cell( - why="call site: prefix of seq_lens_cuda, int32->int64", - spec=dict( - lengths=BATCHES["mixed_p2"], - in_dtype=torch.int32, - out_dtype=torch.int64, - layout="prefix", - ), - ), - Cell( - why="non-contiguous [::2] seq_lens view", - spec=dict( - lengths=(7, 1, 13), - in_dtype=torch.int32, - out_dtype=torch.int32, - layout="strided", - ), + Cell( + why="call site: prefix of seq_lens_cuda, int32->int64", + spec=dict( + lengths=BATCHES["mixed_p2"], + in_dtype=torch.int32, + out_dtype=torch.int64, + layout="prefix", ), - Cell( - why="zero-length rows give repeated offsets", - spec=dict(lengths=(0, 5, 0, 0, 3, 0), in_dtype=torch.int32, out_dtype=torch.int32), + ), + Cell( + why="MTP verify p=8, 32 requests: prefix int32->int64", + spec=dict( + lengths=(8,) * 32, in_dtype=torch.int32, out_dtype=torch.int64, layout="prefix" ), - Cell( - why="empty batch returns [0]", - spec=dict(lengths=(), in_dtype=torch.int32, out_dtype=torch.int32), + ), + Cell( + why="padded_decode_p2: zero-length padded rows, int32->int64", + spec=dict( + lengths=BATCHES["padded_decode_p2"], in_dtype=torch.int32, out_dtype=torch.int64 ), - ) - + tuple( - Cell( - why=f"narrow {_name(d)} input does not limit the accumulator", - spec=dict(lengths=(200, 200, 200), in_dtype=d, out_dtype=torch.int32), - ) - for d in (torch.uint8, torch.int16) - ) + ), ) note = """ @@ -118,9 +85,15 @@ class CuSeqlensFromSeqLens(OpWrapper): `dtype` has no default: `causal_conv1d_fwd` requires int32 and `chunk_gated_delta_rule` keys its autotuner on `cu_seqlens`' dtype. - 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. + 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. + + Not certified: a `seq_lens` dtype other than int32 (the guard admits any + integer dtype), a non-contiguous `seq_lens` view. Hard limits out[-1] accumulated in `dtype` boundary 2^31 - 1 for int32 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 index 5bf4c9257981..8ff846801270 100644 --- 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 @@ -38,9 +38,10 @@ #: 3, hence the 9-column state the mamba cache manager allocates per slot. CKPT_CHANNELS = 10240 CKPT_DILATION = 3 -#: Largest PLE-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 target hands the whole pool over unsliced. +#: 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 @@ -52,7 +53,10 @@ def _scattered(num_seqs: int) -> tuple[int, ...]: def _ckpt(why: str, lengths: tuple[int, ...], **spec) -> Cell: - return Cell(why=why, spec=dict(channels=CKPT_CHANNELS, lengths=lengths, **spec)) + return Cell( + why=why, + spec=dict(channels=CKPT_CHANNELS, lengths=lengths, pool_slots=ENGINE_SLOTS, **spec), + ) class DilatedShortConvPaged(OpWrapper): @@ -68,7 +72,8 @@ class DilatedShortConvPaged(OpWrapper): 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, not a race). + 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 @@ -82,125 +87,21 @@ class DilatedShortConvPaged(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec keys: channels, kernel_size (4), dilation (3), lengths, pool_slots (8), - # slots (None -> arange), initial (default all True; "none" -> None), - # dtype, bias, silu, pad_slot_id, x_form ("channel_major" | "token_major"), - # index_dtype (int32), qsl_dtype (int32), no_indices. + # 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("checkpoint mix 37+1+128, 8-slot pool", (37, 1, 128)), - _ckpt("checkpoint mix 1024+1, 8-slot pool", (1024, 1)), - _ckpt("checkpoint mix 2047+1+1+1, 8-slot pool", (2047, 1, 1, 1)), - ) - + tuple( - # The call site's own form: a bare .t() of the token-major activation, - # query_start_loc_long (int64), int32 slots, the whole pool. - _ckpt( - f"call-site form at the 461-slot pool, {why}", - lengths, - pool_slots=ENGINE_SLOTS, - slots=_scattered(len(lengths)), - x_form="token_major", - qsl_dtype=torch.int64, - ) - for why, lengths in ( - ("37+1+128", (37, 1, 128)), - ("1024+1", (1024, 1)), - ("2047+1+1+1", (2047, 1, 1, 1)), - ("32 one-token rows", (1,) * 32), - ) - ) - + ( - _ckpt( - "MTP verify at 461 slots: a 37-token context + 8 rows of 3 drafts", - (37,) + (3,) * 8, - pool_slots=ENGINE_SLOTS, - slots=_scattered(9), - initial=(False,) + (True,) * 8, - x_form="token_major", - qsl_dtype=torch.int64, - ), - ) - + tuple( - Cell( - why=f"decode batch of {n} one-token rows, 32-slot pool", - spec=dict(channels=256, lengths=(1,) * n, pool_slots=32), - ) - for n in (1, 2, 8, 32) - ) - + tuple( - Cell( - why=f"prefill/ragged {'+'.join(map(str, lengths))}", - spec=dict(channels=256, lengths=lengths), - ) - for lengths in ((1024,), (512, 7, 512), (1, 2, 3, 4, 5), (4096,)) - ) - + ( - Cell( - why="fresh sequences carry zeros, not the slot's contents", - spec=dict(channels=256, lengths=(5, 9, 40), initial=(False, False, False)), - ), - Cell( - why="carry flags mixed True/False/True", - spec=dict(channels=256, lengths=(5, 9, 40), initial=(True, False, True)), - ), - Cell( - why="a zero-token sequence republishes its carry", - spec=dict(channels=256, lengths=(3, 0, 11), slots=(2, 5, 1)), - ), - Cell( - why="state longer/shorter than a sequence: tail and carry mix", - spec=dict(channels=256, lengths=(14, 2), slots=(3, 1)), - ), - Cell( - why="pad_slot_id -1 row next to live rows, slot 0 unnamed", - spec=dict(channels=256, lengths=(6, 4, 17), slots=(2, -1, 3), pool_slots=6), - ), - Cell( - why="pad_slot_id -1 row while slot 0 is itself live", - spec=dict(channels=256, lengths=(6, 4, 17), slots=(0, -1, 3), pool_slots=6), - ), - Cell( - why="two sequences on one slot: the last one's state survives", - spec=dict(channels=256, lengths=(6, 9, 12), slots=(1, 3, 1), pool_slots=5), - ), - Cell(why="bias, prefill-shaped", spec=dict(channels=256, lengths=(9, 9), bias=True)), - Cell( - why="bias, decode-shaped", spec=dict(channels=256, lengths=(1, 1, 1, 1), bias=True) - ), - Cell( - why="silu off: the bare convolution, with bias", - spec=dict(channels=256, lengths=(12, 3), slots=(1, 0), bias=True, silu=False), - ), - ) - + tuple( - Cell( - why=f"{str(dtype).removeprefix('torch.')} {'+'.join(map(str, lengths))}", - spec=dict(channels=256, lengths=lengths, dtype=dtype), - ) - for dtype in (torch.float32, torch.float16) - for lengths in ((16, 1, 5), (1,) * 8) - ) - + tuple( - Cell( - why=f"kernel_size {k}, dilation {d}", - spec=dict(channels=256, lengths=(7, 1, 12), kernel_size=k, dilation=d), - ) - for k in (2, 3, 4, 5) - for d in (1, 2, 3) - ) - + ( - Cell( - why="cache_indices and has_initial_state None", - spec=dict(channels=256, lengths=(3, 4, 5), no_indices=True, initial="none"), - ), - Cell( - why="int64 query_start_loc and cache_indices", - spec=dict( - channels=256, lengths=(3, 4, 5), index_dtype=torch.int64, qsl_dtype=torch.int64 - ), - ), - ) + _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 = """ @@ -211,23 +112,28 @@ class DilatedShortConvPaged(OpWrapper): `[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 at 8, 32 - and 461 slots; at 461 bitwise identical to 8 for the same slot content. + 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). - Kernel sizes 2-5 and dilations 1-3 certified; there is no width dispatch. + Only the checkpoint's kernel size 4 and dilation 3 are certified; there is + no width dispatch. - CUDA-graph capture is certified: no host sync, all batch reads on device, - a captured graph stays correct when the batch tensors' values change. + 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, token counts above 8192, dilation - above 3, kernel above 5, `torch.compile`, a pool dtype unlike the model's. + 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__( 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 index b2f09209e898..434e48ce4b55 100644 --- 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 @@ -20,8 +20,10 @@ ) #: The GDN mixer's packed `[q|k|v]` width for this checkpoint; the target hands -#: the whole row to this op. +#: 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): @@ -36,106 +38,47 @@ class ExtractTransposePrefillSlice(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # `num_rows`/`num_cols` build `src`; `row_stop`/`col_slice` turn it into the - # view the call site passes (a row slice of a split token buffer, or a - # column slice of a wider projection). + # `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, ...] = ( - tuple( + 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, whole 10240-wide row, {n} tokens", + why=f"GDN prefill, the q|k|v column view of the in_proj, {n} tokens", spec=dict( num_rows=n, - num_cols=GDN_MIXED_QKV_WIDTH, + 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 (1, 7, 33, 128, 1024, 4096) - ) - + ( - Cell( - why="mixed batch: prefill half of a split [512, 10240] buffer, a row-slice view", - spec=dict( - num_rows=512, - num_cols=GDN_MIXED_QKV_WIDTH, - row_stop=400, - n=400, - start_col=0, - width=GDN_MIXED_QKV_WIDTH, - dtype=torch.bfloat16, - ), - ), - Cell( - why="fewer prefill rows than the buffer holds: the row crop is honoured", - spec=dict( - num_rows=512, - num_cols=GDN_MIXED_QKV_WIDTH, - n=400, - start_col=0, - width=GDN_MIXED_QKV_WIDTH, - dtype=torch.bfloat16, - ), - ), - Cell( - why="interior column range of a contiguous 16384-wide buffer", - spec=dict( - num_rows=200, - num_cols=16384, - n=150, - start_col=1024, - width=GDN_MIXED_QKV_WIDTH, - dtype=torch.bfloat16, - ), - ), - Cell( - why="column-slice view: row stride wider than the extracted width", - spec=dict( - num_rows=200, - num_cols=16384, - col_slice=(1024, 1024 + GDN_MIXED_QKV_WIDTH), - n=150, - start_col=0, - width=GDN_MIXED_QKV_WIDTH, - dtype=torch.bfloat16, - ), - ), - ) - + tuple( - Cell( - why=f"dtype {str(dtype).removeprefix('torch.')}: elements are moved, never interpreted", - spec=dict(num_rows=64, num_cols=256, n=48, start_col=32, width=128, dtype=dtype), - ) - for dtype in ( - torch.float16, - torch.float32, - torch.float64, - torch.float8_e4m3fn, - torch.int32, - torch.int8, - torch.bool, - ) - ) - + ( - Cell( - why="zero prefill tokens: an empty [width, 0] result", - spec=dict( - num_rows=64, num_cols=256, n=0, start_col=0, width=256, dtype=torch.bfloat16 - ), - ), - Cell( - why="zero width: an empty [0, n] result", - spec=dict( - num_rows=64, num_cols=256, n=1, start_col=0, width=0, 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. @@ -152,6 +95,9 @@ class ExtractTransposePrefillSlice(OpWrapper): 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__( 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/flashinfer_chunk_gated_delta_rule.py new file mode 100644 index 000000000000..e6a503da7c0f --- /dev/null +++ b/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 @@ -0,0 +1,248 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Chunked gated delta rule on FlashInfer's Blackwell kernel: the GDN prefill core. + +Wraps the in-tree module callable +``tensorrt_llm._torch.modules.fla.flashinfer_chunk.chunk_gated_delta_rule`` +(tier-2 sourcing), the adapter in front of the installed FlashInfer's +``flashinfer.chunk_gated_delta_rule`` (CuTe-DSL tcgen05 kernel on sm_100/103). +One call is: ``g -> exp(g)``, a Triton gather of the indexed pool slots into a +packed fp32 ``[N, HV, V, K]`` buffer, the FlashInfer kernel, and a Triton +scatter of the final states back into the same slots. Call-compatible with the +vendored Triton ``chunk_gated_delta_rule`` (``modules/fla/chunk``). + +FlashInfer API: ``flashinfer.chunk_gated_delta_rule`` (``flashinfer.gdn_prefill``), +reached through the adapter above. +""" + +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.flashinfer_chunk import ( + chunk_gated_delta_rule as _fi_chunk_gated_delta_rule, +) + +#: 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). +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 FlashinferChunkGatedDeltaRule(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 (value head `i` reads key head `i // (HV // H)`), 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 on tensor cores (WY form, bf16 operands) -- + the same recurrence as the Triton op, a different sum. `g` is in log space + (the adapter exponentiates it); `scale=None` is `K ** -0.5`; + `use_qk_l2norm_in_kernel` normalizes q/k first (a Triton `l2norm_fwd`). + + One packed row (`q.shape[0] == 1`) split at `cu_seqlens` (required). + `initial_state` is required: a pool `[S, HV, V, K]`; sequence `n` starts + from `initial_state[initial_state_indices[n]]`. With + `inplace_indexed_state_update=True` its final state is written back into + that slot and `final_state` is `None`; otherwise `output_final_state` + returns a fresh `[N, HV, V, K]` in the pool's dtype and the pool is not + written. + + Fusion boundary: gate exponentiation, optional q/k norm, scale, slot + gather, recurrence, output, slot scatter. The conv, q/k/v/g/beta + preparation, slot allocation and the gated norm on `o` are the caller's. + + Returns `(o, final_state)`: `o` `[1, T, HV, V]` in `q.dtype` (`output` + itself when given). + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: geometry (H, HV, K, V); seqlens; `slots`-row pool with `indices` + # (or _engine_slots); output (None | "buffer"). Every cell is + # qwen3.8-flash-next-nvfp4's geometry, in place over an fp32 pool. + CELLS: tuple[Cell, ...] = tuple( + Cell( + why=f"qwen3.8 prefill into output=, 461-slot pool, seqlens {list(s)}", + spec=dict( + geometry=QWEN38_GEOMETRY, + seqlens=s, + slots=ENGINE_SLOTS, + indices=_engine_slots(len(s)), + output="buffer", + ), + ) + for s in ((200, 37, 91), (8192,)) + ) + ( + Cell( + why="qwen3.8 context chunk + 1-token rows [8, 1, 1, 1], 461-slot pool", + spec=dict( + geometry=QWEN38_GEOMETRY, + seqlens=(8, 1, 1, 1), + slots=ENGINE_SLOTS, + indices=_engine_slots(4), + ), + ), + ) + + note = """ + Metadata consumed: none. + + Call site as qwen3.8-flash-next-nvfp4's `_gdn_block` makes it (in place of + the Triton `chunk_gated_delta_rule`): + the same positional arguments, `(q, k, v, g, beta, None, ssm_states, + 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). + + Certified at the 461-slot fp32 pool: three ragged context rows and one + 8192-token row, each into an `output=` buffer, and an 8-token context + chunk with three one-token rows. + + Memory per call: a packed fp32 copy of the named slots in and out (2 x + 3 MiB per sequence), `exp(g)`, 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. No cold-call defect: + nothing in the call is autotuned. + + Loud, from the callable: `head_first=True`, `q.shape[0] != 1`, H_q != H_k, + fp32 q, non-fp32 `g`, `cu_seqlens` or `initial_state` None, in-place + without indices. + + Not checked: indices distinct and in `[0, S)`; q/k unit-norm (by the + caller or the flag). + + Not certified: `state_workspace=`, bf16/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`. + + Hard limits + none 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, + head_first: bool = False, + use_qk_l2norm_in_kernel: bool = False, + output: Optional[torch.Tensor] = None, + state_workspace: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + return _fi_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=head_first, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + output=output, + state_workspace=state_workspace, + ) + + 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, + head_first: bool = False, + use_qk_l2norm_in_kernel: bool = False, + output: Optional[torch.Tensor] = None, + state_workspace: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + ) -> tuple[torch.Tensor, Optional[torch.Tensor], 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 the pool's dtype. `final_state` is `None` on the in-place + path, as the callable returns it. Mutates nothing; `output`, + `head_first` and `state_workspace` are ignored (they choose where + results land).""" + dt = torch.float64 + _, _, num_k_heads, head_k = q.shape + num_v_heads = v.shape[2] + repeats = num_v_heads // num_k_heads + scale = head_k**-0.5 if scale is None else scale + bounds = cu_seqlens.tolist() + slots = initial_state_indices.tolist() + out = torch.zeros(1, q.shape[1], num_v_heads, v.shape[3], dtype=dt, device=q.device) + finals = [] + for n in range(len(bounds) - 1): + state = initial_state[slots[n]].to(dt) + for t in range(bounds[n], bounds[n + 1]): + q_t, k_t = q[0, t].to(dt), k[0, 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[0, t].to(dt).exp()[:, None, None] + u = (v[0, t].to(dt) - (state * k_t[:, None, :]).sum(-1)) * beta[0, t].to(dt)[ + :, None + ] + state = state + u[:, :, None] * k_t[:, None, :] + out[0, t] = (state * q_t[:, None, :]).sum(-1) + finals.append(state) + pool = initial_state.clone() + final_state = None + if inplace_indexed_state_update: + for slot, state in zip(slots, finals): + pool[slot] = state.to(pool.dtype) + elif output_final_state: + final_state = torch.stack(finals).to(initial_state.dtype) + return out.to(q.dtype), final_state, pool + + +flashinfer_chunk_gated_delta_rule = FlashinferChunkGatedDeltaRule() 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 index e81f6f95b7b8..bc996ad6dd0d 100644 --- 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 @@ -51,90 +51,19 @@ class FusedGdnGating(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: rows (or num_decodes x draft); layout "contiguous" | "target_rows" - # (the decode rows of a contiguous [T, 48] `a` projection, past T - rows - # prefill rows) | "ba_column_half" (the stock mixer's [:, 48:] of [T, 96]) - # | "ba_row_and_column" (a row slice of that); a_dtype, param_dtype, - # num_heads, beta, threshold, a_scale, fill (constant a), overlong. + # 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, ...] = ( - tuple( - Cell( - why=f"target verify call: {b} seqs x {p} tokens, contiguous a", - spec=dict(rows=b * p), - ) - for b, p in ((1, 2), (4, 2), (8, 2), (8, 4), (32, 2)) - ) - + ( - Cell( - why="target mixed verify: decode rows of a [T, 48] projection", - spec=dict(rows=16, layout="target_rows"), - ), - Cell( - why="stock mixer: a = ba[:, 48:], row stride 96", - spec=dict(rows=16, layout="ba_column_half"), - ), - Cell( - why="stock mixer mixed batch: ba[48:, 48:], offset row slice", - spec=dict(rows=16, layout="ba_row_and_column"), - ), - ) - + tuple( - Cell(why=f"{n} token rows{note}", spec=dict(rows=n)) - for n, note in ( - (1, ""), - (37, ""), - (512, ""), - (4096, ": prefill-sized"), - ) - ) - + tuple( - Cell( - why=f"a in {str(d).removeprefix('torch.')}", - spec=dict(rows=64, a_dtype=d), - ) - for d in (torch.float16, torch.float32) - ) - + tuple( - Cell( - why=f"A_log/dt_bias in {str(d).removeprefix('torch.')}", - spec=dict(rows=64, param_dtype=d), - ) - for d in (torch.bfloat16, torch.float16) - ) - + tuple( - Cell( - why=f"{h} heads: the 8-wide head block's masked tail", - spec=dict(rows=7, num_heads=h), - ) - for h in (1, 5, 8, 12, 100) - ) - + tuple( - Cell( - why=f"softplus beta {beta}, threshold {threshold:g}{note}", - spec=dict(rows=96, beta=beta, threshold=threshold, a_scale=4.0), - ) - for beta, threshold, note in ( - (1.0, 20.0, ""), - (0.5, 20.0, ""), - (2.0, 20.0, ""), - (1.0, 0.0, ": nearly all on the linear branch"), - (1.0, 1e9, ": all on the log branch"), - ) - ) - + tuple( - Cell( - why=f"a = {value:g}: linear branch; exp(x) is evaluated, inf never leaks", - spec=dict(rows=4, fill=value, a_dtype=torch.float32), - ) - for value in (50.0, 88.0, 200.0, 1e4) - ) - + ( - Cell( - why="A_log/dt_bias longer than num_heads: tail never read", - spec=dict(rows=32, overlong=16), - ), - Cell(why="zero rows: an empty [0, 48] fp32 tensor", spec=dict(rows=0)), - ) + 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 = """ @@ -143,15 +72,18 @@ class FusedGdnGating(OpWrapper): 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`. + `[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 free (consumed in place, bitwise equal to a packed - copy); everything else is addressed without strides. `is_valid` guards the - three violations measured to return a finite wrong `g`: a stride-2 A_log + `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 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 index b97bbdec86f4..e34c3fdb77f1 100644 --- 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 @@ -62,58 +62,22 @@ class FusedGdnPostConv(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: num_prefill, num_decode; layout "contiguous" | "target_mixed" - # (decode a row slice of a [T, 10240] projection, as the target's split - # makes it) | "production" (the stock mixer's views: decode a row slice of - # a column slice of [T, 10240 + 6144], a/b column halves of [T, 96]) | - # "prefill_view" (prefill a `.t()` of a wider buffer, stride (1, 16384)); - # beta_dtype, gate_dtype (A_log/dt_bias), dt_bias_fill. + # 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, ...] = ( - tuple( - Cell(why=f"prefill only, {n} tokens", spec=dict(num_prefill=n, num_decode=0)) - for n in (1, 37, 129, 2048) - ) - + tuple( + *( Cell( - why=f"mixed batch, {p} prefill + {d} decode tokens", + why=f"target in_proj views, {p} prefill + {d} decode tokens", spec=dict(num_prefill=p, num_decode=d), ) - for p, d in ((128, 16), (37, 8), (1, 1), (0, 8)) - ) - + ( - Cell( - why="target mixed call: decode a row slice of the projection", - spec=dict(num_prefill=200, num_decode=32, layout="target_mixed"), - ), - Cell( - why="stock mixer views: strided decode, a/b column halves", - spec=dict(num_prefill=96, num_decode=12, layout="production"), - ), - Cell( - why="prefill a non-contiguous .t() view, read through both strides", - spec=dict(num_prefill=96, num_decode=0, layout="prefill_view"), - ), - Cell( - why="beta_dtype bf16 (the MTP verify-path convention)", - spec=dict(num_prefill=37, num_decode=0, beta_dtype=torch.bfloat16), - ), - Cell( - why="A_log/dt_bias bf16, as stored in the checkpoint", - spec=dict(num_prefill=64, num_decode=8, gate_dtype=torch.bfloat16), - ), - Cell( - why="every head below softplus_threshold", - spec=dict(num_prefill=64, num_decode=0, dt_bias_fill=SOFTPLUS_THRESHOLD - 8.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), - ), - Cell( - why="zero tokens: five empty tensors, no launch", - spec=dict(num_prefill=0, num_decode=0), - ), - ) + 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 = """ @@ -124,13 +88,16 @@ class FusedGdnPostConv(OpWrapper): `(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, `a`/`b` the - `[T, 48]` projections. The stock mixer's verify call passes - `beta_dtype=b.dtype`. + `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 production + 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. @@ -148,7 +115,10 @@ class FusedGdnPostConv(OpWrapper): (`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. + 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. 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py new file mode 100644 index 000000000000..4e6c81fe6fe5 --- /dev/null +++ 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_cached_replay_update.py @@ -0,0 +1,512 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gated DeltaNet speculative-verify recurrence under the engine's cached-replay +state update: a checkpoint state per slot plus a compact history of the tokens +accepted since it, instead of a per-step state buffer. + +Wraps the in-tree module callable +``tensorrt_llm._torch.modules.fla.cached_replay.fused_recurrent_gated_delta_rule_cached_replay_update`` +(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. +""" + +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 +from tensorrt_llm._torch.modules.fla.cached_replay import ( + fused_recurrent_gated_delta_rule_cached_replay_update as _cached_replay_update, +) + +#: This checkpoint's GDN geometry at tp1. +NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM = 16, 48, 128, 128 +#: `[q | k | v]` row width of the post-conv activations a verify step packs. +PACKED_QKV_WIDTH = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240 +#: The cache manager's `replay_history_size`: max(MIN_REPLAY_HISTORY_SIZE=16, T). +HISTORY_SIZE = 16 +#: From this batch size on the engine runs the partitioned replay (work items, no +#: in-kernel commit, one all-layer commit after every layer has run): +#: `CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE` in `modules/fla/cached_replay.py`. +ALL_LAYER_COMMIT_MIN_BATCH = 16 +#: The recurrent pool's slot count of a full-size engine run. +ENGINE_SLOTS = 461 +#: 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 + + +def _cell_spec(path, n, t, **spec) -> dict: + return dict(path=path, n=n, t=t, **spec) + + +def _bf16_cell(why, path, n, t, **spec) -> Cell: + return Cell( + why=f"bf16 state {why}", spec=_cell_spec(path, n, t, state_dtype=torch.bfloat16, **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`), + with `p = prev_num_accepted_tokens[slot]`, `b = cache_buf_idx[slot]`: + + 1. the state the step starts from is rebuilt from the checkpoint + `h0 = ssm_states[slot]` (`[HV, V, K]`, K innermost) and the `p` pending + history entries in half `b` -- update vectors `u_j = old_u[slot, b, j]`, + normalized keys `k_j = old_k[slot, b, j]`, cumulative log-decays + `G_j = old_G[slot, b, :, j]` measured from the checkpoint: + `h = exp(G_{p-1}) h0 + sum_j exp(G_{p-1} - G_j) u_j k_j^T`; + 2. the `T` new tokens run the gated delta rule in order (q/k head + `hv // (HV // H)`, so a v-head group shares one q/k head): + `h *= exp(g_t); u_t = beta_t (v_t - h k_t); h += u_t k_t^T; o_t = h (scale q_t)`, + q and k L2-normalized first (`x / sqrt(sum(x^2) + 1e-6)`) when + `use_qk_l2norm_in_kernel`, `scale=None` meaning `K ** -0.5`; with + `A_log` given, `g = -exp(A_log) * softplus(g + dt_bias)` and + `beta = sigmoid(beta)` are applied to the raw projections inside; + 3. the new tokens' `(u_t, k_t, G_t)` are appended to the history. When + `p + T <= history_size` they go to half `b` at `p .. p+T-1` with `G` + still cumulative from the checkpoint. Otherwise (a *wrap*) they go to + half `1 - b` at `0 .. T-1` with `G` restarted from zero, and the + checkpoint must advance to the state after the `p` pending tokens: + without work items the kernel writes that state into `ssm_states[slot]` + itself; with `replay_work_items` (`use_all_layer_commit=True`) it does + not, and the caller owes one `commit_gdn_cached_replay_history_layers` + over every layer before the next step. + + Nothing about acceptance happens here: the caller advances `p`/`b` after + sampling (`p <- accepted` and `b <- 1 - b` on a wrap, else `p += accepted`, + `accepted` counting the bonus token), so rejected tokens' history entries + are simply never read. + + Fusion boundary: L2 norms, scale, the gating of the raw `a`/`b` + projections, the history rebuild, the recurrence, the history append and + (small batches) the checkpoint commit are inside. The conv, the projections, + the all-layer commit and the `p`/`b` bookkeeping are not. + + Returns `[N, T, HV, V]` in `output`'s dtype, else `q`'s; it *is* `output` + when one is given. `ssm_states`, `old_u`, `old_k`, `old_G` are written in + place and not returned. Rows with a negative slot are skipped entirely: + nothing is read or written for them, their output rows included. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: path ("small": no work items, the kernel commits a wrap itself | + # "all_layer": work items + use_all_layer_commit, the test then runs the + # manager's all-layer commit), n sequences x t tokens per verify step, + # steps (consecutive verify steps, each followed by an acceptance per row), + # accept ("random": 1..t per row, the default | "full": t every step, so from + # a fresh slot at t=8 the second step appends at `pnat + t == 16` without + # wrapping and every second step after wraps), + # state_dtype (the checkpoint pool's dtype: fp32 | bf16, as + # `mamba_ssm_cache_dtype` allocates it; fp32 by default), + # layers (each with its own activations, gates + # and history; one V2-style pool interleaving the layers per slot), slots + # (the pool's slot count), start ("warm": random pending histories in both + # halves, a wrap and an append forced at step 0 | "fresh": pnat 0 and + # 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`, + # 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. + CELLS: tuple[Cell, ...] = ( + Cell( + why="all_layer n=32 T=8 full accept: commit every second step", + spec=_cell_spec("all_layer", 32, 8, start="fresh", accept="full", steps=8), + ), + _bf16_cell( + "small n=4 T=8 full accept: pnat+T==16 appends, then wraps", + "small", + 4, + 8, + start="fresh", + accept="full", + steps=8, + ), + Cell( + why="small n=8 T=3 dup: rows 5..7 share one padding slot", + spec=_cell_spec("small", 8, 3, dup=(5, 6, 7), steps=6), + ), + Cell( + why="all_layer n=16 T=3 L=8 S=461: slot 460's offset is past 2^31 elements", + spec=_cell_spec( + "all_layer", + 16, + 3, + layers=PAST_INT32_LAYERS, + slots=ENGINE_SLOTS, + edge_slots=True, + steps=3, + ), + ), + ) + + note = """ + Call sites. The engine's GDN mixer (`modules/mamba/gdn_mixer.py`) is the + reference caller: the target-verify step of a generation-only batch, and the + decode rows of a mixed context + verify batch -- `N = num_decodes`, `T = + runtime_draft_len + 1` (2..8 here: max_draft_len 1..7). A plain decode step never reaches it. q/k/v + are the `[N, T, 16|48, 128]` views of the post-conv `[N*T, 10240]` rows, plus + `packed_qkv` = those rows for a generation-only batch (a mixed batch passes the + views alone and the op copies them contiguous); `g=a`, `beta=b` are the raw + bf16 `[N, T, 48]` projections, with the layer's fp32 `[48]` `A_log` and + `dt_bias`; `ssm_states` is the layer's pool view `[S, 48, 128, 128]`, fp32 or bf16 as + `mamba_ssm_cache_dtype` allocated it (V2 interleaves the layers inside each slot, so + the slot stride is a multiple of one state); `state_indices` the decode rows' int32 slots; `old_u, old_k, old_G, + old_beta` the layer's `old_x, old_B, old_dt, old_dA_cumsum` from + `mamba_layer_cache(layer)`; `cache_buf_idx`, `prev_num_accepted_tokens`, + `history_size` from `get_replay_state_update_metadata()` (16 here); + `use_qk_l2norm_in_kernel=True`; and exactly when `N >= 16` and the manager's + `use_gdn_cached_replay_all_layer_commit` holds, `replay_work_items = + mamba_metadata.replay_work_items[:N]`, `n_writes = mamba_metadata.replay_n_writes`, + `use_all_layer_commit=True`. That condition has to be the manager's own: + `update_mamba_states` runs the all-layer commit iff it holds, so work items + passed when it does not leave every wrap uncommitted, and work items withheld + when it does commit twice -- both silent at the op and far off the exact + recurrence one step later. + + 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 + presence switches the in-kernel commit off. `old_beta` is never read or written. + + Preconditions not checkable from metadata: + * rows on distinct slots, except padding: two rows on one slot both append to the + same history positions and the appended u is a per-element mixture of the two + rows, so that slot's state and those rows' outputs are unspecified. The rows of + CUDA-graph padding repeat one shared dummy slot (the manager's `_padding_slot`); + the rows on unique slots are exact whatever the padding does, and the dummy slot + is nobody's state; + * `prev_num_accepted_tokens` in `[0, history_size]`, `cache_buf_idx` in `{0, 1}`, + and the pending entries written by earlier calls of this op. A slot fresh from + prefill starts at 0 pending (the manager zeroes the count, the half and the + history); + * after every step, before the next: with work items, + `commit_gdn_cached_replay_history_layers` over every local GDN layer; then the + advance (`_advance_replay_state` in `pyexecutor/kv_cache/mamba_cache_manager.py`) + with the accepted counts including the bonus token. Advancing as if every + draft were accepted is silent at the op and wrong from the next step on; + * 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. + + 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 + paths, `packed_qkv` against separate q/k/v, and `launch_with_pdl` on and off are not + 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 + 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). + + 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 < + T` (a bare assert); and `history_size <= 8`, which does not compile -- the + history tile is `next_pow2(history_size)` wide and feeds a `tl.dot` whose + reduction must be 16, so only 9..16 work (the engine's is 16). + + `is_valid` guards what is accepted and silently wrong: `state_indices` shorter + than N (the trailing rows read past it and append to a slot nobody named); + history buffers whose shape disagrees with `history_size` or the head geometry + (the kernel addresses them by `history_size`, so 16-deep buffers called at 12 + read other slots' entries); history buffers, `prev_num_accepted_tokens` or + `cache_buf_idx` covering fewer slots than `ssm_states` (a slot of the pool past + their end is read and written out of bounds); a strided `prev_num_accepted_tokens` + or `cache_buf_idx` (read at element `slot`). A strided `state_indices` is + 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. + """ + + 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, + ) -> 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, + block_v=block_v, + num_warps=num_warps, + use_l2_state_cache=use_l2_state_cache, + use_l2_streaming_inputs=use_l2_streaming_inputs, + use_l2_shared_inputs=use_l2_shared_inputs, + main_num_stages=main_num_stages, + packed_qkv=packed_qkv, + use_all_layer_commit=use_all_layer_commit, + output=output, + ) + + 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, + ) -> 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.""" + 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() + + 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) + if replay_work_items is None 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, + ) -> 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" + ) + + +fused_recurrent_gated_delta_rule_cached_replay_update = ( + FusedRecurrentGatedDeltaRuleCachedReplayUpdate() +) 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 index a290b44717ab..072f874f6d4e 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_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 @@ -27,11 +27,7 @@ #: 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 -#: The recurrent pool's slot count the cache manager allocates (10 .. 461 over -#: the dumps). The target gathers a batch-scoped copy before this call, so the -#: write-back mode is bracketed here rather than driven. -ENGINE_SLOTS = 461 -#: Pool and buffer rows of the small cells. +#: Default pool and buffer rows of the test's builds. SLOTS = 8 @@ -39,15 +35,11 @@ def _verify(n, p, why, **spec): return Cell(why=why, spec=dict(mode="verify", n=n, p=p, **spec)) -def _write_back(n, p, why, **spec): - return Cell(why=why, spec=dict(mode="write_back", 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 / (||q|| + 1e-6), k / (||k|| + 1e-6) # if use_qk_l2norm_in_kernel + 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, @@ -73,111 +65,43 @@ class FusedRecurrentGatedDeltaRuleUpdate(OpWrapper): ARCHS = frozenset({Arch.SM_103}) # spec: mode ("verify": disable_state_update + buffer with cache_steps = p - # | "write_back" | "both"), n sequences x p tokens (or lengths for the - # varlen form), pool_slots (None: no pool), cache_rows, slots (None: - # arange(n)), pool_dtype / cache_dtype, beta ("fp32" | "bf16": sigmoid of - # a bf16 projection, as the target passes it | "channel"), qkv ("views" + # | "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), output (None | "prefix" of a longer buffer | "fp32"), - # index_dtype, cu_dtype, pool_view ("layer_major"). - CELLS: tuple[Cell, ...] = ( - tuple( - _verify( - n, - p, - f"target verify: n={n} p={p}, n-row gather, 32-row buffer", - pool_slots=n, - cache_rows=ENGINE_VERIFY_ROWS, - beta="bf16", - qkv="contiguous", - ) - for n, p in ((1, 2), (8, 2), (16, 2), (32, 2), (8, 3), (8, 4)) - ) - + ( - _verify( - 4, - 2, - "verify: rows 0/31/17/3 of the 32-row buffer", - pool_slots=ENGINE_VERIFY_ROWS, - cache_rows=ENGINE_VERIFY_ROWS, - slots=(0, 31, 17, 3), - ), - _verify( - 4, - 2, - "verify: -1-padded rows naming 31 and 0", - pool_slots=ENGINE_VERIFY_ROWS, - cache_rows=ENGINE_VERIFY_ROWS, - slots=(31, -1, 0, -1), - ), - ) - + tuple( - _verify(n, p, f"verify P={p}, n={n}: a state cached per token") - for n, p in ((1, 1), (3, 1), (3, 2), (3, 3), (3, 4), (3, 8)) - ) - + tuple(_verify(n, 4, f"verify batch width n={n}, P=4") for n in (1, 2, 4, 8)) - + ( - _verify(4, 3, "verify scattered slots 7/1/4/0", slots=(7, 1, 4, 0)), - _verify(3, 4, "verify slot -1: zero start, no cache row", slots=(2, -1, 5)), - _verify(2, 3, "verify into a bf16 per-step buffer", cache_dtype=torch.bfloat16), - _verify(2, 3, "verify with no pool: every sequence from zero", pool_slots=None), - _write_back( - 3, - 2, - "write-back into a 461-slot pool: slots 0, 460, 230", - pool_slots=ENGINE_SLOTS, - slots=(0, ENGINE_SLOTS - 1, 230), - ), - ) - + tuple( - _write_back(4, p, f"write-back P={p}, scattered slots 7/1/4/0", slots=(7, 1, 4, 0)) - for p in (1, 2, 3, 4, 8, 16) + # 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, ) - + ( - _write_back(3, 4, "write-back into a prefix of a longer output", output="prefix"), - _write_back(2, 1, "write-back P=1 into a prefix of an output", output="prefix"), - _write_back(3, 3, "write-back slot -1: no slot written", slots=(2, -1, 5)), - _write_back(2, 4, "write-back, every slot -1", slots=(-1, -1)), - _write_back(3, 4, "write-back into a bf16 pool", pool_dtype=torch.bfloat16), - _write_back(2, 1, "write-back P=1 into a bf16 pool", pool_dtype=torch.bfloat16), - _write_back(3, 1, "bf16 beta (the mixer's dtype), P=1", slots=(2, 5, 0), beta="bf16"), - _write_back(3, 4, "bf16 beta (the mixer's dtype), P=4", slots=(2, 5, 0), beta="bf16"), - _write_back(2, 3, "per-channel beta [N, P, HV, V]", slots=(1, 2), beta="channel"), - _write_back(2, 3, "fp32 output buffer against bf16 q", slots=(1, 2), output="fp32"), - _write_back(2, 3, "int64 initial_state_indices", slots=(1, 2), index_dtype=torch.int64), - _write_back( - 2, - 3, - "a [layers, S, ...] pool's middle layer", - slots=(1, 4), - pool_view="layer_major", - ), - Cell(why="both sinks at once", spec=dict(mode="both", n=2, p=4)), - _write_back(3, None, "varlen: ragged lengths 1/5/2", lengths=(1, 5, 2)), - _write_back( - 3, - None, - "varlen: uniform lengths, scattered slots", - lengths=(4, 4, 4), - slots=(2, 6, 3), - ), - _write_back(2, None, "varlen: int32 cu_seqlens", lengths=(2, 4), cu_dtype=torch.int32), + 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`), 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]` - fp32 gather of the live slots, `int_idx = intermediate_state_indices[:n]` - (int32 arange), `int_ssm` the layer's `[32, p, 48, 128, 128]` per-step - buffer. 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. + 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 @@ -203,10 +127,17 @@ class FusedRecurrentGatedDeltaRuleUpdate(OpWrapper): 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. - Deterministic; the first call in a fresh process is bitwise the second. + The first call in a fresh process is bitwise the second (measured on the + write-back sink; the kernel has no autotuner). - Not certified: `disable_output_calculation=True` alone (its cached states - were never checked); `K` or `V` other than 128; fp16/fp32 q/k/v; P > 16. + 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; @@ -295,8 +226,8 @@ def reference( 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 / (q_t.square().sum(-1, keepdim=True).sqrt() + 1e-6) - k_t = k_t / (k_t.square().sum(-1, keepdim=True).sqrt() + 1e-6) + 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] 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 index 12dd485c4bfe..19e5cfdc4d9a 100644 --- 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 @@ -19,18 +19,17 @@ fused_sigmoid_gating_delta_rule_update as _fused_sigmoid_gating_delta_rule_update, ) -#: The pool width the bulk of the cells run at, from the old test: a width the -#: engine never allocates, kept because its cells were certified there. -SLOTS = 24 -#: Largest slot count the cache manager was seen to allocate the recurrent pool -#: at for this checkpoint (10 .. 461 over twelve runtime-value dumps, batch never -#: above 32). The target passes the layer's whole pool unsliced. +#: 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 included.""" - return ((0, ENGINE_SLOTS - 1) + tuple(i * 13 + 4 for i in range(n_seqs)))[:n_seqs] + """`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): @@ -42,7 +41,7 @@ class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper): (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 / (|q| + 1e-6), k / (|k| + 1e-6) (only with use_qk_l2norm_in_kernel) + 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` @@ -61,130 +60,72 @@ class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: seq_lens, slots (pool rows; None for no pool), indices (None -> - # arange), pool_dtype (fp32), l2norm (True), normalize_qk (caller-normalized - # q/k when l2norm is off), output ("none" | "prefix": a length-prefix slice - # of a longer buffer, as the stock mixer passes), ab ("packed": column - # halves of a [T, 96] projection | "separate": two contiguous [T, 48], as - # the target's two matmuls make them), layout ("varlen" | "layer_view": - # pool[1] of a [3, slots, ...] cache | "batched": cu_seqlens None). + # 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, ...] = ( - tuple( - Cell( - why=f"target decode call, N={n}, 461-slot pool", - spec=dict( - seq_lens=(1,) * n, slots=ENGINE_SLOTS, indices=_scattered(n), ab="separate" - ), - ) - for n in (1, 4, 16, 32) - ) - + ( - Cell( - why="varlen [3, 5, 1, 7] at the 461-slot pool", - spec=dict(seq_lens=(3, 5, 1, 7), slots=ENGINE_SLOTS, indices=_scattered(4)), - ), - Cell( - why="CUDA-graph 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", - ), - ), - ) - + tuple( - Cell( - why=f"decode N={n}, one token per sequence, 24 slots", - spec=dict(seq_lens=(1,) * n, slots=SLOTS), - ) - for n in (1, 2, 4, 8, 16) - ) - + tuple( - Cell( - why=f"output= a length-prefix slice of a longer buffer, N={n}", - spec=dict(seq_lens=(1,) * n, slots=SLOTS, output="prefix"), - ) - for n in (1, 4, 16) - ) - + ( - Cell( - why="scattered, unsorted slots [17, 3, 0, 11]", - spec=dict(seq_lens=(1,) * 4, slots=SLOTS, indices=(17, 3, 0, 11)), - ), - Cell( - why="scattered slots [23, 5]", - spec=dict(seq_lens=(1, 1), slots=SLOTS, indices=(23, 5)), - ), - Cell( - why="negative slot indices [2, -1, 5, -1]", - spec=dict(seq_lens=(1,) * 4, slots=SLOTS, indices=(2, -1, 5, -1)), - ), - Cell( - why="every slot index negative", - spec=dict(seq_lens=(1, 1), slots=SLOTS, indices=(-1, -1)), - ), - ) - + tuple( - Cell( - why=f"varlen multi-token sequences {list(s)}", - spec=dict(seq_lens=s, slots=SLOTS), - ) - for s in ((3, 5), (2, 1, 7, 1), (64, 33)) - ) - + ( - Cell( - why="l2norm off, raw single-token q/k (one step cannot amplify)", - spec=dict(seq_lens=(1,) * 4, slots=SLOTS, l2norm=False), - ), - ) - + tuple( - Cell( - why=f"l2norm off, caller-normalized q/k, {list(s)}", - spec=dict(seq_lens=s, slots=SLOTS, l2norm=False, normalize_qk=True), - ) - for s in ((3, 5), (2, 1, 7, 1)) - ) - + tuple( - Cell( - why=f"bf16 pool on the Triton branch (T != N), {list(s)}", - spec=dict(seq_lens=s, slots=SLOTS, pool_dtype=torch.bfloat16), - ) - for s in ((3, 5), (2, 6)) - ) - + ( - Cell( - why="initial_state_source None: zero state, no write-back", - spec=dict(seq_lens=(1,) * 4, slots=None), - ), - Cell( - why="pool[layer] of a [layers, slots, HV, V, K] cache", - spec=dict(seq_lens=(1,) * 4, slots=SLOTS, layout="layer_view"), + 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="cu_seqlens None: batched [N, T, ...] layout, N=3, T=4", - spec=dict(seq_lens=(4, 4, 4), slots=SLOTS, layout="batched"), + ), + 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): `(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 reshaped - column views of the conv output `[T, 10240]` (read through `stride(1)`, no - repack), `a`/`b` two contiguous `[T, 48]` projections, the layer's whole - fp32 pool `[slots, 48, 128, 128]`, int32 slots, int64 `cu_seqlens`, one - token per sequence. Certified at 24 and 461 slots; at 461 output and state - are bitwise identical to 24 for the same slot content. + 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 when the pool is - bf16, `cu_seqlens` is given and `T == N` (head dims 128, supported arch); - everything else -- every fp32-pool call, i.e. this target's - `--mamba_ssm_cache_dtype=float32` -- runs the vendored Triton kernel. Only - the Triton branch is certified; the test proves the fp32 pool never - reaches FlashInfer and that the witness is not blind. + 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`. @@ -207,8 +148,11 @@ class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper): not measured: N > 1365 at this geometry). Not certified: the FlashInfer branch, K or V other than 128, fp16/fp32 - operands, `softplus_beta`/`softplus_threshold` other than 1.0/20.0, - `B > 1` with `cu_seqlens`. + 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. @@ -301,8 +245,8 @@ def reference( ) beta = torch.sigmoid(b3[row, t].float()) if use_qk_l2norm_in_kernel: - q_t = q_t / (q_t.pow(2).sum(-1, keepdim=True).sqrt() + 1e-6) - k_t = k_t / (k_t.pow(2).sum(-1, keepdim=True).sqrt() + 1e-6) + 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, :] @@ -338,6 +282,16 @@ def is_valid( 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_decode_gated_norm_pdl.py b/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 new file mode 100644 index 000000000000..e9179d141741 --- /dev/null +++ b/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 @@ -0,0 +1,509 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""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. +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 +before `griddepcontrol.wait`, while the input projection ahead of this kernel +streams its weight. + +A v-head is split over a cluster of `CLUSTER` CTAs of `WARPS` warps; each warp owns +`V / (CLUSTER * WARPS)` value rows of the head's `[V, K]` state and lane `l` the +`K / 32` key columns from `l * K / 32` of each. Small CTAs fit beside the input +projection's last wave, so their state reads start before it ends. The gated RMS +norm (per v-head, 128 wide) needs the head's whole output: every warp writes the +sum of squares of its rows into each CTA of the cluster through distributed shared +memory, and every CTA then normalizes and gates its own rows. +""" + +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 +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. +CLUSTER = 8 +WARPS = 4 + + +@dsl_user_op +def _st_shared_cluster_f32(ptr, val, *, loc=None, ip=None): + """`st.shared::cluster.f32` to a pointer from `cute.arch.map_dsmem_ptr`.""" + llvm.inline_asm( + None, + [ + cutlass.Int32(ptr.toint(loc=loc, ip=ip)).ir_value(loc=loc, ip=ip), + cutlass.Float32(val).ir_value(loc=loc, ip=ip), + ], + "st.shared::cluster.f32 [$0], $1;", + "r,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +class _GdnKernel: + def __init__(self, *, k_heads, v_heads, kdim, vdim, pdl): + assert vdim % (CLUSTER * WARPS) == 0 and CLUSTER * WARPS <= 32 + assert kdim % 32 == 0 and v_heads % k_heads == 0 + self.k_heads, self.v_heads, self.kdim, self.vdim, self.pdl = ( + k_heads, + v_heads, + kdim, + vdim, + pdl, + ) + self.rows_per_cta = vdim // CLUSTER + self.rows_per_warp = vdim // (CLUSTER * WARPS) + self.per_lane = kdim // 32 + self.k_off = k_heads * kdim + self.v_off = 2 * k_heads * kdim + self.z_off = self.v_off + v_heads * vdim + self.a_off = self.z_off + v_heads * vdim + self.b_off = self.a_off + v_heads + + @cute.jit + def __call__( + self, + mProj, + mPool, + mSlot, + mALog, + mDt, + mW, + mOut, + eps: cutlass.Float32, + stream: cuda_drv.CUstream, + ): + rows = cute.size(mProj, mode=[0]) + ld_state = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cutlass.Float32, + num_bits_per_copy=32 * self.per_lane, + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, + ) + st_state = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=32 * self.per_lane + ) + ld_act = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + cutlass.BFloat16, + num_bits_per_copy=16 * self.per_lane, + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, + ) + self.kernel( + mProj, mPool, mSlot, mALog, mDt, mW, mOut, eps, ld_state, st_state, ld_act + ).launch( + grid=[self.v_heads * CLUSTER, rows, 1], + block=[WARPS * 32, 1, 1], + cluster=[CLUSTER, 1, 1], + smem=32 * 4 + 16, + stream=stream, + use_pdl=self.pdl, + # <= 64 registers a thread: the CTAs fit beside the producer's last wave. + min_blocks_per_mp=max(1, 8 * 4 // WARPS), + ) + + @cute.kernel + def kernel( + self, + mProj, + mPool, + mSlot, + mALog, + mDt, + mW, + mOut, + eps: cutlass.Float32, + ld_state: cute.CopyAtom, + st_state: cute.CopyAtom, + ld_act: cute.CopyAtom, + ): + bidx, n, _ = cute.arch.block_idx() + warp = cute.arch.warp_idx() + lane = cute.arch.lane_idx() + rpw: cutlass.Constexpr = self.rows_per_warp + pl: cutlass.Constexpr = self.per_lane + kdim: cutlass.Constexpr = self.kdim + vdim: cutlass.Constexpr = self.vdim + parts: cutlass.Constexpr = CLUSTER * WARPS + hv = bidx // CLUSTER + part = bidx % CLUSTER # the CTA's rank in its cluster: clusters tile the grid's x + kh = hv // (self.v_heads // self.k_heads) + v0 = part * self.rows_per_cta + warp * rpw + k0 = lane * pl + ss_parts = cutlass.utils.SmemAllocator().allocate_tensor( + cutlass.Float32, cute.make_layout((32,), stride=(1,)), byte_alignment=16 + ) + if const_expr(CLUSTER > 1): + # Every CTA of the cluster must have started before another writes into its + # shared memory; this arrival is waited on just before those writes. + cute.arch.cluster_arrive_relaxed() + + # 1. Static and step-invariant operands, before the wait: the gates' constants, + # the norm weight of the warp's rows, the slot and the warp's state rows. + a_log_h = mALog[hv] + dtb = mDt[hv] + wn = cute.make_rmem_tensor(cute.make_layout((rpw,), stride=(1,)), cutlass.Float32) + for r in cutlass.range_constexpr(rpw): + wn[r] = mW[v0 + r].to(cutlass.Float32) + slot = mSlot[n] + # 64-bit: the cache manager coalesces every GDN layer's state of a slot into one + # page, so a layer's slot stride is ~28M elements and slot * stride overflows int32. + slot64 = cutlass.Int64(slot) + # (slot, hv, v, (pl, lane)): lane l's 16-byte vector of a state row. + pool_vec = cute.logical_divide(mPool, (None, None, None, pl)) + h = cute.make_rmem_tensor(cute.make_layout((rpw, pl), stride=(pl, 1)), cutlass.Float32) + h.fill(0.0) + if slot >= 0: + for r in cutlass.range_constexpr(rpw): + cute.copy(ld_state, pool_vec[slot64, hv, v0 + r, (None, lane)], h[r, None]) + if const_expr(self.pdl): + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + + # 2. The projection's q, k, v, z, a, b for this row and head. + # (n, (pl, cols / pl)): 8-byte vectors of the row; every offset below is a multiple of pl. + proj_vec = cute.logical_divide(mProj, (None, pl)) + raw = cute.make_rmem_tensor(cute.make_layout((2, pl), stride=(pl, 1)), cutlass.BFloat16) + cute.copy(ld_act, proj_vec[n, (None, (kh * kdim + k0) // pl)], raw[0, None]) + cute.copy(ld_act, proj_vec[n, (None, (self.k_off + kh * kdim + k0) // pl)], raw[1, None]) + qn = cute.make_rmem_tensor(cute.make_layout((pl,), stride=(1,)), cutlass.Float32) + kn = cute.make_rmem_tensor(cute.make_layout((pl,), stride=(1,)), cutlass.Float32) + zv = cute.make_rmem_tensor(cute.make_layout((rpw,), stride=(1,)), cutlass.Float32) + vv = cute.make_rmem_tensor(cute.make_layout((rpw,), stride=(1,)), cutlass.Float32) + for e in cutlass.range_constexpr(pl): + qn[e] = raw[0, e].to(cutlass.Float32) + kn[e] = raw[1, e].to(cutlass.Float32) + for r in cutlass.range_constexpr(rpw): + vv[r] = mProj[n, self.v_off + hv * vdim + v0 + r].to(cutlass.Float32) + zv[r] = mProj[n, self.z_off + hv * vdim + v0 + r].to(cutlass.Float32) + a = mProj[n, self.a_off + hv].to(cutlass.Float32) + b = mProj[n, self.b_off + hv].to(cutlass.Float32) + + # q, k L2 norm (eps 1e-6 inside the root), the 1/sqrt(K) scale on q, the gates. + sq = cutlass.Float32(0.0) + sk = cutlass.Float32(0.0) + for e in cutlass.range_constexpr(pl): + sq = qn[e] * qn[e] + sq + sk = kn[e] * kn[e] + sk + sq = cute.arch.warp_reduction_sum(sq) + sk = cute.arch.warp_reduction_sum(sk) + rq = cute.math.rsqrt(sq + 1e-6) * (1.0 / cute.math.sqrt(cutlass.Float32(kdim))) + rk = cute.math.rsqrt(sk + 1e-6) + for e in cutlass.range_constexpr(pl): + qn[e] = qn[e] * rq + kn[e] = kn[e] * rk + x = a + dtb + sp = x + if x <= 20.0: + sp = cute.math.log1p(cute.math.exp(x)) + decay = cute.math.exp(-cute.math.exp(a_log_h) * sp) + beta = 1.0 / (1.0 + cute.math.exp(-b)) + + # 3. The delta rule, one value row at a time; the new rows go straight back. + ov = cute.make_rmem_tensor(cute.make_layout((rpw,), stride=(1,)), cutlass.Float32) + ssw = cutlass.Float32(0.0) + for r in cutlass.range_constexpr(rpw): + hk = cutlass.Float32(0.0) + for e in cutlass.range_constexpr(pl): + hk = h[r, e] * kn[e] + hk + hk = cute.arch.warp_reduction_sum(hk) + u = beta * (vv[r] - decay * hk) + o = cutlass.Float32(0.0) + for e in cutlass.range_constexpr(pl): + h[r, e] = decay * h[r, e] + u * kn[e] + o = h[r, e] * qn[e] + o + o = cute.arch.warp_reduction_sum(o) + if slot >= 0: + cute.copy(st_state, h[r, None], pool_vec[slot64, hv, v0 + r, (None, lane)]) + else: + o = cutlass.Float32(0.0) + # The recurrence's output is a bf16 tensor in the unfused composition. + ov[r] = o.to(cutlass.BFloat16).to(cutlass.Float32) + ssw = ov[r] * ov[r] + ssw + + # 4. The head's sum of squares: warp `w` of CTA `c` writes its rows' share into + # slot `c * WARPS + w` of every CTA of the cluster (lane `i` into CTA `i`). + if const_expr(CLUSTER > 1): + cute.arch.cluster_wait() + if lane < CLUSTER: + _st_shared_cluster_f32( + cute.arch.map_dsmem_ptr(ss_parts.iterator + (part * WARPS + warp), lane), ssw + ) + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + else: + if lane == 0: + ss_parts[warp] = ssw + cute.arch.sync_threads() + ss = cutlass.Float32(0.0) + if lane < parts: + ss = ss_parts[lane] + ss = cute.arch.warp_reduction_sum(ss) + rstd = cute.math.rsqrt(ss / vdim + eps) + + # 5. Gated RMS norm of the warp's rows: lane `r` writes row `r`. + o_r = ov[0] + z_r = zv[0] + w_r = wn[0] + for r in cutlass.range_constexpr(1, rpw): + o_r = cutlass.Float32(cutlass.select_(lane == r, ov[r], o_r)) + z_r = cutlass.Float32(cutlass.select_(lane == r, zv[r], z_r)) + w_r = cutlass.Float32(cutlass.select_(lane == r, wn[r], w_r)) + y = o_r * rstd * w_r * (1.0 / (1.0 + cute.math.exp(-z_r))) + if lane < rpw: + mOut[n, hv * vdim + v0 + lane] = y.to(cutlass.BFloat16) + + +def _dl(t: torch.Tensor, align: int = 4): + return from_dlpack(t, assumed_align=align) + + +_COMPILED: dict = {} + + +def _operands(proj, pool, slots, A_log, dt_bias, norm_weight, out, eps, stream): + return ( + _dl(proj, 8), + _dl(pool, 16), + _dl(slots), + _dl(A_log), + _dl(dt_bias), + _dl(norm_weight, 2), + _dl(out, 2), + cutlass.Float32(eps), + stream, + ) + + +def launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, pdl=True): + """The launch behind `GdnDecodeGatedNormPdl.__call__`; updates `pool` in place and + returns the gated-normed output `[rows, HV * V]` bf16.""" + rows = proj.shape[0] + v_heads, vdim = pool.shape[1], pool.shape[2] + out = torch.empty((rows, v_heads * vdim), dtype=torch.bfloat16, device=proj.device) + if rows == 0: + return out + pdl = bool(pdl and get_env_enable_pdl()) + stream = cuda_drv.CUstream(torch.cuda.current_stream(proj.device).cuda_stream) + operands = (proj, pool, slots, A_log, dt_bias, norm_weight, out, eps, stream) + key = ( + rows, + k_heads, + v_heads, + kdim, + vdim, + proj.stride(0), + tuple(pool.shape), + pool.stride(), + pdl, + ) + 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 state cache). + kern = _GdnKernel(k_heads=k_heads, v_heads=v_heads, kdim=kdim, vdim=vdim, pdl=pdl) + fn = _COMPILED[key] = cute_ext.compile(kern, *_operands(*operands)) + fn(*_operands(*operands)) + return out + + +def _recurrence_reference(q, k, v, A_log, a, dt_bias, b, pool, slots): + """`(o, pool_after)` of one gated-delta-rule decode step per row in fp64, rounded once: + `o` to bf16, the state to the pool's dtype. Row `n` reads and writes `pool[slots[n]]` + (a negative slot skips the row: zero output, no state access); q and k are + L2-normalized, the decay is `exp(-exp(A_log) * softplus(a + dt_bias))`, `beta = + sigmoid(b)`, the scale `K^-0.5`. `q` / `k` `[n, 1, H, K]`, `v` `[n, 1, HV, V]`, `a` / + `b` `[n, 1, HV]`, `pool` `[slots, HV, V, K]`. Mutates nothing.""" + dt = torch.float64 + batch, _, num_k_heads, k_dim = q.shape + hv = v.shape[2] + rep = hv // num_k_heads + scale = k_dim**-0.5 + reads = slots.tolist() + pool_after = pool.clone() + o = torch.zeros(batch, 1, hv, v.shape[3], dtype=dt, device=q.device) + neg_a = -torch.exp(A_log.to(dt)) + for n in range(batch): + if reads[n] < 0: + continue + q_n = q[n, 0].to(dt).repeat_interleave(rep, dim=0) + k_n = k[n, 0].to(dt).repeat_interleave(rep, dim=0) + q_n = q_n * torch.rsqrt(q_n.pow(2).sum(-1, keepdim=True) + 1e-6) + k_n = k_n * torch.rsqrt(k_n.pow(2).sum(-1, keepdim=True) + 1e-6) + x = a[n, 0].to(dt) + dt_bias.to(dt) + softplus = torch.where(x <= 20.0, torch.log1p(torch.exp(x)), x) + beta = torch.sigmoid(b[n, 0].to(dt)) + h = pool[reads[n]].to(dt) * torch.exp(neg_a * softplus)[:, None, None] + u = (v[n, 0].to(dt) - torch.einsum("hvk,hk->hv", h, k_n)) * beta[:, None] + h = h + u[:, :, None] * k_n[:, None, :] + o[n, 0] = torch.einsum("hvk,hk->hv", h, q_n * scale) + pool_after[reads[n]] = h.to(pool.dtype) + return o.to(torch.bfloat16), pool_after + + +class GdnDecodeGatedNormPdl(OpWrapper): + """A single-token Gated-DeltaNet step per row over the fp32 slot pool, then the + sigmoid-gated RMS norm of its output per v-head. With `proj` the decode input + projection's output row `[q (H*K) | k (H*K) | v (HV*V) | z (HV*V) | a (HV) | b (HV)]` + (after the conv step: q, k, v are the conv's SiLU outputs):: + + o, pool = gated_delta_rule_decode_pretranspose(q, k, v, A_log, a, dt_bias, b, + use_qk_l2norm=True, pool, slots) + out = rms_norm_gated_token_major(o (bf16), z, norm_weight, eps, "sigmoid") + + Returns `out` `[rows, HV * V]` bf16 and updates `pool[slots[n]]` in place; a + negative slot skips the row's state (no read or write) and its output is zero. + `o` is rounded to bf16 between the two, as the composition stores it. + + Fusion boundary: both gates, the q/k L2 norm, the recurrence, the indexed state + read and write, the gated RMS norm. The projection and the conv are the caller's. + """ + + 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). + 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)), + Cell( + 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), + ), + ) + + note = """ + Metadata consumed: `slots` (the cache manager's state indices for the step). + + Preconditions: bf16 `proj` `[rows, 2*H*K + 2*HV*V + 2*HV]` with unit column stride; + the pool fp32 `[slots, HV, V, K]` with K innermost (any outer strides); `A_log`, + `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 + 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. + + Launch: `HV * 8` CTAs a row, clusters of 8 (one v-head), 128 threads, <= 64 + registers. Numerics: fp32 inside; `o` rounded to bf16 before the norm, the output + 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`. + """ + + def __call__( + self, + proj: torch.Tensor, + pool: torch.Tensor, + slots: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + k_heads: int, + kdim: int, + pdl: bool = True, + ) -> torch.Tensor: + return launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, pdl) + + def reference( + self, + proj: torch.Tensor, + pool: torch.Tensor, + slots: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + k_heads: int, + kdim: int, + pdl: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """`(out, pool_after)`: the fp64 recurrence (`_recurrence_reference`; o rounded to + bf16, the pool to fp32), then the gated norm entry's fp32 reference.""" + rows = proj.shape[0] + v_heads, vdim = pool.shape[1], pool.shape[2] + q, k, v, z, a, b = torch.split( + proj, + [k_heads * kdim, k_heads * kdim, v_heads * vdim, v_heads * vdim, v_heads, v_heads], + 1, + ) + o, pool_after = _recurrence_reference( + q.reshape(rows, 1, k_heads, kdim), + k.reshape(rows, 1, k_heads, kdim), + v.reshape(rows, 1, v_heads, vdim), + A_log, + a.reshape(rows, 1, v_heads), + dt_bias, + b.reshape(rows, 1, v_heads), + pool, + slots, + ) + out = _gated_norm.reference( + o.reshape(-1, vdim), z.reshape(rows, v_heads, vdim), norm_weight, eps, None, "sigmoid" + ) + return out.reshape(rows, v_heads * vdim), pool_after + + def is_valid( + self, + proj: torch.Tensor, + pool: torch.Tensor, + slots: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + k_heads: int, + kdim: int, + pdl: bool = True, + ) -> None: + rows = proj.shape[0] + v_heads, vdim, kd = pool.shape[1], pool.shape[2], pool.shape[3] + assert rows <= MAX_ROWS, f"at most {MAX_ROWS} rows" + assert proj.dtype == torch.bfloat16 and proj.stride(1) == 1, "bf16 proj, unit column stride" + assert proj.shape[1] == 2 * k_heads * kdim + 2 * v_heads * vdim + 2 * v_heads, "proj width" + assert pool.dtype == torch.float32 and pool.stride(3) == 1 and kd == kdim, ( + "fp32 pool [S, HV, V, K]" + ) + assert vdim % (CLUSTER * WARPS) == 0 and kdim % 32 == 0, "V % (CLUSTER * WARPS), K % 32" + assert v_heads % k_heads == 0, "HV a multiple of H" + assert slots.dtype == torch.int32 and slots.shape == (rows,), "int32 slots [rows]" + assert A_log.dtype == torch.float32 and dt_bias.dtype == torch.float32, ( + "fp32 gate constants" + ) + assert norm_weight.dtype == torch.bfloat16 and norm_weight.shape == (vdim,), ( + "bf16 norm weight [V]" + ) + + +gdn_decode_gated_norm_pdl = GdnDecodeGatedNormPdl() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_decode_rows_gated_norm.py new file mode 100644 index 000000000000..6ec2bd7904fb --- /dev/null +++ b/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 @@ -0,0 +1,689 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One Gated-DeltaNet decode step per row and the gated RMS norm on its output, streamed +over the fp32 state pool: the decode step at batch sizes past the cluster kernel's. + +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, read +and written once -- so a CTA owns one whole state 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 512-byte state row. The state is loaded before +`griddepcontrol.wait` (nothing in the step writes it ahead of this kernel). The two dot +products a value row takes over K are reduced across the warp by a transposed +reduction (RPW rows in RPW - 1 + log2(32 / RPW) shuffles, not 5 RPW), which leaves row +j's sum in lanes [j * 32 / RPW, (j + 1) * 32 / RPW); `u` goes back to every lane by +index shuffles. The CTA holds the head's whole output, so the gated RMS norm runs in +place. +""" + +from __future__ import annotations + +import math + +import cutlass +import cutlass.cute as cute +import torch +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.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 (a 16-byte 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. +ENGINE_SLOTS = 461 + + +def _iv(x, loc, ip): + return x.ir_value(loc=loc, ip=ip) + + +def _asm(ret, args, text, constraints, loc, ip, side_effects=False): + return llvm.inline_asm( + ret, + [_iv(a, loc, ip) for a in args], + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _asm_n(n, ty, args, text, constraints, loc, ip, side_effects=False): + res = _asm( + llvm.StructType.get_literal([ty] * n), args, text, constraints, loc, ip, side_effects + ) + return [llvm.extractvalue(ty, res, [i], loc=loc, ip=ip) for i in range(n)] + + +@dsl_user_op +def _ld_state(addr, pred, *, loc=None, ip=None): + """Four fp32 of a state row (streaming), or zeros where `pred` is 0.""" + r = _asm_n( + 4, + T.f32(), + [cutlass.Int64(addr), cutlass.Int32(pred)], + "{ .reg .pred p; setp.ne.b32 p, $5, 0; mov.b32 $0, 0; mov.b32 $1, 0; mov.b32 $2, 0; " + "mov.b32 $3, 0; @p ld.global.cs.v4.f32 {$0, $1, $2, $3}, [$4]; }", + "=f,=f,=f,=f,l,r", + loc, + ip, + side_effects=True, + ) + return [cutlass.Float32(v) for v in r] + + +@dsl_user_op +def _st_state(addr, pred, a, b, c, d, *, loc=None, ip=None): + """Four fp32 of a state row (streaming) where `pred` is not 0.""" + _asm( + T.i32(), + [cutlass.Int64(addr), cutlass.Int32(pred)] + [cutlass.Float32(v) for v in (a, b, c, d)], + "{ .reg .pred p; setp.ne.b32 p, $2, 0; @p st.global.cs.v4.f32 [$1], {$3, $4, $5, $6}; " + "mov.u32 $0, 0; }", + "=r,l,r,f,f,f,f", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _ld_bf16x4(addr, *, loc=None, ip=None): + """Four bf16 (8 bytes) as fp32.""" + r = _asm_n( + 4, + T.f32(), + [cutlass.Int64(addr)], + "{ .reg .b32 w0, w1; ld.global.v2.b32 {w0, w1}, [$4]; shl.b32 $0, w0, 16; " + "and.b32 $1, w0, 0xFFFF0000; shl.b32 $2, w1, 16; and.b32 $3, w1, 0xFFFF0000; }", + "=f,=f,=f,=f,l", + loc, + ip, + side_effects=True, + ) + return [cutlass.Float32(v) for v in r] + + +@dsl_user_op +def _ld_bf16(addr, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Int64(addr)], + "{ .reg .b16 h; ld.global.b16 h, [$1]; cvt.f32.bf16 $0, h; }", + "=f,l", + loc, + ip, + side_effects=True, + ) + ) + + +@dsl_user_op +def _st_bf16(addr, pred, x, *, loc=None, ip=None): + """`x` rounded to bf16 where `pred` is not 0.""" + _asm( + T.i32(), + [cutlass.Int64(addr), cutlass.Int32(pred), cutlass.Float32(x)], + "{ .reg .pred p; .reg .b16 h; setp.ne.b32 p, $2, 0; cvt.rn.bf16.f32 h, $3; " + "@p st.global.b16 [$1], h; mov.u32 $0, 0; }", + "=r,l,r,f", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _round_bf16(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + T.f32(), + [cutlass.Float32(x)], + "{ .reg .b16 h; cvt.rn.bf16.f32 h, $1; cvt.f32.bf16 $0, h; }", + "=f,f", + loc, + ip, + ) + ) + + +def _rows_reduce(vals, lane): + """Sum each of `vals` (a lane's partial of one row each) over the warp. Each halving + level (offset 16 first) trades half the rows with the partner lane, so row j's sum + lands in lanes [j * 32 / n, (j + 1) * 32 / n); full butterflies finish those lanes.""" + vals = list(vals) + off = 16 + while len(vals) > 1: + half = len(vals) // 2 + upper = (lane & off) != 0 + kept = [] + for i in range(half): + lo, hi = vals[i], vals[i + half] + recv = cute.arch.shuffle_sync_bfly(cutlass.select_(upper, lo, hi), offset=off) + kept.append(cutlass.select_(upper, hi, lo) + recv) + vals = kept + off //= 2 + v = vals[0] + while off >= 1: + v = v + cute.arch.shuffle_sync_bfly(v, offset=off) + off //= 2 + return v + + +class _GdnRowsKernel: + """CTA (v-head, row): the row's whole state of that head. `norm`: the gated RMS norm on + the head's output; else the recurrence's output, rounded to bf16, is stored as is.""" + + def __init__(self, k_heads, v_heads, kdim, vdim, pdl, norm=True): + assert kdim == 32 * _PL and vdim % _WARPS == 0 + self.h, self.hv, self.k, self.v, self.pdl = k_heads, v_heads, kdim, vdim, pdl + self.norm = norm + self.rpw = vdim // _WARPS + assert self.rpw in (4, 8, 16, 32) + # Lanes that hold one row's reduced sums. + self.lpr = 32 // self.rpw + self.k_off = k_heads * kdim + self.v_off = 2 * k_heads * kdim + self.z_off = self.v_off + v_heads * vdim + self.a_off = self.z_off + v_heads * vdim + self.b_off = self.a_off + v_heads + + @cute.jit + def __call__( + self, + proj_t: cute.Tensor, + pool_ptr: cute.Pointer, + slot_ptr: cute.Pointer, + a_log_ptr: cute.Pointer, + dt_ptr: cute.Pointer, + w_ptr: cute.Pointer, + out_ptr: cute.Pointer, + rows: cutlass.Int32, + s_proj: cutlass.Int32, + s_slot: cutlass.Int64, + s_head: cutlass.Int32, + s_row: cutlass.Int32, + s_out: cutlass.Int32, + eps: cutlass.Float32, + stream, + ): + slots = cute.make_tensor(slot_ptr, cute.make_layout((rows,), stride=(1,))) + a_log_t = cute.make_tensor(a_log_ptr, cute.make_layout((self.hv,), stride=(1,))) + dt = cute.make_tensor(dt_ptr, cute.make_layout((self.hv,), stride=(1,))) + # `proj` arrives as a tensor: the TVM-FFI entry takes its stream from a GPU tensor. + self.kernel( + proj_t.iterator.toint(), + pool_ptr.toint(), + slots, + a_log_t, + dt, + w_ptr.toint(), + out_ptr.toint(), + s_proj, + s_slot, + s_head, + s_row, + s_out, + eps, + ).launch( + grid=[self.hv, rows, 1], + block=[32 * _WARPS, 1, 1], + smem=_WARPS * 4 + 64, + stream=stream, + use_pdl=self.pdl, + ) + + @cute.kernel + def kernel( + self, + proj_base: cutlass.Int64, + pool_base: cutlass.Int64, + mSlot: cute.Tensor, + mALog: cute.Tensor, + mDt: cute.Tensor, + w_base: cutlass.Int64, + out_base: cutlass.Int64, + s_proj: cutlass.Int32, + s_slot: cutlass.Int64, + s_head: cutlass.Int32, + s_row: cutlass.Int32, + s_out: cutlass.Int32, + eps: cutlass.Float32, + ): + tidx, _, _ = cute.arch.thread_idx() + hv, row, _ = cute.arch.block_idx() + warp = tidx // 32 + lane = tidx % 32 + kh = hv // (self.hv // self.h) + v0 = warp * self.rpw + jl = lane // self.lpr # the value row whose sums this lane ends up holding + red = cutlass.utils.SmemAllocator().allocate_tensor( + cutlass.Float32, cute.make_layout((_WARPS,), stride=(1,)), byte_alignment=16 + ) + + # 1. Before the wait: the warp's state rows (a negative slot reads none: zero state), + # the gate constants and the norm weight of this lane's row. + slot = mSlot[row] + live = cutlass.Int32(0) + if slot >= 0: + live = cutlass.Int32(1) + # 64-bit: a coalesced cache page puts every GDN layer's state of a slot in one + # page, so `slot * s_slot` passes int32. + base = ( + pool_base + + ( + cutlass.Int64(cutlass.max(slot, cutlass.Int32(0))) * s_slot + + cutlass.Int64(hv * s_head + v0 * s_row + lane * _PL) + ) + * 4 + ) + h = [_ld_state(base + cutlass.Int64(j * s_row * 4), live) for j in range(self.rpw)] + a_log = mALog[hv] + dtb = mDt[hv] + w_l = cutlass.Float32(0.0) + if cutlass.const_expr(self.norm): + w_l = _ld_bf16(w_base + cutlass.Int64(v0 + jl) * 2) + if cutlass.const_expr(self.pdl): + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() + + # 2. The row's q and k (this lane's four columns), v and z (this lane's value + # row), a and b; the q/k L2 norms (eps 1e-6 inside the root, 1/sqrt(K) on q) and + # the gates (softplus beta 1, threshold 20). + prow = proj_base + cutlass.Int64(row) * s_proj * 2 + q4 = _ld_bf16x4(prow + cutlass.Int64(kh * self.k + lane * _PL) * 2) + k4 = _ld_bf16x4(prow + cutlass.Int64(self.k_off + kh * self.k + lane * _PL) * 2) + v_l = _ld_bf16(prow + cutlass.Int64(self.v_off + hv * self.v + v0 + jl) * 2) + z_l = cutlass.Float32(0.0) + if cutlass.const_expr(self.norm): + z_l = _ld_bf16(prow + cutlass.Int64(self.z_off + hv * self.v + v0 + jl) * 2) + a = _ld_bf16(prow + cutlass.Int64(self.a_off + hv) * 2) + b = _ld_bf16(prow + cutlass.Int64(self.b_off + hv) * 2) + sq = cutlass.Float32(0.0) + sk = cutlass.Float32(0.0) + for e in cutlass.range_constexpr(_PL): + sq = q4[e] * q4[e] + sq + sk = k4[e] * k4[e] + sk + sq = cute.arch.warp_reduction_sum(sq) + sk = cute.arch.warp_reduction_sum(sk) + rq = cute.math.rsqrt(sq + 1e-6) * (1.0 / math.sqrt(self.k)) + rk = cute.math.rsqrt(sk + 1e-6) + qn = [q4[e] * rq for e in range(_PL)] + kn = [k4[e] * rk for e in range(_PL)] + x = a + dtb + sp = x + if x <= 20.0: + sp = cute.math.log1p(cute.math.exp(x)) + decay = cute.math.exp(-cute.math.exp(a_log) * sp) + beta = 1.0 / (1.0 + cute.math.exp(-b)) + + # 3. Per value row: hk = h.k, u = beta (v - decay hk); h = decay h + u k; o = h.q. + part = [] + for j in cutlass.range_constexpr(self.rpw): + acc = cutlass.Float32(0.0) + for e in cutlass.range_constexpr(_PL): + acc = h[j][e] * kn[e] + acc + part.append(acc) + hk = _rows_reduce(part, lane) + u_l = beta * (v_l - decay * hk) + part = [] + for j in cutlass.range_constexpr(self.rpw): + u = cute.arch.shuffle_sync(u_l, j * self.lpr) + acc = cutlass.Float32(0.0) + for e in cutlass.range_constexpr(_PL): + h[j][e] = u * kn[e] + decay * h[j][e] + acc = h[j][e] * qn[e] + acc + part.append(acc) + _st_state(base + cutlass.Int64(j * s_row * 4), live, *h[j]) + o = _rows_reduce(part, lane) + + # 4. The gated RMS norm over the head's V outputs, each rounded to bf16 first (the + # recurrence's output dtype in the composition): y = o rstd w sigmoid(z). + o_b = _round_bf16(o) * cutlass.Float32(live) + own = cutlass.Int32(0) + if lane % self.lpr == 0: + own = cutlass.Int32(1) + y = o_b + if cutlass.const_expr(self.norm): + ss = cute.arch.warp_reduction_sum(o_b * o_b * cutlass.Float32(own)) + if lane == 0: + red[warp] = ss + cute.arch.sync_threads() + tot = red[0] + for i in cutlass.range_constexpr(1, _WARPS): + tot = tot + red[i] + rstd = cute.math.rsqrt(tot / self.v + eps) + y = o_b * rstd * w_l * (1.0 / (1.0 + cute.math.exp(-z_l))) + _st_bf16(out_base + (cutlass.Int64(row) * s_out + hv * self.v + v0 + jl) * 2, own, y) + + +_COMPILED: dict = {} + + +def launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, pdl=True, out=None): + """The launch behind `GdnDecodeRowsGatedNorm.__call__` and, with `norm_weight=None`, + `GdnDecodeRows.__call__`; updates `pool` in place and returns the gated-normed output + (or the recurrence's) `[rows, HV * V]` bf16, written into `out` when given.""" + rows = proj.shape[0] + v_heads, vdim = pool.shape[1], pool.shape[2] + norm = norm_weight is not None + if out is None: + out = torch.empty((rows, v_heads * vdim), dtype=torch.bfloat16, device=proj.device) + if rows == 0: + return out + pdl = bool(pdl and get_env_enable_pdl()) + key = (k_heads, v_heads, kdim, vdim, pdl, norm) + ptrs = ( + proj, + pool.data_ptr(), + slots.data_ptr(), + A_log.data_ptr(), + dt_bias.data_ptr(), + (norm_weight if norm else A_log).data_ptr(), + out.data_ptr(), + ) + fn = _COMPILED.get(key) + if fn is None: + gm = cute.AddressSpace.gmem + cargs = ( + from_dlpack(proj, assumed_align=8).mark_layout_dynamic(leading_dim=1), + make_ptr(cutlass.Float32, ptrs[1], gm, assumed_align=16), + make_ptr(cutlass.Int32, ptrs[2], gm, assumed_align=4), + make_ptr(cutlass.Float32, ptrs[3], gm, assumed_align=4), + make_ptr(cutlass.Float32, ptrs[4], gm, assumed_align=4), + make_ptr(cutlass.BFloat16, ptrs[5], gm, assumed_align=2), + make_ptr(cutlass.BFloat16, ptrs[6], gm, assumed_align=2), + ) + fn = cute.compile( + _GdnRowsKernel(k_heads, v_heads, kdim, vdim, pdl, norm), + *cargs, + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Int64(1), + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Float32(1e-6), + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--opt-level 2 --enable-tvm-ffi", + ) + _COMPILED[key] = fn + fn( + *ptrs, + rows, + proj.stride(0), + pool.stride(0), + pool.stride(1), + pool.stride(2), + out.stride(0), + float(eps) if norm else 0.0, + ) + return out + + +class GdnDecodeRowsGatedNorm(OpWrapper): + """A single-token Gated-DeltaNet step per row over the fp32 slot pool, then the + sigmoid-gated RMS norm of its output per v-head. With `proj` the decode input + projection's output row `[q (H*K) | k (H*K) | v (HV*V) | z (HV*V) | a (HV) | b (HV)]` + (after the conv step: q, k, v are the conv's SiLU outputs):: + + o, pool = gated_delta_rule_decode_pretranspose(q, k, v, A_log, a, dt_bias, b, + use_qk_l2norm=True, pool, slots) + out = rms_norm_gated_token_major(o (bf16), z, norm_weight, eps, "sigmoid") + + Returns `out` `[rows, HV * V]` bf16 and updates `pool[slots[n]]` in place; a + negative slot skips the row's state (no read or write) and its output is zero. + `o` is rounded to bf16 between the two, as the composition stores it. + + Fusion boundary: both gates, the q/k L2 norm, the recurrence, the indexed state + read and write, the gated RMS norm. The projection and the conv are the caller's. + """ + + 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). + 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)), + Cell( + 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), + ), + ) + + note = """ + Metadata consumed: `slots` (the cache manager's state indices for the step). + + Preconditions: bf16 `proj` `[rows, 2*H*K + 2*HV*V + 2*HV]` with unit column stride + and 8-byte aligned rows; the pool fp32 `[slots, HV, V, K]` with K innermost (any + outer strides, multiples of 4) and 16-byte aligned; `K == 128`, `V % 8 == 0` with + `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 + 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). GB300, graph replay, 461-slot pool, us (FlashInfer's + recurrence + the gated norm entry -> this): 128 rows 132.2 -> 118.8 (6.78 TB/s + against an in-place `mul_` over the same bytes at 6.83), 64 rows 70.7 -> 62.7, 16 + rows 19.1 -> 14.9, 4 rows 8.5 -> 8.4. + + 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 + rounded once. Inside torch's default band for bf16 / fp32 of the composition's + float64 reference, not bitwise FlashInfer's kernel. + + 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). Under graph replay the op reads the + projection after `griddepcontrol.wait` (16 rows). Not certified: `pdl=False`. + """ + + def __call__( + self, + proj: torch.Tensor, + pool: torch.Tensor, + slots: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + k_heads: int, + kdim: int, + pdl: bool = True, + ) -> torch.Tensor: + return launch(proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, pdl) + + def reference( + self, + proj: torch.Tensor, + pool: torch.Tensor, + slots: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + k_heads: int, + kdim: int, + pdl: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """`(out, pool_after)`: `gdn_decode_gated_norm_pdl`'s reference, the same contract -- + the fp64 recurrence (o rounded to bf16, the pool to fp32), then the gated norm + entry's fp32 reference.""" + return _one_launch.reference( + proj, pool, slots, A_log, dt_bias, norm_weight, eps, k_heads, kdim, pdl + ) + + def is_valid( + self, + proj: torch.Tensor, + pool: torch.Tensor, + slots: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + k_heads: int, + kdim: int, + pdl: bool = True, + ) -> None: + rows = proj.shape[0] + v_heads, vdim, kd = pool.shape[1], pool.shape[2], pool.shape[3] + assert proj.dtype == torch.bfloat16 and proj.stride(1) == 1, "bf16 proj, unit column stride" + assert proj.shape[1] == 2 * k_heads * kdim + 2 * v_heads * vdim + 2 * v_heads, "proj width" + assert proj.data_ptr() % 8 == 0 and proj.stride(0) % 4 == 0, "8-byte aligned proj rows" + assert pool.dtype == torch.float32 and pool.stride(3) == 1 and kd == kdim, ( + "fp32 pool [S, HV, V, K]" + ) + assert pool.data_ptr() % 16 == 0 and all(s % 4 == 0 for s in pool.stride()[:3]), ( + "16-byte aligned pool rows" + ) + assert kdim == 32 * _PL, f"K == {32 * _PL}" + assert vdim % _WARPS == 0 and vdim // _WARPS in (4, 8, 16, 32), "V / 8 in {4, 8, 16, 32}" + assert v_heads % k_heads == 0, "HV a multiple of H" + assert slots.dtype == torch.int32 and slots.shape == (rows,), "int32 slots [rows]" + assert slots.is_contiguous(), "contiguous slots" + assert A_log.dtype == torch.float32 and dt_bias.dtype == torch.float32, ( + "fp32 gate constants" + ) + assert A_log.is_contiguous() and dt_bias.is_contiguous(), "contiguous gate constants" + assert norm_weight.dtype == torch.bfloat16 and norm_weight.shape == (vdim,), ( + "bf16 norm weight [V]" + ) + + +class GdnDecodeRows(OpWrapper): + """`GdnDecodeRowsGatedNorm` without the gated norm: the recurrence's output `o` (rounded + to bf16) written into `out` `[rows, HV * V]` (unit column stride, any row stride) -- + a mixed step's generation rows, whose gated norm runs with the context rows'. The same + kernel, operands and preconditions otherwise (`proj`'s z columns are not read). + + Fusion boundary: both gates, the q/k L2 norm, the recurrence, the indexed state read and + write. Returns `out`. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="mixed step: 121 generation rows into a row range of the step output", + spec=dict(rows=121, out="row_range"), + ), + Cell(why="3 rows, one pad slot (-1)", spec=dict(rows=3, pad=True, out="row_range")), + Cell( + why="2 rows, the cache manager's coalesced pool: slot stride 36 layers", + spec=dict(rows=2, pool="coalesced", out="row_range"), + ), + ) + + note = """ + Metadata consumed: `slots`. Preconditions and launch: `GdnDecodeRowsGatedNorm`'s, plus + a bf16 `out` `[rows, HV * V]` with unit column stride. The pool rows and `slots` are + read before `griddepcontrol.wait`: in a mixed step nothing queued ahead of the call + 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. + """ + + def __call__( + self, + proj: torch.Tensor, + pool: torch.Tensor, + slots: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + k_heads: int, + kdim: int, + out: torch.Tensor, + pdl: bool = True, + ) -> torch.Tensor: + return launch(proj, pool, slots, A_log, dt_bias, None, 0.0, k_heads, kdim, pdl, out=out) + + def reference( + self, + proj: torch.Tensor, + pool: torch.Tensor, + slots: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + k_heads: int, + kdim: int, + out: torch.Tensor, + pdl: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """`(o, pool_after)`: the fp64 recurrence, o rounded to bf16 (the pool to fp32).""" + rows = proj.shape[0] + v_heads, vdim = pool.shape[1], pool.shape[2] + q, k, v, _, a, b = torch.split( + proj, + [k_heads * kdim, k_heads * kdim, v_heads * vdim, v_heads * vdim, v_heads, v_heads], + 1, + ) + o, pool_after = _recurrence_reference( + q.reshape(rows, 1, k_heads, kdim), + k.reshape(rows, 1, k_heads, kdim), + v.reshape(rows, 1, v_heads, vdim), + A_log, + a.reshape(rows, 1, v_heads), + dt_bias, + b.reshape(rows, 1, v_heads), + pool, + slots, + ) + return o.reshape(rows, v_heads * vdim), pool_after + + def is_valid( + self, + proj: torch.Tensor, + pool: torch.Tensor, + slots: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + k_heads: int, + kdim: int, + out: torch.Tensor, + pdl: bool = True, + ) -> None: + v_heads, vdim = pool.shape[1], pool.shape[2] + gdn_decode_rows_gated_norm.is_valid( + proj, + pool, + slots, + A_log, + dt_bias, + torch.empty(vdim, dtype=torch.bfloat16), + 0.0, + k_heads, + kdim, + ) + assert out.dtype == torch.bfloat16 and out.dim() == 2, "bf16 out [rows, HV * V]" + assert out.shape == (proj.shape[0], v_heads * vdim) and out.stride(1) == 1, ( + "out [rows, HV * V], unit column stride" + ) + + +gdn_decode_rows_gated_norm = GdnDecodeRowsGatedNorm() +gdn_decode_rows = GdnDecodeRows() 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 new file mode 100644 index 000000000000..4110469820d5 --- /dev/null +++ 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 @@ -0,0 +1,432 @@ +# 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/gdn_prefill_conv_post.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py new file mode 100644 index 000000000000..a24cbbf3865a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/gdn_prefill_conv_post.py @@ -0,0 +1,1043 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gated-DeltaNet prefill preparation in one pass: the causal depthwise conv1d +of the context rows, read token-major straight off the input projection, with +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. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda +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.kernels.blackwell.utils import make_ptr + +#: 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 +#: 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). +ENGINE_SLOTS = 461 +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 + + +@triton.jit +def _gdn_prefill_conv_publish_kernel( + x_ptr, + stride_xt, + state_ptr, + stride_ss, + stride_sc, + stride_sw, + qsl_ptr, + slot_ptr, + init_ptr, + conv_dim, + pad_slot_id, + BLOCK_C: tl.constexpr, +): + """One program: one sequence x BLOCK_C channels of its new 3-column window. + + New column j holds position `e - 3 + j`: the raw input when it lies inside + the sequence, else the carried column `L + j` (or 0 without a carry). All + three are loaded before any is stored, and each channel's three columns + live in one thread, so a short sequence's shift reads the old window.""" + n = tl.program_id(0) + c = tl.program_id(1) * BLOCK_C + tl.arange(0, BLOCK_C) + c_mask = c < conv_dim + slot = tl.load(slot_ptr + n) + if slot != pad_slot_id: + s0 = tl.load(qsl_ptr + n) + e = tl.load(qsl_ptr + n + 1) + carry = tl.load(init_ptr + n) != 0 + base = state_ptr + slot.to(tl.int64) * stride_ss + c.to(tl.int64) * stride_sc + p0 = e - 3 + p1 = e - 2 + p2 = e - 1 + x0 = tl.load(x_ptr + p0.to(tl.int64) * stride_xt + c, mask=c_mask & (p0 >= s0), other=0.0) + x1 = tl.load(x_ptr + p1.to(tl.int64) * stride_xt + c, mask=c_mask & (p1 >= s0), other=0.0) + x2 = tl.load(x_ptr + p2.to(tl.int64) * stride_xt + c, mask=c_mask & (p2 >= s0), other=0.0) + s_0 = tl.load(base + (3 + p0 - s0) * stride_sw, mask=c_mask & (p0 < s0) & carry, other=0.0) + s_1 = tl.load(base + (3 + p1 - s0) * stride_sw, mask=c_mask & (p1 < s0) & carry, other=0.0) + s_2 = tl.load(base + (3 + p2 - s0) * stride_sw, mask=c_mask & (p2 < s0) & carry, other=0.0) + n0 = tl.where(p0 >= s0, x0, s_0) + n1 = tl.where(p1 >= s0, x1, s_1) + n2 = tl.where(p2 >= s0, x2, s_2) + tl.store(base, n0, mask=c_mask) + tl.store(base + stride_sw, n1, mask=c_mask) + tl.store(base + 2 * stride_sw, n2, mask=c_mask) + + +# -------------------------------------------------------------------------------------- +# The conv pass in CuTe DSL. A CTA is one 128-channel head x `_TT` tokens, 128 threads: +# it first copies its `_TT + 3`-row tile (3 history rows) into shared memory with +# 16-byte `cp.async` (rows outside `[0, T)` zero-filled), so the bytes in flight cost no +# registers; then a thread (4 channels x `_TPT` consecutive tokens) slides the width-4 +# window over shared-memory rows, every element converted to fp32 once. The conv runs +# as `fma.rn.f32x2` on channel pairs (first tap from zero, then in tap order: per element +# an `fma.rn.f32` chain); the SiLU is `acc * rcp(1 + 2^(-acc * log2e))` +# with `ex2.approx.ftz` and `rcp.approx.ftz`, rounded in pairs by `cvt.rn.bf16x2.f32`. +# -------------------------------------------------------------------------------------- + +_VEC = 4 +_TPT = 16 +_RG = 4 +_TT = _TPT * _RG +_TILE_ROWS = _TT + 3 +_THREADS = 128 +_CHUNKS = _TILE_ROWS * 16 + + +def _iv(x, loc, ip): + return x.ir_value(loc=loc, ip=ip) + + +def _asm1(ret, args, text, constraints, loc, ip, side_effects=False): + return llvm.inline_asm( + ret, + [_iv(a, loc, ip) for a in args], + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _asm2(ty, args, text, constraints, loc, ip, side_effects=False): + res = _asm1( + llvm.StructType.get_literal([ty, ty]), args, text, constraints, loc, ip, side_effects + ) + return ( + llvm.extractvalue(ty, res, [0], loc=loc, ip=ip), + llvm.extractvalue(ty, res, [1], loc=loc, ip=ip), + ) + + +@dsl_user_op +def _ffma2(a0, a1, b0, b1, c0, c1, *, loc=None, ip=None): + """`fma.rn.f32x2`: two fp32 FMAs, each rounded once.""" + r = _asm2( + T.f32(), + [cutlass.Float32(v) for v in (a0, a1, b0, b1, c0, c1)], + "{ .reg .b64 a, b, c, d; mov.b64 a, {$2, $3}; mov.b64 b, {$4, $5}; " + "mov.b64 c, {$6, $7}; fma.rn.f32x2 d, a, b, c; mov.b64 {$0, $1}, d; }", + "=f,=f,f,f,f,f,f,f", + loc, + ip, + ) + return cutlass.Float32(r[0]), cutlass.Float32(r[1]) + + +@dsl_user_op +def _fmul2(a0, a1, b0, b1, *, loc=None, ip=None): + """`mul.rn.f32x2`.""" + r = _asm2( + T.f32(), + [cutlass.Float32(v) for v in (a0, a1, b0, b1)], + "{ .reg .b64 a, b, d; mov.b64 a, {$2, $3}; mov.b64 b, {$4, $5}; " + "mul.rn.f32x2 d, a, b; mov.b64 {$0, $1}, d; }", + "=f,=f,f,f,f,f", + loc, + ip, + ) + return cutlass.Float32(r[0]), cutlass.Float32(r[1]) + + +@dsl_user_op +def _silu2(a0, a1, *, loc=None, ip=None): + """`acc * rcp(1 + 2^(acc * -log2e))` for two values: `mul.rn.f32x2`, `ex2.approx.ftz`, + `add.rn.f32x2`, `rcp.approx.ftz`, `mul.rn`. For `1 + e` in [1, 2^126] this equals a + Triton `acc / (1 + exp(-acc))` (`div.full.f32`); past it (acc < -87.3) it gives -0 + where that gives a value below 1e-36.""" + r = _asm2( + T.f32(), + [cutlass.Float32(a0), cutlass.Float32(a1)], + "{ .reg .b64 a, k, t, o, d; .reg .f32 c, t0, t1, d0, d1; " + "mov.f32 c, 0fBFB8AA3B; mov.b64 k, {c, c}; mov.f32 c, 0f3F800000; mov.b64 o, {c, c}; " + "mov.b64 a, {$2, $3}; mul.rn.f32x2 t, a, k; mov.b64 {t0, t1}, t; " + "ex2.approx.ftz.f32 t0, t0; ex2.approx.ftz.f32 t1, t1; mov.b64 t, {t0, t1}; " + "add.rn.f32x2 d, t, o; mov.b64 {d0, d1}, d; " + "rcp.approx.ftz.f32 d0, d0; rcp.approx.ftz.f32 d1, d1; " + "mul.rn.f32 $0, $2, d0; mul.rn.f32 $1, $3, d1; }", + "=f,=f,f,f", + loc, + ip, + ) + return cutlass.Float32(r[0]), cutlass.Float32(r[1]) + + +@dsl_user_op +def _fma_rn(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm1( + T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _inv_norm(ss, eps, *, loc=None, ip=None): + """1 / sqrt(ss + eps): `add.f32`, `sqrt.approx.ftz.f32`, `div.full.f32` (a Triton + `1.0 / tl.sqrt(...)`).""" + return cutlass.Float32( + _asm1( + T.f32(), + [cutlass.Float32(ss), cutlass.Float32(eps)], + "{ .reg .f32 t; add.f32 t, $1, $2; sqrt.approx.ftz.f32 t, t; " + "div.full.f32 $0, 0f3F800000, t; }", + "=f,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16x2(lo, hi, *, loc=None, ip=None): + """`cvt.rn.bf16x2.f32`: `lo` in the low half (the lower address).""" + return cutlass.Int32( + _asm1( + T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + loc, + ip, + ) + ) + + +@dsl_user_op +def _unpack2(w, *, loc=None, ip=None): + """A bf16x2 word as `(low, high)` fp32.""" + r = _asm2( + T.f32(), + [cutlass.Int32(w)], + "{ .reg .b32 l, h; shl.b32 l, $2, 16; and.b32 h, $2, 0xFFFF0000; " + "mov.b32 $0, l; mov.b32 $1, h; }", + "=f,=f,r", + loc, + ip, + ) + return cutlass.Float32(r[0]), cutlass.Float32(r[1]) + + +@dsl_user_op +def _st_v2(addr, a, b, *, loc=None, ip=None): + _asm1( + T.i32(), + [cutlass.Int64(addr), cutlass.Int32(a), cutlass.Int32(b)], + "{ st.global.v2.b32 [$1], {$2, $3}; mov.u32 $0, 0; }", + "=r,l,r,r", + loc, + ip, + side_effects=True, + ) + + +@dsl_user_op +def _cp_async16(smem_addr, gmem_addr, src_bytes, *, loc=None, ip=None): + """16 bytes global -> shared, the rest of the 16 zero-filled past `src_bytes`.""" + _asm1( + T.i32(), + [cutlass.Int32(smem_addr), cutlass.Int64(gmem_addr), cutlass.Int32(src_bytes)], + "{ cp.async.cg.shared.global [$1], [$2], 16, $3; mov.u32 $0, 0; }", + "=r,r,l,r", + 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) + + +@dsl_user_op +def _lds_v2(addr, *, loc=None, ip=None): + r = _asm2( + T.i32(), + [cutlass.Int32(addr)], + "ld.shared.v2.b32 {$0, $1}, [$2];", + "=r,=r,r", + loc, + ip, + side_effects=True, + ) + return cutlass.Int32(r[0]), cutlass.Int32(r[1]) + + +class _ConvPassKernel: + """q / k / v of the conv pass; the gates and the publish are separate launches.""" + + def __init__(self, num_k_heads: int, num_v_heads: int, head_dim: int, x_stride: int): + assert head_dim == _VEC * 32 + self.hk, self.hv, self.d = num_k_heads, num_v_heads, head_dim + self.heads = 2 * num_k_heads + num_v_heads + self.qkv = self.heads * head_dim + self.sx = x_stride + + @cute.jit + def __call__( + self, + x_ptr: cute.Pointer, + w_ptr: cute.Pointer, + st_ptr: cute.Pointer, + qsl_ptr: cute.Pointer, + slot_ptr: cute.Pointer, + init_ptr: cute.Pointer, + q_t: cute.Tensor, + k_ptr: cute.Pointer, + v_ptr: cute.Pointer, + num_tokens: cutlass.Int32, + num_seqs: cutlass.Int32, + search_steps: cutlass.Int32, + pad_slot_id: cutlass.Int32, + eps: cutlass.Float32, + w_s0: cutlass.Int32, + w_s1: cutlass.Int32, + st_s0: cutlass.Int32, + st_s1: cutlass.Int32, + st_s2: cutlass.Int32, + 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)) + ) + 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,))) + blocks = (num_tokens + _TT - 1) // _TT + # `q` arrives as a tensor: the TVM-FFI entry takes its stream from a GPU tensor. + self.kernel( + x_ptr.toint(), + weight, + pool, + qsl, + slots, + init, + q_t.iterator.toint(), + k_ptr.toint(), + v_ptr.toint(), + num_tokens, + num_seqs, + search_steps, + pad_slot_id, + eps, + ).launch( + grid=[self.heads, blocks, 1], + block=[_THREADS, 1, 1], + smem=_TILE_ROWS * 128 * 2 + 128, + stream=stream, + ) + + @cute.kernel + def kernel( + self, + x_base: cutlass.Int64, + mW: cute.Tensor, + mState: cute.Tensor, + mQsl: cute.Tensor, + mSlot: cute.Tensor, + mInit: cute.Tensor, + q_base: cutlass.Int64, + k_base: cutlass.Int64, + v_base: cutlass.Int64, + num_tokens: cutlass.Int32, + num_seqs: cutlass.Int32, + search_steps: cutlass.Int32, + pad_slot_id: cutlass.Int32, + eps: cutlass.Float32, + ): + tidx, _, _ = cute.arch.thread_idx() + head, tb, _ = cute.arch.block_idx() + lane = tidx % 32 + rg = tidx // 32 + t0 = tb * _TT + + smem = cutlass.utils.SmemAllocator() + tile = smem.allocate_tensor( + cutlass.BFloat16, cute.make_layout((_TILE_ROWS * 128,), stride=(1,)), byte_alignment=16 + ) + tile_base = cutlass.Int32(tile.iterator.toint()) + + # 1. The tile: rows t0-3 .. t0+TT-1 of this head's 128 channels, zero-filled outside [0, T). + for it in cutlass.range_constexpr((_CHUNKS + _THREADS - 1) // _THREADS): + chunk = tidx + it * _THREADS + if chunk < _CHUNKS: + r = chunk // 16 + cc = chunk - r * 16 + row = t0 - 3 + r + valid = cutlass.Int32(0) + if row >= 0: + if row < num_tokens: + valid = cutlass.Int32(16) + rowc = cutlass.min(cutlass.max(row, cutlass.Int32(0)), num_tokens - 1) + _cp_async16( + tile_base + (r * 128 + cc * 8) * 2, + x_base + (cutlass.Int64(rowc) * self.sx + head * 128 + cc * 8) * 2, + valid, + ) + + # 2. While the copies land: the weights, the first token's sequence, the output row. + c0 = head * 128 + lane * _VEC + w = cute.make_rmem_tensor(cute.make_layout((_VEC, 4), stride=(4, 1)), cutlass.Float32) + for i in cutlass.range_constexpr(_VEC * 4): + w[i // 4, i % 4] = mW[c0 + i // 4, i % 4].to(cutlass.Float32) + tf = t0 + rg * _TPT + lo = cutlass.Int32(0) + hi = num_seqs + for _ in cutlass.range(search_steps, unroll=1): + mid = (lo + hi) // 2 + if mQsl[mid] <= tf: + lo = mid + else: + hi = mid + s0f = mQsl[lo] + ef = mQsl[lo + 1] + # One uniform branch a thread: its tokens and their windows inside one sequence. + fast = cutlass.Int32(0) + if tf - s0f >= 3: + if tf + _TPT <= ef: + if tf + _TPT <= num_tokens: + fast = cutlass.Int32(1) + is_qk = head < 2 * self.hk + o_base = v_base + cutlass.Int64((head - 2 * self.hk) * 128 + lane * _VEC) * 2 + o_stride = cutlass.Int64(self.hv * 128 * 2) + if head < self.hk: + o_base = q_base + cutlass.Int64(head * 128 + lane * _VEC) * 2 + o_stride = cutlass.Int64(self.hk * 128 * 2) + else: + if is_qk: + o_base = k_base + cutlass.Int64((head - self.hk) * 128 + lane * _VEC) * 2 + o_stride = cutlass.Int64(self.hk * 128 * 2) + + _cp_async_wait_all() + cute.arch.sync_threads() + + # 3. The window over tile rows rg*TPT .. (tile row r is token t0 - 3 + r). + win = cute.make_rmem_tensor(cute.make_layout((4, _VEC), stride=(_VEC, 1)), cutlass.Float32) + for i in cutlass.range_constexpr(3): + lw, hw = _lds_v2(tile_base + ((rg * _TPT + i) * 128 + lane * _VEC) * 2) + win[i, 0], win[i, 1] = _unpack2(lw) + win[i, 2], win[i, 3] = _unpack2(hw) + acc = cute.make_rmem_tensor(cute.make_layout((_VEC,), stride=(1,)), cutlass.Float32) + for j in cutlass.range_constexpr(_TPT): + t = tf + j + lw, hw = _lds_v2(tile_base + ((rg * _TPT + j + 3) * 128 + lane * _VEC) * 2) + win[(j + 3) % 4, 0], win[(j + 3) % 4, 1] = _unpack2(lw) + win[(j + 3) % 4, 2], win[(j + 3) % 4, 3] = _unpack2(hw) + if fast == 1: + for p in cutlass.range_constexpr(_VEC // 2): + a0, a1 = _ffma2( + win[j % 4, 2 * p], + win[j % 4, 2 * p + 1], + w[2 * p, 0], + w[2 * p + 1, 0], + cutlass.Float32(0.0), + cutlass.Float32(0.0), + ) + for kk in cutlass.range_constexpr(1, 4): + a0, a1 = _ffma2( + win[(j + kk) % 4, 2 * p], + win[(j + kk) % 4, 2 * p + 1], + w[2 * p, kk], + w[2 * p + 1, kk], + a0, + a1, + ) + acc[2 * p] = a0 + acc[2 * p + 1] = a1 + else: + # The token's sequence; a tap before its start reads the carried + # column W-1-(s0-pos), or 0 without a carry. + sq = lo + s0 = s0f + if t >= ef: + lo2 = cutlass.Int32(0) + hi2 = num_seqs + for _ in cutlass.range(search_steps, unroll=1): + mid = (lo2 + hi2) // 2 + if mQsl[mid] <= t: + lo2 = mid + else: + hi2 = mid + sq = lo2 + s0 = mQsl[lo2] + slot = mSlot[sq] + init = cutlass.Int32(mInit[sq]) + for v in cutlass.range_constexpr(_VEC): + a = cutlass.Float32(0.0) + for kk in cutlass.range_constexpr(4): + pos = t - 3 + kk + xv = win[(j + kk) % 4, v] + if pos < s0: + xv = cutlass.Float32(0.0) + if init != 0: + if slot != pad_slot_id: + xv = mState[slot, c0 + v, 3 + pos - s0].to(cutlass.Float32) + a = _fma_rn(xv, w[v, kk], a) + acc[v] = a + y0, y1 = _silu2(acc[0], acc[1]) + y2, y3 = _silu2(acc[2], acc[3]) + p01 = _bf16x2(y0, y1) + p23 = _bf16x2(y2, y3) + if is_qk: + # L2 norm of the rounded head row: fp32 squares, a 32-lane butterfly. + f0, f1 = _unpack2(p01) + f2, f3 = _unpack2(p23) + s_a, s_b = _ffma2(f0, f1, f0, f1, cutlass.Float32(0.0), cutlass.Float32(0.0)) + s_a, s_b = _ffma2(f2, f3, f2, f3, s_a, s_b) + ss = s_a + s_b + for off in cutlass.range_constexpr(5): + ss = ss + cute.arch.shuffle_sync_bfly(ss, offset=1 << off) + inv = _inv_norm(ss, eps) + n0, n1 = _fmul2(f0, f1, inv, inv) + n2, n3 = _fmul2(f2, f3, inv, inv) + p01 = _bf16x2(n0, n1) + p23 = _bf16x2(n2, n3) + if t < num_tokens: + _st_v2(o_base + cutlass.Int64(t) * o_stride, p01, p23) + + +@triton.jit +def _gdn_gates_kernel( + a_ptr, + stride_at, + stride_ah, + b_ptr, + stride_bt, + stride_bh, + A_log_ptr, + dt_bias_ptr, + g_ptr, + beta_ptr, + num_tokens, + softplus_threshold, + NUM_V_HEADS: tl.constexpr, + HP: tl.constexpr, + BLOCK_T: tl.constexpr, +): + """The decay/beta gates over `[T, NUM_V_HEADS]`, one row-coalesced tile a program.""" + t = (tl.program_id(0) * BLOCK_T + tl.arange(0, BLOCK_T))[:, None] + h = tl.arange(0, HP)[None, :] + m = (t < num_tokens) & (h < NUM_V_HEADS) + t64 = t.to(tl.int64) + a_val = tl.load(a_ptr + t64 * stride_at + h * stride_ah, mask=m, other=0.0) + b_val = tl.load(b_ptr + t64 * stride_bt + h * stride_bh, mask=m, other=0.0) + A_log = tl.load(A_log_ptr + h, mask=h < NUM_V_HEADS, other=0.0).to(tl.float32) + dt_bias = tl.load(dt_bias_ptr + h, mask=h < NUM_V_HEADS, other=0.0).to(tl.float32) + gate_input = a_val.to(tl.float32) + dt_bias + softplus = tl.where( + gate_input <= softplus_threshold, tl.log(1.0 + tl.exp(gate_input)), gate_input + ) + g_off = t64 * NUM_V_HEADS + h + tl.store(g_ptr + g_off, -tl.exp(A_log) * softplus, mask=m) + tl.store(beta_ptr + g_off, tl.sigmoid(b_val.to(tl.float32)), mask=m) + + +_CONV_PASS: dict = {} + + +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.""" + 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 + ) + + +def _cute_conv_pass( + x, weight, conv_states, query_start_loc, cache_indices, has_initial_state, q, k, v, + num_k_heads, num_v_heads, head_k_dim, pad_slot_id, l2_norm_eps, +): # fmt: skip + num_tokens = x.shape[0] + num_seqs = query_start_loc.shape[0] - 1 + key = (num_k_heads, num_v_heads, head_k_dim, x.stride(0)) + q2 = q.view(num_tokens, -1) + ptrs = ( + x.data_ptr(), + weight.data_ptr(), + conv_states.data_ptr(), + query_start_loc.data_ptr(), + cache_indices.data_ptr(), + has_initial_state.data_ptr(), + q2, + k.data_ptr(), + v.data_ptr(), + ) + fn = _CONV_PASS.get(key) + if fn is None: + gm = cute.AddressSpace.gmem + types = ( + (cutlass.BFloat16, 16), + (cutlass.BFloat16, 2), + (cutlass.BFloat16, 2), + (cutlass.Int32, 4), + (cutlass.Int32, 4), + (cutlass.Uint8, 1), + None, + (cutlass.BFloat16, 16), + (cutlass.BFloat16, 16), + ) + cptrs = [ + from_dlpack(p, assumed_align=16).mark_layout_dynamic(leading_dim=1) + if spec is None + else make_ptr(spec[0], p, gm, assumed_align=spec[1]) + for spec, p in zip(types, ptrs) + ] + fn = cute.compile( + _ConvPassKernel(num_k_heads, num_v_heads, head_k_dim, x.stride(0)), + *cptrs, + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Int32(1), + cutlass.Int32(-1), + cutlass.Float32(1e-6), + *(cutlass.Int32(1),) * 5, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--opt-level 2 --enable-tvm-ffi", + ) + _CONV_PASS[key] = fn + fn( + *ptrs, + num_tokens, + num_seqs, + max(num_seqs - 1, 0).bit_length() + 1, + pad_slot_id, + l2_norm_eps, + *weight.stride(), + *conv_states.stride(), + ) + + +def _launch( + x, + weight, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + a, + b, + A_log, + dt_bias, + num_k_heads, + head_k_dim, + num_v_heads, + head_v_dim, + pad_slot_id, + l2_norm_eps, + softplus_threshold, +): + assert head_k_dim == head_v_dim and weight.shape[-1] == CONV_WIDTH, ( + "one program per 128-channel head: head_k_dim must equal head_v_dim, and the " + f"conv width must be {CONV_WIDTH}" + ) + num_tokens = x.shape[0] + num_seqs = query_start_loc.shape[0] - 1 + conv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim + dev, dt = x.device, x.dtype + q = torch.empty((1, num_tokens, num_k_heads, head_k_dim), dtype=dt, device=dev) + k = torch.empty_like(q) + v = torch.empty((1, num_tokens, num_v_heads, head_v_dim), dtype=dt, device=dev) + g = torch.empty((1, num_tokens, num_v_heads), dtype=torch.float32, device=dev) + beta = torch.empty((1, num_tokens, num_v_heads), dtype=torch.float32, device=dev) + if num_seqs == 0: + return q, k, v, g, beta + if num_tokens > 0: + assert _cute_conv_pass_ok(x, weight, conv_states, num_k_heads, head_k_dim, num_v_heads), ( + "the conv pass takes bf16 x with 128-channel heads, 16-byte aligned rows and " + f"int32 weight/pool strides; got {x.dtype}, head_k_dim {head_k_dim}, x at " + f"{x.data_ptr() % 16} mod 16 with row stride {x.stride(0)}" + ) + _cute_conv_pass( + x, weight, conv_states, query_start_loc, cache_indices, has_initial_state, q, k, v, + num_k_heads, num_v_heads, head_k_dim, pad_slot_id, l2_norm_eps, + ) # fmt: skip + _gdn_gates_kernel[(triton.cdiv(num_tokens, 64),)]( + a, + a.stride(0), + a.stride(1), + b, + b.stride(0), + b.stride(1), + A_log, + dt_bias, + g, + beta, + num_tokens, + softplus_threshold, + NUM_V_HEADS=num_v_heads, + HP=triton.next_power_of_2(num_v_heads), + BLOCK_T=64, + num_warps=4, + ) + _gdn_prefill_conv_publish_kernel[(num_seqs, triton.cdiv(conv_dim, _PUBLISH_BLOCK_C))]( + x, + x.stride(0), + conv_states, + *conv_states.stride(), + query_start_loc, + cache_indices, + has_initial_state, + conv_dim, + pad_slot_id, + BLOCK_C=_PUBLISH_BLOCK_C, + num_warps=4, + ) + return q, k, v, g, beta + + +def _engine_slots(num_seqs: int) -> tuple[int, ...]: + """`num_seqs` distinct slots of the engine pool, the last and slot 0 first.""" + rest = tuple((i * 37 + 4) % (ENGINE_SLOTS - 2) + 1 for i in range(max(num_seqs - 2, 0))) + slots = ((ENGINE_SLOTS - 1, 0) + rest)[:num_seqs] + assert len(set(slots)) == num_seqs + return slots + + +class GdnPrefillConvPost(OpWrapper): + """Causal conv + SiLU of the context rows, then the GDN q/k/v/g/beta. + + `x` is `[T, >= qkv_dim]` token-major with unit channel stride (any row + stride: the target passes the q|k|v column view of its `[T, 16480]` input + projection), packed `[q | k | v]`, `qkv_dim = 2 * num_k_heads * head_k_dim + + num_v_heads * head_v_dim`. For sequence `s` over rows `[b, e)` of + `query_start_loc` and `slot = cache_indices[s]`: + + xpad = cat([conv_states[slot] if has_initial_state[s] else 0, x[b:e]]) + y[t] = bf16(silu(sum_k weight[:, k] * xpad[t + k])) (fp32 taps, in order) + q, k = y / sqrt(sum(y^2) + l2_norm_eps) (per head, fp32, rounded once) + v = y (bit-exact) + g = -exp(A_log) * softplus(a + dt_bias) (fp32) + beta = sigmoid(b) (fp32) + conv_states[slot] = xpad[e - b : e - b + W - 1] (raw input, exact) + + -- the composition `extract_transpose_prefill_slice` -> `causal_conv1d_fwd` + (SiLU, no bias) -> `fused_gdn_post_conv(prefill, None, ...)`, in one pass + over the activation plus a per-sequence window publish. + + Fusion boundary: the conv-state read, conv, SiLU, split, q/k norm, gates + and the state publish are inside. The projection, the generation rows' + conv and the recurrence are the caller's (`use_qk_l2norm_in_kernel=False` + downstream). + + Returns `(q, k, v, g, beta)`: fresh contiguous `[1, T, num_k_heads, + head_k_dim]` x2, `[1, T, num_v_heads, head_v_dim]` in `x.dtype`, and + `[1, T, num_v_heads]` fp32 twice. `conv_states` is advanced in place for + every sequence whose slot is not `pad_slot_id`; nothing else is written. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # 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. + CELLS: tuple[Cell, ...] = ( + Cell( + why="target call, engine 461-slot pool: 37+1+128 tokens at slots 460, 0, 230", + spec=dict( + lengths=(37, 1, 128), + pool_slots=ENGINE_SLOTS, + slots=(ENGINE_SLOTS - 1, 0, ENGINE_SLOTS // 2), + initial=(True, False, True), + ), + ), + Cell( + why="mixed step's context half: first 8065 rows of a 8192-row projection", + spec=dict( + lengths=(8065,), + pool_slots=ENGINE_SLOTS, + initial=(False,), + layout="rows", + extra_rows=127, + ), + ), + Cell( + why="128 short contexts (the batch ceiling), some shorter than the window", + spec=dict( + lengths=tuple(1 + (i * 7) % 64 for i in range(128)), + pool_slots=ENGINE_SLOTS, + initial=tuple(i % 3 != 0 for i in range(128)), + ), + ), + ) + + note = """ + Metadata consumed: none. Batch layout (`query_start_loc`, int32 `[N + 1]`, + `query_start_loc[0] == 0`, `query_start_loc[N] == x.shape[0]`), slot map + (`cache_indices`, int32 `[N]`) and carry flags (`has_initial_state`, bool + `[N]`) are explicit. + + Call sites (`PrefillTarget._gdn_prefill_prep`, from `_gdn_block`): the + context rows of every prefill-carrying step -- prefill-only `(mixed_qkv, + a, b, qsl_all, si_all, hi_all)`, mixed and MTP-verify `(mq_p, a_p, b_p, + qsl_p, si_p, hi_p)` -- with the layer's `[10240, 4]` conv weight, its + whole conv pool unsliced, `pad_slot_id` -1, `l2_norm_eps` 1e-6 and + `softplus_threshold` 20. `x`, `a`, `b` are column (and row) views, row + stride 16480, of the one input-projection output; the outputs feed + `flashinfer_chunk_gated_delta_rule` (the Triton chunk on the verify path) + with `use_qk_l2norm_in_kernel=False`. + + 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 + 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 `[64, 48]` tile a + program. 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). + + 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 + (`fma.rn.f32x2` on channel pairs). SiLU: `acc * rcp(1 + 2^(-acc * + log2e))` with `ex2.approx.ftz` and `rcp.approx.ftz` -- a Triton + `acc / (1 + exp(-acc))`'s `ex2.approx` / `div.full` values wherever + `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 + one bf16 ulp on 1e-6..1e-4 of elements, g/beta within the fp32 band, conv + pool bitwise. + + Certified over the 461-slot pool, seen as one layer of a two-layer buffer: + three contexts of 37, 1 and 128 tokens; a mixed step's 8065-token context + half, the first rows of an 8192-row projection; and 128 short contexts (1 to + 64 tokens, some shorter than the window). + + 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 + context row a real slot; a pad context is not certified. + + Preconditions (`is_valid`): `x` 2-D, unit channel stride, at least + `qkv_dim` wide; `a`/`b` at least `T` rows (read through both strides); + A_log/dt_bias contiguous with `num_v_heads` entries; `weight` `[qkv_dim, + 4]` in `x`'s dtype (any strides); `conv_states` `[slots, qkv_dim, 3]` in + `x`'s dtype (any strides); int32 `query_start_loc`/`cache_indices`, bool + `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. + + 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. + + 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. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor, + has_initial_state: 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, + pad_slot_id: int = -1, + l2_norm_eps: float = 1e-6, + softplus_threshold: float = 20.0, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return _launch( + x, + weight, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + a, + b, + A_log, + dt_bias, + num_k_heads, + head_k_dim, + num_v_heads, + head_v_dim, + pad_slot_id, + l2_norm_eps, + softplus_threshold, + ) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor, + has_initial_state: 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, + pad_slot_id: int = -1, + l2_norm_eps: float = 1e-6, + softplus_threshold: float = 20.0, + ) -> tuple[torch.Tensor, ...]: + """`(q, k, v, g, beta, conv_states as the call leaves them)`; mutates + nothing. Conv taps in fp32 in tap order and the SiLU in fp32, rounded + to `x.dtype` once; the q/k norm in fp32 from that rounded value, + rounded once; `g` in fp64 with the stabilised softplus, rounded to + fp32; `beta` the fp32 sigmoid.""" + qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim + width = weight.shape[-1] + n_tok = x.shape[0] + w = weight.float() + y = torch.zeros(n_tok, qkv_dim, dtype=x.dtype, device=x.device) + states = conv_states.clone() + bounds = query_start_loc.tolist() + slots = cache_indices.tolist() + carry = has_initial_state.tolist() + for s in range(len(bounds) - 1): + lo, hi, slot = bounds[s], bounds[s + 1], slots[s] + tokens = x[lo:hi, :qkv_dim].float() + if slot != pad_slot_id and carry[s]: + prefix = conv_states[slot].float().t() + else: + prefix = tokens.new_zeros(width - 1, qkv_dim) + padded = torch.cat([prefix, tokens], dim=0) + length = hi - lo + acc = torch.zeros_like(tokens) + for kk in range(width): + acc = acc + padded[kk : kk + length] * w[:, kk] + y[lo:hi] = torch.nn.functional.silu(acc).to(x.dtype) + if slot != pad_slot_id: + states[slot] = padded[length : length + width - 1].t().to(states.dtype) + kd = num_k_heads * head_k_dim + + def l2norm(part): + part = part.reshape(n_tok, num_k_heads, head_k_dim).float() + return (part / torch.sqrt(part.pow(2).sum(-1, keepdim=True) + l2_norm_eps)).to(x.dtype) + + q, k = l2norm(y[:, :kd]), l2norm(y[:, kd : 2 * kd]) + v = y[:, 2 * kd :].reshape(n_tok, num_v_heads, head_v_dim) + gx = a[:n_tok, :num_v_heads].double() + dt_bias[:num_v_heads].double() + softplus = torch.where(gx <= softplus_threshold, torch.log1p(torch.exp(gx)), gx) + g = (-torch.exp(A_log[:num_v_heads].double()) * softplus).float() + beta = torch.sigmoid(b[:n_tok, :num_v_heads].float()) + return (*(t.unsqueeze(0) for t in (q, k, v, g, beta)), states) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor, + has_initial_state: 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, + pad_slot_id: int = -1, + l2_norm_eps: float = 1e-6, + softplus_threshold: float = 20.0, + ) -> None: + # Every address below is computed from the arguments' own extents and + # strides only where stated; the rest is `ptr + index`, unmasked. + qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim + n_tok = x.shape[0] + assert x.dim() == 2 and x.stride(1) == 1 and x.shape[1] >= qkv_dim, ( + f"x must be [T, >= {qkv_dim}] with unit channel stride; got {tuple(x.shape)} " + f"strides {x.stride()} -- the channel axis is read as `ptr + c`" + ) + assert weight.dtype == x.dtype and weight.shape[0] == qkv_dim, ( + f"weight must be [{qkv_dim}, 4] in x's dtype; got {tuple(weight.shape)} {weight.dtype}" + ) + assert conv_states.dtype == x.dtype and tuple(conv_states.shape[1:]) == (qkv_dim, 3), ( + f"conv_states must be [slots, {qkv_dim}, 3] in x's dtype; got " + f"{tuple(conv_states.shape)} {conv_states.dtype}" + ) + assert a.shape[0] >= n_tok and b.shape[0] >= n_tok, ( + f"a and b must have at least T={n_tok} rows; the token axis is unmasked" + ) + assert A_log.stride(0) == 1 and dt_bias.stride(0) == 1, ( + "A_log and dt_bias must be contiguous 1-D; they are read as `ptr + head`" + ) + assert A_log.numel() >= num_v_heads and dt_bias.numel() >= num_v_heads, ( + f"A_log and dt_bias must hold num_v_heads={num_v_heads} entries; the load is unmasked" + ) + assert query_start_loc.dtype == torch.int32 and cache_indices.dtype == torch.int32, ( + "query_start_loc and cache_indices must be int32" + ) + assert has_initial_state.dtype == torch.bool, "has_initial_state must be bool" + n_seq = query_start_loc.shape[0] - 1 + assert cache_indices.shape[0] >= n_seq and has_initial_state.shape[0] >= n_seq, ( + f"cache_indices and has_initial_state must cover the {n_seq} sequences" + ) + + +gdn_prefill_conv_post = GdnPrefillConvPost() 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 index 175b7f695dfd..8fe4119dbaaa 100644 --- 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 @@ -50,74 +50,15 @@ class PackGdnDecodeQkv(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: num_tokens, or (num_decodes, draft) for a verify batch; layout - # "contiguous" | "conv_dim_major" (the conv's [B, C, P] output transposed - # and flattened, as the stock mixer builds it; a view for B == 1) | - # "row_slice_of_projection" | "column_stride_2" | "transposed_view" | - # "wide" (trailing columns past qkv_dim); dtype; heads. - CELLS: tuple[Cell, ...] = ( - tuple( - Cell(why=f"{n} tokens, contiguous{note}", spec=dict(num_tokens=n)) - for n, note in ( - (1, ""), - (2, ""), - (5, ""), - (15, ": one short of the 16-token block"), - (16, ": exactly one 16-token block"), - (17, ": one past the 16-token block"), - (32, ": two full blocks"), - (127, ": a partial final block"), - ) + # 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), ) - + 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 ((1, 2), (8, 2), (32, 2), (8, 4)) - ) - + tuple( - Cell( - why=f"stock mixer verify: {b} x {p} off the [B, C, P] conv output", - spec=dict(num_decodes=b, draft=p, layout="conv_dim_major"), - ) - for b, p in ((1, 1), (4, 1), (1, 2), (5, 2), (3, 3), (8, 4), (16, 2), (2, 5)) - ) - + ( - Cell( - why="row slice of a [T, qkv_dim + z] projection, stride (16384, 1)", - spec=dict(num_tokens=24, layout="row_slice_of_projection"), - ), - Cell( - why="column-stride-2 view, stride (20480, 2)", - spec=dict(num_tokens=24, layout="column_stride_2"), - ), - Cell( - why="transposed dim-major view, stride (1, 24)", - spec=dict(num_tokens=24, layout="transposed_view"), - ), - Cell( - why="6144 trailing columns past qkv_dim are dropped", - spec=dict(num_tokens=24, layout="wide"), - ), - ) - + tuple( - Cell( - why=f"dtype {str(dtype).removeprefix('torch.')} is carried, not chosen", - spec=dict(num_tokens=11, dtype=dtype), - ) - for dtype in (torch.bfloat16, torch.float16, torch.float32) - ) - + tuple( - Cell(why=f"geometry {heads}{note}", spec=dict(num_tokens=13, heads=heads)) - for heads, note in ( - ((16, 128, 32, 128), ""), - ((4, 64, 8, 64), ""), - ((2, 96, 6, 80), ": head dims off a power of two"), - ) - ) - + (Cell(why="zero tokens: three empty tensors, no launch", spec=dict(num_tokens=0)),) - ) + 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 @@ -125,7 +66,9 @@ class PackGdnDecodeQkv(OpWrapper): 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. + 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` 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 index 5416d7bc8808..b61b97d3bb67 100644 --- 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 @@ -22,9 +22,10 @@ #: [conv_dim, d_conv - 1] = [10240, 3] convolution window. REC_SHAPE = (48, 128, 128) CONV_SHAPE = (10240, 3) -#: Largest slot count the cache manager was seen to allocate both pools at for -#: this checkpoint (10 .. 461 over twelve runtime-value dumps, batch never above -#: 32). The target hands both pools over unsliced. +#: 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 @@ -58,115 +59,32 @@ class ResetRecurrentStateRows(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: slots (pool rows), indices, flags, index_dtype (int32), flag_dtype - # (bool), rec_dtype (fp32), conv (True), rec_shape/conv_shape (checkpoint), - # layout: "plain" | "row_padding" (stride(0) > row) | "layer_view" - # (pool[1] of a [3, slots, ...] allocation); `indices="perm"` is a random - # permutation of every slot with random flags. + # 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="checkpoint geometry prefill batch, mixed flags", - spec=dict(slots=6, indices=(3, 0, 5, 1), flags=(False, True, False, True)), - ), - Cell( - why="engine pool 461, 4 requests, slot 460 reset", - spec=_engine_case(4, flip=False), - ), Cell( why="engine pool 461, 32 requests, slot 460 reset", - spec=_engine_case(32, flip=False), - ), - Cell( - why="engine pool 461, 32 requests, slot 0 reset", spec=_engine_case(32, flip=True), ), Cell( - why="every request continues: pools bitwise unchanged", - spec=dict(slots=6, indices=(0, 1, 2, 3, 4, 5), flags=(True,) * 6), - ), - Cell( - why="one named slot; the other five untouched", - spec=dict(slots=6, indices=(2,), flags=(False,)), - ), - ) - CELLS += tuple( - Cell( - why=f"negative and out-of-range slots skipped, {str(d)[6:]}", - spec=dict(slots=6, indices=(-1, 6, 23, -99999, 2), flags=(False,) * 5, index_dtype=d), - ) - for d in (torch.int32, torch.int64) - ) - CELLS += tuple( - Cell( - why=f"index dtype {str(d)[6:]} widened after the load", - spec=dict(slots=4, indices=(1,), flags=(False,), index_dtype=d), - ) - for d in (torch.int8, torch.int16, torch.int64, torch.uint8) - ) - CELLS += tuple( - Cell( - why=f"flag dtype {str(d)[6:]}: a nonzero test, not a low bit", - spec=dict( - slots=4, - indices=(0, 1, 2, 3), - flags=(0, 1, 1, 1) if d is torch.bool else (0, 1, 2, 3), - flag_dtype=d, - ), - ) - for d in (torch.bool, torch.uint8, torch.int32, torch.int64, torch.float32) - ) - CELLS += ( - Cell( - why="conv_states omitted: that pool is not touched", - spec=dict(slots=6, indices=(4,), flags=(False,), conv=False), - ), - ) - CELLS += tuple( - Cell( - why=f"slot named twice, flags {flags}", - spec=dict(slots=6, indices=(2, 2), flags=flags), - ) - for flags in ((False, True), (True, False), (True, True)) - ) - CELLS += ( - Cell( - why="stride(0) > row size: the padding lane untouched", - spec=dict(slots=6, indices=(1,), flags=(False,), conv=False, layout="row_padding"), + 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"), ), - Cell( - why="empty batch launches nothing", - spec=dict(slots=6, indices=(), flags=()), - ), - ) - CELLS += tuple( - Cell( - why=f"recurrent pool {str(d)[6:]} (--mamba_ssm_cache_dtype)", - spec=dict(slots=3, indices=(1,), flags=(False,), rec_dtype=d), - ) - for d in (torch.bfloat16, torch.float16) - ) - CELLS += ( - Cell( - why="conv row wider than the recurrent row", - spec=dict(slots=3, indices=(1,), flags=(False,), rec_shape=(8,), conv_shape=(5000,)), - ), - Cell( - why="64-slot random permutation: slot-indexed, not position-indexed", - spec=dict(slots=64, indices="perm", flags="random"), - ), ) 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 3..64 - and 461 slots; at 461 the result is bitwise identical to the 64-slot - result for the same slot content. + `[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 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 index 7f9d3b609836..16b691b4008a 100644 --- 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 @@ -47,63 +47,26 @@ class PatchEmbedConv3d(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # spec: num_patches, dtype; `in_channels`, `kernel`, `embed_dim` (default - # the checkpoint's); `bias` (default True); `layout` "flat" (default), - # "five_d" ([N, C, kt, kh, kw]), "column_slice" (a [N, 2*row] buffer's first - # half) or "transposed" (column-major rows). fp32 cells run with - # `torch.backends.cudnn.allow_tf32 = False` -- see note. - CELLS: tuple[Cell, ...] = ( - tuple( - Cell( - why=f"checkpoint geometry, {n} patches bf16", - spec=dict(num_patches=n, dtype=torch.bfloat16), - ) - for n in (1, 2, 7, 64, 1024, 4096) - ) - + ( - Cell( - why="explicit [N, C, kt, kh, kw] rows", - spec=dict(num_patches=64, dtype=torch.bfloat16, layout="five_d"), - ), - Cell( - why="bias=None (the Qwen2-VL form)", - spec=dict(num_patches=128, dtype=torch.bfloat16, bias=False), - ), - Cell(why="fp16", spec=dict(num_patches=256, dtype=torch.float16)), - Cell( - why="fp32 with cudnn TF32 off", - spec=dict(num_patches=256, dtype=torch.float32), - ), - ) - + tuple( - Cell( - why=f"geometry read from weight: C={c} kernel={k} out={o}", - spec=dict( - num_patches=33, - dtype=torch.bfloat16, - in_channels=c, - kernel=k, - embed_dim=o, - ), - ) - for c, k, o in ((1, (1, 4, 4), 8), (4, (2, 8, 8), 64), (3, (1, 16, 16), 128)) - ) - + ( - Cell( - why="column slice of a wider pixel buffer", - spec=dict(num_patches=8, dtype=torch.bfloat16, layout="column_slice"), - ), - Cell( - why="column-major pixel buffer (row elements strided)", - spec=dict(num_patches=8, dtype=torch.bfloat16, layout="transposed"), - ), + # 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 are just more rows; there is no batch axis. + 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 @@ -114,7 +77,14 @@ class PatchEmbedConv3d(OpWrapper): `[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 certified: float8 operands; patch counts past 4096. + 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__( 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 index ba26e3250f11..b616a48a86de 100644 --- 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 @@ -28,22 +28,6 @@ HIDDEN = 1152 MERGE = 2 -#: Grids the processor produces: square, both non-square orientations, the -#: exact-fit grid, a downsample, upsamples, video (t > 1) and a batch. -CHECKPOINT_GRIDS = ( - ((1, 2, 2),), - ((1, 4, 4),), - ((1, 4, 8),), - ((1, 8, 4),), - ((1, 48, 48),), - ((1, 62, 90),), - ((1, 64, 96),), - ((2, 4, 6),), - ((3, 2, 8),), - ((1, 4, 8), (2, 6, 4), (1, 48, 48)), - ((1, 128, 256),), -) - 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.""" @@ -75,58 +59,37 @@ class PosEmbedInterpolate(OpWrapper): # spec: grid (tuple of (t, h, w)), merge, dtype, rows (table rows), hidden. CELLS: tuple[Cell, ...] = ( - tuple( - Cell( - why=f"checkpoint table bf16, grid {list(map(list, grid))}", - spec=dict( - grid=grid, merge=MERGE, dtype=torch.bfloat16, rows=NUM_GRID**2, hidden=HIDDEN - ), - ) - for grid in CHECKPOINT_GRIDS - ) - + ( - 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="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, - ), + ), + 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, ), - ) - + tuple( - Cell( - why=f"{str(dtype)[6:]} table, grid {list(map(list, grid))}", - spec=dict(grid=grid, merge=MERGE, dtype=dtype, rows=NUM_GRID**2, hidden=64), - ) - for dtype in (torch.float16, torch.float32) - for grid in (((1, 6, 10),), ((1, 62, 90),), ((2, 4, 8), (1, 48, 48))) - ) - + tuple( - Cell( - why=f"spatial_merge_size {merge}", - spec=dict( - grid=((1, 8, 12), (1, 4, 4), (2, 8, 8)), - merge=merge, - dtype=torch.bfloat16, - rows=NUM_GRID**2, - hidden=64, - ), - ) - for merge in (1, 4) - ) + ), + # 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 = """ @@ -147,7 +110,11 @@ class PosEmbedInterpolate(OpWrapper): 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. - Not certified: a CPU table (it runs and returns on CPU). + 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__( 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 index e839617d3405..1df6b9943f24 100644 --- 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 @@ -9,9 +9,21 @@ 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: `[up ; gate]`, interleave, 32-row block shuffle of weight - and scale bytes, then 128x4 scale swizzle (FC2: shuffle only); +* 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). @@ -23,9 +35,9 @@ 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`, and -`mtp.layers.0.self_attn.indexer.*` always (the draft layer runs dense attention). -MTP experts are plain row-major FP8 128x128: `[up ; gate]` concat, bf16 +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 @@ -40,7 +52,7 @@ from cuda.bindings import runtime as cudart from torch import nn -from tensorrt_llm._utils import CUASSERT, prefer_pinned +from tensorrt_llm._utils import CUASSERT #: `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 @@ -103,7 +115,7 @@ def sources(self, d: SimpleNamespace, i: int | None, prefix: str) -> list[tuple[ # ----------------------------------------------------------- host residency -- # # Allocates the one PLE n-gram table. The table stays a `cuda` tensor of the same -# shape/dtype. Pinned, device-mapped pages (not a host gather) because the gather's +# 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. @@ -123,10 +135,45 @@ def memory_kind(ptr: int) -> int: return int(attributes.type) -def _host_resident_empty(rows: int, cols: int, dtype, device) -> tuple[torch.Tensor, torch.Tensor]: +#: `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 + + +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) -> None: + 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 + ) + ) + # 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 pinned CPU allocation that owns the pages. The + 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. """ @@ -137,14 +184,8 @@ def _host_resident_empty(rows: int, cols: int, dtype, device) -> tuple[torch.Ten ) assert device.type == "cuda", f"host offload needs a cuda device to map into; got {device}" - # torch's pinned allocator rounds up to a power of two (47.7 GiB -> 64 GiB resident), - # and keeps the block cached after release. - host = torch.empty(rows * cols, dtype=torch.uint8, pin_memory=prefer_pinned()) - assert host.is_pinned(), ( - "host offload maps pinned memory into the device, which is unavailable " - "here (e.g. Confidential Compute)" - ) - (device_pointer,) = CUASSERT(cudart.cudaHostGetDevicePointer(host.data_ptr(), 0)) + host = _RegisteredHostBuffer(rows * cols) + (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. @@ -176,15 +217,6 @@ def _host_resident_empty(rows: int, cols: int, dtype, device) -> tuple[torch.Ten _MTP = "mtp" _MTP_LAYER = "mtp.layers.0" -# The three tensors of the draft layer that are never loaded, with their reason -# in the module docstring: its own QSA index projection and the two per-head -# norms that go with it. The draft layer runs dense attention. -_MTP_NOT_LOADED = ( - f"{_MTP_LAYER}.self_attn.indexer.index_qk_proj.weight", - f"{_MTP_LAYER}.self_attn.indexer.q_layernorm.weight", - f"{_MTP_LAYER}.self_attn.indexer.k_layernorm.weight", -) - # NVFP4 block size: one e4m3 scale per 16 contiguous elements along K. SF_VEC = 16 @@ -196,7 +228,15 @@ def _host_resident_empty(rows: int, cols: int, dtype, device) -> tuple[torch.Ten # 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", "ple_table", "mtp_fc1_w", "mtp_fc2_w") +_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: @@ -221,8 +261,25 @@ def _interleave_perm(rows: int, device) -> torch.Tensor: return p.to(device) -def _fc1_perm(rows: int, device) -> torch.Tensor: - """Interleave then 32-row block shuffle, composed into one gather.""" +#: 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)] @@ -233,41 +290,89 @@ def _swizzle(x: torch.Tensor) -> torch.Tensor: 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 gate|up weight bytes: **gate** rows first, the half - order `flashinfer_silu_and_mul` reads.""" - return torch.cat(t, dim=0).contiguous() + """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)).contiguous() + 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.""" + """(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)) + 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).""" + """`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.""" + """`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 - s = torch.index_select(s, 0, _block32_perm(rows, s.device)) + 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 @@ -309,6 +414,24 @@ def _gdn_layers(d): 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) @@ -440,9 +563,34 @@ def _mixer(tag, prefix, layers): ) -def _gated_attention(tag, src, 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.""" + 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", @@ -462,6 +610,11 @@ def _gated_attention(tag, src, layers): 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), @@ -473,17 +626,46 @@ def _gated_attention(tag, src, 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, and its per-token gate.""" + """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, d.hidden), + 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, ) @@ -497,12 +679,6 @@ def _shared_expert(tag, prefix, layers): src=lambda d, i: [(f"{prefix(i)}.shared_expert.down_proj.weight", None)], layers=layers, ), - W( - f"{tag}sh_gate", - shape=lambda d: (1, d.hidden), - src=lambda d, i: [(f"{prefix(i)}.shared_expert_gate.weight", None)], - layers=layers, - ), ) @@ -528,51 +704,24 @@ class Qwen4ExpTp1Weights: ) ), # --- the gated GQA layers and their QSA index branch --- - *_gated_attention("", lambda s: f"{{p}}.self_attn.{s}", _attn_layers), - # One replicated, unquantized projection producing [index_heads q | 1 k], and - # two per-head gemma norms. - W( - "idx_w", - shape=lambda d: (d.index_width, d.hidden), - src="{p}.self_attn.indexer.index_qk_proj.weight", - layers=_attn_layers, - ), - W( - "idx_qn", - shape=lambda d: (d.index_head_dim,), - src="{p}.self_attn.indexer.q_layernorm.weight", - layers=_attn_layers, - ), - W( - "idx_kn", - shape=lambda d: (d.index_head_dim,), - src="{p}.self_attn.indexer.k_layernorm.weight", - layers=_attn_layers, - ), + # 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 --- - W( - "in_qkv_w", - shape=lambda d: (d.conv_dim, d.hidden), - src="{p}.linear_attn.in_proj_qkv.weight", - layers=_gdn_layers, - ), - W( - "in_z_w", - shape=lambda d: (d.value_dim, d.hidden), - src="{p}.linear_attn.in_proj_z.weight", - layers=_gdn_layers, - ), - # The two per-head scalar projections. - W( - "in_a_w", - shape=lambda d: (d.num_v_heads, d.hidden), - src="{p}.linear_attn.in_proj_a.weight", - layers=_gdn_layers, - ), - W( - "in_b_w", - shape=lambda d: (d.num_v_heads, d.hidden), - src="{p}.linear_attn.in_proj_b.weight", + # 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( @@ -637,6 +786,39 @@ class Qwen4ExpTp1Weights: 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 @@ -872,6 +1054,9 @@ class Qwen4ExpTp1Weights: ) ), *_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), @@ -946,10 +1131,15 @@ def dims(self, core) -> SimpleNamespace: 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, @@ -969,17 +1159,16 @@ def dims(self, core) -> SimpleNamespace: 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, exactly the three draft - indexer tensors. Named explicitly rather than by prefix so that a fourth unloaded - `mtp.*` key fails the coverage assert instead of hiding.""" + """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: - for key in _MTP_NOT_LOADED: - assert key in weights, f"checkpoint key missing: {key}" - return set(_MTP_NOT_LOADED) + return set() return {k for k in weights if k.startswith(f"{_MTP}.")} def declare(self, d: SimpleNamespace) -> nn.ParameterDict: @@ -1082,9 +1271,7 @@ def fill(param: nn.Parameter, ckpt_key, index, transform) -> None: torch.cuda.synchronize() leftover = set(weights.keys()) - consumed expected = self.expected_unconsumed(core, weights) - what = ( - "the draft layer's three QSA indexer tensors" if core.mtp_enabled else "the MTP module" - ) + 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]}, " 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 267a7c3b7ff6..ef2f31f01234 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 @@ -12,12 +12,19 @@ from __future__ import annotations -from typing import Any, Optional, Tuple +from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Tuple + +import torch from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig +from tensorrt_llm.quantization.mode import QuantAlgo from ..._router_index import NoTarget +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + # (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) @@ -27,6 +34,105 @@ # list this target lacks. _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. +_MAX_DRAFT_LEN = 7 +_MAX_LEGACY_DRAFT_LEN = 3 +_MAX_CUTEDSL_DRAFT_LEN = 3 + +# Recurrent-state cache dtypes the target serves (`_SSM_STATE_DTYPES` in modeling.py). +_SSM_STATE_DTYPES = (torch.float32, torch.bfloat16) + + +class EngineChoices(NamedTuple): + """Engine choices that fix which code the 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 + moe_backend: str # `moe_config.backend` as requested, before AUTO resolves + mamba_ssm_cache_dtype: torch.dtype # the recurrent-state cache's dtype + gdn_replay: bool # TRTLLM_USE_GDN_REPLAY, which the engine reads for eligible MTP workloads + prefill_cuda_graph_backend: str # `prefill_cuda_graph_backend`: disabled, piecewise, breakable + + +def engine_choices(model_config: ModelConfig, llm_args: TorchLlmArgs) -> EngineChoices: + # Imported on the claimed branch only, as the target is. + from tensorrt_llm._torch.pyexecutor.model_loader import effective_mamba_ssm_cache_dtype + from tensorrt_llm._torch.utils import is_gdn_replay_enabled + + kv_cache_config = llm_args.kv_cache_config + return EngineChoices( + torch_dtype=model_config.torch_dtype, + moe_backend=model_config.moe_backend_requested, + mamba_ssm_cache_dtype=effective_mamba_ssm_cache_dtype( + model_config, kv_cache_config.mamba_ssm_cache_dtype + ), + gdn_replay=is_gdn_replay_enabled(), + prefill_cuda_graph_backend=str(llm_args.prefill_cuda_graph_backend.value), + ) + + +def _check_engine( + spec_config: Optional[Any], kv_cache_quant_algo: Optional[Any], engine: Optional[EngineChoices] +) -> None: + """Decline what the target's constructor or first forward would refuse. + + Without this ``auto`` claims the deployment and then fails mid-build instead of + falling back to the built-in implementation. + """ + 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 + ): + raise NoTarget( + "the draft length changes at runtime (draft_len_schedule, max_concurrency or " + "acceptance_rate_threshold), but this target's verify step runs at the one " + "width max_draft_len + 1" + ) + if kv_cache_quant_algo not in (None, QuantAlgo.FP8, QuantAlgo.NVFP4): + raise NoTarget( + f"a {kv_cache_quant_algo} KV cache; the target stores K/V as the checkpoint's bf16, " + "as FP8 or as NVFP4" + ) + if engine is None: + return + 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: + raise NoTarget( + f"moe_config.backend={engine.moe_backend!r}; the target runs the trunk's routed " + f"experts on {', '.join(_MOE_BACKENDS)}" + ) + if engine.mamba_ssm_cache_dtype not 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)" + ) + if engine.prefill_cuda_graph_backend == "piecewise": + raise NoTarget( + "prefill_cuda_graph_backend='piecewise' makes the engine torch.compile the model, " + "which the target does not implement; 'breakable' runs the prefill step as a " + "breakable CUDA graph" + ) + if spec_config is None: + return + draft = spec_config.max_draft_len + if not engine.gdn_replay and draft > _MAX_LEGACY_DRAFT_LEN: + 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" + ) + if str(engine.moe_backend).upper() == "CUTEDSL" and draft > _MAX_CUTEDSL_DRAFT_LEN: + raise NoTarget( + f"max_draft_len={draft} with moe_config.backend=CUTEDSL: the CuteDSL expert " + f"pipeline is certified at max_draft_len up to {_MAX_CUTEDSL_DRAFT_LEN}" + ) + def route( pretrained_config: Any, @@ -34,6 +140,8 @@ def route( compute_capability: Tuple[int, int], spec_config: Optional[Any], kv_cache_quant_algo: Optional[Any], + model_config: Optional[ModelConfig] = None, + llm_args: Optional[TorchLlmArgs] = None, ) -> type: """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits. @@ -41,7 +149,8 @@ def route( (see ``_router_index._kv_cache_quant_algo``), not the checkpoint's alone. The target is imported where it is returned, so declining costs no import - of modeling code. + of modeling code. With ``model_config`` and ``llm_args`` it also declines the + engine choices (`EngineChoices`) the target's constructor would refuse. """ text = getattr(pretrained_config, "text_config", None) vision = getattr(pretrained_config, "vision_config", None) @@ -73,17 +182,16 @@ def route( ) # The target implements the one-model MTP-Eagle loop over the checkpoint's - # single MTP module, and no other draft path. A draft longer than three - # needs predicted_tokens_per_seq past the 1..4 thop_attention certifies, - # which the target asserts; declining here falls back instead of crashing. - # The draft must also be a linear chain over that one module; the target - # asserts each of these at construction, so declining here is what lets - # ``auto`` fall back instead of crashing. + # single MTP module, and no other draft path. A draft longer than seven needs + # more verify rows per request than the GDN cached replay and thop_attention + # are certified at. The draft must also be a linear chain over that one + # module; the target asserts each of these at construction, so declining here + # is what lets ``auto`` fall back instead of crashing. mtp_eagle_or_none = spec_config is None or ( isinstance(spec_config, MTPDecodingConfig) and spec_config.spec_dec_mode.is_mtp_eagle_one_model() and getattr(spec_config, "num_nextn_predict_layers", None) in (None, 1) - and 1 <= spec_config.max_draft_len <= 3 + and 1 <= spec_config.max_draft_len <= _MAX_DRAFT_LEN and not getattr(spec_config, "use_dynamic_tree", False) and getattr(spec_config, "eagle_choices", None) is None and spec_config.max_total_draft_tokens == spec_config.max_draft_len @@ -96,8 +204,6 @@ def route( and vshape == _VISION_SHAPE and mapping.world_size == 1 and mtp_eagle_or_none - # The target runs the checkpoint's bf16 KV pool only; it asserts so at construction. - and kv_cache_quant_algo is None ): # One line on purpose: the claims suite reads the routed targets off # `from ..modeling import ` lines. @@ -107,6 +213,12 @@ def route( # isort: on # fmt: on + engine = ( + engine_choices(model_config, llm_args) + if model_config is not None and llm_args is not None + else None + ) + _check_engine(spec_config, kv_cache_quant_algo, engine) return ModelingV2Qwen38FlashNextNvfp4Sm103Tp1 raise NoTarget( @@ -116,6 +228,5 @@ def route( 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 " - "one-model MTP-Eagle (max_draft_len 1..3) or no speculative decoding, and a " - "bf16 KV cache" + f"one-model MTP-Eagle (max_draft_len 1..{_MAX_DRAFT_LEN}) or no speculative decoding" ) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py index a5ddcaa70614..0d5273ebef92 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py @@ -11,6 +11,7 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm._utils import TensorWrapper, binding_to_torch_dtype, convert_to_torch_tensor +from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, DataRole, PageIndexMode @@ -31,10 +32,22 @@ _SHARED_POSITION_ROLE_INDEX = 0 _INDEX_K_ELEMENT_BYTES = torch.empty((), dtype=QSA_INDEX_K_CACHE_DTYPE).element_size() _POSITION_ELEMENT_BYTES = torch.empty((), dtype=QSA_POSITION_CACHE_DTYPE).element_size() +# NVFP4 packs two e2m1 values per byte and keeps one e4m3 scale byte per +# 16-element block. +_NVFP4_VALUES_PER_BYTE = 2 +_NVFP4_BLOCK_SIZE = 16 +# Granularity by which the packed side-plane page grows until its size cannot +# be a main-cache page; keeps every slot's planes 16-byte aligned. +_SIDE_PLANES_SIZE_STEP = 16 # Per-token RoPE/mRoPE coordinates used when raw index keys are compressed. # They are request state shared by all local QSA layers. QSA_INDEX_POSITION = DataRole("qsa_index_position") +# Every local layer's index-K page followed by the position page, as one +# buffer per lifecycle slot. Registered instead of the per-layer INDEX_KEY and +# the position role when either page could be the size of a K/V or block-scale +# page (see `_extra_buffers_per_layer`). +QSA_SIDE_PLANES = DataRole("qsa_side_planes") class QSAMambaHybridCacheManagerV2(MambaHybridCacheManagerV2): @@ -47,6 +60,17 @@ class QSAMambaHybridCacheManagerV2(MambaHybridCacheManagerV2): Index geometry is resolved from the same serving/checkpoint configuration as the indexer. A local fallback could otherwise give the indexer and its paged cache different shapes or page strides. + + V2 storage coalesces the buffers of one lifecycle by page size alone. The + attention op addresses a layer's K/V pages and its NVFP4 block-scale pages + with one block-offset table, so a side page that lands in the K/V pool + (NVFP4 at 32 tokens per block: index-K 8192 B == K == V) changes the K/V + slot stride but not the block-scale one, and any side page in the K/V pool + breaks the K/V pool mapping. When a side page could be the size of a main + page, the side planes are therefore registered as one packed per-slot buffer + whose size can be no main-cache page, and this manager serves their + per-layer views itself. Both the choice and the packed size are the same on + every TP and attention-DP layout, so disaggregated peers agree on the roles. """ def __init__( @@ -76,6 +100,9 @@ def __init__( layer_idx for layer_idx, enabled in enumerate(layer_mask) if enabled ] self.qsa_position_layer_id: Optional[int] = None + # Bytes of one packed side-plane page, or None while the side planes + # are registered as per-layer roles. Set by `_extra_buffers_per_layer`. + self.qsa_side_planes_size: Optional[int] = None super().__init__(*args, layer_mask=layer_mask, **kwargs) def get_disagg_role_mapper_kinds(self) -> Dict[DataRole, MapperKind]: @@ -86,6 +113,7 @@ def get_disagg_role_mapper_kinds(self) -> Dict[DataRole, MapperKind]: return { **super().get_disagg_role_mapper_kinds(), QSA_INDEX_POSITION: MapperKind.REPLICATED, + QSA_SIDE_PLANES: MapperKind.REPLICATED, } @cached_property @@ -110,37 +138,118 @@ def _extra_buffers_per_layer( *, tokens_per_block: int, ) -> dict[int, list[BufferConfig]]: - """Register per-layer index K and lifecycle-aligned position pages.""" + """Register per-layer index K and lifecycle-aligned position pages. + + Both are registered as their own roles unless either page could have + the size of a K/V or block-scale page at some split of the KV heads over + ranks (`_main_cache_page_sizes`), which V2 would coalesce into that pool. + They are then packed, in local sparse-layer order with the position page + last, into one `QSA_SIDE_PLANES` buffer on the first local sparse layer, + grown until its size is none of those page sizes. Neither the choice nor + the size depends on the KV heads this rank holds, so peers of different + TP or attention-DP layouts register the same transfer roles. + """ index_size = ( self.qsa_index_kv_heads * self.qsa_index_dim * _INDEX_K_ELEMENT_BYTES * tokens_per_block ) + position_size = QSA_POSITION_COORDINATE_AXES * _POSITION_ELEMENT_BYTES * tokens_per_block local_sparse_layers = self.qsa_local_sparse_layer_ids # Coordinates are request-wide, so all local indexers use one view. # Register the position role once; duplicating it on every sparse layer # wastes one three-axis int32 page per layer without adding state. self.qsa_position_layer_id = next(iter(local_sparse_layers), None) + self.qsa_side_planes_size = None + if self.qsa_position_layer_id is None: + return {} + position_local_idx = self.layer_offsets[self.qsa_position_layer_id] + main_page_sizes = self._main_cache_page_sizes(tokens_per_block) + if index_size in main_page_sizes or position_size in main_page_sizes: + packed_size = len(local_sparse_layers) * index_size + position_size + while packed_size in main_page_sizes: + packed_size += _SIDE_PLANES_SIZE_STEP + self.qsa_side_planes_size = packed_size + return {position_local_idx: [BufferConfig(role=QSA_SIDE_PLANES, size=packed_size)]} result = { self.layer_offsets[layer_id]: [BufferConfig(role=Role.INDEX_KEY, size=index_size)] for layer_id in local_sparse_layers } - if self.qsa_position_layer_id is not None: - local_idx = self.layer_offsets[self.qsa_position_layer_id] - result[local_idx].append( - BufferConfig( - role=QSA_INDEX_POSITION, - size=( - QSA_POSITION_COORDINATE_AXES * _POSITION_ELEMENT_BYTES * tokens_per_block - ), - ) - ) + result[position_local_idx].append(BufferConfig(role=QSA_INDEX_POSITION, size=position_size)) return result + def _main_cache_page_sizes(self, tokens_per_block: int) -> set[int]: + """Every size, in bytes, a local layer's K/V or block-scale page can have + at any split of that layer's KV heads over ranks: one head per rank up to + all of them. The set is the same on every TP and attention-DP layout.""" + global_layer = {local: layer for layer, local in self.layer_offsets.items()} + sizes = set() + for local_idx in range(self.num_local_layers): + heads = self.num_kv_heads_per_layer[local_idx] + total = ( + self.num_kv_heads + if isinstance(self.num_kv_heads, int) + else self.num_kv_heads[global_layer[local_idx]] + ) + for role in self._get_buffer_roles_for_layer(local_idx): + page = self.get_layer_bytes_per_token(local_idx, role) * tokens_per_block + if page: + sizes.update(page // heads * h for h in range(1, total + 1)) + return sizes + + def _side_planes(self) -> torch.Tensor: + """uint8 ``[slots, qsa_side_planes_size]`` over the packed side-plane pool.""" + local_idx = self.layer_offsets[self.qsa_position_layer_id] + page_stride = self.impl.get_page_stride(local_idx, QSA_SIDE_PLANES) + if page_stride != self.qsa_side_planes_size: + raise RuntimeError( + "QSA side-plane page stride mismatch: " + f"expected {self.qsa_side_planes_size}, got {page_stride}" + ) + converter = self.impl.get_page_index_converter(local_idx, QSA_SIDE_PLANES) + if ( + int(converter.scale) != 1 + or int(converter.layer_offset) != 0 + or int(converter.expansion) != _IDENTITY_PAGE_INDEX_EXPANSION + ): + raise RuntimeError( + "QSA side planes share their V2 pool with another buffer: " + f"scale={int(converter.scale)}, layer_offset={int(converter.layer_offset)}, " + f"expansion={int(converter.expansion)}" + ) + addr = self.impl.get_mem_pool_base_address(local_idx, QSA_SIDE_PLANES, PageIndexMode.SHARED) + num_slots = self.impl.get_page_index_upper_bound(local_idx, QSA_SIDE_PLANES) + return convert_to_torch_tensor( + TensorWrapper(addr, torch.uint8, [num_slots, self.qsa_side_planes_size]) + ) + def get_index_k_buffer( self, layer_idx: int, kv_layout: str = "NHD", ) -> Optional[torch.Tensor]: - """Return the index-K view using the indexer's resolved geometry.""" + """Return the index-K view using the indexer's resolved geometry. + + ``[slots, tokens_per_block, heads, dim]`` (NHD) or ``[slots, heads, + tokens_per_block, dim]`` (HND), indexed by lifecycle slot. + """ + if self.qsa_side_planes_size is not None: + if kv_layout not in ("NHD", "HND"): + raise ValueError(f"Unsupported kv_layout: {kv_layout}") + if layer_idx not in self.qsa_local_sparse_layer_ids: + return None + ordinal = self.qsa_local_sparse_layer_ids.index(layer_idx) + index_size = ( + self.qsa_index_kv_heads + * self.qsa_index_dim + * _INDEX_K_ELEMENT_BYTES + * self.tokens_per_block + ) + planes = self._side_planes() + plane = planes[:, ordinal * index_size : (ordinal + 1) * index_size] + if kv_layout == "NHD": + page_shape = [self.tokens_per_block, self.qsa_index_kv_heads, self.qsa_index_dim] + else: + page_shape = [self.qsa_index_kv_heads, self.tokens_per_block, self.qsa_index_dim] + return plane.view(QSA_INDEX_K_CACHE_DTYPE).view(planes.shape[0], *page_shape) return super().get_index_k_buffer( layer_idx, num_heads=self.qsa_index_kv_heads, @@ -158,11 +267,62 @@ def get_buffers( QSA tables store lifecycle slot IDs. Preserving V2's coalesced stride is therefore required even when several layers share one physical pool. + + ``[slots, 2, heads, tokens_per_block, width]`` (HND) or ``[slots, 2, + tokens_per_block, heads, width]`` (NHD), K then V. Half, bf16 and fp8 + pools are viewed in their element type at ``width = head_dim``. An + NVFP4 pool is viewed as uint8 at ``width = head_dim // 2``: two e2m1 + values per byte, element ``2i`` in the low nibble; its block scales are + `get_block_scale_buffers`. Any other format is the parent's layout. + """ + if self.dtype in QSA_SPARSE_KV_CACHE_DTYPES: + return self._paired_role_view( + layer_idx, + (Role.KEY, Role.VALUE), + binding_to_torch_dtype(self.dtype), + 1, + kv_layout, + ) + if self.dtype == DataType.NVFP4: + return self._paired_role_view( + layer_idx, (Role.KEY, Role.VALUE), torch.uint8, _NVFP4_VALUES_PER_BYTE, kv_layout + ) + return super().get_buffers(layer_idx, kv_layout) + + def get_block_scale_buffers( + self, + layer_idx: int, + kv_layout: str = "NHD", + ) -> Optional[torch.Tensor]: + """Return an NVFP4 pool's e4m3 block scales for this layer, K then V. + + uint8 ``[slots, 2, heads, tokens_per_block, head_dim // 16]`` (HND) or + ``[slots, 2, tokens_per_block, heads, head_dim // 16]`` (NHD) with the + same slot indexing as `get_buffers`: one scale byte per 16 values. The + attention op writes K's scales token-major within a head page, and V's + in groups of four tokens, ``[tokens_per_block / 4, head_dim / 16, 4]`` + per head; the view is the page as bytes and leaves that order to the + reader. None for a layer this rank does not hold. """ - if self.dtype not in QSA_SPARSE_KV_CACHE_DTYPES: - # The parent owns packed data and scale-page layouts such as - # NVFP4. QSA returns to the regular backend for these formats. - return super().get_buffers(layer_idx, kv_layout) + if self.dtype != DataType.NVFP4: + raise ValueError(f"block scales exist only for an NVFP4 pool, not {self.dtype}") + return self._paired_role_view( + layer_idx, + (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE), + torch.uint8, + _NVFP4_BLOCK_SIZE, + kv_layout, + ) + + def _paired_role_view( + self, + layer_idx: int, + roles: tuple[DataRole, DataRole], + dtype: torch.dtype, + values_per_element: int, + kv_layout: str, + ) -> Optional[torch.Tensor]: + """Slot-indexed view over a layer's adjacent K-side and V-side pages.""" if kv_layout not in ("NHD", "HND"): raise ValueError(f"Unsupported kv_layout: {kv_layout}") if layer_idx not in self.layer_offsets: @@ -170,21 +330,22 @@ def get_buffers( if self.kv_cache_type == CacheTypeCpp.SELFKONLY: raise NotImplementedError("QSA sparse attention requires both K and V cache buffers") + key_role, value_role = roles layer_offset = self.layer_offsets[layer_idx] - addr_key = self.impl.get_mem_pool_base_address(layer_offset, Role.KEY, PageIndexMode.SHARED) + addr_key = self.impl.get_mem_pool_base_address(layer_offset, key_role, PageIndexMode.SHARED) addr_value = self.impl.get_mem_pool_base_address( - layer_offset, Role.VALUE, PageIndexMode.SHARED + layer_offset, value_role, PageIndexMode.SHARED ) - key_stride = self.impl.get_page_stride(layer_offset, Role.KEY) - value_stride = self.impl.get_page_stride(layer_offset, Role.VALUE) + key_stride = self.impl.get_page_stride(layer_offset, key_role) + value_stride = self.impl.get_page_stride(layer_offset, value_role) if addr_key + key_stride != addr_value or key_stride != value_stride: raise RuntimeError( "QSA K/V buffers are not adjacent equal-sized V2 roles: " - f"layer={layer_idx}, key_stride={key_stride}, " - f"value_stride={value_stride}" + f"layer={layer_idx}, roles={key_role}/{value_role}, " + f"key_stride={key_stride}, value_stride={value_stride}" ) - converter = self.impl.get_page_index_converter(layer_offset, Role.KEY) + converter = self.impl.get_page_index_converter(layer_offset, key_role) scale = int(converter.scale) if scale < QSA_MAIN_KV_ROLES: raise RuntimeError( @@ -193,7 +354,7 @@ def get_buffers( if int(converter.expansion) != _IDENTITY_PAGE_INDEX_EXPANSION: raise RuntimeError("QSA does not support expanded V2 page indices") layer_offset_pages = int(converter.layer_offset) - page_upper = self.impl.get_page_index_upper_bound(layer_offset, Role.KEY) + page_upper = self.impl.get_page_index_upper_bound(layer_offset, key_role) num_slots_total = page_upper + layer_offset_pages if num_slots_total % scale != 0: raise RuntimeError( @@ -202,17 +363,23 @@ def get_buffers( ) num_slots = num_slots_total // scale - torch_dtype = binding_to_torch_dtype(self.dtype) head_dim = self.head_dim_per_layer[layer_offset] num_heads = self.num_kv_heads_per_layer[layer_offset] + width = head_dim // values_per_element if kv_layout == "NHD": - page_shape = [self.tokens_per_block, num_heads, head_dim] + page_shape = [self.tokens_per_block, num_heads, width] else: - page_shape = [num_heads, self.tokens_per_block, head_dim] + page_shape = [num_heads, self.tokens_per_block, width] + page_bytes = self.tokens_per_block * num_heads * width * dtype.itemsize + if page_bytes != key_stride: + raise RuntimeError( + f"QSA {key_role} page is {key_stride} B, but {page_shape} {dtype} is " + f"{page_bytes} B: layer={layer_idx}" + ) full_view = convert_to_torch_tensor( TensorWrapper( addr_key, - torch_dtype, + dtype, [num_slots, scale, *page_shape], ) ) @@ -220,9 +387,27 @@ def get_buffers( @torch.compiler.disable def get_qsa_position_buffer(self) -> Optional[torch.Tensor]: - """Return per-token three-axis RoPE/mRoPE position coordinates.""" + """Return per-token three-axis RoPE/mRoPE position coordinates. + + int32 ``[slots, tokens_per_block, 3]``, indexed by lifecycle slot. + """ if self.qsa_position_layer_id is None: return None + if self.qsa_side_planes_size is not None: + position_size = ( + QSA_POSITION_COORDINATE_AXES * _POSITION_ELEMENT_BYTES * self.tokens_per_block + ) + planes = self._side_planes() + start = len(self.qsa_local_sparse_layer_ids) * ( + self.qsa_index_kv_heads + * self.qsa_index_dim + * _INDEX_K_ELEMENT_BYTES + * self.tokens_per_block + ) + plane = planes[:, start : start + position_size] + return plane.view(QSA_POSITION_CACHE_DTYPE).view( + planes.shape[0], self.tokens_per_block, QSA_POSITION_COORDINATE_AXES + ) local_idx = self.layer_offsets[self.qsa_position_layer_id] addr = self.impl.get_mem_pool_base_address( local_idx, @@ -269,6 +454,30 @@ def get_qsa_position_buffer(self) -> Optional[torch.Tensor]: ) return full_view[:, _SHARED_POSITION_ROLE_INDEX] + def _build_pool_mapping_tensors(self): + """Refuse an NVFP4 layout whose K/V and block-scale slots differ. + + The attention op reads a layer's K/V pages and their block scales + through one block-offset table: a page index scaled for the K/V pool + addresses the scale pool too. A layer whose two page-index scales + differ would read and write other slots' scales without any error. + """ + if self.dtype == DataType.NVFP4: + for layer_id in self.pp_layers: + local_idx = self.layer_offsets[layer_id] + if self._is_local_mamba_layer(local_idx): + continue + data_scale = self.impl.get_page_index_scale(local_idx, Role.KEY) + scale_scale = self.impl.get_page_index_scale(local_idx, Role.KEY_BLOCK_SCALE) + if data_scale != scale_scale: + raise RuntimeError( + "NVFP4 K/V and block-scale pools hold a different number of pages " + f"per lifecycle slot (layer {layer_id}: {data_scale} vs {scale_scale}); " + "another buffer was coalesced into one of them, and the attention op " + "would address the block scales of other slots" + ) + return super()._build_pool_mapping_tensors() + def get_qsa_attention_pool_layout(self) -> tuple[int, int]: """Return the shared attention-pool index and V2 page-index scale. @@ -307,4 +516,4 @@ def get_qsa_attention_pool_layout(self) -> tuple[int, int]: return pool_id, scale -__all__ = ["QSAMambaHybridCacheManagerV2"] +__all__ = ["QSA_INDEX_POSITION", "QSA_SIDE_PLANES", "QSAMambaHybridCacheManagerV2"] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/kernels.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/kernels.py index 74834c64fd48..358b5f424b43 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/kernels.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/kernels.py @@ -1119,7 +1119,24 @@ def _qsa_paged_index_scores_kernel( BLOCK_N: tl.constexpr, TILES_PER_PROGRAM: tl.constexpr, ONLY_VISIBLE_BLOCKS: tl.constexpr, + USE_PDL: tl.constexpr = False, + PREFETCH: tl.constexpr = False, ): + if PREFETCH: + # One tile a program: its pages and keys are read before the wait (see + # ``triton_qsa_paged_index_scores``), q after it. + tl.static_assert(USE_PDL and TILES_PER_PROGRAM == 1 and ONLY_VISIBLE_BLOCKS) + _qsa_paged_index_scores_prefetch_tile( + q, index_cache, block_table, query_positions, request_indices, output, score_scale, + q_stride_row, q_stride_head, q_stride_dim, cache_stride_page, cache_stride_token, + cache_stride_dim, block_table_stride_request, block_table_stride_page, + output_stride_row, output_stride_block, NUM_INDEX_HEADS, INDEX_HEAD_DIM, BLOCK_D, + TOKENS_PER_BLOCK, COMPRESS_RATIO, MAX_COMPRESSED_BLOCKS, BLOCK_H, BLOCK_N, + ) # fmt: skip + return + if USE_PDL: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() row = tl.program_id(0) request = tl.load(request_indices + row).to(tl.int64) query_position = tl.load(query_positions + row) @@ -1188,6 +1205,95 @@ def _qsa_paged_index_scores_kernel( ) +@triton.jit +def _qsa_paged_index_scores_prefetch_tile( + q, + index_cache, + block_table, + query_positions, + request_indices, + output, + score_scale, + 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, + BLOCK_H: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """One generation row's tile of visible compressed blocks under PDL, its static + operands read before ``gdc_wait``: the row's metadata, the tile's pages and every + key but the one a decode step's own prologue may have just written -- the last + visible block's, compressed at its anchor when the row's position ends a group -- + which is read again after the wait, with q.""" + tl.extra.cuda.gdc_launch_dependents() + row = tl.program_id(0) + request = tl.load(request_indices + row).to(tl.int64) + query_position = tl.load(query_positions + row) + visible_blocks = (query_position + 1) // COMPRESS_RATIO + first_tile = tl.program_id(1) + if first_tile * BLOCK_N >= visible_blocks: + tl.extra.cuda.gdc_wait() + return + dim_offsets = tl.arange(0, BLOCK_D) + head_offsets = tl.arange(0, BLOCK_H) + block_columns = first_tile * BLOCK_N + tl.arange(0, BLOCK_N) + valid_blocks = (block_columns < MAX_COMPRESSED_BLOCKS) & (block_columns < visible_blocks) + 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) + key_ptrs = ( + index_cache + + physical_pages[:, None] * cache_stride_page + + token_in_page[:, None] * cache_stride_token + + dim_offsets[None, :] * cache_stride_dim + ) + key_mask = valid_blocks[:, None] & (dim_offsets[None, :] < INDEX_HEAD_DIM) + keys = tl.load(key_ptrs, mask=key_mask, other=0.0, eviction_policy="evict_first") + tl.extra.cuda.gdc_wait() + fresh = (block_columns == visible_blocks - 1) & ((query_position + 1) % COMPRESS_RATIO == 0) + keys = tl.where( + fresh[:, None], tl.load(key_ptrs, mask=key_mask & fresh[:, None], other=0.0), keys + ) + query_values = tl.load( + q + + row * q_stride_row + + dim_offsets[:, None] * q_stride_dim + + head_offsets[None, :] * q_stride_head, + mask=(dim_offsets[:, None] < INDEX_HEAD_DIM) & (head_offsets[None, :] < NUM_INDEX_HEADS), + other=0.0, + ) + per_head_scores = tl.dot(keys, query_values, out_dtype=tl.float32) + scores = tl.sum( + tl.where(head_offsets[None, :] < NUM_INDEX_HEADS, tl.maximum(per_head_scores, 0.0), 0.0), + axis=1, + ) + tl.store( + output + row * output_stride_row + block_columns * output_stride_block, + scores * score_scale, + mask=valid_blocks, + ) + + @triton.jit def _qsa_paged_index_scores_qtile_kernel( q, @@ -1248,6 +1354,11 @@ def _qsa_paged_index_scores_qtile_kernel( other=0.0, ) column_offsets = tl.arange(0, BLOCK_N) + # Each tile row's (lane's) visible bound, indexed by lane instead of by column. + lane_rows = row_base + tl.arange(0, QUERY_TILE) + lane_present = lane_rows < num_rows + lane_positions = tl.load(query_positions + lane_rows, mask=lane_present, other=0) + lane_visible = tl.where(lane_present, (lane_positions + 1) // COMPRESS_RATIO, 0) # A gather is only shareable across rows that page through the same block # table, so score one request at a time. Context rows arrive packed per @@ -1266,6 +1377,10 @@ def _qsa_paged_index_scores_qtile_kernel( request = tl.max(tl.where(dot_columns == leader, requests, 0)) group = pending & (requests == request) pending = pending & (requests != request) + # The lanes whose columns are in this group (none on a pass past the last group). + lane_group = ( + tl.max(tl.reshape(group.to(tl.int32), [QUERY_TILE, NUM_INDEX_HEADS]), axis=1) > 0 + ) group_visible = tl.max(tl.where(group, visible_blocks, 0)) last_tile = tl.minimum( first_tile + TILES_PER_PROGRAM, @@ -1300,37 +1415,34 @@ def _qsa_paged_index_scores_qtile_kernel( tl.dot(keys, query_values, out_dtype=tl.float32), 0.0, ) - for lane in tl.static_range(QUERY_TILE): - row_columns = group & (column_row == lane) - scored = tl.max(row_columns.to(tl.int32)) > 0 - row_visible = tl.max(tl.where(row_columns, visible_blocks, 0)) - scores = tl.sum( - tl.where(row_columns[None, :], positive_scores, 0.0), - axis=1, + # 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 + ) + output_ptrs = ( + output + + lane_rows[None, :] * output_stride_row + + block_columns[:, None] * output_stride_block + ) + row_valid = block_columns[:, None] < lane_visible[None, :] + if ONLY_VISIBLE_BLOCKS: + tl.store( + output_ptrs, + scores * score_scale, + mask=lane_group[None, :] & row_valid, ) - output_ptrs = ( - output - + (row_base + lane) * output_stride_row - + block_columns * output_stride_block + else: + tl.store( + output_ptrs, + tl.where(row_valid, scores * score_scale, -float("inf")), + mask=lane_group[None, :] & (block_columns[:, None] < MAX_COMPRESSED_BLOCKS), ) - row_valid = block_columns < row_visible - if ONLY_VISIBLE_BLOCKS: - tl.store( - output_ptrs, - scores * score_scale, - mask=scored & row_valid, - ) - else: - tl.store( - output_ptrs, - tl.where(row_valid, scores * score_scale, -float("inf")), - mask=scored & (block_columns < MAX_COMPRESSED_BLOCKS), - ) - - -# ``tl.dot`` needs at least 16 columns on SM90+, so a narrower model dimension -# leaves the rest of every tensor-core tile computed and discarded. -_INDEX_SCORE_TILE_WIDTH = 16 + + +# Dot columns of a context-row index-scoring tile: whole query rows of +# ``num_index_heads`` columns each. A wide tile serves more rows with each +# gathered key tile (``tl.dot`` needs at least 16 columns on SM90+). +_INDEX_SCORE_TILE_WIDTH = 128 def _qsa_index_scores_query_tile(num_index_heads: int) -> int: @@ -1358,6 +1470,8 @@ def triton_qsa_paged_index_scores( compress_ratio: int, only_visible_blocks: bool = False, context_rows: bool = False, + use_pdl: bool = False, + prefetch: bool = False, ) -> torch.Tensor: """Score compressed keys directly in the paged side cache. @@ -1368,6 +1482,17 @@ def triton_qsa_paged_index_scores( ``context_rows`` declares that consecutive rows mostly share a request, so one gather of compressed keys can serve a tile of query rows. Generation batches hold one row per request and would only lose parallelism. + + ``use_pdl`` launches the generation-row kernel with programmatic dependent + launch: it waits for its predecessor before its first read and releases its + dependents at once, so every dependent must itself wait before reading the + scores. + + ``prefetch`` (with ``use_pdl``, ``only_visible_blocks`` and generation-sized + batches) reads the pages and compressed keys before the wait: sound only when the + predecessor writes no compressed key but, at a row whose position ends a group, the + last visible block's -- the decode step's own K/V-and-index prologue -- which is + read again after the wait. """ if q.ndim != 3: raise ValueError("QSA index queries must have shape [rows, heads, dim]") @@ -1401,10 +1526,12 @@ def triton_qsa_paged_index_scores( ) if rows == 0: return output - # Decode favors many narrow score tiles; large prefill batches reuse each - # query across a deeper column range to reduce launch and gather overhead. + # A decode batch of a few rows favors many narrow score tiles; past that, + # each program's fixed cost (the query load, the page lookups) is spread + # over a wider tile. Large prefill batches also reuse each query across a + # deeper column range to reduce launch and gather overhead. is_prefill_sized = rows > 256 - block_n = 64 if is_prefill_sized else 8 + block_n = 64 if rows > 4 else 8 tiles_per_program = 8 if is_prefill_sized else 1 query_tile = 1 if context_rows and is_prefill_sized: @@ -1412,7 +1539,7 @@ def triton_qsa_paged_index_scores( if query_tile > 1: # Row packing reduces the grid, so each program walks a deeper column # range to retain enough work per launch. - tiles_per_program = 32 + tiles_per_program = 8 grid = ( triton.cdiv(rows, query_tile), triton.cdiv(max_compressed_blocks, block_n * tiles_per_program), @@ -1488,7 +1615,10 @@ def triton_qsa_paged_index_scores( BLOCK_N=block_n, TILES_PER_PROGRAM=tiles_per_program, ONLY_VISIBLE_BLOCKS=only_visible_blocks, + USE_PDL=use_pdl, + PREFETCH=prefetch and use_pdl and tiles_per_program == 1 and only_visible_blocks, num_warps=2, + launch_pdl=use_pdl, ) return output @@ -2178,6 +2308,7 @@ def _qsa_mtp_reuse_topk_kernel( visible_ptr, out_ptr, BLOCK_TOPK: tl.constexpr, + DISPLACE_OLDEST: tl.constexpr, ): """Rebuild a query row's selection from its request's captured row. @@ -2185,6 +2316,10 @@ def _qsa_mtp_reuse_topk_kernel( groups that completed afterwards are written on top of it. They are always larger than every captured ID, because the capture's Top-K searched only ``[0, captured_visible)``, so the row never gains a duplicate. + + A full row has no room for them: they take the slots of the row's last IDs, + or with ``DISPLACE_OLDEST`` the survivors shift to the front and they take + the slots after them, which for an ascending row displaces its smallest IDs. """ row = tl.program_id(0) request = tl.load(request_ptr + row).to(tl.int32) @@ -2197,8 +2332,16 @@ def _qsa_mtp_reuse_topk_kernel( # Valid captured IDs form a prefix; the rest is -1 padding. filled = tl.sum((values >= 0).to(tl.int32)) fresh = tl.minimum(tl.maximum(current - captured, 0), BLOCK_TOPK) - # Append after the prefix when there is room, otherwise displace its tail. + # Append after the prefix when there is room, otherwise displace its tail, or its head + # (the survivors shifted to the front) with DISPLACE_OLDEST. start = tl.minimum(filled, BLOCK_TOPK - fresh) + if DISPLACE_OLDEST: + source = slots + (filled - start) + values = tl.load( + shared_ptr + request.to(tl.int64) * BLOCK_TOPK + source, + mask=source < BLOCK_TOPK, + other=-1, + ) offset = slots - start is_fresh = (offset >= 0) & (offset < fresh) values = tl.where(is_fresh, (captured + offset).to(values.dtype), values) @@ -2212,8 +2355,13 @@ def triton_qsa_mtp_reuse_topk( request_indices: torch.Tensor, visible_blocks: torch.Tensor, output: torch.Tensor, + displace_oldest: bool = False, ) -> torch.Tensor: - """Expand captured per-request selections back to per-query-row selections.""" + """Expand captured per-request selections back to per-query-row selections. + + ``displace_oldest`` is for captured rows in ascending order: a full row gives up + its first IDs, rather than its last, to the groups completed since the capture. + """ rows, block_topk = output.shape if shared_indices.shape[1] != block_topk: raise ValueError( @@ -2239,6 +2387,7 @@ def triton_qsa_mtp_reuse_topk( visible_blocks, output, BLOCK_TOPK=block_topk, + DISPLACE_OLDEST=displace_oldest, num_warps=4, ) return output diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py index fe1a441e3584..4ccfb3ce2928 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py @@ -53,6 +53,7 @@ def __post_init__(self) -> None: self.qsa_shared_topk_indices: Optional[torch.Tensor] = None self.qsa_shared_topk_visible_blocks: Optional[torch.Tensor] = None self.qsa_shared_topk_captured_slots: set[int] = set() + self.qsa_draft_block_table: Optional[torch.Tensor] = None if not isinstance(self.kv_cache_manager, QSAMambaHybridCacheManagerV2): raise TypeError("QSA sparse attention requires QSAMambaHybridCacheManagerV2") if not isinstance(self.sparse_metadata_params, QSASparseMetadataParams): @@ -80,6 +81,7 @@ def __post_init__(self) -> None: self._qsa_attention_pool_id, self._qsa_page_index_scale = ( self.kv_cache_manager.get_qsa_attention_pool_layout() ) + self._init_qsa_draft_block_table() # Packed-row mapping shared by compression, selection, and sparse GQA. self.qsa_req_idx_per_token = self.get_empty( @@ -164,7 +166,11 @@ def __post_init__(self) -> None: # compressed groups were complete when the row was captured, so the # reuse can tell which groups formed afterwards. if self.sparse_metadata_params.mtp_index_share: + # Slots are numbered by the bound manager; the draft loop binds the + # separate draft manager. num_slots = len(self.kv_cache_manager.qsa_shared_topk_slots) + if self.qsa_draft_block_table is not None: + num_slots = max(num_slots, len(self.draft_kv_cache_manager.qsa_shared_topk_slots)) self.qsa_shared_topk_indices = self.get_empty( buffers, ( @@ -184,6 +190,32 @@ def __post_init__(self) -> None: capture_graph=capture_graph, ) + def _init_qsa_draft_block_table(self) -> None: + """Allocate the page table of a separate one-model draft KV cache. + + The draft layers' pages live in the draft manager's pools, which the + target table does not address. This table is its own graph-stable + buffer, so a captured graph reads the target table in the target + forward and this one in the draft loop; ``prepare()`` refreshes both + before every replay. + """ + draft_manager = self.draft_kv_cache_manager + if ( + not isinstance(draft_manager, QSAMambaHybridCacheManagerV2) + or draft_manager.qsa_position_layer_id is None + ): + return + self.qsa_draft_block_table = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences, draft_manager.max_blocks_per_seq), + cache_name="qsa_draft_block_table", + dtype=torch.int32, + capture_graph=self.is_cuda_graph, + ) + self._qsa_draft_attention_pool_id, self._qsa_draft_page_index_scale = ( + draft_manager.get_qsa_attention_pool_layout() + ) + def set_skip_topk(self, skip: bool) -> None: """Mark a draft step that reuses the captured selection.""" self.qsa_indexer_skip_topk = skip @@ -245,21 +277,45 @@ def _refresh_qsa_token_mapping(self) -> None: ((logical + 1) // self.sparse_metadata_params.compress_ratio).to(torch.int32) ) - def _refresh_qsa_block_table(self) -> None: - """Decode V2 attention offsets into raw page IDs used by QSA kernels.""" + def _unscale_qsa_block_table( + self, + block_offsets: torch.Tensor, + pool_id: int, + page_index_scale: int, + block_table: torch.Tensor, + ) -> None: + """Decode one manager's V2 attention offsets into raw page IDs.""" num_seqs = self.num_seqs if num_seqs <= 0: return - scaled = self.kv_cache_block_offsets[ - self._qsa_attention_pool_id, - :num_seqs, - QSA_KEY_ROLE_INDEX, - :, - ] triton_qsa_unscale_block_table( - scaled_block_table=scaled, - block_table=self.qsa_block_table[:num_seqs], - page_index_scale=self._qsa_page_index_scale, + scaled_block_table=block_offsets[pool_id, :num_seqs, QSA_KEY_ROLE_INDEX, :], + block_table=block_table[:num_seqs], + page_index_scale=page_index_scale, + ) + + def _refresh_qsa_block_table(self) -> None: + """Decode the bound manager's offsets into the table QSA kernels read. + + While ``prepare_for_draft_forward()`` is in effect every operand here is + the draft manager's, so this rebuilds the draft table. + """ + self._unscale_qsa_block_table( + self.kv_cache_block_offsets, + self._qsa_attention_pool_id, + self._qsa_page_index_scale, + self.qsa_block_table, + ) + + def _refresh_qsa_draft_block_table(self) -> None: + """Decode the separate draft manager's offsets into the draft table.""" + if self.qsa_draft_block_table is None: + return + self._unscale_qsa_block_table( + self.draft_kv_cache_block_offsets, + self._qsa_draft_attention_pool_id, + self._qsa_draft_page_index_scale, + self.qsa_draft_block_table, ) def _refresh_qsa_speculative_snapshot_state(self) -> None: @@ -276,8 +332,29 @@ def _refresh_qsa_device_state(self) -> None: if not self.qsa_has_local_layers: return self._refresh_qsa_block_table() + self._refresh_qsa_draft_block_table() self._refresh_qsa_token_mapping() + def prepare_for_draft_forward(self) -> None: + """Point the QSA page table at the separate draft manager's pages. + + The caller has already swapped ``kv_cache_manager`` and + ``kv_cache_block_offsets`` to the draft manager's. The draft table was + decoded in ``prepare()``, so this only rebinds fields and launches + nothing; ``restore_draft_swaps()`` restores the target bindings. The + row mapping is manager-independent and stays as it is. + """ + if self.qsa_draft_block_table is None: + return + if self.kv_cache_manager is not self.draft_kv_cache_manager: + raise RuntimeError( + "QSA draft forward expects the metadata's draft KV cache manager " + f"to be bound; got {type(self.kv_cache_manager).__name__}" + ) + self.swap_for_draft("qsa_block_table", self.qsa_draft_block_table) + self.swap_for_draft("_qsa_attention_pool_id", self._qsa_draft_attention_pool_id) + self.swap_for_draft("_qsa_page_index_scale", self._qsa_draft_page_index_scale) + def prepare(self) -> None: """Refresh host decisions, then rebuild graph-stable device mappings.""" super().prepare() diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 84edf64cfd06..9983032e9b68 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -284,6 +284,8 @@ class ModelConfig(Generic[TConfig]): # disables the optimization; hardware support is resolved by the backend. skip_correction_threshold: float = 0.0 moe_backend: str = 'CUTLASS' # options can be CUTLASS, TRTLLM + # The backend the user requested, before AUTO resolution. + moe_backend_requested: str = 'AUTO' # IF true, disables FC2+finalize fusion in CUTLASS MoE backend moe_disable_finalize_fusion: bool = False # If true, use low precision combine in MoE operations (only for NVFP4 quantization) @@ -1526,6 +1528,7 @@ def _recursive_update_config(config: transformers.PretrainedConfig, architecture, quant_config=quant_config, ) + kwargs['moe_backend_requested'] = requested_moe_backend # V4.1 names the routed experts exactly as V4 does and encodes MXFP4 the # same way (I8 storage), so the detector and setter apply verbatim. diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py index ac57b3ef45eb..4e1c6efaf56d 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py @@ -292,6 +292,9 @@ class HostInputs: n_gen: int = 0 has_any_multimodal_request: bool = False + # Whether a verify group or generation request of this pass carries real + # MRoPE metadata; without one, the delta read slots are dropped. + has_gen_mrope_delta: bool = False beam_width: int = 1 num_tokens: int = 0 @@ -1242,6 +1245,30 @@ def _helix_pack_extend(request, group: int) -> int: host.helix_pack_extend = _helix_pack_extend host.has_cp_helix = _has_cp_helix + @staticmethod + def _decode_mrope_position_delta(request: LlmRequest) -> torch.Tensor | None: + """The MRoPE position delta of a decode-phase request, or None. + + A request that ran its context phase on this worker already holds the + delta in ``py_mrope_position_delta``. Otherwise (for example the first + step of a generation-only request) it is read from the request's + ``mrope_config``, moved to the device once, and cached on the request. + None means the request carries no MRoPE metadata (a text-only prompt), + whose delta is zero. + """ + mrope_position_delta = getattr(request, "py_mrope_position_delta", None) + if mrope_position_delta is None and request.py_multimodal_data: + mrope_config = request.py_multimodal_data.get("mrope_config") or {} + mrope_position_delta = mrope_config.get("mrope_position_deltas") + if mrope_position_delta is not None: + if mrope_position_delta.device.type == "cpu": + mrope_position_delta = maybe_pin_memory(mrope_position_delta).to( + device="cuda", dtype=torch.int32, non_blocking=True + ) + mrope_config["mrope_position_deltas"] = mrope_position_delta + request.py_mrope_position_delta = mrope_position_delta + return mrope_position_delta + def _pack_extend_requests( self, host: HostInputs, @@ -1273,8 +1300,24 @@ def _pack_extend_requests( _has_cp_helix = host.has_cp_helix _helix_local_len_host = host.helix_local_len_host _helix_pack_extend = host.helix_pack_extend + mrope_position_ids = host.lists.mrope_position_ids + mrope_delta_write_seq_slots = host.lists.mrope_delta_write_seq_slots + mrope_delta_read_seq_slots = host.lists.mrope_delta_read_seq_slots + multimodal_params_list = host.lists.multimodal_params_list + _use_mrope = host.use_mrope + mrope_dummy_seq_slot = host.mrope_dummy_seq_slot + has_gen_mrope_delta = host.has_gen_mrope_delta runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 + # Under extend_ctx the attention kernel runs extend requests as chunked + # context rows: their rotary comes from position_ids alone, and they take + # no entry in the generation-indexed MRoPE delta read slots. + extend_requests_are_context = ( + _use_mrope + and enable_spec_decode + and spec_config.spec_dec_mode.extend_ctx(self._config.attention_backend) + and spec_config.is_linear_tree + ) for request in extend_requests: is_promoted_context = request.py_request_id in promoted_context_request_ids if getattr(request, "py_needs_onehot_draft_probs", False): @@ -1406,8 +1449,61 @@ def _pack_extend_requests( else: prompt_lengths.append(request.py_prompt_len) + if _use_mrope: + mrope_position_delta = self._decode_mrope_position_delta(request) + if mrope_position_delta is not None: + # Every row of a verify group lies past the prompt, so it is + # text: its three MRoPE coordinates are equal, the scalar + # position plus the request's delta. The overlap scheduler's + # accepted-token correction later adds to all three axes. + num_request_tokens = sequence_lengths[-1] + token_start = len(position_ids) - num_request_tokens + first_position = position_ids[token_start] + extend_mrope_position_ids = ( + torch.arange( + first_position, + first_position + num_request_tokens, + dtype=torch.int32, + device=mrope_position_delta.device, + ) + + mrope_position_delta.reshape(1, 1, 1) + ).expand(3, 1, num_request_tokens) + mrope_position_ids.append( + (token_start, len(position_ids), extend_mrope_position_ids) + ) + if ( + request.py_seq_slot is not None + and not request.is_dummy + and getattr(request, "py_mrope_delta_cache_slot", None) + != request.py_seq_slot + ): + multimodal_params_list.append( + MultimodalParams( + multimodal_data={ + "mrope_config": {"mrope_position_deltas": mrope_position_delta} + } + ) + ) + mrope_delta_write_seq_slots.append(request.py_seq_slot) + request.py_mrope_delta_cache_slot = request.py_seq_slot + if not extend_requests_are_context: + # The kernel indexes deltas by generation request, so every + # verify group takes one read slot, in batch order; a dummy, + # unslotted, or text-only request reads the zero slot. + if mrope_position_delta is not None: + has_gen_mrope_delta = True + mrope_delta_read_seq_slots.append( + request.py_seq_slot + if mrope_position_delta is not None + and not request.is_dummy + and request.py_seq_slot is not None + else mrope_dummy_seq_slot + ) + extra_inputs.add_generation_request(request) + host.has_gen_mrope_delta = has_gen_mrope_delta + def _pack_generation_requests( self, host: HostInputs, @@ -1439,9 +1535,9 @@ def _pack_generation_requests( _use_mrope = host.use_mrope _has_cp_helix = host.has_cp_helix mrope_dummy_seq_slot = host.mrope_dummy_seq_slot - # Whether any generation request in this batch carries real MRoPE - # metadata; see the post-loop cleanup below. - has_gen_mrope_delta = False + # Whether any generation request in this batch, verify groups included, + # carries real MRoPE metadata; see the post-loop cleanup below. + has_gen_mrope_delta = host.has_gen_mrope_delta # Cache invariant method result to avoid repeated calls per-request _n_gen = len(generation_requests) @@ -1572,17 +1668,7 @@ def _pack_generation_requests( gather_ids.append(len(position_ids) - 1) if _use_mrope: - mrope_position_delta = getattr(request, "py_mrope_position_delta", None) - if mrope_position_delta is None and request.py_multimodal_data: - mrope_config = request.py_multimodal_data.get("mrope_config") or {} - mrope_position_delta = mrope_config.get("mrope_position_deltas") - if mrope_position_delta is not None: - if mrope_position_delta.device.type == "cpu": - mrope_position_delta = maybe_pin_memory(mrope_position_delta).to( - device="cuda", dtype=torch.int32, non_blocking=True - ) - mrope_config["mrope_position_deltas"] = mrope_position_delta - request.py_mrope_position_delta = mrope_position_delta + mrope_position_delta = self._decode_mrope_position_delta(request) if mrope_position_delta is not None: has_gen_mrope_delta = True # NOTE: Expanding position_ids to 3D tensor who is using mrope diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 9f2f4e12a855..c53cb02cd55e 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -133,16 +133,22 @@ def _validate_and_adjust_mamba_snapshot_config(config: ModelConfig, kv_cache_config.enable_block_reuse = False +def effective_mamba_ssm_cache_dtype(config: ModelConfig, + mamba_ssm_cache_dtype: str) -> torch.dtype: + """The recurrent-state cache dtype the engine runs with.""" + if mamba_ssm_cache_dtype == "auto": + return resolve_auto_ssm_cache_dtype(config.pretrained_config, + config.torch_dtype) + return str_dtype_to_torch(mamba_ssm_cache_dtype) + + def validate_and_set_mamba_ssm_cache_dtype( config: ModelConfig, mamba_ssm_cache_dtype: str, mamba_ssm_stochastic_rounding: bool = False, mamba_ssm_philox_rounds: int = 10) -> None: - if mamba_ssm_cache_dtype == "auto": - mamba_ssm_cache_dtype = resolve_auto_ssm_cache_dtype( - config.pretrained_config, config.torch_dtype) - else: - mamba_ssm_cache_dtype = str_dtype_to_torch(mamba_ssm_cache_dtype) + mamba_ssm_cache_dtype = effective_mamba_ssm_cache_dtype( + config, mamba_ssm_cache_dtype) validate_kimi_kda_state_dtype(config.pretrained_config, mamba_ssm_cache_dtype, mamba_ssm_stochastic_rounding) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 9455105c8029..56c095e69cc1 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -6496,6 +6496,13 @@ def _validate_request(self, request: LlmRequest): # Perform sampler-specific validation self.sampler.validate_request(request) + # Perform model-specific validation, for models that define the hook. + # An executor without a model engine has no hook to call. + model = getattr(getattr(self, "model_engine", None), "model", None) + validate = getattr(model, "validate_request", None) + if validate is not None: + validate(request) + # Check if request has enough budget self._validate_request_budget(request) 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 f7764bd3fa60..949051ff1800 100644 --- a/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py +++ b/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py @@ -19,16 +19,35 @@ and report the built-in's numbers as the target's. """ +import random +import re +import time + import pytest +from PIL import Image, ImageDraw +from transformers import AutoProcessor, AutoTokenizer 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 + modeling as target_modeling, +) +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner import BreakableCUDAGraphRunner from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.executor.utils import RequestError +from tensorrt_llm.inputs import VideoData +from tensorrt_llm.inputs.utils import ( + ConversationMessage, + MultimodalDataTracker, + apply_chat_template, + apply_mm_placeholders, +) from tensorrt_llm.llmapi import ( CudaGraphConfig, KvCacheConfig, MoeConfig, MTPDecodingConfig, + QSASparseAttentionConfig, SamplingParams, ) @@ -46,11 +65,6 @@ get_sm_version() != 103, reason="modeling_v2 targets in this batch are certified on sm_103 only" ) -# Speculative cases need the legacy Gated-DeltaNet verify path: the engine's -# default (replay) commits through a callable the target's closed op vocabulary -# has no entry for, and the target asserts the legacy path is selected. -_GDN_LEGACY = {"TRTLLM_USE_GDN_REPLAY": "0"} - # Check the target's step contract on the first forward, against metadata the # engine prepared; the flag has to reach the worker rank. _VALIDATE = {MODELING_V2_VALIDATE_ENV: "1"} @@ -67,6 +81,266 @@ def _require_target(monkeypatch: pytest.MonkeyPatch, **env: str) -> dict: return dict(modeling_v2="require", env_overrides=dict(env)) +def _target_core(llm: LLM): + """The target's core module. + + Reachable only with the executor in this process (``TLLM_WORKER_USE_SINGLE_PROCESS=1``), + where the iteration stats that ``compute_acceptance_length`` reads are not collected. + """ + return llm._executor.engine.model_engine.model.model + + +def _count_prefill_graph_replays(monkeypatch: pytest.MonkeyPatch) -> dict[int, int]: + """Replays of the breakable prefill graph by bucket, filled as the engine runs. + + Needs the executor in this process. Without it a case that sets + ``prefill_cuda_graph_backend="breakable"`` could pass on the eager path. + """ + replays: dict[int, int] = {} + replay = BreakableCUDAGraphRunner.replay + + def counting_replay(self, num_tokens): + replays[num_tokens] = replays.get(num_tokens, 0) + 1 + return replay(self, num_tokens) + + monkeypatch.setattr(BreakableCUDAGraphRunner, "replay", counting_replay) + return replays + + +def _acceptance_length_from_requests(outputs, draft_len: int) -> float: + """Mean acceptance length from the requests' own speculative-decoding metrics. + + Needs ``SamplingParams(return_perf_metrics=True)``. A draft loop yields one target token + plus the drafts it accepted. + """ + accepted = drafted = 0 + for output in outputs: + spec = output.outputs[0].request_perf_metrics.speculative_decoding + accepted += spec.total_accepted_draft_tokens + drafted += spec.total_draft_tokens + assert drafted > 0, "no request drafted a token" + return 1.0 + accepted * draft_len / drafted + + +# ---------------------------------------------------------------- long prompts --- +# Prompts past the QSA indexer's 2048-token budget, so every decode step of the +# target and of its MTP draft layer takes the sparse attention path. + +_PASSKEY_LENGTHS = (3000, 4000, 5000, 6000) +_PASSKEY_LENGTHS_LONG = (3000, 6000, 9000, 12000) +# 33 prompts, each past the indexer's budget, 32 of which one batch carries. +_PASSKEY_DEPTHS = (0.15, 0.5, 0.85) +_PASSKEY_MAX_TOKENS = 160 + +_COMPONENTS = ( + "scheduler", + "allocator", + "tokenizer", + "router", + "profiler", + "exporter", + "planner", + "balancer", + "resolver", + "dispatcher", +) +_ACTIONS = ( + "reads", + "rewrites", + "caches", + "validates", + "forwards", + "compacts", + "audits", + "replays", + "indexes", + "batches", +) +_OBJECTS = ( + "pending requests", + "block tables", + "shard manifests", + "routing hints", + "retry budgets", + "lease records", + "checkpoint deltas", + "quota counters", +) +_CLAUSES = ( + "before the next flush", + "whenever the queue drains", + "after every shard reports", + "once the lease expires", + "while the window is open", + "until the barrier clears", +) +_STATES = ("consistent", "bounded", "ordered", "reproducible", "idle") +_TOPICS = ("rules", "limits", "guarantees", "hooks", "settings") + + +def _section(rng: random.Random, number: int) -> str: + heading = f"## Section {number}: {rng.choice(_COMPONENTS).capitalize()} {rng.choice(_TOPICS)}" + sentences = [ + f"The {rng.choice(_COMPONENTS)} {rng.choice(_ACTIONS)} the {rng.choice(_OBJECTS)} " + f"{rng.choice(_CLAUSES)}, so the {rng.choice(_OBJECTS)} stay {rng.choice(_STATES)}." + for _ in range(6) + ] + return heading + "\n\n" + " ".join(sentences) + + +def _passkey_prompts(tokenizer, lengths=_PASSKEY_LENGTHS) -> list[tuple[str, str]]: + """``(chat prompt, passkey)`` pairs over ``lengths`` x ``_PASSKEY_DEPTHS``. + + Each prompt is a generated document of about that many tokens with one pass key sentence + at that relative depth, the same bytes in every run. + """ + rng = random.Random(20261014) + question = ( + "Based on the document above, first state the secret pass key on a line of its own " + "in the form 'Pass key: '. Then list, one per line, the section headings " + "that appear in the document, in order." + ) + prompts = [] + for length in lengths: + for depth in _PASSKEY_DEPTHS: + key = str(rng.randint(100000, 999999)) + sections, tokens = [], 0 + while tokens < length - 150: + sections.append(_section(rng, len(sections) + 1)) + tokens += len(tokenizer.encode(sections[-1])) + needle = ( + f"IMPORTANT NOTE FOR THE READER: the secret pass key of this document is {key}. " + f"Remember the pass key {key} exactly." + ) + position = max(1, min(len(sections) - 1, round(depth * len(sections)))) + body = sections[:position] + [needle] + sections[position:] + text = "Document:\n\n" + "\n\n".join(body) + "\n\n" + question + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": text}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + prompts.append((prompt, key)) + return prompts + + +# ------------------------------------------------------------------ word problems --- + +_WORD_PROBLEM_SUFFIX = "Solve it step by step and finish with 'The answer is .'" +_WORD_PROBLEMS = ( + ( + "A bakery bakes 7 trays of 12 muffins and sells all but 9. How many muffins did it sell?", + "75", + ), + ( + "Tom has 3 boxes of 24 pencils and gives away 19 pencils. How many pencils does he " + "have left?", + "53", + ), + ( + "A train travels 60 miles per hour for 2.5 hours, then 40 miles per hour for 3 hours. " + "How many miles does it travel in total?", + "270", + ), + ( + "Sara buys 5 notebooks at $3 each and 4 pens at $2 each and pays with a $50 bill. " + "How many dollars of change does she get?", + "27", + ), + ("A rectangle is 14 cm long and 9 cm wide. What is its perimeter in cm?", "46"), + ( + "125 students attend; 40% play soccer and one fifth of the others play tennis. " + "How many students play tennis?", + "15", + ), +) + + +# ---------------------------------------------------------------- media inputs --- + +_MODEL_TYPE = "qwen4_exp" +_RED, _BLUE, _GREEN = (220, 30, 30), (30, 60, 220), (30, 170, 60) +# name -> (square colour, direction, colour word, direction phrases an answer may use) +_VIDEOS = { + "red_rightward": (_RED, "right", "red", ("left to right", "to the right", "rightward")), + "blue_leftward": (_BLUE, "left", "blue", ("right to left", "to the left", "leftward")), + "green_downward": (_GREEN, "down", "green", ("top to bottom", "downward", "down")), +} +_VIDEO_QUESTION = ( + "What colour is the moving square, and in which direction does it move across the frame? " + "Answer in one short sentence." +) +_MIXED_QUESTION = ( + "The image shows a square and the video shows a moving square. Answer in exactly this " + "format: 'Image: . Video: .'" +) + + +def _synthetic_video(name: str, num_frames: int = 16, size: int = 256) -> VideoData: + """A square of the named colour crossing a plain background in the named direction. + + The frames are built in memory, so no video decoder is involved. + """ + color, direction, _, _ = _VIDEOS[name] + side, margin = 48, 8 + frames = [] + for i in range(num_frames): + t = i / (num_frames - 1) + travel = margin + t * (size - side - 2 * margin) + centre = (size - side) // 2 + x, y = { + "right": (travel, centre), + "left": (size - side - travel, centre), + "down": (centre, travel), + }[direction] + frame = Image.new("RGB", (size, size), (235, 235, 235)) + ImageDraw.Draw(frame).rectangle((x, y, x + side, y + side), fill=color) + frames.append(frame) + return VideoData( + frames=frames, + metadata=dict( + total_num_frames=num_frames, + fps=2.0, + duration=num_frames / 2.0, + frames_indices=list(range(num_frames)), + ), + ) + + +def _synthetic_image(color: tuple, size: int = 256) -> Image.Image: + """A square of ``color`` on a plain background.""" + image = Image.new("RGB", (size, size), (235, 235, 235)) + quarter = size // 4 + ImageDraw.Draw(image).rectangle((quarter, quarter, size - quarter, size - quarter), fill=color) + return image + + +def _media_request(tokenizer, processor, text: str, items: list[tuple[str, object]]) -> dict: + """One request carrying ``items``, ``(modality, data)`` pairs in prompt order. + + It holds the chat prompt with its placeholders, the media and the ``mm_item_order`` + manifest that splices several modalities back in prompt order. + """ + tracker = MultimodalDataTracker(_MODEL_TYPE) + for modality, data in items: + tracker.add_data(modality, data) + counts = tracker.placeholder_counts() + message = ConversationMessage(role="user", content=text, media=[]) + apply_mm_placeholders(_MODEL_TYPE, message, counts, tracker) + prompt = apply_chat_template( + model_type=_MODEL_TYPE, + tokenizer=tokenizer, + processor=processor, + conversation=[message], + add_generation_prompt=True, + mm_placeholder_counts=[counts], + chat_template_kwargs=dict(enable_thinking=False), + ) + media, _ = tracker.retrieve_all_sync() + 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.""" @@ -88,13 +362,35 @@ class TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1(LlmapiAccuracyTestHarness): GSM8K_MAX_OUTPUT_LEN = 512 # `mamba_ssm_cache_dtype` is float32 because the checkpoint declares - # `mamba_ssm_dtype: float32` (the engine default is bfloat16). `dtype` is - # left at the checkpoint's bf16 KV, which is what the anchors were measured at. + # `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( free_gpu_memory_fraction=0.6, enable_block_reuse=True, mamba_ssm_cache_dtype="float32", ) + KV_CACHE_FP8_BF16_STATE = KvCacheConfig( + dtype="fp8", + free_gpu_memory_fraction=0.6, + enable_block_reuse=True, + mamba_ssm_cache_dtype="bfloat16", + ) + # 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( + 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( + 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 @@ -110,11 +406,131 @@ class TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1(LlmapiAccuracyTestHarness): moe_config=MOE, ) - # One-model MTP-Eagle with the checkpoint's single MTP module. + # The breakable prefill CUDA graph is opt-in; these buckets cover the prompts of + # `test_breakable_prefill_matches_eager`. + 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. MTP1 = MTPDecodingConfig(max_draft_len=1) + MTP3 = MTPDecodingConfig(max_draft_len=3) + MTP7 = MTPDecodingConfig(max_draft_len=7) - # Acceptance-length anchor, keyed by target and variant. + # 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" + ) + ACCEPTANCE_KEY_MTP7_LONG_CONTEXT = ( + "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::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``. + + 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``. + """ + 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. + + 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``. + """ + 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)), + ) + + def _video_checks(self, llm): + """Synthetic videos through ``llm``. + + Three videos of a moving square (frames built in memory, since the container has no video + decoder) must each be answered with the right colour and direction. A request mixing an + image and a video, in either order, must reach the forward: its ``mm_item_order`` + manifest has to survive the input processor, the request and the admission check, which + refuses a mixed request without a usable one. A frame past the vision tower's patch + 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``. + """ + tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) + processor = AutoProcessor.from_pretrained( + self.MODEL_PATH, use_fast=True, trust_remote_code=True + ) + greedy = SamplingParams(max_tokens=96, temperature=0) + + def ask(text, items): + request = _media_request(tokenizer, processor, text, items) + return llm.generate_async(request, sampling_params=greedy).result().outputs[0].text + + for name, (_, _, colour, directions) in _VIDEOS.items(): + answer = ask(_VIDEO_QUESTION, [("video", _synthetic_video(name))]).lower() + assert colour in answer and any(d in answer for d in directions), (name, answer) + + # Each label carries its own item's colour. + image = _synthetic_image((250, 200, 0)) + for items, video_colour in ( + ([("video", _synthetic_video("blue_leftward")), ("image", image)], "blue"), + ([("image", image), ("video", _synthetic_video("red_rightward"))], "red"), + ): + answer = ask(_MIXED_QUESTION, items).lower() + assert re.search(r"image:\W*yellow", answer), answer + assert re.search(rf"video:\W*{video_colour}", answer), answer + + with pytest.raises(RequestError, match="frames of 192x192 = 36864 patches"): + ask( + "What shape and colour is in the image?", + [("image", _synthetic_image(_RED, 3072))], + ) + answer = ask(_VIDEO_QUESTION, [("video", _synthetic_video("red_rightward"))]).lower() + assert "red" in answer and "right" in answer, answer @skip_not_sm103 def test_gsm8k(self, mocker, monkeypatch): @@ -130,61 +546,243 @@ def test_gsm8k(self, mocker, monkeypatch): task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) @skip_not_sm103 - def test_mtp1_gsm8k_and_acceptance(self, mocker, monkeypatch): - """MTP with one draft: gsm8k against the MTP reference, plus acceptance length. - - A miscomputed draft layer only makes generation slower (every draft is - rejected, the text stays correct), so acceptance length is the only - detector. The floor in references/acceptance_length.yaml is a collapse - tripwire. The draft layer runs dense attention and the PLE pools have - no speculative-state rollback; no gate here sees either. - """ - mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN) + def test_mtp1_gsm8k_mmmu_and_video(self, mocker, monkeypatch): + """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 + makes generation slower (every draft rejected, the text still correct), so the + acceptance length of the same run is the only detector; its floor is a collapse + tripwire. + - MMMU against the MTP row of references/mmmu.yaml: the vision tower and the mrope path. + - Videos (`_video_checks`): the only case that sends a video. + + 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) with LLM( self.MODEL_PATH, speculative_config=self.MTP1, enable_iter_perf_stats=True, - **self.ENGINE, - **_require_target(monkeypatch, **_GDN_LEGACY, **_VALIDATE), + **engine, + **_require_target(monkeypatch, **_VALIDATE), ) as llm: - task = GSM8K(self.MODEL_NAME) - task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) - acceptance_length = compute_acceptance_length(llm) - print(f"[AL] mtp1 acceptance_length = {acceptance_length:.3f}") - assert_acceptance_length(self.ACCEPTANCE_KEY, acceptance_length) + self._gsm8k_and_acceptance(mocker, llm, self.ACCEPTANCE_KEY, "mtp1") + self._mmmu(llm) + self._video_checks(llm) @skip_not_sm103 - def test_mtp1_mmmu(self, monkeypatch): - """The vision tower and mrope path under MTP, against the MTP row of references/mmmu.yaml. + def test_mtp3_cutedsl_fp8_kv_bf16_state(self, mocker, monkeypatch): + """gsm8k and 3k-6k-token pass keys with three drafts over the opt-in choices. - Block reuse is off here (on in the cases above) because the vision path - requires it. + One engine runs gsm8k with an acceptance-length floor, then the pass keys. The trunk's + routed experts run on CuteDSL next to the draft layer's FP8 block-scale experts; the + KV cache is FP8 (unit scales) and the recurrent state bf16 (what ``auto`` + resolves to for this checkpoint), rounded at every decode token, prefill call and + cached-replay commit. gsm8k reads the FP8-KV MTP row of references/gsm8k.yaml; its + prompts stay on the dense attention path. The pass keys run the FP8 sparse store and + GQA kernels of the target and of the draft layer (IndexShare off, so the draft layer + selects on every step). Three drafts give the verify step four rows per request. + Acceptance length detects a trunk hidden state or a drifted state the draft layer + cannot use, which accuracy does not. """ - engine = dict(self.ENGINE) - engine["kv_cache_config"] = KvCacheConfig( - free_gpu_memory_fraction=0.6, - enable_block_reuse=False, - mamba_ssm_cache_dtype="float32", + engine = dict( + self.ENGINE, + moe_config=MoeConfig(backend="CUTEDSL"), + kv_cache_config=self.KV_CACHE_FP8_BF16_STATE, + sparse_attention_config=QSASparseAttentionConfig(index_share_for_mtp_iteration=False), ) + with LLM( + self.MODEL_PATH, + speculative_config=self.MTP3, + enable_iter_perf_stats=True, + **engine, + **_require_target(monkeypatch, **_VALIDATE), + ) as llm: + self._gsm8k_and_acceptance(mocker, llm, self.ACCEPTANCE_KEY_OPT_IN, "mtp3 opt-in") + self._passkeys( + llm, + _PASSKEY_LENGTHS, + self.MTP3.max_draft_len, + self.ACCEPTANCE_KEY_OPT_IN_LONG_CONTEXT, + "mtp3 opt-in long context", + ) - # 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|>", - ) + 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 LLM( self.MODEL_PATH, - speculative_config=self.MTP1, + speculative_config=speculative_config, **engine, - **_require_target(monkeypatch, **_GDN_LEGACY, **_VALIDATE), + **_require_target(monkeypatch, **_VALIDATE), ) as llm: - task = MMMU(self.MODEL_NAME) - task.evaluate( - llm, - sampling_params=sampling_params, - extra_evaluator_kwargs=dict(chat_template_kwargs=dict(enable_thinking=False)), + 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. + """ + self._long_context_passkey( + monkeypatch, + self.MTP3, + self.ENGINE, + True, + _PASSKEY_LENGTHS, + self.ACCEPTANCE_KEY_LONG_CONTEXT, + "mtp3 long context", + ) + + @skip_not_sm103 + def test_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 + cut at the edges of the buckets 128 and 1024 (one token under, at, one over), at the + top bucket 3072 and one past it, and a 5000-token prompt prefilled in two chunks. A + batch of word problems and a 5000-token prompt then prefills while the earlier ones + decode, so replayed steps carry generation rows beside the context rows. + + A prompt that fills its bucket runs the same kernels on the same shapes as the eager + step, so its tokens must be identical. A padded prompt changes the row count every + GEMM sees, which moves the logits by what a different batch size does (about 0.2 nats + at most here); it must still agree on its first four tokens. The same engine given + the same prompts twice must repeat itself exactly. Replays are counted, since a + replay that never happens passes everything above. + """ + monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") + replays = _count_prefill_graph_replays(monkeypatch) + tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) + document = tokenizer.encode(_passkey_prompts(tokenizer, (6000,))[0][0]) + lengths = (127, 128, 129, 1023, 1024, 1025, 3072, 3073, 5000) + filled = (128, 1024, 3072) + # Word problems, answered in full, and one prompt of five thousand tokens prefilled in + # chunks while they decode. A word problem's first token can be a near tie between two + # openings, so its answer is checked, not its tokens. + problems = [ + tokenizer.apply_chat_template( + [{"role": "user", "content": f"{question} {_WORD_PROBLEM_SUFFIX}"}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, ) + for question, _ in _WORD_PROBLEMS + ] + greedy = SamplingParams(max_tokens=24, temperature=0) + answering = SamplingParams(max_tokens=384, temperature=0) + + def run(engine): + with LLM( + self.MODEL_PATH, + speculative_config=self.MTP3, + **engine, + **_require_target(monkeypatch, **_VALIDATE), + ) as llm: + single = [ + llm.generate({"prompt_token_ids": document[:n]}, greedy, use_tqdm=False) + .outputs[0] + .token_ids + for n in lengths + ] + again = [ + llm.generate({"prompt_token_ids": document[:n]}, greedy, use_tqdm=False) + .outputs[0] + .token_ids + for n in lengths + ] + # The word problems are decoding when the long prompt arrives: their short + # prefills are done well inside the pause, and a step-by-step answer is + # still being written after it. + futures = [llm.generate_async(prompt, answering) for prompt in problems] + time.sleep(0.2) + assert not all(future.finished for future in futures), ( + "every word problem finished before the long prompt arrived" + ) + long_future = llm.generate_async({"prompt_token_ids": document[:5000]}, greedy) + answers = [future.result().outputs[0].text for future in futures] + 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)) + 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) + ) + assert not replays, replays + + for n, got, want in zip(lengths, single, eager_single): + head = len(want) if n in filled else 4 + assert got[:head] == want[:head], (n, got, want) + for texts in (answers, eager_answers): + wrong = [ + (question, answer, text[-120:]) + for (question, answer), text in zip(_WORD_PROBLEMS, texts) + if f"answer is {answer}" not in text.replace("**", "") + ] + assert not wrong, f"wrong or missing answers: {wrong}" + assert long_tokens[:4] == eager_long_tokens[:4], (long_tokens, eager_long_tokens) + + @skip_not_sm103 + def test_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 + with eight rows per request in the verify step, the cached replay kernel's ceiling + (``T = 8``), and draft decode steps that reuse the first step's selection up to six + positions past it, where the groups completed since the capture displace two of its + ids. The NVFP4 pool runs the target's and the draft layer's NVFP4 sparse store and GQA + kernels, and the context rows' dequantized scratch. The bf16 state is rounded at every + cached-replay checkpoint commit and once per prefill chunk, so the rounding adds up + with the length of the prompt and of the answer; pass keys stated back over 3k-12k + 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._long_context_passkey( + monkeypatch, + self.MTP7, + engine, + True, + _PASSKEY_LENGTHS_LONG, + self.ACCEPTANCE_KEY_MTP7_LONG_CONTEXT, + "mtp7 nvfp4 kv bf16 state long context", + ) diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml index 45a750e9f47f..5f921108d166 100644 --- a/tests/integration/defs/accuracy/references/acceptance_length.yaml +++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml @@ -123,21 +123,25 @@ ModelingV2DeepseekR10528Nvfp4Sm10xDefault::mtp3: TestGLM53FlashFP8::test_mtp: ref_al: 3.384 min_al: 3.215 -# Read by TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1:: -# test_mtp1_gsm8k_and_acceptance. Keyed by target and variant for the same -# reason as the ModelingV2DeepseekR10528Nvfp4Sm10xDefault::mtp3 entry above. +# Read by the acceptance-length cases of TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1 +# (test_qwen4_exp.py). Keyed by target 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. # -# max_draft_len is 1 on this checkpoint (`mtp_num_hidden_layers: 1`, one module -# replayed once per step), so acceptance length lives in [1.0, 2.0]: 1.0 is a -# draft path that is rejected every time -- the exact failure a miscomputed MTP -# layer produces, with the text still correct and every accuracy gate green -- -# and 2.0 is every draft accepted. +# The checkpoint ships one MTP module (`mtp_num_hidden_layers: 1`), replayed max_draft_len +# times per step, so acceptance length lives in [1.0, 1 + max_draft_len]: 1.0 is a draft +# path that is rejected every time -- the exact failure a miscomputed MTP layer produces, +# with the text still correct and every accuracy gate green -- and 1 + max_draft_len is +# 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: # 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. - # There is no full-gsm8k reading yet: the accuracy runs were made without - # enable_iter_perf_stats, so they recorded no specDecodingStats. + # There is no full-gsm8k STOCK reading yet; modeling_v2's own over full gsm8k + # (1319 samples) is 1.971. ref_al: 1.7403 # Hand-set as a collapse tripwire, NOT the automatic 95% of ref_al -- 1.65 # would be a floor derived from a two-prompt smoke applied to a 1319-sample @@ -145,12 +149,48 @@ ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1: # the variant's whole range: a collapsed or miscomputed draft path scores 1.0, # and anything at or above 1.5 is accepting at least half its drafts. # - # This unit has a disclosed deviation the gate is deliberately sized around: - # the draft layer runs dense attention rather than QSA, because - # `QSAAttentionMetadata` has no `prepare_for_draft_forward`. That costs - # acceptance on long prompts and costs nothing on short ones, so a floor set - # tight against a short-prompt reading would fail on the benchmark for a - # reason already written down. Replace both numbers with a full-gsm8k STOCK - # reading (TRTLLM_POPULATE_ACCEPTANCE_LENGTH=1 on the stock leg) when one is + # Replace both numbers with a full-gsm8k STOCK reading + # (TRTLLM_POPULATE_ACCEPTANCE_LENGTH=1 on the stock leg) when one is # 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 +# CuteDSL over an fp8 KV cache and a bf16 recurrent state. +ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::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 + # 95% of ref_al, above the two-draft reading (~2.8): a third draft position that is never + # 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: + # The same prompts with the draft layer selecting on every step (IndexShare off): 4.000, + # equal to the IndexShare reading. + ref_al: 4.000 + # Hand-set, 2.5% below ref_al rather than the automatic 95% (3.8): a draft loop whose + # reused selection is deliberately wrong (the oldest 512 groups of every row) scores + # 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: + # 3.962 here, against 4.000 over bf16 KV. + ref_al: 3.962 + # 2.5% below ref_al, as for mtp3_long_context. + min_al: 3.86 + +# test_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: + ref_al: 6.867 + # 2.5% below ref_al, as for mtp3_long_context. + min_al: 6.69 diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index 4413b6b0f811..a2593fb2a263 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -299,6 +299,18 @@ Qwen/Qwen3.8-Flash-Next: - quant_algo: MIXED_PRECISION spec_dec_algo: MTP accuracy: 95.57 + # FP8 KV cache over the checkpoint's bf16 KV, measured on the built-in implementation: + # TP1 on GB300, MTP draft length 1, TRTLLM MoE, fp32 Mamba SSM cache. + - quant_algo: MIXED_PRECISION + kv_cache_quant_algo: FP8 + spec_dec_algo: MTP + accuracy: 95.94 + # NVFP4 KV cache over the checkpoint's bf16 KV, measured on the built-in implementation: + # TP1 on GB300, MTP draft length 1, TRTLLM MoE, fp32 Mamba SSM cache. + - quant_algo: MIXED_PRECISION + kv_cache_quant_algo: NVFP4 + spec_dec_algo: MTP + accuracy: 95.26 moonshotai/Kimi-K2.5: - quant_algo: NVFP4 kv_cache_quant_algo: FP8 diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index fd4554680089..43e69d30002d 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -135,15 +135,18 @@ l0_b300: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_chunked_prefill[quant_dtype=fp8-kv_cache_reuse=True-fp8kv=True-overlap_scheduler=True] # Cover nvbugs 6084445 - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=True] - # TIMEOUT measured: 348 cases in 5m34s on one GB300, so 30 is 5x headroom. + # TIMEOUT measured: 3122 cases in 13m54s on one GB300 (the cpu_only files run on the CPU + # stage instead), so 30 is 2.2x headroom. - unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm TIMEOUT (30) - accuracy/modeling_v2/test_gpt_oss.py::TestModelingV2GptOss120bSm10xDefault::test_gsm8k[tp1] - # qwen4_exp target: gsm8k; gsm8k with MTP plus the acceptance-length floor, the - # only gate that sees a miscomputed draft layer; mmmu with MTP, the only gate - # that reaches the vision tower. + # 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_and_acceptance - - accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_mtp1_mmmu + - 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 - condition: ranges: system_gpu_count: @@ -164,3 +167,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 diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 726469ed58ac..87c73a700e44 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -68,13 +68,23 @@ 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 rest of - # unittest/_torch/modeling_v2 needs a GPU. + # 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. - unittest/_torch/modeling_v2/test_modeling_v2_claims.py - unittest/_torch/modeling_v2/test_modeling_v2_routing.py - unittest/_torch/modeling_v2/test_modeling_v2_deepseek_default_shards.py - unittest/_torch/modeling_v2/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/test_modeling_v2_qwen4_exp_multimodal_embeds.py + - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py + - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py + - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py + - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_kv_quant.py + - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_thop_attention_calls.py + - unittest/_torch/modeling_v2/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/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py new file mode 100644 index 000000000000..18abace17337 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_cache_manager_layout.py @@ -0,0 +1,281 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The QSA hybrid cache manager's V2 page layout, on real managers. + +The attention op addresses a layer's K/V page as ``pool + (scale * slot + 2 * +offset) * page`` and its NVFP4 block-scale page with the same block index in +the scale pool. These tests build the managers a Qwen3.8-Flash-Next engine +builds -- the 48-layer trunk (12 QSA layers among 36 recurrent ones) and the +one-layer MTP draft manager -- and check that every view the manager serves is +where that formula lands, for bf16, fp8 and NVFP4 pools. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.qsa import QSAMambaHybridCacheManagerV2 +from tensorrt_llm._torch.attention.backends.sparse.qsa.cache_manager import ( + QSA_INDEX_POSITION, + QSA_SIDE_PLANES, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.llmapi.llm_args import KvCacheConfig, QSASparseAttentionConfig +from tensorrt_llm.mapping import Mapping + +_CHECKPOINT = SimpleNamespace( + indexer_n_heads=4, + indexer_kv_heads=1, + indexer_head_dim=128, + indexer_budget=2048, + indexer_compress_ratio=4, +) +_TRUNK_LAYERS = 48 +_QSA_LAYERS = list(range(3, _TRUNK_LAYERS, 4)) +_DRAFT_LAYER = _TRUNK_LAYERS +_KV_HEADS, _HEAD_DIM, _TPB = 2, 256, 32 +_INDEX_PAGE = 1 * 128 * 2 * _TPB +_POSITION_PAGE = 3 * 4 * _TPB + + +def _manager( + dtype: DataType, *, draft: bool, tp_size: int = 1, attention_dp: bool = False +) -> QSAMambaHybridCacheManagerV2: + """The trunk or the separate MTP draft manager, as the engine builds them (rank 0 of + `tp_size`, with attention data parallelism when `attention_dp`).""" + if draft: + attention_mask = [False] * _TRUNK_LAYERS + [True] + mamba_mask = [False] * (_TRUNK_LAYERS + 1) + else: + attention_mask = [layer in _QSA_LAYERS for layer in range(_TRUNK_LAYERS)] + mamba_mask = [not attention for attention in attention_mask] + return QSAMambaHybridCacheManagerV2( + mamba_d_state=8, + mamba_d_conv=4, + mamba_num_heads=4, + mamba_n_groups=1, + mamba_head_dim=8, + mamba_num_layers=sum(mamba_mask), + mamba_layer_mask=mamba_mask, + mamba_cache_dtype=torch.bfloat16, + mamba_ssm_cache_dtype=torch.float32, + kv_cache_config=KvCacheConfig(max_tokens=4096, enable_block_reuse=False), + kv_cache_type=CacheTypeCpp.SELF, + num_layers=sum(attention_mask), + num_kv_heads=_KV_HEADS, + head_dim=_HEAD_DIM, + tokens_per_block=_TPB, + max_seq_len=1024, + max_batch_size=4, + mapping=Mapping( + world_size=tp_size, + rank=0, + tp_size=tp_size, + pp_size=1, + enable_attention_dp=attention_dp, + ), + dtype=dtype, + layer_mask=attention_mask, + is_draft=draft, + vocab_size=1024, + conv_state_layout="q_k_v", + sparse_attention_config=QSASparseAttentionConfig(), + pretrained_config=_CHECKPOINT, + ) + + +def _layers(draft: bool) -> list[int]: + return [_DRAFT_LAYER] if draft else _QSA_LAYERS + + +def _kv_page_bytes(dtype: DataType) -> int: + element_bits = {DataType.BF16: 16, DataType.FP8: 8, DataType.NVFP4: 4}[dtype] + return _KV_HEADS * _TPB * _HEAD_DIM * element_bits // 8 + + +def _thop_addresses(cm, layer: int, slot: int, page_bytes: int, plane: int) -> tuple[int, int]: + """Byte addresses of this layer's K and V page for `slot`, as the attention + op forms them from the manager's pool pointers, mapping and page-index + scale (`plane` 0 = data, 1 = NVFP4 block scales).""" + local = cm.layer_offsets[layer] + pool, offset = cm.kv_cache_pool_mapping[local].tolist() + pointers = cm.kv_cache_pool_pointers + base = int(pointers[pool, 0, plane]) if pointers.dim() == 3 else int(pointers[pool, 0]) + key_entry = int(cm.index_scales[pool]) * slot + value_entry = key_entry + int(cm.kv_offset[pool]) + shift = offset * cm.kv_factor * page_bytes + return base + key_entry * page_bytes + shift, base + value_entry * page_bytes + shift + + +def _slots_to_check(view: torch.Tensor) -> list[int]: + return [0, 1, view.shape[0] // 2, view.shape[0] - 1] + + +@pytest.fixture(scope="module", autouse=True) +def _require_cuda(): + assert torch.cuda.is_available(), "the managers allocate their pools on the GPU" + + +@pytest.mark.parametrize("draft", [False, True], ids=["trunk", "draft"]) +@pytest.mark.parametrize("dtype", [DataType.BF16, DataType.FP8, DataType.NVFP4]) +def test_views_are_where_the_attention_op_addresses(dtype: DataType, draft: bool) -> None: + """Every K/V (and NVFP4 block-scale) view of every QSA layer starts, at + every checked slot, at the address the attention op computes for it.""" + cm = _manager(dtype, draft=draft) + try: + page = _kv_page_bytes(dtype) + for layer in _layers(draft): + kv = cm.get_buffers(layer, kv_layout="HND") + assert kv.shape[1:] == ( + 2, + _KV_HEADS, + _TPB, + _HEAD_DIM // (2 if dtype == DataType.NVFP4 else 1), + ) + for slot in _slots_to_check(kv): + key, value = _thop_addresses(cm, layer, slot, page, plane=0) + assert kv[slot, 0].data_ptr() == key, (layer, slot) + assert kv[slot, 1].data_ptr() == value, (layer, slot) + if dtype != DataType.NVFP4: + continue + scales = cm.get_block_scale_buffers(layer, kv_layout="HND") + assert scales.dtype == torch.uint8 + assert scales.shape == (kv.shape[0], 2, _KV_HEADS, _TPB, _HEAD_DIM // 16) + for slot in _slots_to_check(scales): + key, value = _thop_addresses(cm, layer, slot, page // 8, plane=1) + assert scales[slot, 0].data_ptr() == key, (layer, slot) + assert scales[slot, 1].data_ptr() == value, (layer, slot) + finally: + cm.shutdown() + + +@pytest.mark.parametrize("draft", [False, True], ids=["trunk", "draft"]) +@pytest.mark.parametrize("dtype", [DataType.FP8, DataType.NVFP4]) +def test_side_planes_pack_when_a_page_could_collide(dtype: DataType, draft: bool) -> None: + """An index page could be a K/V page here -- FP8 at one KV head a rank (1 x 256 x 32 + tokens), NVFP4 at both -- so the index and position planes are packed into a buffer + of their own instead of landing in the K/V pool, and under NVFP4 the K/V and + block-scale pools hold the same number of pages per slot.""" + cm = _manager(dtype, draft=draft) + try: + layers = _layers(draft) + assert cm.qsa_side_planes_size == len(layers) * _INDEX_PAGE + _POSITION_PAGE + for ordinal, layer in enumerate(layers): + local = cm.layer_offsets[layer] + data_scale = cm.impl.get_page_index_scale(local, Role.KEY) + if dtype == DataType.NVFP4: + assert data_scale == cm.impl.get_page_index_scale(local, Role.KEY_BLOCK_SCALE) + assert data_scale == 2 * len(layers) + assert cm.kv_cache_pool_mapping[local, 1] == ordinal + planes_owner = cm.layer_offsets[layers[0]] + assert cm.impl.get_page_index_scale(planes_owner, QSA_SIDE_PLANES) == 1 + finally: + cm.shutdown() + + +@pytest.mark.parametrize("draft", [False, True], ids=["trunk", "draft"]) +def test_bf16_keeps_per_layer_side_roles(draft: bool) -> None: + """A bf16 K/V page (at least 1 x 256 x 2 bytes x 32 tokens) can never be an index or + position page, so the side planes stay the per-layer INDEX_KEY role plus one position + role, served by the same views the generic accessor builds.""" + cm = _manager(DataType.BF16, draft=draft) + try: + layers = _layers(draft) + assert cm.qsa_side_planes_size is None + for layer in layers: + local = cm.layer_offsets[layer] + assert cm.impl.get_page_index_scale(local, Role.INDEX_KEY) == len(layers) + assert cm.impl.get_page_index_scale(local, Role.KEY) == 2 * len(layers) + generic = MambaHybridCacheManagerV2.get_index_k_buffer( + cm, layer, num_heads=1, head_dim=128, dtype=torch.bfloat16 + ) + served = cm.get_index_k_buffer(layer) + assert (served.data_ptr(), served.shape, served.stride()) == ( + generic.data_ptr(), + generic.shape, + generic.stride(), + ) + owner = cm.layer_offsets[layers[0]] + assert cm.impl.get_page_index_scale(owner, QSA_INDEX_POSITION) == 1 + assert cm.get_qsa_position_buffer().stride(0) == _TPB * 3 + finally: + cm.shutdown() + + +@pytest.mark.parametrize("draft", [False, True], ids=["trunk", "draft"]) +@pytest.mark.parametrize("dtype", [DataType.BF16, DataType.FP8, DataType.NVFP4]) +def test_side_and_main_views_do_not_alias(dtype: DataType, draft: bool) -> None: + """One slot written through every view reads back unchanged through each: + no K/V, scale, index or position page overlaps another.""" + cm = _manager(dtype, draft=draft) + try: + views = [] + for layer in _layers(draft): + kv = cm.get_buffers(layer, kv_layout="HND") + views += [kv[:, 0], kv[:, 1]] + if dtype == DataType.NVFP4: + scales = cm.get_block_scale_buffers(layer, kv_layout="HND") + views += [scales[:, 0], scales[:, 1]] + views.append(cm.get_index_k_buffer(layer)) + views.append(cm.get_qsa_position_buffer()) + slots = {view.shape[0] for view in views} + assert len(slots) == 1, f"views disagree on the slot count: {sorted(slots)}" + slot = slots.pop() - 1 + as_bytes = [view[slot].view(torch.uint8) for view in views] + for i, raw in enumerate(as_bytes): + raw.fill_(i + 1) + torch.cuda.synchronize() + for i, raw in enumerate(as_bytes): + assert bool((raw == i + 1).all()), f"view {i} was overwritten by another view" + finally: + cm.shutdown() + + +def _logical_page(view: torch.Tensor) -> tuple: + """What a peer copies of one side view per slot: the page's shape, dtype and bytes, and + whether they are contiguous. Where the page sits in its pool (the slot stride, other + buffers coalesced beside it) is physical and may differ between layouts.""" + page = view[0] + return tuple(page.shape), page.dtype, page.numel() * page.element_size(), page.is_contiguous() + + +#: (tp_size, attention_dp): the layouts a disaggregated pair can mix over one global +#: geometry (2 KV heads): TP1 and ADP2 hold both heads per rank, TP2 one. +_LAYOUTS = {"tp1": (1, False), "tp2": (2, False), "adp2": (2, True)} + + +@pytest.mark.parametrize("dtype", [DataType.BF16, DataType.FP8, DataType.NVFP4]) +def test_transfer_views_do_not_depend_on_the_kv_head_split(dtype: DataType) -> None: + """Disaggregated peers compare the replicated side roles each layer registers and copy + their pages byte for byte, so roles, mapper kinds and every side view's page (shape, + dtype, bytes) must not depend on how many KV heads a rank holds; where pages sit in + their pools may. At TP2 an FP8 K/V page (1 head x 256 x 32 tokens) is the size of an + index page; under NVFP4 the same holds at TP1 and ADP2 (2 heads x 128 bytes x 32 + tokens). Both pack on every layout; bf16, which can never collide, on none.""" + seen = {} + for name, (tp_size, attention_dp) in _LAYOUTS.items(): + cm = _manager(dtype, draft=False, tp_size=tp_size, attention_dp=attention_dp) + try: + buffers = cm._extra_buffers_per_layer(tokens_per_block=_TPB) + roles = { + layer: sorted( + (str(b.role), b.size) for b in buffers.get(cm.layer_offsets[layer], []) + ) + for layer in _QSA_LAYERS + } + views = {layer: cm.get_index_k_buffer(layer) for layer in _QSA_LAYERS} + views["position"] = cm.get_qsa_position_buffer() + logical = {key: _logical_page(view) for key, view in views.items()} + seen[name] = (roles, logical, cm.get_disagg_role_mapper_kinds()) + finally: + cm.shutdown() + assert seen["tp1"] == seen["tp2"] == seen["adp2"] + roles = {role for layer in seen["tp1"][0].values() for role, _ in layer} + if dtype == DataType.BF16: + assert roles == {str(Role.INDEX_KEY), str(QSA_INDEX_POSITION)} + else: + assert roles == {str(QSA_SIDE_PLANES)} diff --git a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py index b766051d0ef0..59308c3778db 100644 --- a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py +++ b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py @@ -11,6 +11,7 @@ QSASparseMetadataParams, QSASparseParams, ) +from tensorrt_llm._torch.attention.backends.sparse.qsa.cache_manager import QSA_SIDE_PLANES from tensorrt_llm._torch.attention.backends.sparse.qsa.indexer import ( QSAIndexer, _position_coordinates, @@ -37,6 +38,11 @@ select_qsa_paged_tokens, select_qsa_tokens, ) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.speculative.interface import ( + prepare_attn_metadata_for_draft_replay, + restore_attn_metadata_after_draft_replay, +) from tensorrt_llm.runtime.kv_cache_manager_v2 import PageIndexMode @@ -219,6 +225,33 @@ def test_qsa_mtp_reuse_topk_appends_groups_completed_since_capture() -> None: assert row.tolist()[: len(valid)] == valid +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_qsa_mtp_reuse_topk_displace_oldest_shifts_a_full_row() -> None: + """With `displace_oldest` a full row gives up its first ids and the rest move to the + front, so an ascending row stays ascending and keeps its newest groups.""" + shared = torch.tensor( + [[1, 2, 4, 6, 7, 9, 11, 12], [1, 2, 4, 6, 7, 9, -1, -1], [3, 5, 6, 8, 9, 10, 11, -1]], + device="cuda", + dtype=torch.int32, + ) + out = torch.full((3, 8), -99, device="cuda", dtype=torch.int32) + + triton_qsa_mtp_reuse_topk( + shared_indices=shared, + shared_visible_blocks=torch.tensor([13, 10, 12], device="cuda", dtype=torch.int32), + request_indices=torch.tensor([0, 1, 2], device="cuda", dtype=torch.int32), + visible_blocks=torch.tensor([15, 11, 14], device="cuda", dtype=torch.int32), + output=out, + displace_oldest=True, + ) + + assert out.tolist() == [ + [4, 6, 7, 9, 11, 12, 13, 14], # full, 2 new: the 2 oldest go + [1, 2, 4, 6, 7, 9, 10, -1], # room: appended + [5, 6, 8, 9, 10, 11, 12, 13], # 7 ids, 2 new: the oldest goes + ] + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") def test_qsa_mtp_reuse_topk_is_verbatim_without_new_groups() -> None: """A draft step that completes no group reuses the captured row unchanged.""" @@ -744,7 +777,7 @@ def test_qsa_paged_index_scores_match_reference(rows: int, head_dim: int) -> Non @pytest.mark.parametrize( "num_index_heads, expected", - [(1, 16), (2, 8), (4, 4), (8, 2), (16, 1), (3, 1), (12, 1), (24, 1), (0, 1)], + [(1, 128), (2, 64), (4, 32), (8, 16), (16, 8), (128, 1), (3, 1), (12, 1), (24, 1), (0, 1)], ) def test_qsa_index_scores_query_tile_divides_the_dot_width( num_index_heads: int, expected: int @@ -752,7 +785,7 @@ def test_qsa_index_scores_query_tile_divides_the_dot_width( query_tile = _qsa_index_scores_query_tile(num_index_heads) assert query_tile == expected if query_tile > 1: - assert query_tile * num_index_heads == 16 + assert query_tile * num_index_heads == 128 def _paged_index_scores_reference( @@ -791,7 +824,7 @@ def test_qsa_paged_index_scores_query_tile_matches_row_wise( head_dim = 16 # The query tile only replaces the row-wise kernel above the prefill row # boundary, so keep the launch there or the comparison is vacuous. - assert _qsa_index_scores_query_tile(num_heads) == 4 + assert _qsa_index_scores_query_tile(num_heads) == 32 if rows_per_request == "many": # Request boundaries deliberately land inside a query tile. lengths = [101, 157, 42] @@ -1456,12 +1489,39 @@ def test_qsa_side_buffers_use_exact_geometry_and_one_position_role() -> None: manager.qsa_index_kv_heads = 1 manager.qsa_sparse_layer_ids = [7, 11] manager.layer_offsets = {7: 2, 11: 4} + # bf16 K/V pages of 2 heads x 128 dims: no side page shares their size. + manager._main_cache_page_sizes = lambda tokens_per_block: {2 * 128 * 2 * tokens_per_block} buffers = manager._extra_buffers_per_layer(tokens_per_block=128) assert buffers[2][0].size == 128 * 2 * 128 assert len(buffers[2]) == 2 assert len(buffers[4]) == 1 + assert manager.qsa_side_planes_size is None + + +def test_qsa_side_buffers_pack_when_a_page_size_matches_the_main_cache() -> None: + """An index page the size of a K/V page would coalesce into the K/V pool, + so both sparse layers' index pages and the position page become one + buffer on the first layer, sized past every main-cache page size.""" + manager = object.__new__(QSAMambaHybridCacheManagerV2) + manager.qsa_index_dim = 128 + manager.qsa_index_kv_heads = 1 + manager.qsa_sparse_layer_ids = [7, 11] + manager.layer_offsets = {7: 2, 11: 4} + index_page, position_page = 128 * 2 * 32, 3 * 4 * 32 + packed = 2 * index_page + position_page + # NVFP4 at 32 tokens per block: K == V == index page; the scale page, and + # a page that happens to equal the packed size, force it to grow once. + manager._main_cache_page_sizes = lambda tokens_per_block: {index_page, 1024, packed} + + buffers = manager._extra_buffers_per_layer(tokens_per_block=32) + + assert list(buffers) == [2] + (packed_buffer,) = buffers[2] + assert packed_buffer.role == QSA_SIDE_PLANES + assert packed_buffer.size == packed + 16 + assert manager.qsa_side_planes_size == packed + 16 def test_qsa_position_buffer_keeps_shared_index_mode_outside_dynamo() -> None: @@ -1478,6 +1538,7 @@ def get_page_stride(self, layer_idx, role): manager = object.__new__(QSAMambaHybridCacheManagerV2) manager.qsa_position_layer_id = 7 + manager.qsa_side_planes_size = None manager.layer_offsets = {7: 2} manager.impl = _Impl() manager.tokens_per_block = 128 @@ -1608,3 +1669,241 @@ def test_qsa_multi_token_generation_captures_speculative_cache_snapshot() -> Non assert pending["token_ordinals"].tolist() == [0, 1, 2, 3] torch.testing.assert_close(pending["index_values"], index_cache[pages, within]) torch.testing.assert_close(pending["position_values"], position_cache[pages, within]) + + +class _QsaPageConverter: + def __init__(self, scale: int) -> None: + self.scale = scale + self.expansion = 1 + + +class _QsaPageImpl: + def __init__(self, scale: int) -> None: + self._scale = scale + + def get_page_index_converter(self, layer_idx, role): + del layer_idx, role + return _QsaPageConverter(self._scale) + + +def _qsa_layout_manager( + *, sparse_layer_ids, pool_id, scale, num_pools, max_blocks_per_seq +) -> QSAMambaHybridCacheManagerV2: + """A V2 QSA manager reduced to the page layout the metadata reads.""" + manager = object.__new__(QSAMambaHybridCacheManagerV2) + manager.qsa_sparse_layer_ids = list(sparse_layer_ids) + manager.layer_offsets = {layer: local for local, layer in enumerate(sparse_layer_ids)} + manager.qsa_position_layer_id = sparse_layer_ids[0] + manager.layer_to_pool_mapping_dict = {local: pool_id for local in range(len(sparse_layer_ids))} + manager.num_attention_op_pools = num_pools + manager.max_blocks_per_seq = max_blocks_per_seq + manager.impl = _QsaPageImpl(scale) + manager.host_kv_cache_block_offsets = torch.zeros( + (num_pools, 4, 2, max_blocks_per_seq), dtype=torch.int32 + ) + return manager + + +# Target: two sparse layers in pool 1 at scale 6. Draft: one MTP layer in pool 0 at +# scale 4, a narrower table, and slot numbers the target never uses. +_TARGET_LAYOUT = dict( + sparse_layer_ids=[3, 7], pool_id=1, scale=6, num_pools=2, max_blocks_per_seq=8 +) +_DRAFT_LAYOUT = dict(sparse_layer_ids=[48], pool_id=0, scale=4, num_pools=1, max_blocks_per_seq=4) + + +def _qsa_metadata_with_draft( + monkeypatch, + *, + mtp_index_share: bool = False, + draft_layout: dict | None = _DRAFT_LAYOUT, + draft_manager=None, +): + """QSA metadata over a target manager and, unless `draft_layout` is None, a separate + draft manager (`draft_manager` replaces the one `draft_layout` describes).""" + monkeypatch.setattr(TrtllmAttentionMetadata, "__post_init__", lambda self: None) + metadata = object.__new__(QSAAttentionMetadata) + metadata.kv_cache_manager = _qsa_layout_manager(**_TARGET_LAYOUT) + if draft_manager is None and draft_layout is not None: + draft_manager = _qsa_layout_manager(**draft_layout) + metadata.draft_kv_cache_manager = draft_manager + metadata.sparse_metadata_params = QSASparseMetadataParams( + token_topk=8, compress_ratio=4, mtp_index_share=mtp_index_share + ) + metadata.max_num_sequences = 4 + metadata.max_num_tokens = 16 + metadata.draft_replay_swapped_attrs = {} + metadata.__post_init__() + return metadata + + +def _set_qsa_offsets(metadata, num_seqs: int, generation: int): + """Write fresh target and draft attention offsets; return the raw slots each + manager's table must decode to.""" + slots = [] + for layout, base in ((_TARGET_LAYOUT, 0), (_DRAFT_LAYOUT, 1000)): + width = layout["max_blocks_per_seq"] + table = torch.arange(num_seqs * width, dtype=torch.int32, device="cuda") + table = table.reshape(num_seqs, width) + base + 100 * generation + table[:, -1] = -1 + # V2 offsets: the key role of the layer's pool holds slot * scale, unallocated + # entries stay -1, and every other pool and role is noise. + offsets = torch.full( + (layout["num_pools"], num_seqs, 2, width), 7777, dtype=torch.int32, device="cuda" + ) + offsets[layout["pool_id"], :, 0, :] = torch.where(table < 0, table, table * layout["scale"]) + slots.append(table) + if layout is _TARGET_LAYOUT: + metadata.kv_cache_block_offsets = offsets + else: + metadata.draft_kv_cache_block_offsets = offsets + return slots + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_qsa_draft_block_table_exists_only_for_a_separate_draft_manager(monkeypatch) -> None: + metadata = _qsa_metadata_with_draft(monkeypatch) + + # Each table is sized, and its pool and page scale taken, from its own manager. + assert metadata.qsa_block_table.shape == (4, _TARGET_LAYOUT["max_blocks_per_seq"]) + assert metadata.qsa_draft_block_table.shape == (4, _DRAFT_LAYOUT["max_blocks_per_seq"]) + assert metadata.qsa_draft_block_table.data_ptr() != metadata.qsa_block_table.data_ptr() + assert (metadata._qsa_attention_pool_id, metadata._qsa_page_index_scale) == (1, 6) + assert (metadata._qsa_draft_attention_pool_id, metadata._qsa_draft_page_index_scale) == (0, 4) + + target_only = _qsa_metadata_with_draft(monkeypatch, draft_layout=None) + assert target_only.qsa_draft_block_table is None + assert target_only.prepare_for_draft_forward() is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("kind", ("not_qsa", "no_local_layer")) +def test_qsa_draft_table_needs_a_qsa_draft_manager_with_local_layers(monkeypatch, kind) -> None: + if kind == "not_qsa": + draft_manager = SimpleNamespace(qsa_position_layer_id=0) + else: + draft_manager = _qsa_layout_manager(**_DRAFT_LAYOUT) + draft_manager.qsa_position_layer_id = None + metadata = _qsa_metadata_with_draft(monkeypatch, draft_manager=draft_manager) + metadata._seq_lens = torch.ones(2, dtype=torch.int32) + + assert metadata.qsa_draft_block_table is None + # The target's own manager is bound, as outside a draft loop, and the hook and the + # refresh both leave the target bindings alone. + metadata.prepare_for_draft_forward() + metadata._refresh_qsa_draft_block_table() + assert metadata.draft_replay_swapped_attrs == {} + assert (metadata._qsa_attention_pool_id, metadata._qsa_page_index_scale) == (1, 6) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize( + "draft_layout, slots", + ( + (dict(_DRAFT_LAYOUT, sparse_layer_ids=[48, 49, 50]), 3), # draft wider than target + (_DRAFT_LAYOUT, 2), # target wider than draft + (None, 2), # no separate draft manager + ), + ids=("draft_wider", "target_wider", "no_draft"), +) +def test_qsa_shared_topk_slots_cover_the_wider_manager(monkeypatch, draft_layout, slots) -> None: + """The shared selection buffers are addressed by whichever manager is bound.""" + metadata = _qsa_metadata_with_draft( + monkeypatch, mtp_index_share=True, draft_layout=draft_layout + ) + + assert metadata.qsa_shared_topk_indices.shape[0] == slots + assert metadata.qsa_shared_topk_visible_blocks.shape[0] == slots + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_qsa_draft_forward_addresses_the_draft_manager_pages(monkeypatch) -> None: + """prepare() decodes both tables; the draft swap makes `qsa_block_table` the draft one + and the restore puts the target's table, pool and scale back.""" + num_seqs = 3 + metadata = _qsa_metadata_with_draft(monkeypatch) + metadata._seq_lens = torch.ones(num_seqs, dtype=torch.int32) + target_slots, draft_slots = _set_qsa_offsets(metadata, num_seqs, generation=0) + target_table = metadata.qsa_block_table + + metadata._refresh_qsa_block_table() + metadata._refresh_qsa_draft_block_table() + torch.testing.assert_close(target_table[:num_seqs], target_slots) + torch.testing.assert_close(metadata.qsa_draft_block_table[:num_seqs], draft_slots) + + draft_manager = metadata.draft_kv_cache_manager + saved = prepare_attn_metadata_for_draft_replay(metadata, draft_manager) + assert saved is not None + assert metadata.kv_cache_manager is draft_manager + assert metadata.qsa_block_table is metadata.qsa_draft_block_table + # Every operand of the table refresh is the draft manager's while swapped, so a + # refresh rebuilds the draft table and leaves the target's alone. + metadata._refresh_qsa_block_table() + torch.testing.assert_close(metadata.qsa_block_table[:num_seqs], draft_slots) + torch.testing.assert_close(target_table[:num_seqs], target_slots) + + restore_attn_metadata_after_draft_replay(metadata, saved) + assert metadata.kv_cache_manager is not draft_manager + assert metadata.qsa_block_table is target_table + assert (metadata._qsa_attention_pool_id, metadata._qsa_page_index_scale) == (1, 6) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_qsa_prepare_decodes_the_target_and_the_draft_table(monkeypatch) -> None: + """`prepare()` is what refreshes both tables before every step and replay: a second + generation of offsets reaches the target table and the draft table through it.""" + num_seqs = 3 + metadata = _qsa_metadata_with_draft(monkeypatch) + metadata._seq_lens = torch.ones(num_seqs, dtype=torch.int32) + # The parent's prepare and the row mapping are not under test here. + monkeypatch.setattr(TrtllmAttentionMetadata, "prepare", lambda self: None) + monkeypatch.setattr(QSAAttentionMetadata, "_refresh_qsa_token_mapping", lambda self: None) + monkeypatch.setattr( + QSAAttentionMetadata, "_refresh_qsa_speculative_snapshot_state", lambda self: None + ) + for generation in (0, 1): + target_slots, draft_slots = _set_qsa_offsets(metadata, num_seqs, generation) + metadata.prepare() + torch.testing.assert_close(metadata.qsa_block_table[:num_seqs], target_slots) + torch.testing.assert_close(metadata.qsa_draft_block_table[:num_seqs], draft_slots) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_qsa_draft_forward_refuses_an_unswapped_manager(monkeypatch) -> None: + metadata = _qsa_metadata_with_draft(monkeypatch) + + with pytest.raises(RuntimeError, match="draft KV cache manager"): + metadata.prepare_for_draft_forward() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_qsa_draft_table_replays_the_refreshed_draft_pages(monkeypatch) -> None: + """A graph captured across the draft swap reads each table by address: the trunk + segment the target table, the draft segment the draft one, and a refresh outside + capture (as prepare() does before every replay) is what the replay sees.""" + num_seqs = 3 + metadata = _qsa_metadata_with_draft(monkeypatch) + metadata._seq_lens = torch.ones(num_seqs, dtype=torch.int32) + _set_qsa_offsets(metadata, num_seqs, generation=0) + metadata._refresh_qsa_block_table() + metadata._refresh_qsa_draft_block_table() + draft_manager = metadata.draft_kv_cache_manager + target_seen = torch.zeros_like(metadata.qsa_block_table[:num_seqs]) + draft_seen = torch.zeros_like(metadata.qsa_draft_block_table[:num_seqs]) + + graph = torch.cuda.CUDAGraph() + torch.cuda.synchronize() + with torch.cuda.graph(graph): + target_seen.copy_(metadata.qsa_block_table[:num_seqs]) + saved = prepare_attn_metadata_for_draft_replay(metadata, draft_manager) + draft_seen.copy_(metadata.qsa_block_table[:num_seqs]) + restore_attn_metadata_after_draft_replay(metadata, saved) + + target_slots, draft_slots = _set_qsa_offsets(metadata, num_seqs, generation=1) + metadata._refresh_qsa_block_table() + metadata._refresh_qsa_draft_block_table() + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close(target_seen, target_slots) + torch.testing.assert_close(draft_seen, draft_slots) diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index aa38d11c9390..f4478bdb41ff 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -529,6 +529,93 @@ def test_model_without_remote_tail_does_not_inspect_request() -> None: PyExecutor._configure_csa2_remote_tail(executor, types.SimpleNamespace()) +def _make_admission_executor(model, requests): + """An executor whose `_fetch_and_activate_new_requests` and `_validate_request` are the + real ones, over a model engine holding `model`; everything else they touch is a stub.""" + executor = types.SimpleNamespace( + model_engine=types.SimpleNamespace(model=model), + _is_kv_manager_v2=False, + kv_cache_manager=None, + kv_cache_transceiver=None, + max_beam_width=1, + sampler=Mock(), + waiting_queue=object(), + active_requests=[], + _mm_encoder_item_scheduling_enabled=False, + _fetch_new_requests=Mock(return_value=requests), + _configure_csa2_remote_tail=Mock(), + _validate_token_id_range=Mock(), + _validate_request_budget=Mock(), + _handle_errors=Mock(), + ) + executor._validate_request = types.MethodType(PyExecutor._validate_request, executor) + return executor + + +def _make_admission_request(request_id: int) -> types.SimpleNamespace: + return types.SimpleNamespace( + request_id=request_id, + py_disaggregated_params=None, + sampling_config=None, + kv_cache_retention_config=None, + ) + + +def test_model_validation_fails_only_the_request_it_rejects() -> None: + requests = [_make_admission_request(i) for i in range(3)] + + def validate_request(request) -> None: + if request.request_id == 1: + raise ValueError("frame too large") + + model = types.SimpleNamespace(validate_request=Mock(side_effect=validate_request)) + executor = _make_admission_executor(model, requests) + + admitted = PyExecutor._fetch_and_activate_new_requests(executor) + + assert admitted == [requests[0], requests[2]] + assert executor.active_requests == [requests[0], requests[2]] + executor._handle_errors.assert_called_once_with( + "frame too large", requests=[requests[1]], charge_budget=False + ) + assert [call.args[0] for call in model.validate_request.call_args_list] == requests + # The checks after the model's still run for the requests it admitted. + assert executor._validate_request_budget.call_count == 2 + + +def test_model_validation_runs_after_the_sampler_and_before_the_budget() -> None: + order = [] + model = types.SimpleNamespace(validate_request=lambda request: order.append("model")) + executor = _make_admission_executor(model, []) + executor.sampler.validate_request.side_effect = lambda request: order.append("sampler") + executor._validate_request_budget.side_effect = lambda request: order.append("budget") + + PyExecutor._validate_request(executor, _make_admission_request(0)) + + assert order == ["sampler", "model", "budget"] + + +@pytest.mark.parametrize( + "model", + [ + types.SimpleNamespace(), + Mock(spec=DecoderModelForCausalLM), + types.SimpleNamespace(validate_request=None), + None, + ], + ids=["no_attribute", "spec_mock", "none_attribute", "no_model"], +) +def test_models_without_the_validation_hook_are_unaffected(model) -> None: + requests = [_make_admission_request(i) for i in range(2)] + executor = _make_admission_executor(model, requests) + + admitted = PyExecutor._fetch_and_activate_new_requests(executor) + + assert admitted == requests + assert executor.active_requests == requests + executor._handle_errors.assert_not_called() + + @pytest.mark.parametrize("mode", ["standalone", "generation", "disabled"]) def test_full_model_requests_without_remote_handoff_are_unchanged(mode: str) -> None: executor, request = _make_remote_tail_executor_and_request(context_only=False) diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 3a7293bd7e9d..f7cc60f0a500 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -59,6 +59,7 @@ from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata from tensorrt_llm._torch.attention.backends.trtllm import \ TrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.vanilla import VanillaAttention from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._torch.speculative.interface import \ INVALID_PROMPT_LOOKAHEAD_TOKEN @@ -67,7 +68,8 @@ from tensorrt_llm._torch.utils import set_per_request_prefill_cuda_graph_flag from tensorrt_llm.inputs.registry import (BaseMultimodalDummyInputsBuilder, BaseMultimodalInputProcessor) -from tensorrt_llm.llmapi import (CudaGraphConfig, SADecodingConfig, +from tensorrt_llm.llmapi import (CudaGraphConfig, MTPDecodingConfig, + NGramDecodingConfig, SADecodingConfig, SamplingParams) from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.mapping import CpType, Mapping @@ -2898,11 +2900,15 @@ def test_prepare_inputs_with_helix_parallelism(self) -> None: actual_seq_lens = attn_metadata.seq_lens.cpu().tolist() self.assertEqual(actual_seq_lens, expected_seq_lens) - def _setup_mrope_engine(self, max_num_tokens: int = 32): + def _setup_mrope_engine(self, + max_num_tokens: int = 32, + spec_config: DecodingBaseConfig | None = None): """Build a DummyModelEngine that takes the MRoPE path, plus its KV cache manager and attention metadata.""" llm_args = TorchLlmArgs(model="dummy") - model_engine = DummyModelEngine(llm_args, dtype=torch.half) + model_engine = DummyModelEngine(llm_args, + dtype=torch.half, + spec_config=spec_config) model_engine.model.model_config.pretrained_config.rope_scaling = { "type": "mrope" } @@ -2931,6 +2937,362 @@ def _setup_mrope_engine(self, max_num_tokens: int = 32): _resize_input_buffers(model_engine._runner, max_num_tokens) return model_engine, kv_cache_manager, attn_metadata + @staticmethod + def _make_mrope_gen_request(num_tokens: int, req_id: int, seq_slot: int, + delta): + """Generation request carrying an MRoPE delta when `delta` is not None, + and no multimodal data at all otherwise (a text-only prompt).""" + request = _create_request(num_tokens, req_id) + request.py_prompt_len = num_tokens + request.py_batch_idx = None + request.py_seq_slot = seq_slot + request.sampling_config.beam_width = 1 + if delta is None: + request.py_multimodal_data = {} + else: + request.py_multimodal_data = { + "mrope_config": { + "mrope_position_deltas": + torch.tensor([[delta]], dtype=torch.int32) + }, + } + return request + + def _make_mrope_spec_batch(self, previous_batch: bool): + """Two multimodal verify groups around a text-only one, plus a CUDA graph + padding dummy, each carrying one draft token. + + Returns the scheduled batch and the multimodal requests.""" + # (num_tokens, req_id, seq_slot, delta); the middle request is text-only. + requests = [ + self._make_mrope_gen_request(4, 1, 0, 10), + self._make_mrope_gen_request(5, 2, 1, None), + self._make_mrope_gen_request(6, 3, 2, 20), + ] + for request in requests: + request.py_draft_tokens = [7] + if previous_batch: + # The overlap scheduler's previous step ran these requests at + # their current slots, which also seeded the delta cache. + request.py_batch_idx = request.py_seq_slot + if request.py_multimodal_data: + request.py_mrope_delta_cache_slot = request.py_seq_slot + # The padding dummy carries the zero delta that `add_dummy_requests` + # attaches on MRoPE models. + dummy = self._make_mrope_gen_request(2, 4, 3, 0) + dummy.is_cuda_graph_dummy = True + dummy.py_draft_tokens = [0] + + scheduled_requests = ScheduledRequests() + scheduled_requests.context_requests_last_chunk = [] + scheduled_requests.generation_requests = requests + [dummy] + return scheduled_requests, [requests[0], requests[2]] + + def test_prepare_inputs_spec_decode_adds_mrope_delta(self) -> None: + """One-model speculative verify groups get the request's MRoPE delta. + + Every token a generation step feeds is text after the prompt's vision + spans, so its three coordinates are `seq_index + delta`, as for an + ordinary generation request. Read slots stay dense, one per verify + group, and a generation-only request seeds the delta cache. + """ + model_engine, kv_cache_manager, attn_metadata = self._setup_mrope_engine( + spec_config=SADecodingConfig(max_draft_len=1)) + scheduled_requests, multimodal_requests = self._make_mrope_spec_batch( + previous_batch=False) + spec_metadata = Mock( + _force_non_greedy_for_capture=False, + context_prompt_lookahead_tokens=None, + ) + + result = model_engine._runner._inputs.prepare_inputs( + scheduled_requests=scheduled_requests, + kv_cache_manager=kv_cache_manager, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + enable_spec_decode=True, + runtime_draft_len=1, + is_dummy=False).kwargs + + # Each verify group is [last token, draft token] at past_seen + 0/1; + # the text-only request and the zero-delta dummy keep the scalar + # positions. + position_ids = result["position_ids"] + self.assertEqual(tuple(position_ids.shape), (3, 1, 8)) + expected = torch.tensor([[[13, 14, 4, 5, 25, 26, 1, 2]]] * 3, + dtype=torch.int32, + device='cuda') + torch.testing.assert_close(position_ids, expected, atol=0, rtol=0) + # Slot 32 is the reserved zero slot (max_num_tokens * pp_size). + self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), + [0, 32, 2, 32]) + self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), + [0, 2]) + self.assertEqual([ + request.py_mrope_delta_cache_slot for request in multimodal_requests + ], [0, 2]) + kv_cache_manager.shutdown() + + def test_prepare_inputs_spec_decode_overlap_keeps_mrope_delta(self) -> None: + """Under the overlap scheduler the accepted-token correction applied in + `preprocess_inputs` lands on top of the delta, on all three axes.""" + model_engine, kv_cache_manager, attn_metadata = self._setup_mrope_engine( + spec_config=SADecodingConfig(max_draft_len=1)) + scheduled_requests, _ = self._make_mrope_spec_batch(previous_batch=True) + spec_metadata = Mock( + _force_non_greedy_for_capture=False, + context_prompt_lookahead_tokens=None, + ) + # The previous step accepted 2, 1 and 2 tokens in slots 0, 1 and 2. + overlap_state = SampleStateTensorsSpec( + new_tokens=torch.zeros((2, 4, 1), dtype=torch.int32, device="cuda"), + new_tokens_lens=torch.tensor([2, 1, 2, 1], + dtype=torch.int32, + device="cuda"), + next_draft_tokens=torch.zeros((4, 1), + dtype=torch.int32, + device="cuda"), + ) + + inputs = model_engine._runner._inputs.prepare_inputs( + scheduled_requests=scheduled_requests, + kv_cache_manager=kv_cache_manager, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + new_tensors_device=overlap_state, + enable_spec_decode=True, + runtime_draft_len=1, + is_dummy=False).kwargs + + # The host packs positions from the stale base (past_seen + delta). + stale = torch.tensor([[[13, 14, 4, 5, 25, 26, 1, 2]]] * 3, + dtype=torch.int32, + device='cuda') + torch.testing.assert_close(inputs["position_ids"], + stale, + atol=0, + rtol=0) + self.assertEqual(inputs["mrope_delta_read_seq_slots"].cpu().tolist(), + [0, 32, 2, 32]) + # The delta cache already holds both multimodal requests' deltas. + self.assertNotIn("mrope_delta_write_seq_slots", inputs) + + inputs = model_engine._runner._inputs.preprocess_inputs( + inputs, enable_spec_decode=True, runtime_draft_len=1) + corrected = torch.tensor([[[15, 16, 5, 6, 27, 28, 1, 2]]] * 3, + dtype=torch.int32, + device='cuda') + torch.testing.assert_close(inputs["position_ids"], + corrected, + atol=0, + rtol=0) + kv_cache_manager.shutdown() + + @staticmethod + def _make_mrope_ctx_request(num_tokens: int, req_id: int, seq_slot: int, + delta: int): + """Single-chunk context request carrying MRoPE metadata. + + Its position on axis `a` of token `i` is `100 + 1000 * a + i`, so a + spliced span says which axis and token it came from.""" + request = _create_request(num_tokens, req_id) + request.py_prompt_len = num_tokens + request.py_batch_idx = None + request.py_seq_slot = seq_slot + request.context_current_position = 0 + request.context_chunk_size = num_tokens + request.sampling_config.beam_width = 1 + positions = (100 + 1000 * torch.arange(3).view(3, 1, 1) + + torch.arange(num_tokens).view(1, 1, -1)).to(torch.int32) + request.py_multimodal_data = { + "mrope_config": { + "mrope_position_ids": positions, + "mrope_position_deltas": torch.tensor([[delta]], + dtype=torch.int32), + }, + } + return request + + @staticmethod + def _prepare_mrope_spec_inputs(model_engine, kv_cache_manager, + attn_metadata, scheduled_requests): + spec_metadata = Mock( + _force_non_greedy_for_capture=False, + context_prompt_lookahead_tokens=None, + ) + result = model_engine._runner._inputs.prepare_inputs( + scheduled_requests=scheduled_requests, + kv_cache_manager=kv_cache_manager, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + enable_spec_decode=True, + runtime_draft_len=1, + is_dummy=False).kwargs + return result + + def _assert_mrope_positions(self, result, rows) -> None: + """`rows` is one list for all three axes, or one list per axis.""" + if not isinstance(rows[0], list): + rows = [rows] * 3 + position_ids = result["position_ids"] + self.assertEqual(tuple(position_ids.shape), (3, 1, len(rows[0]))) + expected = torch.tensor([rows[0:1], rows[1:2], rows[2:3]], + dtype=torch.int32, + device='cuda') + torch.testing.assert_close(position_ids, expected, atol=0, rtol=0) + + def test_prepare_inputs_spec_decode_mrope_plain_rows_follow_verify_groups( + self) -> None: + """Verify groups are packed ahead of plain generation requests, and the + MRoPE delta read slots follow that packing rather than scheduler order. + + The attention kernel indexes the deltas by generation batch index: row + i of the packed generation batch must read the delta of the request + packed i-th, so a verify group's slot comes before every plain + request's, whatever order the scheduler listed them in. + """ + model_engine, kv_cache_manager, attn_metadata = self._setup_mrope_engine( + spec_config=MTPDecodingConfig(max_draft_len=1)) + # `verify_*` carry a draft token, `plain_*` do not; `*_mm` have an + # MRoPE delta, `*_text` none. + plain_mm = self._make_mrope_gen_request(6, 1, 0, 20) + verify_mm = self._make_mrope_gen_request(4, 2, 1, 10) + plain_text = self._make_mrope_gen_request(5, 3, 2, None) + verify_text = self._make_mrope_gen_request(3, 4, 3, None) + verify_mm.py_draft_tokens = [7] + verify_text.py_draft_tokens = [7] + + scheduled_requests = ScheduledRequests() + scheduled_requests.context_requests_last_chunk = [] + scheduled_requests.generation_requests = [ + plain_mm, verify_mm, plain_text, verify_text + ] + result = self._prepare_mrope_spec_inputs(model_engine, kv_cache_manager, + attn_metadata, + scheduled_requests) + + # Packed order: verify_mm, verify_text, plain_mm, plain_text. The + # verify groups hold two rows each, the plain requests one. + self._assert_mrope_positions(result, [13, 14, 2, 3, 25, 4]) + # Slot 32 is the reserved zero slot (max_num_tokens * pp_size) read by + # the text-only requests. + self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), + [1, 32, 0, 32]) + self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), + [1, 0]) + kv_cache_manager.shutdown() + + def test_prepare_inputs_spec_decode_mrope_context_takes_no_read_slot( + self) -> None: + """A context request packs ahead of the verify groups with its own + MRoPE positions and takes no generation read slot, so the verify + groups' slots start the list.""" + model_engine, kv_cache_manager, attn_metadata = self._setup_mrope_engine( + spec_config=MTPDecodingConfig(max_draft_len=1)) + # A context request with multimodal content takes the placeholder-token + # split, which reads the vocabulary size. + model_engine.model.config.vocab_size = 1000 + context = self._make_mrope_ctx_request(5, 1, 0, 30) + verify_mm = self._make_mrope_gen_request(4, 2, 1, 10) + verify_mm.py_draft_tokens = [7] + plain_mm = self._make_mrope_gen_request(6, 3, 2, 20) + + scheduled_requests = ScheduledRequests() + scheduled_requests.context_requests_last_chunk = [context] + scheduled_requests.generation_requests = [plain_mm, verify_mm] + result = self._prepare_mrope_spec_inputs(model_engine, kv_cache_manager, + attn_metadata, + scheduled_requests) + + # The context span keeps the processor's per-axis positions; the + # generation rows after it are text at position + delta on all axes. + self._assert_mrope_positions(result, [ + [100, 101, 102, 103, 104, 13, 14, 25], + [1100, 1101, 1102, 1103, 1104, 13, 14, 25], + [2100, 2101, 2102, 2103, 2104, 13, 14, 25], + ]) + self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), + [1, 2]) + # The context request seeds the delta cache first, then the verify + # group and the plain request seed theirs. + self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), + [0, 1, 2]) + self.assertEqual(context.py_mrope_delta_cache_slot, 0) + kv_cache_manager.shutdown() + + def _setup_extend_as_context_engine(self): + """A two-model draft mode on an attention backend that runs verify groups + as chunked context rows. + + `extend_ctx` is False for every one-model mode and for the TRTLLM + backend below sm_120, so this takes a drafter-based mode and swaps in + the vanilla backend. + """ + model_engine, kv_cache_manager, attn_metadata = self._setup_mrope_engine( + spec_config=NGramDecodingConfig(max_draft_len=1)) + _set_runner_config(model_engine._runner, + attention_backend=VanillaAttention, + disable_overlap_scheduler=True) + spec_config = model_engine._runner._config.spec_config + self.assertTrue(spec_config.spec_dec_mode.extend_ctx(VanillaAttention)) + self.assertTrue(spec_config.is_linear_tree) + return model_engine, kv_cache_manager, attn_metadata + + def test_prepare_inputs_mrope_verify_as_context_takes_no_read_slot( + self) -> None: + """When the kernel runs verify groups as chunked context rows, their + rotary comes from `position_ids` alone and only the plain generation + requests take read slots.""" + model_engine, kv_cache_manager, attn_metadata = ( + self._setup_extend_as_context_engine()) + verify_mm = self._make_mrope_gen_request(4, 1, 0, 10) + verify_text = self._make_mrope_gen_request(3, 2, 1, None) + plain_mm = self._make_mrope_gen_request(6, 3, 2, 20) + plain_text = self._make_mrope_gen_request(5, 4, 3, None) + verify_mm.py_draft_tokens = [7] + verify_text.py_draft_tokens = [7] + + scheduled_requests = ScheduledRequests() + scheduled_requests.context_requests_last_chunk = [] + scheduled_requests.generation_requests = [ + verify_mm, verify_text, plain_mm, plain_text + ] + result = self._prepare_mrope_spec_inputs(model_engine, kv_cache_manager, + attn_metadata, + scheduled_requests) + + self._assert_mrope_positions(result, [13, 14, 2, 3, 25, 4]) + self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), + [2, 32]) + # The verify group still seeds the delta cache for the plain steps + # that follow it. + self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), + [0, 2]) + kv_cache_manager.shutdown() + + def test_prepare_inputs_mrope_verify_as_context_alone_emits_no_read_slots( + self) -> None: + """A batch whose only delta carriers are verify groups run as context + rows has nothing for the generation-indexed read.""" + model_engine, kv_cache_manager, attn_metadata = ( + self._setup_extend_as_context_engine()) + verify_mm = self._make_mrope_gen_request(4, 1, 0, 10) + verify_mm.py_draft_tokens = [7] + plain_text = self._make_mrope_gen_request(5, 2, 1, None) + + scheduled_requests = ScheduledRequests() + scheduled_requests.context_requests_last_chunk = [] + scheduled_requests.generation_requests = [verify_mm, plain_text] + result = self._prepare_mrope_spec_inputs(model_engine, kv_cache_manager, + attn_metadata, + scheduled_requests) + + self._assert_mrope_positions(result, [13, 14, 4]) + self.assertNotIn("mrope_delta_read_seq_slots", result) + self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), + [0]) + kv_cache_manager.shutdown() + @skip_pre_blackwell_unittest def test_ced_warmup_passes_dummy_context_requests(self) -> None: for warmup in (False, True): diff --git a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py index ee5b97703635..6d8d1b727a55 100644 --- a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py +++ b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py @@ -365,7 +365,7 @@ def test_qsa_cache_manager_requires_sparse_config() -> None: QSAMambaHybridCacheManagerV2(layer_mask=[True]) -@pytest.mark.parametrize("dtype", (DataType.NVFP4, DataType.FLOAT)) +@pytest.mark.parametrize("dtype", (DataType.FLOAT,)) def test_qsa_cache_manager_delegates_regular_kv_layout( monkeypatch: pytest.MonkeyPatch, dtype: DataType ) -> None: 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 b3e446dd7b23..5c196b997868 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 @@ -175,7 +175,8 @@ within-block mask read out one-hot at P = 2, 3, 4 over a bf16 pool at 24/2/256 and 24/4/256 and an fp8 one at 24/4/256 (bottom-right causal, later siblings bitwise zero, mask_type inert, a page-crossing block, - consecutive steps), with the full-block model as the measured rival. + consecutive steps), with the full-block model as the measured rival; the + entry's cells gate P up to 8 (the target's max_draft_len 7) against `reference`. 12. qwen4_exp's full-attention cell, GQA 24/2/256 over a bf16 pool on the paged-context path: the cached-prefix rival, the sparse / skip-softmax group measured bitwise inert, and local_layer_idx over its 48-row mapping @@ -183,11 +184,16 @@ real 48-layer manager's identity rows, and a bitwise A/B of a single-layer pool against mapping row 47. 13. The no-KV-cache (encoder-only) mode qwen4_exp's vision tower runs: all - three cache-addressing tensors None, one context sequence per frame, MHA - 16/16 at head_size 72 and 128, the merged-segment rival, the cache + three cache-addressing tensors None, one context sequence per frame (an + image, or one frame group of a video), MHA 16/16 at head_size 72 and 128, + the merged-segment rival, the cache arguments that are inert there, the workspace high-water mark over the mmmu serving shapes, and the two boundaries (update_kv_cache=False is rejected; head_size 72 has no generation phase). +14. The same full-attention cell over an NVFP4 pool (e2m1 data and e4m3 block + scales in a second pool, unit fp32 [3] second-level scales, bf16 output): + the pool and its block scales bitwise as the append writes them, and the + output against `reference` over the NVFP4 values (the cells). The cells at the end drive the entry's own `reference` and `_compare` over the configurations the targets run. """ @@ -241,6 +247,7 @@ MASK_PADDING = 0 # AttentionMaskType.padding: no mask (bidirectional) QUANT_MODE_FP8_KV_CACHE = 128 # QuantMode.FP8_KV_CACHE (sole bit set) +QUANT_MODE_NVFP4_KV_CACHE = 8192 # QuantMode.NVFP4_KV_CACHE (sole bit set) # fp8-e4m3 KV-pool decode tolerance. The reference dequantizes the mirrored # e4m3 pool content and puts q through the same e4m3 round-trip the kernel @@ -291,6 +298,88 @@ FP8_MLA_RTOL = 2**-4 +# NVFP4: one e4m3 block scale per this many values; e2m1 magnitudes by code +# and the round-to-nearest-even steps between them (inclusive where the upper +# neighbour is the even code). +NVFP4_BLOCK = 16 +E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +E2M1_STEPS = ((0.25, False), (0.75, True), (1.25, False), (1.75, True), (2.5, False)) +E2M1_STEPS += ((3.5, True), (5.0, False)) + + +def _nvfp4_quantize(rows: torch.Tensor, global_scale: torch.Tensor): + """[n, D] rows -> (packed e2m1 bytes [n, D/2], e4m3 scale bytes [n, D/16]): + the attention op's append, in torch (`cvt_warp_fp16_to_fp4`).""" + n, d = rows.shape + x = rows.float().reshape(n, d // NVFP4_BLOCK, NVFP4_BLOCK) + block_max = x.abs().amax(-1) + sixth = torch.tensor(6.0, dtype=torch.float32, device=rows.device).reciprocal() + code = (global_scale * (block_max * sixth)).to(torch.float8_e4m3fn) + out_scale = torch.where( + block_max != 0, + (code.float() * global_scale.reciprocal()).reciprocal(), + torch.zeros_like(block_max), + ) + scaled = x * out_scale[..., None] + magnitude = scaled.abs() + codes = torch.zeros(scaled.shape, dtype=torch.int32, device=rows.device) + for threshold, inclusive in E2M1_STEPS: + codes += (magnitude >= threshold if inclusive else magnitude > threshold).to(torch.int32) + codes |= torch.signbit(scaled).to(torch.int32) << 3 + # 0 x inf, in a block whose scale underflows to zero: the conversion makes +6. + codes = torch.where(torch.isnan(scaled), torch.full_like(codes, 0x7), codes) + codes = codes.reshape(n, d // 2, 2) + return (codes[..., 0] | (codes[..., 1] << 4)).to(torch.uint8), code.view(torch.uint8) + + +def _nvfp4_dequantize(packed: torch.Tensor, scales: torch.Tensor, global_scale) -> torch.Tensor: + """Inverse of `_nvfp4_quantize`: fp32 [n, D].""" + lut = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=packed.device) + lut = torch.cat([lut, -lut]) + raw = packed.long() + values = torch.stack([lut[raw & 0xF], lut[raw >> 4]], -1).flatten(-2) + scale = scales.view(torch.float8_e4m3fn).float().repeat_interleave(NVFP4_BLOCK, -1) + return values * scale / global_scale + + +def _nvfp4_roundtrip(kv: torch.Tensor, global_scale) -> torch.Tensor: + """[T, Hkv, D] bf16 -> the bf16 values an NVFP4 pool returns for them.""" + t, h, d = kv.shape + packed, scales = _nvfp4_quantize(kv.reshape(t * h, d), global_scale) + return _nvfp4_dequantize(packed, scales, global_scale).reshape(t, h, d).to(kv.dtype) + + +def _check_nvfp4_cache(env, request_id: int) -> None: + """An NVFP4 pool holds each token's e2m1 bytes and its block scales -- K's + token-major, V's in four-token groups -- bitwise as the append writes them.""" + tpb, hkv, d = env.tokens_per_block, env.num_kv_heads, env.head_dim + groups = d // NVFP4_BLOCK + for role, history in ((0, env.k_history[request_id]), (1, env.v_history[request_id])): + seq = torch.cat(history) + total = seq.shape[0] + packed, scales = _nvfp4_quantize( + seq.reshape(total * hkv, d), env.kv_scale_orig_quant[1 + role] + ) + packed = packed.reshape(total, hkv, d // 2) + scales = scales.reshape(total, hkv, groups) + for j, page in enumerate(env.pages[request_id]): + n = min(tpb, total - j * tpb) + if n <= 0: + break + rows = slice(j * tpb, j * tpb + n) + got = env.kv_cache[page, role, :, :n].permute(1, 0, 2) + assert torch.equal(got, packed[rows]), f"request {request_id} role {role} page {j}" + sf_page = env.kv_block_scales[page, role] + if role == 0: + got_sf = sf_page[:, :n].permute(1, 0, 2) + else: + grouped = sf_page.reshape(hkv, tpb // 4, groups, 4).permute(0, 1, 3, 2) + got_sf = grouped.reshape(hkv, tpb, groups)[:, :n].permute(1, 0, 2) + assert torch.equal(got_sf, scales[rows]), ( + f"request {request_id} role {role} page {j}: block scales" + ) + + def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: """A softmax chain, so the band is the longest in this catalog. @@ -378,26 +467,46 @@ def __init__( self.page_ring = page_ring self.rings: dict[int, List[int]] = {} - # Single-layer, single-pool paged KV cache (HND layout). + # Single-layer, single-pool paged KV cache (HND layout). An NVFP4 pool + # holds e2m1 pairs, and its e4m3 block scales live in a second pool + # whose pointer is the third axis of the pool pointers. + self.nvfp4 = bool(quant_mode & QUANT_MODE_NVFP4_KV_CACHE) self.kv_cache = torch.zeros( num_blocks, 2, num_kv_heads, tokens_per_block, - head_dim, - dtype=pool_dtype, + head_dim // 2 if self.nvfp4 else head_dim, + dtype=torch.uint8 if self.nvfp4 else pool_dtype, device="cuda", ) - self.pool_pointers = torch.zeros(1, 2, dtype=torch.int64, device="cpu") - self.pool_pointers[0, 0] = self.kv_cache.data_ptr() + if self.nvfp4: + self.kv_block_scales = torch.zeros( + num_blocks, + 2, + num_kv_heads, + tokens_per_block, + head_dim // NVFP4_BLOCK, + dtype=torch.uint8, + device="cuda", + ) + self.pool_pointers = torch.zeros(1, 2, 2, dtype=torch.int64, device="cpu") + self.pool_pointers[0, 0, 0] = self.kv_cache.data_ptr() + self.pool_pointers[0, 0, 1] = self.kv_block_scales.data_ptr() + else: + self.pool_pointers = torch.zeros(1, 2, dtype=torch.int64, device="cpu") + self.pool_pointers[0, 0] = self.kv_cache.data_ptr() # KV-cache scaling factor s: dequant = quantized * s. Mirrors the # production construction (quant_orig = s, orig_quant = 1/s, both - # fp32 [1] CUDA); None = unquantized pool, scale args passed as None. + # fp32 [1] CUDA, [3] -- q, K, V -- over an NVFP4 pool); None = + # unquantized pool, scale args passed as None. if kv_scaling_factor is None: self.kv_scale_quant_orig: Optional[torch.Tensor] = None self.kv_scale_orig_quant: Optional[torch.Tensor] = None else: - factor = torch.full((1,), kv_scaling_factor, dtype=torch.float32, device="cuda") + factor = torch.full( + (3 if self.nvfp4 else 1,), kv_scaling_factor, dtype=torch.float32, device="cuda" + ) self.kv_scale_quant_orig = factor self.kv_scale_orig_quant = 1.0 / factor self.pool_mapping = torch.zeros(1, 2, dtype=torch.int32, device="cpu") @@ -702,7 +811,11 @@ def cache_pages_content(self, request_id: int) -> Tuple[torch.Tensor, ...]: def check_cache(self, request_id: int) -> None: """The paged cache must hold exactly the expected rows, bitwise (the - bf16 K/V fed so far, or their scaled e4m3 casts for an fp8 pool).""" + bf16 K/V fed so far, or their scaled e4m3 casts for an fp8 pool, or + their NVFP4 bytes and block scales for an NVFP4 pool).""" + if self.nvfp4: + _check_nvfp4_cache(self, request_id) + return got_k, got_v, exp_k, exp_v = self.cache_pages_content(request_id) assert torch.equal(got_k.view(torch.uint8), exp_k.view(torch.uint8)), ( f"K cache mismatch: request {request_id}" @@ -5426,9 +5539,13 @@ def _op_reference( causal: bool = True, e4m3: bool = False, layer: Optional[int] = None, + nvfp4: bool = False, ) -> torch.Tensor: """`op.reference` over the history the env mirrored for these requests. + nvfp4=True reads K and V as an NVFP4 pool returns them (the history + through `_nvfp4_roundtrip`) before the e4m3 arithmetic `e4m3` selects. + Call it after the op call, so each history ends with this call's K/V: `reference`'s bottom-right alignment then puts this call's q rows at the end of it. causal=False on a generation call is the full-block rival -- @@ -5438,6 +5555,9 @@ def _op_reference( """ hq, hkv, d = env.num_heads, env.num_kv_heads, env.head_dim ks, vs = zip(*(_history(env, rid, layer) for rid in request_ids)) + if nvfp4: + ks = tuple(_nvfp4_roundtrip(k, env.kv_scale_orig_quant[1]) for k in ks) + vs = tuple(_nvfp4_roundtrip(v, env.kv_scale_orig_quant[2]) for v in vs) return op.reference( qkv[:, : hq * d], torch.cat(ks).reshape(-1, hkv * d), @@ -5930,8 +6050,8 @@ def repeat_pools(alias: bool) -> List[torch.Tensor]: # (its draft chain), token-major. On sm_103 a linear draft leaves every # spec_decoding_* tensor None and is_spec_decoding_enabled False, so the # within-block mask has to come from predicted_tokens_per_seq alone. The cells -# gate P = 1..4 against `op.reference`; these read the mask out directly and -# measure the rival. +# gate P = 1..8 against `op.reference`; these read the mask out directly and +# measure the rival at P = 2..4. # One configuration per pool the cells certify P over: qwen4_exp's 24/2/256 # over bf16, and 24/4/256 over bf16 and over fp8 at s = 1.0. The weight @@ -5942,6 +6062,7 @@ def repeat_pools(alias: bool) -> List[torch.Tensor]: ("bf16 24/2/256", 2, torch.bfloat16, 0, None, SWA_WEIGHT_RTOL), ("bf16 24/4/256", 4, torch.bfloat16, 0, None, SWA_WEIGHT_RTOL), ("fp8 24/4/256", 4, torch.float8_e4m3fn, QUANT_MODE_FP8_KV_CACHE, 1.0, 2**-4), + ("fp8 24/2/256", 2, torch.float8_e4m3fn, QUANT_MODE_FP8_KV_CACHE, 1.0, 2**-4), ] @@ -6261,7 +6382,11 @@ def again(mask_type: int = MASK_CAUSAL, **overrides) -> torch.Tensor: def _qwen_layer_pools( - attn_blocks: int = 48, other_blocks: int = 32, max_seq_len: int = 256, max_batch: int = 4 + attn_blocks: int = 48, + other_blocks: int = 32, + max_seq_len: int = 256, + max_batch: int = 4, + **pool_kw, ) -> Tuple[_MultiLayerPagedAttnEnv, _MultiLayerPagedAttnEnv]: """The two pools of the hybrid mapping, at the full-attention cell. @@ -6269,7 +6394,8 @@ def _qwen_layer_pools( rows name a real, distinct pool base and the untouched assertion has somewhere to look. Both managers take the same max_seq_len and max_batch, which makes their block-offset tables concatenable into the per-pool table - the op indexes by the mapping's pool column. + the op indexes by the mapping's pool column. `pool_kw` (dtype, quant_mode, + kv_scaling_factor) goes to both. """ kw = dict( num_heads=QWEN_HQ, @@ -6278,6 +6404,7 @@ def _qwen_layer_pools( tokens_per_block=QWEN_TPB, max_seq_len=max_seq_len, max_batch=max_batch, + **pool_kw, ) attn = _MultiLayerPagedAttnEnv(num_layers=len(QWEN_ATTN_LAYERS), num_blocks=attn_blocks, **kw) other = _MultiLayerPagedAttnEnv( @@ -6971,13 +7098,14 @@ def _cell_paged(cell, arm) -> None: A first, uncertified context call fills each sequence's `cached` tokens; the certified call then appends `new` rows per sequence (P for a generation row) and is gated against `op.reference` over the history it should have - attended. Over an fp8 pool every operand that call reads went through + attended. Over an fp8 or NVFP4 pool every operand that call reads went through e4m3, except on the packed context path, which never reads the pool. """ spec = cell.spec arm(op) _seed(cell) fp8 = bool(spec["quant_mode"] & QUANT_MODE_FP8_KV_CACHE) + nvfp4 = bool(spec["quant_mode"] & QUANT_MODE_NVFP4_KV_CACHE) env = _PagedAttnEnv( num_heads=spec["num_heads"], num_kv_heads=spec["num_kv_heads"], @@ -6988,7 +7116,7 @@ def _cell_paged(cell, arm) -> None: max_blocks_per_seq=16, pool_dtype=torch.float8_e4m3fn if fp8 else torch.bfloat16, quant_mode=spec["quant_mode"], - kv_scaling_factor=spec["kv_scale"] if fp8 else None, + kv_scaling_factor=spec["kv_scale"] if fp8 or nvfp4 else None, ) cached, new, nc = spec["cached"], spec["new"], spec["num_contexts"] p, paged = spec["predicted_tokens_per_seq"], spec["use_paged_context_fmha"] @@ -7023,9 +7151,10 @@ def _cell_paged(cell, arm) -> None: # One e4m3 flag covers the whole call: the packed path's context rows read # bf16, its generation rows e4m3, so a cell never mixes the two there. assert not (fp8 and not paged and 0 < nc < len(rids)) - e4m3 = fp8 and (paged or nc == 0) + assert not (nvfp4 and not paged), "an NVFP4 cell runs the paged-context path" + e4m3 = (fp8 or nvfp4) and (paged or nc == 0) causal = spec["mask_type"] == MASK_CAUSAL - _compare(out, _op_reference(env, qkv, list(new), rids, causal=causal, e4m3=e4m3)) + _compare(out, _op_reference(env, qkv, list(new), rids, causal=causal, e4m3=e4m3, nvfp4=nvfp4)) for rid in rids: env.check_cache(rid) env.check_unwritten_pool_zero() @@ -7046,8 +7175,14 @@ def _cell_layered(cell, arm) -> None: row, p = spec["local_layer_idx"], spec["predicted_tokens_per_seq"] cached, new, nc = spec["cached"], spec["new"], spec["num_contexts"] assert all(cached), "the history call drives every sequence of the cell" + fp8 = bool(spec["quant_mode"] & QUANT_MODE_FP8_KV_CACHE) + pool_kw = ( + dict(dtype=DataType.FP8, quant_mode=spec["quant_mode"], kv_scaling_factor=spec["kv_scale"]) + if fp8 + else {} + ) if spec["pool_layout"] == "hybrid48": - env, other = _qwen_layer_pools(max_batch=spec["max_num_requests"]) + env, other = _qwen_layer_pools(max_batch=spec["max_num_requests"], **pool_kw) layer = QWEN_ATTN_LAYERS.index(row) else: env = _MultiLayerPagedAttnEnv( @@ -7059,6 +7194,7 @@ def _cell_layered(cell, arm) -> None: num_blocks=24, max_batch=spec["max_num_requests"], max_seq_len=128, + **pool_kw, ) other, layer = None, row @@ -7104,7 +7240,7 @@ def tables() -> Dict: predicted_tokens_per_seq=p, **tables(), ) - _compare(out, _op_reference(env, qkv, list(new), rids, layer=layer)) + _compare(out, _op_reference(env, qkv, list(new), rids, layer=layer, e4m3=fp8)) env.check_caches(rids, layers=[layer]) _assert_only_this_layer_moved(env, snapshot, layer, cell.why) for before, after in zip(other_before, [] if other is None else other.layer_views): diff --git a/tests/unittest/_torch/modeling_v2/graph/test_modeling_v2_eager_on_graph.py b/tests/unittest/_torch/modeling_v2/graph/test_modeling_v2_eager_on_graph.py new file mode 100644 index 000000000000..33e34b6db83b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/graph/test_modeling_v2_eager_on_graph.py @@ -0,0 +1,285 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the eager_on_graph catalog entry. + +Each cell captures (or, for a passthrough cell, eagerly runs) a body of `breaks` breaks with +row-wise ops between them, the shape of a breakable prefill body, and replays it at the cell's +live rows. Every break's `inner` reads the step's rows and scale from live state, mixes rows +`[:n]` of its operand across the batch, writes a fixed output allocated inside the body and +zeroes its rows `[n:]`. The replayed body must equal the same body run eagerly at the same +live state, bit for bit; the hazards the entry's note states are each shown by a test below. +""" + +import ctypes +import logging +import zlib +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.graph.eager_on_graph import ( + eager_on_graph as op, +) +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import ( + BreakableCUDAGraph, + BreakableCUDAGraphCapture, +) + +__extra_import_path__ = [".."] +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "eager_on_graph requires a CUDA device" + +_BCG_LOGGER = "tensorrt_llm._torch.pyexecutor.breakable_cuda_graph.breakable_cuda_graph" + +#: `cudaError_t` values the pending-error checks compare against. +_CUDA_SUCCESS = 0 +_CUDA_ERROR_ILLEGAL_STATE = 401 + + +def _loaded_cudart() -> ctypes.CDLL: + """The CUDA runtime library torch already loaded into this process. Its + per-thread last error is the one torch's launch checks read, which + `cuda.bindings.runtime` (a separate runtime instance) does not see.""" + soname = f"/libcudart.so.{torch.version.cuda.split('.')[0]}" + with open("/proc/self/maps") as maps: + paths = sorted({line.split()[-1] for line in maps if soname in line}) + assert len(paths) == 1, f"expected exactly one loaded {soname[1:]}, found {paths}" + return ctypes.CDLL(paths[0]) + + +_CUDART = _loaded_cudart() + + +def _take_pending_cuda_error() -> int: + """Read and clear the CUDA runtime's pending last error on this thread.""" + return int(_CUDART.cudaGetLastError()) + + +@pytest.fixture(autouse=True) +def _leaves_cuda_clean(): + """Every test here hands the process on as it found it: no stream + capturing and no CUDA runtime error pending. A pending error is reported + by whichever CUDA launch comes next -- in a shared pytest process, the + first call of the next test file -- so it is taken here, after every test + and whether or not the test passed, and a test that left one fails here + instead of there.""" + yield + pending = _take_pending_cuda_error() + torch.cuda.synchronize() + assert not torch.cuda.is_current_stream_capturing(), "a test left the stream capturing" + assert pending == _CUDA_SUCCESS, f"the test left CUDA runtime error {pending} pending" + + +class _Live: + """The per-step state a break reads, refreshed before every step.""" + + def __init__(self, rows: int, scale: float) -> None: + self.rows = rows + self.scale = scale + + +def _inner(live: _Live): + """A break's eager work: rows `[:n]` centred over the live batch, scaled by + the live scale and offset by the static layer index, into `out` in place; + rows `[n:]` zeroed. Returns None, as an in-place break does.""" + + def inner(src: torch.Tensor, out: torch.Tensor, layer: int) -> None: + n = live.rows + rows = src[:n].float() + mixed = (rows - rows.mean(0, keepdim=True)) * live.scale + layer + out[:n].copy_(mixed.to(out.dtype)) + out[n:].zero_() + + return inner + + +def _body(x: torch.Tensor, breaks: int, call) -> tuple[torch.Tensor, list[torch.Tensor]]: + """Row-wise ops around `breaks` breaks; each break's output is allocated in + the body, as a captured body allocates it from the graph pool.""" + h = torch.sin(x) * 2.0 + outs = [] + for layer in range(breaks): + out = torch.empty_like(h) + call(h, out, layer) + outs.append(out) + h = torch.tanh(out) + h * 0.5 + return h, outs + + +def _capture(x_static: torch.Tensor, body) -> tuple: + """Two eager warmups, then the capture, on a side stream as the engine runs it.""" + graph = BreakableCUDAGraph() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(2): + body(x_static) + with BreakableCUDAGraphCapture(graph, pool=torch.cuda.graph_pool_handle(), stream=stream): + h, outs = body(x_static) + torch.cuda.current_stream().wait_stream(stream) + return graph, h, outs + + +@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 + tokens, width, breaks = spec["tokens"], spec["width"], spec["breaks"] + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + live = _Live(rows=tokens, scale=1.0) + inner = _inner(live) + + def broken(*operands): + return op(*operands, inner=inner) + + def eager(*operands): + return op.reference(*operands, inner=inner) + + x_static = torch.randn([tokens, width], generator=gen, device="cuda").to(torch.bfloat16) + graph = None + if spec["capture"]: + graph, h_static, outs_static = _capture(x_static, lambda x: _body(x, breaks, broken)) + assert graph.num_breaks == breaks and graph.num_segments == breaks + 1 + assert graph.num_segments == breaks + 1 + + for step, rows in enumerate(spec["live_rows"]): + assert 0 < rows <= tokens + live.rows, live.scale = rows, 0.25 + 0.5 * step + x_new = torch.randn([tokens, width], generator=gen, device="cuda").to(torch.bfloat16) + if graph is not None: + x_static.copy_(x_new) + graph.replay() + got, got_outs = h_static.clone(), [o.clone() for o in outs_static] + else: + got, got_outs = _body(x_new, breaks, broken) + want, want_outs = _body(x_new, breaks, eager) + assert len(got_outs) == len(want_outs) == breaks + op.compare(got, want) + for got_out, want_out in zip(got_outs, want_outs): + op.compare(got_out, want_out) + assert torch.count_nonzero(got_out[rows:]) == 0, "padding rows left unwritten" + + +def _single_break_graph(inner, x: torch.Tensor, *extra): + """Capture `inner(h, out, *extra)` as one break between two segments that + read `x` and copy `out` into the returned result. Every tensor the graph + touches is allocated inside the capture and kept by the caller.""" + graph = BreakableCUDAGraph() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + with BreakableCUDAGraphCapture(graph, pool=torch.cuda.graph_pool_handle(), stream=stream): + h = x * 1.0 + out = torch.empty_like(h) + op(h, out, *extra, inner=inner) + result = out + 0.0 + torch.cuda.current_stream().wait_stream(stream) + return graph, result, (h, out) + + +def test_a_scalar_argument_is_frozen_at_capture() -> None: + """A per-step count handed in as an argument replays at its capture-time value.""" + + def inner(src, out, n): + out[:n].copy_(src[:n] * 2) + out[n:].zero_() + + x = torch.randn(64, 8, device="cuda") + graph, result, _keep = _single_break_graph(inner, x, 64) + # The step now carries 17 rows, but nothing can hand the replay that number. + graph.replay() + torch.cuda.synchronize() + assert torch.equal(result, x * 2), "the replay must run at the capture-time n" + stale = torch.zeros_like(x) + stale[:17] = x[:17] * 2 + assert not torch.equal(result, stale) + + +def test_a_tensor_argument_is_read_at_its_capture_address() -> None: + """A tensor argument is re-read at its capture-time address: a new tensor the + step computed is never seen, and the captured one is read as it now is.""" + + def inner(src, out, scale): + torch.mul(src, scale, out=out) + + x = torch.randn(32, 8, device="cuda") + captured = torch.full((1,), 1.0, device="cuda") + graph, result, _keep = _single_break_graph(inner, x, captured) + fresh = torch.full((1,), 2.0, device="cuda") # what a per-step computation would make + assert fresh.data_ptr() != captured.data_ptr() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(result, x * 1.0), "replay must read the capture-time tensor" + captured.fill_(3.0) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(result, x * 3.0), "replay must read the capture-time address as it is now" + + +def test_a_wait_on_a_non_capturing_stream_is_dropped(caplog) -> None: + """Inside the inner's own capture-time run, a wait on a stream that is not + capturing is dropped, with a warning naming the call site.""" + side = torch.cuda.Stream() + + def inner(src, out, layer): + torch.cuda.current_stream().wait_stream(side) + out.copy_(src) + + with caplog.at_level(logging.WARNING, logger=_BCG_LOGGER): + _single_break_graph(inner, torch.randn(8, 8, device="cuda"), 0) + assert any("Dropping a wait" in record.message for record in caplog.records) + + +def test_an_object_argument_raises_under_capture() -> None: + """Any argument the replay cannot hold weakly is refused at capture, after + the inner already ran once. Outside capture the same call is a passthrough. + The refusal is a break raising under capture, so it also leaves the error + `test_a_raising_break_leaves_an_error_pending` shows, taken here.""" + ran = [] + + def inner(src, out, config): + ran.append(config.tag) + out.copy_(src) + + config = SimpleNamespace(tag="step") + x = torch.randn(8, 8, device="cuda") + out = torch.empty_like(x) + op(x, out, config, inner=inner) + assert torch.equal(out, x) and ran == ["step"] + + with pytest.raises((TypeError, RuntimeError)) as info: + _single_break_graph(inner, x, config) + chain = [] + error = info.value + while error is not None: + chain.append(error) + error = error.__context__ + assert any(isinstance(e, TypeError) and "Invalid type" in str(e) for e in chain) + assert ran == ["step", "step"] + assert _take_pending_cuda_error() == _CUDA_ERROR_ILLEGAL_STATE + torch.cuda.synchronize() + + +def test_a_raising_break_leaves_an_error_pending() -> None: + """A break that raises under capture has already ended the captured + segment, and the capture's exit ends it again: the exit raises + RuntimeError over the break's own error and leaves cudaErrorIllegalState + pending, which the process's next CUDA launch reports as its own failure + -- once, since that launch's check clears it.""" + + def inner(src, out, layer): + out.copy_(src) + raise ValueError("the break's own failure") + + with pytest.raises(RuntimeError) as info: + _single_break_graph(inner, torch.randn(8, 8, device="cuda"), 0) + assert isinstance(info.value.__context__, ValueError), info.value.__context__ + torch.cuda.synchronize() + + assert _CUDART.cudaPeekAtLastError() == _CUDA_ERROR_ILLEGAL_STATE + with pytest.raises(torch.AcceleratorError, match="cannot be performed in the present state"): + torch.ones(1, device="cuda").add_(1) + assert _take_pending_cuda_error() == _CUDA_SUCCESS + assert torch.equal(torch.ones(1, device="cuda").add_(1).cpu(), torch.full((1,), 2.0)) 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 index 973a12dc0501..1cfb3de32ee3 100644 --- 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 @@ -22,131 +22,51 @@ def _seed(text: str) -> int: def _build(spec, seed): - """`(buffer, out, region, gate_buffer, gate)` for one spec. + """`(out, gate_buffer, gate)` for one spec. - `out` is `region(buffer)`: the whole buffer, a padded window of it, or every - other row. `gate` is a view of `gate_buffer` -- for the checkpoint layout the - second half of each `[T, heads, 2 * D]` q_proj head. + `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, dt=dtype): - return torch.randn(*shape, generator=g, device="cuda").to(dt) - - layout = spec.get("out", "contiguous") - if layout == "window": - buffer = randn(t + 2, hidden + 2048) - - def region(x): - return x[:t, :hidden] - elif layout == "every_other_row": - buffer = randn(2 * t, hidden) - - def region(x): - return x[::2] - else: - buffer = randn(t, hidden) - - def region(x): - return x + def randn(*shape): + return torch.randn(*shape, generator=g, device="cuda").to(dtype) - gate_layout = spec["gate"] - if gate_layout == "interleaved": + out = randn(t, hidden) + if spec["gate"] == "interleaved": gate_buffer = randn(t, heads, 2 * d) gate = gate_buffer[:, :, d:] - elif gate_layout == "packed": - gate_buffer = gate = randn(t, heads, d) else: - gate_buffer = gate = randn(t, hidden, dt=spec.get("gate_dtype", dtype)) - return buffer, region(buffer), region, gate_buffer, gate + 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) - buffer, out, region, gate_buffer, gate = _build(cell.spec, seed=_seed(cell.why)) - if cell.spec["gate"] == "interleaved" and cell.spec["tokens"] > 1: - assert gate.stride() == ( - cell.spec["heads"] * 2 * cell.spec["head_dim"], - 2 * cell.spec["head_dim"], - 1, - ) - buffer_before, gate_before = buffer.clone(), gate_buffer.clone() + 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) - # In place: everything outside the [T, hidden] window, and the gate, bitwise unchanged. - touched = torch.zeros_like(buffer, dtype=torch.bool) - region(touched).fill_(True) - assert torch.equal(buffer[~touched], buffer_before[~touched]) assert torch.equal(gate_buffer, gate_before) -def test_the_gate_half_is_read_through_its_own_strides() -> None: - """The strided view and its compacted copy agree bitwise; the q half of the - same buffer does not -- so the cells can tell the two halves apart.""" - spec = dict(op.CELLS[1].spec) - _, out, _, q_and_gate, gate = _build(spec, seed=7) - q_half = q_and_gate[:, :, : spec["head_dim"]] - - from_strided = op(out.clone(), gate) - from_compacted = op(out.clone(), gate.contiguous()) - from_wrong_half = op(out.clone(), q_half) - - assert torch.equal(from_strided, from_compacted) - torch.testing.assert_close(from_strided, op.reference(out, gate)) - assert (from_wrong_half.float() - from_strided.float()).abs().max().item() > 0.5 - - -@pytest.mark.parametrize( - "case", - [ - "output_3d", - "output_last_stride", - "gate_numel_mismatch", - "gate_token_mismatch", - "gate_2d_shape_mismatch", - "gate_1d", - "gate_4d", - "gate_last_stride", - ], -) -def test_metadata_misuse_raises(case: str) -> None: - """Every layout violation is loud, which is why the entry has no `is_valid`.""" +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) - - def randn(*shape): - return torch.randn(*shape, generator=g, device="cuda").to(torch.bfloat16) - - attention_output, gate = randn(4, 128), randn(4, 2, 64) - if case == "output_3d": - attention_output = randn(4, 2, 64) - elif case == "output_last_stride": - attention_output = randn(4, 256)[:, ::2] - elif case == "gate_numel_mismatch": - gate = randn(4, 2, 32) - elif case == "gate_token_mismatch": - gate = randn(3, 2, 64) - elif case == "gate_2d_shape_mismatch": - gate = randn(4, 64) - elif case == "gate_1d": - gate = randn(512) - elif case == "gate_4d": - gate = randn(4, 2, 2, 32) - elif case == "gate_last_stride": - gate = randn(4, 256)[:, ::2] - + 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) - - -def test_a_cpu_tensor_is_rejected_by_triton() -> None: - attention_output = torch.randn(4, 128, dtype=torch.bfloat16) - gate = torch.randn(4, 128, dtype=torch.bfloat16) - with pytest.raises(ValueError, match="Pointer argument cannot be accessed from Triton"): - op(attention_output, gate) 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_erf.py index ab0b602f1189..5255d4197ee8 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_erf.py @@ -67,18 +67,7 @@ def _build(spec, seed): ) bias = FC1_BIAS_STD * torch.randn(MERGER_WIDTH, generator=g, device="cuda") return (grouped @ weight.to(dtype).float().t() + bias.to(dtype).float()).to(dtype) - layout = spec.get("layout", "contiguous") lo, hi = spec["lo"], spec["hi"] - if layout == "column_slice": - rows, width = shape - x = _uniform((rows, 2 * width), lo, hi, dtype, seed)[:, :width] - assert not x.is_contiguous() and x.stride(-1) == 1 - return x - if layout == "transposed": - rows, width = shape - x = _uniform((width, rows), lo, hi, dtype, seed).t() - assert x.stride(-1) != 1 - return x return _uniform(shape, lo, hi, dtype, seed) @@ -89,15 +78,8 @@ def test_certified_cells(cell, certifying): # noqa: F811 before = x.clone() out = op(x) assert out.shape == x.shape and out.dtype == x.dtype and out.device == x.device - if x.numel(): - assert out.data_ptr() != x.data_ptr(), "the entry is out-of-place" + assert out.data_ptr() != x.data_ptr(), "the entry is out-of-place" assert torch.equal(before, x), "the entry must not mutate its input" - if cell.spec.get("layout") == "transposed": - assert out.stride() == x.stride(), "the output keeps the input's memory format" - assert torch.equal(out, op(x.contiguous())) - if cell.spec.get("layout") == "column_slice": - assert out.is_contiguous() - assert torch.equal(out, op(x.contiguous())) torch.testing.assert_close(out, op.reference(x)) @@ -157,75 +139,3 @@ def test_identity_is_the_erf_form_and_not_tanh() -> None: print(f" bf16: correct {correct_b:.6e}; tanh variant {tanh_b:.6e} ({tanh_b / correct_b:.3f}x)") assert tanh_b < 4.0 * correct_b, (tanh_b, correct_b) torch.testing.assert_close(outb.to(torch.bfloat16), op.reference(xb)) - - -def test_peak_divergence_between_the_two_gelus() -> None: - """Where and how far apart the two GELUs are: the gap bounds every claim - about which one an implementation is.""" - xs = torch.linspace(-8.0, 8.0, 2_000_001, device="cuda", dtype=torch.float64) - gap = (_erf64(xs) - _tanh64(xs)).abs() - i = int(gap.argmax().item()) - print(f"\n max|erf - tanh| = {gap[i].item():.6e} at x = {xs[i].item():.4f}") - assert 4.0e-4 < gap[i].item() < 1e-3 - - -def test_extremes_and_non_finite() -> None: - """Saturation, signed-zero underflow, no overflow at finfo.max, -inf -> nan.""" - for dtype in (torch.bfloat16, torch.float16, torch.float32): - finfo = torch.finfo(dtype) - probe = torch.tensor( - [0.0, -0.0, 1e-3, -1e-3, 20.0, -20.0, 1000.0, -1000.0, finfo.max, -finfo.max], - device="cuda", - dtype=dtype, - ) - out = op(probe) - torch.testing.assert_close(out, op.reference(probe)) - assert out[4].item() == 20.0 and out[6].item() == 1000.0 - assert out[8].item() == finfo.max, "no overflow at the top of the dtype's range" - for index in (5, 7, 9): - assert out[index].item() == 0.0 and math.copysign(1.0, out[index].item()) < 0 - - nonfinite = torch.tensor( - [float("inf"), float("-inf"), float("nan")], device="cuda", dtype=dtype - ) - got = op(nonfinite) - assert got[0].item() == float("inf") - assert math.isnan(got[1].item()), "gelu_erf(-inf) is nan" - assert math.isnan(got[2].item()) - ref = _erf64(nonfinite) - assert ref[0].item() == float("inf"), "the closed form agrees: the function, not the kernel" - assert math.isnan(ref[1].item()) and math.isnan(ref[2].item()) - - -def test_rank_does_not_change_the_arithmetic() -> None: - """The one engine op computing this function rejects non-2-D input; this one does not.""" - x = _uniform((2, 37, MERGER_WIDTH), -16.0, 16.0, torch.bfloat16, 26) - assert torch.equal(op(x).reshape(-1, MERGER_WIDTH), op(x.reshape(-1, MERGER_WIDTH))) - - -def test_rows_are_independent() -> None: - """The 4-to-1 grouping in front of the merger reorders rows; the activation - must not care. Permuting rows permutes the result, bitwise.""" - x = _uniform((140, MERGER_WIDTH), -16.0, 16.0, torch.bfloat16, 27) - order = torch.randperm(140, generator=torch.Generator().manual_seed(0)).to("cuda") - assert torch.equal(op(x[order].contiguous()), op(x)[order]) - - -def test_cuda_graph_capture_and_replay() -> None: - x = _uniform((64, MERGER_WIDTH), -16.0, 16.0, torch.bfloat16, 28) - - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - for _ in range(3): - op(x) - torch.cuda.current_stream().wait_stream(warmup) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured = op(x) - - x.copy_(_uniform((64, MERGER_WIDTH), -16.0, 16.0, torch.bfloat16, 29)) - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(captured, op.reference(x)) 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_tanh.py index aa52e550fce0..b09168720ce5 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_tanh.py @@ -58,18 +58,7 @@ def _build(spec, seed): ).to(dtype) bias = (FC1_BIAS_STD * torch.randn(shape[-1], generator=g, device="cuda")).to(dtype) return (hidden.float() @ weight.float().t() + bias.float()).to(dtype) - layout = spec.get("layout", "contiguous") lo, hi = spec["lo"], spec["hi"] - if layout == "column_slice": - rows, width = shape - x = _uniform((rows, 2 * width), lo, hi, dtype, seed)[:, :width] - assert not x.is_contiguous() and x.stride(-1) == 1 - return x - if layout == "transposed": - rows, width = shape - x = _uniform((width, rows), lo, hi, dtype, seed).t() - assert x.stride(-1) != 1 - return x return _uniform(shape, lo, hi, dtype, seed) @@ -80,15 +69,8 @@ def test_certified_cells(cell, certifying): # noqa: F811 before = x.clone() out = op(x) assert out.shape == x.shape and out.dtype == x.dtype and out.device == x.device - if x.numel(): - assert out.data_ptr() != x.data_ptr(), "the entry is out-of-place" + assert out.data_ptr() != x.data_ptr(), "the entry is out-of-place" assert torch.equal(before, x), "the entry must not mutate its input" - if cell.spec.get("layout") == "transposed": - assert out.stride() == x.stride(), "the output keeps the input's memory format" - assert torch.equal(out, op(x.contiguous())) - if cell.spec.get("layout") == "column_slice": - assert out.is_contiguous() - assert torch.equal(out, op(x.contiguous())) torch.testing.assert_close(out, op.reference(x)) @@ -149,63 +131,3 @@ def test_identity_is_the_tanh_form_and_not_erf() -> None: print(f" bf16: correct {correct_b:.6e}; erf variant {erf_b:.6e} ({erf_b / correct_b:.3f}x)") assert erf_b < 4.0 * correct_b, (erf_b, correct_b) torch.testing.assert_close(outb.to(torch.bfloat16), op.reference(xb)) - - -def test_extremes_and_non_finite() -> None: - """Saturation, signed-zero underflow, no overflow at finfo.max, -inf -> nan.""" - for dtype in (torch.bfloat16, torch.float16, torch.float32): - finfo = torch.finfo(dtype) - probe = torch.tensor( - [0.0, -0.0, 1e-3, -1e-3, 20.0, -20.0, 1000.0, -1000.0, finfo.max, -finfo.max], - device="cuda", - dtype=dtype, - ) - out = op(probe) - torch.testing.assert_close(out, op.reference(probe)) - assert out[4].item() == 20.0 and out[6].item() == 1000.0 - assert out[8].item() == finfo.max, "x**3 is not evaluated in the input dtype" - for index in (5, 7, 9): - assert out[index].item() == 0.0 and math.copysign(1.0, out[index].item()) < 0 - - nonfinite = torch.tensor( - [float("inf"), float("-inf"), float("nan")], device="cuda", dtype=dtype - ) - got = op(nonfinite) - assert got[0].item() == float("inf") - assert math.isnan(got[1].item()), "gelu_tanh(-inf) is nan" - assert math.isnan(got[2].item()) - ref = _tanh64(nonfinite) - assert ref[0].item() == float("inf"), "the closed form agrees: the function, not the kernel" - assert math.isnan(ref[1].item()) and math.isnan(ref[2].item()) - - -def test_rank_does_not_change_the_arithmetic() -> None: - x = _uniform((2, 37, INTERMEDIATE_SIZE), -16.0, 16.0, torch.bfloat16, 26) - assert torch.equal(op(x).reshape(-1, INTERMEDIATE_SIZE), op(x.reshape(-1, INTERMEDIATE_SIZE))) - - -def test_rows_are_independent() -> None: - """Permuting rows permutes the result, bitwise.""" - x = _uniform((128, INTERMEDIATE_SIZE), -16.0, 16.0, torch.bfloat16, 27) - order = torch.randperm(128, generator=torch.Generator().manual_seed(0)).to("cuda") - assert torch.equal(op(x[order].contiguous()), op(x)[order]) - - -def test_cuda_graph_capture_and_replay() -> None: - x = _uniform((64, INTERMEDIATE_SIZE), -16.0, 16.0, torch.bfloat16, 28) - - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - for _ in range(3): - op(x) - torch.cuda.current_stream().wait_stream(warmup) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured = op(x) - - x.copy_(_uniform((64, INTERMEDIATE_SIZE), -16.0, 16.0, torch.bfloat16, 29)) - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(captured, op.reference(x)) 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_silu_and_mul_rows.py new file mode 100644 index 000000000000..e5291a35eb15 --- /dev/null +++ b/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 @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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.""" + +import zlib + +import pytest +import torch + +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 + silu_and_mul_rows as op, +) + +__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" + +INTER = 640 + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + rows, width = spec["rows"], 2 * INTER + 8 + x = (torch.randn(rows, width, generator=g, device="cuda") * 3.0).to(torch.bfloat16) + if rows: + x[0, :4] = torch.tensor([0.0, -0.0, 88.0, -100.0], device="cuda").to(torch.bfloat16) + return x + + +@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) + x = _build(cell.spec, zlib.crc32(cell.why.encode())) + out = op(x, INTER) + assert out.shape == (x.shape[0], INTER) and out.is_contiguous() + 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/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 new file mode 100644 index 000000000000..0561d585eeae --- /dev/null +++ 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 @@ -0,0 +1,167 @@ +# 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_expand_block_indices.py b/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 index 3c09c66fc05f..0c83dcf90793 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_expand_block_indices.py @@ -31,8 +31,6 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: DEV = "cuda" CR, TOPK = 4, 2048 -BLOCK_TOPK = TOPK // CR -FINAL_TOPK = TOPK + CR - 1 ENGINE_MAX_SEQ_LEN = 8192 @@ -64,7 +62,12 @@ def _rows(batch) -> tuple[list[int], list[int]]: elif kind == "serving": return _serving_rows(batch[1], batch[2]) else: - positions = list(batch[1]) + 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 + return positions, lengths return positions, [p + 1 for p in positions] @@ -89,23 +92,9 @@ def _packed_selection(positions: torch.Tensor, compress_ratio: int, token_topk: def _build(spec, seed: int): positions, lengths = _rows(spec["batch"]) cr, topk = spec["compress_ratio"], spec["token_topk"] - block_topk = topk // cr query_positions = torch.tensor(positions, dtype=spec["position_dtype"], device=DEV) sequence_lengths = torch.tensor(lengths, dtype=spec["length_dtype"], device=DEV) blocks = _packed_selection(query_positions, cr, topk, seed) - rows = blocks.shape[0] - junk = torch.randint(0, 64, (rows + 37, 2 * block_topk), dtype=torch.int32, device=DEV) - if spec["layout"] == "row_slice": - # The engine's buffer is max_num_tokens tall; rows past this step hold - # whatever the last step left there. - tall = junk[:, :block_topk].clone() - tall[:rows] = blocks - blocks = tall[:rows] - elif spec["layout"] == "column_slice": - wide = junk[:rows].clone() - wide[:, :block_topk] = blocks - blocks = wide[:, :block_topk] - assert not blocks.is_contiguous() and blocks.stride(1) == 1 return (blocks, query_positions, sequence_lengths), dict(compress_ratio=cr, token_topk=topk) @@ -126,87 +115,11 @@ def test_certified_cells(cell, certifying): # noqa: F811 ) -def _expand(blocks, positions, lengths, *, compress_ratio=CR, token_topk=TOPK): - """One call, gated on the entry's own reference, returned on the host.""" - args = ( - blocks.to(DEV), - torch.as_tensor(positions, dtype=torch.int64).to(DEV), - torch.as_tensor(lengths, dtype=torch.int32).to(DEV), - ) - geometry = dict(compress_ratio=compress_ratio, token_topk=token_topk) - out = op(*args, **geometry) - _compare(out, op.reference(*args, **geometry)) - return out.cpu() - - -def test_worked_example_small_geometry() -> None: - """Two rows written out by hand, end to end: the reference agrees too.""" - # Row 0: groups 1 and 0 -> tokens 4..7 then 0..3; position 9 leaves {8, 9}. - # Row 1: group 0 -> 0..3; position 5 leaves {4, 5}; the rest is padding. - blocks = torch.tensor([[1, 0], [0, -1]], dtype=torch.int32) - out = _expand(blocks, [9, 5], [10, 6], compress_ratio=4, token_topk=8) - assert out.tolist() == [ - [4, 5, 6, 7, 0, 1, 2, 3, 8, 9, -1], - [0, 1, 2, 3, 4, 5, -1, -1, -1, -1, -1], - ] - - -def test_launch_shapes_agree_on_shared_rows() -> None: - """The same rows expand identically either side of the 128-row switch.""" - positions = torch.tensor([4096 + 13 * i for i in range(128)], device=DEV) - blocks = _packed_selection(positions, CR, TOPK, seed=7) - whole_row = _expand(blocks, positions, positions + 1) - column_split = _expand(blocks[:127], positions[:127], positions[:127] + 1) - assert torch.equal(whole_row[:127], column_split) - - -@pytest.mark.parametrize( - "position, expected_tail", - [(4095, 0), (4096, 1), (4097, 2), (4098, 3), (4099, 0), (7, 0), (3, 0), (0, 1)], -) -def test_tail_width_tracks_the_query_position(position: int, expected_tail: int) -> None: - """`(position + 1) % cr` tokens are appended right after the selection, ascending.""" - blocks = _packed_selection(torch.tensor([position], device=DEV), CR, TOPK, seed=position + 1) - out = _expand(blocks, [position], [position + 1]) - selected = int((blocks[0] >= 0).sum()) * CR - assert int((out[0] >= 0).sum()) == selected + expected_tail - tail_start = (position + 1) // CR * CR - assert out[0, selected : selected + expected_tail].tolist() == list( - range(tail_start, tail_start + expected_tail) - ) - assert bool((out[0, selected + expected_tail :] == -1).all()) - - -def test_selected_groups_are_bounded_by_sequence_length_not_query_position() -> None: - """The causal bound on the selection is the producer's, not this call's. - - Both inputs are illegal for a QSA producer; they pin what the call does. - """ - # Group 8 is past position 7 but inside length 64: expanded verbatim. - out = _expand(torch.tensor([[8, 0]], dtype=torch.int32), [7], [64], token_topk=8) - assert out.tolist() == [[32, 33, 34, 35, 0, 1, 2, 3, -1, -1, -1]] - # Group 2 is incomplete at length 10: -1 holes in place, then its real - # tokens repeated in the tail. - out = _expand(torch.tensor([[0, 1, 2, -1]], dtype=torch.int32), [9], [10], token_topk=16) - assert out.tolist() == [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, -1, -1, 8, 9, -1, -1, -1, -1, -1]] - - -def test_unpacked_prefix_is_silently_dropped() -> None: - """A -1 before a valid group loses that group, with no error. - - Prefix packing is a value property, so no metadata guard can check it. - """ - blocks = torch.tensor([[-1, 3]], dtype=torch.int32, device=DEV) - positions = torch.tensor([15], device=DEV) - lengths = torch.tensor([16], dtype=torch.int32, device=DEV) - out = op(blocks, positions, lengths, compress_ratio=4, token_topk=8) - # Group 3 is tokens 12..15; every one of them is gone. - assert out.cpu().tolist() == [[-1] * 11] - - def test_cuda_graph_capture_and_replay() -> None: - """Capturable, and a replay recomputes from whatever the buffers now hold.""" - positions = torch.tensor([4096 + 7 * i for i in range(4)], device=DEV) + """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 + 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) @@ -233,78 +146,11 @@ def call(): _compare(captured, op.reference(blocks, positions, lengths, compress_ratio=CR, token_topk=TOPK)) -def test_output_is_fresh_and_inputs_are_untouched() -> None: - """Every column is written, two calls do not alias, and nothing is in place.""" - positions = torch.tensor([4096, 97, 4095], device=DEV) - lengths = (positions + 1).to(torch.int32) - blocks = _packed_selection(positions, CR, TOPK, seed=41) - before = (blocks.clone(), positions.clone(), lengths.clone()) - outputs = [op(blocks, positions, lengths, compress_ratio=CR, token_topk=TOPK) for _ in range(5)] - assert len({out.data_ptr() for out in outputs}) == len(outputs) - # The output comes from torch.empty: a skipped column would differ between - # calls or land outside the legal range. - for out in outputs[1:]: - assert torch.equal(outputs[0], out) - first = outputs[0] - assert bool(((first == -1) | ((first >= 0) & (first < int(lengths.max())))).all()) - for tensor, saved in zip((blocks, positions, lengths), before): - assert torch.equal(tensor, saved) - - -def test_op_rejects_malformed_arguments() -> None: - """Every metadata violation is the callable's own loud error: why is_valid is empty.""" +def test_op_rejects_a_width_that_does_not_match_the_geometry() -> None: + """A selection buffer whose width is not `token_topk // compress_ratio` is the + callable's own loud error, not a silent mis-expansion: why is_valid is empty.""" positions = torch.tensor([4096, 4097], device=DEV) lengths = torch.tensor([4097, 4098], dtype=torch.int32, device=DEV) blocks = _packed_selection(positions, CR, TOPK, seed=17) - geometry = dict(compress_ratio=CR, token_topk=TOPK) - cases = [ - (r"must have shape \[rows, block_topk\]", (blocks.reshape(-1), positions, lengths), {}), - (r"must have shape \[rows, block_topk\]", (blocks.unsqueeze(0), positions, lengths), {}), - ("divisible by compress_ratio >= 2", (blocks, positions, lengths), dict(compress_ratio=1)), - ("divisible by compress_ratio >= 2", (blocks, positions, lengths), dict(token_topk=0)), - ("divisible by compress_ratio >= 2", (blocks, positions, lengths), dict(token_topk=2049)), - ("width does not match", (blocks, positions, lengths), dict(token_topk=1024)), - ("one entry per row", (blocks, torch.zeros(3, dtype=torch.int64, device=DEV), lengths), {}), - ("must use int32 storage", (blocks.long(), positions, lengths), {}), - ("must use integer storage", (blocks, positions.to(torch.int16), lengths), {}), - ("must use integer storage", (blocks, positions, lengths.float()), {}), - ( - "contiguous along their last dimension", - ( - torch.zeros((2, 2 * BLOCK_TOPK), dtype=torch.int32, device=DEV)[:, ::2], - positions, - lengths, - ), - {}, - ), - ( - "metadata must be contiguous", - (blocks, torch.zeros(4, dtype=torch.int64, device=DEV)[::2], lengths), - {}, - ), - ("requires CUDA tensors", (blocks.cpu(), positions, lengths), {}), - ("requires CUDA tensors", (blocks, positions.cpu(), lengths), {}), - ] - for message, args, override in cases: - with pytest.raises(ValueError, match=message): - op(*args, **{**geometry, **override}) - - -def test_discriminates_plausible_wrong_variants() -> None: - """The bitwise gate separates this layout from the ones a rewrite would produce.""" - positions = [4096, 97, 4095] - blocks = _packed_selection(torch.tensor(positions, device=DEV), CR, TOPK, seed=23).cpu() - out = _expand(blocks, positions, [p + 1 for p in positions]) - first_free = (blocks >= 0).sum(dim=1) * CR - - dropped_tail = out.clone() - tail_at_fixed_column = out.clone() - for row, position in enumerate(positions): - dropped_tail[row, first_free[row] :] = -1 - tail_at_fixed_column[row, first_free[row] : TOPK] = -1 - tail_start = (position + 1) // CR * CR - for offset in range(min(position + 1 - tail_start, CR - 1)): - tail_at_fixed_column[row, TOPK + offset] = tail_start + offset - assert not torch.equal(out, dropped_tail) - assert not torch.equal(out, tail_at_fixed_column) - assert not torch.equal(out, torch.sort(out, dim=1).values) + with pytest.raises(ValueError, match="width does not match"): + op(blocks, positions, lengths, compress_ratio=CR, token_topk=1024) 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_mtp_reuse_topk.py new file mode 100644 index 000000000000..1f8e34122860 --- /dev/null +++ b/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 @@ -0,0 +1,248 @@ +# 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. + +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 +sit in the engine's buffers -- one slot of a `[slots, max_num_sequences, K]` +shared buffer holding stale rows elsewhere, and row slices of the +`max_num_tokens`-tall request map, visible-count and Top-K buffers. Every gate +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.attention.qsa_mtp_reuse_topk import ( # noqa: E501 + qsa_mtp_reuse_topk as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "qsa_mtp_reuse_topk requires a CUDA device" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Bitwise: these are group ids, and one wrong entry attends to a different + group. There is no rounding for a tolerance to absorb.""" + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + torch.testing.assert_close(out, ref, rtol=0, atol=0) + + +DEV = "cuda" +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. +ANCHORS = (2, 3, 402, 2046, 2050, 8188, 2051, 2047, 1022, 4095, 5998, 7001, 37, 2044, 3001, 6143) + + +def _captured_row(visible: int, block_topk: int, gen: torch.Generator) -> torch.Tensor: + """A legal capture: `min(K, visible)` distinct groups below `visible`, then -1.""" + row = torch.full((block_topk,), -1, dtype=torch.int32) + if visible > 0: + take = min(visible, block_topk) + row[:take] = torch.randperm(visible, generator=gen)[:take].to(torch.int32) + return row + + +def _requests(contexts, drafted, gen) -> 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] + 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) + out.append(((position + 1) // CR, (position + d + 1) // CR)) + return out + + +def _build(spec, seed: int, requests: int | None = None) -> dict: + """Turn a cell spec into the op's keyword arguments, plus the buffers behind them. + + Returns `dict(kwargs=..., buffers=...)`; `buffers` holds the full tensors the + arguments are views of, so a test can check that nothing outside the view moved. + """ + gen = torch.Generator().manual_seed(seed) + block_topk = spec["block_topk"] + contexts = spec["contexts"] + if requests is not None: + contexts = (contexts[0], requests) + pairs = _requests(contexts, spec["drafted"], gen) + num_requests = len(pairs) + per_request = spec["rows_per_request"] + rows = num_requests * per_request + num_slots, slot = spec["slots"] + shared_rows = max(spec["shared_rows"], num_requests) + meta_dtype = spec["metadata_dtype"] + + # Stale rows everywhere: other layers' slots and requests past this batch. + shared = torch.randint(-1, 3000, (num_slots, shared_rows, block_topk), generator=gen) + shared = shared.to(torch.int32) + shared_vis = torch.randint(0, 3000, (num_slots, shared_rows), generator=gen).to(meta_dtype) + for q, (captured, _) in enumerate(pairs): + shared[slot, q] = _captured_row(captured, block_topk, gen) + shared_vis[slot, q] = captured + + request, visible = [], [] + for q in range(num_requests): + captured, current = pairs[q] + for j in range(per_request): + request.append(q) + # A later row of the same request sees at least as many groups. + visible.append(current + j // CR) + request = torch.tensor(request, dtype=meta_dtype).reshape(rows) + visible = torch.tensor(visible, dtype=meta_dtype).reshape(rows) + + height = max(ENGINE_MAX_NUM_TOKENS, rows + 37) + tall_out = torch.randint(-1, 3000, (height, block_topk), generator=gen).to(torch.int32) + tall_req = torch.randint(0, shared_rows, (height,), generator=gen).to(meta_dtype) + tall_vis = torch.randint(0, 3000, (height,), generator=gen).to(meta_dtype) + tall_req[:rows], tall_vis[:rows] = request, visible + buffers = dict( + shared=shared.to(DEV), + shared_vis=shared_vis.to(DEV), + request=tall_req.to(DEV), + visible=tall_vis.to(DEV), + output=tall_out.to(DEV), + ) + kwargs = dict( + shared_indices=buffers["shared"][slot], + shared_visible_blocks=buffers["shared_vis"][slot], + request_indices=buffers["request"][:rows], + visible_blocks=buffers["visible"][:rows], + output=buffers["output"][:rows], + ) + return dict(kwargs=kwargs, buffers=buffers, rows=rows) + + +def _assert_legal_selection(out: torch.Tensor, kwargs: dict) -> None: + """What the producer's contract buys: no duplicate group, none past the horizon.""" + if out.shape[0] == 0: + return + ids = out.long() + ordered = torch.sort(ids, dim=1).values + assert not bool(((ordered[:, 1:] == ordered[:, :-1]) & (ordered[:, 1:] >= 0)).any()), ( + "a group id appears twice in a row" + ) + captured = kwargs["shared_visible_blocks"].long()[kwargs["request_indices"].long()] + horizon = torch.maximum(captured, kwargs["visible_blocks"].long()) + assert bool((ids < horizon[:, None]).all()), "a group id reaches past the row's horizon" + assert bool((ids >= -1).all()) + + +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"] + before = {name: tensor.clone() for name, tensor in buffers.items()} + out = op(**kwargs) + torch.cuda.synchronize() + assert out.data_ptr() == kwargs["output"].data_ptr(), "the op returns the buffer it wrote" + assert out.dtype == torch.int32 and tuple(out.shape) == (rows, spec["block_topk"]) + _compare(out, op.reference(**kwargs)) + _assert_legal_selection(out, kwargs) + for name, tensor in buffers.items(): + if name == "output": + assert torch.equal(tensor[rows:], before[name][rows:]), "rows past the view moved" + else: + assert torch.equal(tensor, before[name]), f"{name} was written" + return out + + +def _check_graph(spec, seed: int) -> None: + """At each request count: eager, capture, then replay on new data.""" + for requests in spec["graph_batches"]: + first = _build(spec, seed + requests, requests) + kwargs = first["kwargs"] + eager = op(**kwargs).clone() + _compare(eager, op.reference(**kwargs)) + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(**kwargs) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(**kwargs) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(captured, eager), f"{requests} requests: replay != eager" + + # The mixed batch's first requests sit at fixed positions, so the new + # data comes from full rows: every request's captured row changes. + refill = {**spec, "contexts": ("full", requests)} + fresh = _build(refill, seed + 1000 + requests)["kwargs"] + for name, tensor in kwargs.items(): + if name != "output": + tensor.copy_(fresh[name]) + graph.replay() + torch.cuda.synchronize() + expected = op.reference(**kwargs) + assert not torch.equal(expected, eager), f"{requests} requests: new data changes nothing" + _compare(captured, expected) + rerun = op(**{**kwargs, "output": torch.empty_like(kwargs["output"])}) + assert torch.equal(captured, rerun), f"{requests} requests: replay != eager on new data" + del graph + + +@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) + seed = zlib.crc32(cell.why.encode()) + if cell.spec["graph_batches"]: + _check_graph(cell.spec, seed) + else: + _check_cell(cell.spec, seed) + + +def _spec_of(prefix: str) -> dict: + (cell,) = [c for c in op.CELLS if c.why.startswith(prefix)] + return cell.spec + + +def _guard_case(): + """Legal arguments, then the whole `[slots, seqs]` visible-count buffer passed where + one slot's view belongs. + + Returns (intended, kwargs): `intended` is the reference on the legal + arguments -- what a caller passing the bad layout meant to get. + """ + spec = _spec_of("decode MTP") + kwargs = _build({**spec, "slots": (3, 2)}, seed=13, requests=6)["kwargs"] + intended = op.reference(**kwargs) + whole = torch.zeros((3, kwargs["shared_visible_blocks"].numel()), dtype=torch.int32) + whole = whole.to(DEV) + whole[2] = kwargs["shared_visible_blocks"] + kwargs["shared_visible_blocks"] = whole + return intended, kwargs + + +def test_guard_refuses_a_whole_multi_slot_count_buffer() -> None: + """Unguarded the op reads slot 0's counts and returns a plausible wrong answer; + `is_valid` refuses it.""" + intended, kwargs = _guard_case() + out = op(**kwargs) + torch.cuda.synchronize() + assert not torch.equal(out, intended), "the op got it right after all" + _, kwargs = _guard_case() + with ( + validating(op), + pytest.raises(ValueError, match="shared_visible_blocks must be one slot's"), + ): + op(**kwargs) 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 index b2ce01b5cbb0..8498ad6cded2 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_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 @@ -3,8 +3,8 @@ """GPU test for the qsa_paged_index_scores catalog entry. The paged state is real: a `[pages, tokens_per_block, 1, dim]` index cache -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 +(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. @@ -37,9 +37,8 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: 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) land more than 1000x outside - (asserted at 64 and 2080 columns; measured 3e5x to 4e6x). Columns the - reference marks NaN are unspecified and not read. + 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) @@ -51,8 +50,7 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: #: 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 = 16 -CAPTURE_PAGES = 260 +DOT_TILE_WIDTH = 128 class _LaunchRecorder: @@ -101,7 +99,7 @@ def grid(self): def _case(rows, requests, **overrides): - """A narrow-table spec for the non-cell tests, same shape as the entry's own.""" + """A narrow-table spec for the non-cell tests, same geometry as the entry's own.""" spec = dict( layout="random", rows=rows, @@ -130,37 +128,17 @@ def _rows_and_requests(spec, gen): def pages_for(seq_len): return math.ceil(seq_len / tpb) - if layout in ("random", "shuffled", "past_table", "arange"): + if layout == "random": rows, requests = spec["rows"], spec["requests"] - if layout == "shuffled": - req = torch.randint(0, requests, (rows,), generator=gen, device=DEV) - else: - req = torch.arange(rows, device=DEV) * requests // max(rows, 1) - if layout == "past_table": - positions = torch.full((rows,), pages * tpb * 3, device=DEV) - elif layout == "arange": - positions = torch.arange(rows, device=DEV) - else: - positions = torch.randint(0, pages * tpb, (rows,), generator=gen, device=DEV) + 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 - if layout == "far": - rows = spec["rows"] - positions = torch.full((rows,), pages * tpb - 1, device=DEV) - return positions, torch.zeros(rows, dtype=torch.long, device=DEV), [pages] - segments = [] # (first position, row count) per request, in row order + # (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"]] - elif layout == "mixed": - max_position = pages * tpb - 1 - contexts = spec["context_lens"] - segments = [ - (0 if i == 0 else (907 * i) % (max_position - n + 2), n) for i, n in enumerate(contexts) - ] - # MTP: each generation request carries two rows at consecutive positions. - segments += [(100 + (499 * g) % 8000, 2) for g in range(spec["requests"] - len(contexts))] else: raise AssertionError(f"unknown layout {layout!r}") positions = torch.cat([torch.arange(s, s + n, device=DEV) for s, n in segments]) @@ -186,18 +164,25 @@ def _build(spec, seed): for slot, count in zip(slot_of.tolist(), mapped): block_table[slot, :count] = page_ids[cursor : cursor + count] cursor += count - index_cache = torch.randn( - pool_pages, spec["tokens_per_block"], 1, head_dim, generator=gen, device=DEV - ).to(dtype) + 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] - elif spec["q_layout"] == "dim_slice": - wide = torch.zeros(rows, heads, 2 * head_dim, dtype=dtype, device=DEV) - wide[:, :, :head_dim] = torch.randn(rows, heads, head_dim, generator=gen, device=DEV) - q = wide[:, :, :head_dim] 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 @@ -247,307 +232,53 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 _compare(out, op.reference(**args)) -@pytest.mark.parametrize("context_rows,branch", [(False, "row_wise"), (True, "query_tile")]) -def test_capture_width_needs_more_than_one_column_program(context_rows, branch) -> None: - """Why the 260-page table is an axis of its own, not a bigger number. - - Both kernels walk the score columns in a grid dimension of their own. At 8 - pages (64 columns) one program covers every column; at 260 pages (2080) it - does not, so only the wide cells ever run a later column program. - """ - grids = {} - for tag, pages in (("narrow", 8), ("wide", CAPTURE_PAGES)): - args = _build(_case(300, 4, pages=pages, context_rows=context_rows), seed=8500 + pages) - with _watch_launches() as watcher: - out = op(**args) - torch.cuda.synchronize() - assert watcher.branch == branch - grids[tag] = watcher.grid - _compare(out, op.reference(**args)) - assert grids["narrow"][1] == 1, grids - assert grids["wide"][1] > 1, grids - - -def _capture_args(seed, only_visible_blocks=False): - spec = _case( - 0, - 1, - layout="capture", - lengths=(8191,), - slots=8, - pages=CAPTURE_PAGES, - only_visible_blocks=only_visible_blocks, - context_rows=True, - ) - return _build(spec, seed) - - -@pytest.mark.parametrize("shape", ["512rows_narrow", "8191rows_capture"]) -def test_the_two_kernels_agree(shape) -> None: - """One entry, two kernels: no further apart than fp32 reordering. - - Measured 9.5e-7 apart (one fp32 ulp at the score's magnitude) at 320 rows - and again at 8191, so the packed kernel does not drift with batch size. - """ - if shape == "512rows_narrow": - args = _build(_case(512, 6), seed=404) - else: - args = _capture_args(seed=8700) - visible = _visible(args) - row_wise = op(**{**args, "context_rows": False}) - query_tile = op(**{**args, "context_rows": True}) - torch.testing.assert_close(row_wise[visible], query_tile[visible]) - - -@pytest.mark.parametrize("shape", ["300rows_row_wise", "300rows_query_tile", "8191rows_capture"]) -def test_results_are_bitwise_reproducible(shape) -> None: - """Three identical calls, bitwise identical, in the padded mode so the - whole tensor is written.""" - if shape == "8191rows_capture": - args = _capture_args(seed=8800) - else: - args = _build(_case(300, 5, context_rows=shape.endswith("tile")), seed=505) - runs = [op(**args).clone() for _ in range(3)] - torch.cuda.synchronize() - assert torch.equal(runs[0], runs[1]) and torch.equal(runs[1], runs[2]) - - -@pytest.mark.parametrize( - "shape", ["64rows_row_wise", "300rows_row_wise", "300rows_query_tile", "8191rows_capture"] -) -def test_only_visible_blocks_leaves_the_tail_unspecified(shape) -> None: - """`only_visible_blocks=True` writes the visible columns and nothing else. - - The output is a fresh `torch.empty`, and the caching allocator hands back - the block freed just before: poisoning one of the right size makes the - unwritten tail observable. Asserted in the allocator-independent form: the - tail is NOT the `-inf` the padded mode writes. - """ - if shape == "8191rows_capture": - args = _capture_args(seed=8900, only_visible_blocks=True) - else: - rows = int(shape.split("rows")[0]) - spec = _case(rows, 4, only_visible_blocks=True, context_rows=shape.endswith("tile")) - args = _build(spec, seed=606 + rows) - expected = op.reference(**args) - poison = torch.full(expected.shape, 98765.0, dtype=torch.float32, device=DEV) - del poison - out = op(**args) - torch.cuda.synchronize() - _compare(out, expected) - tail = torch.isnan(expected) - assert bool(tail.any()), "this case must have unscored columns to say anything" - assert not bool(torch.isneginf(out[tail]).all()), "the tail must not be the -inf padding" - - -def test_both_modes_agree_on_the_visible_columns() -> None: - """The flag decides what lands outside the causal bound, nothing inside.""" - args = _build(_case(96, 4), seed=707) - padded = op(**{**args, "only_visible_blocks": False}) - unspecified = op(**{**args, "only_visible_blocks": True}) +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) - assert torch.equal(padded[visible], unspecified[visible]) - - -def test_capture_shape_survives_cuda_graph_capture_and_replay() -> None: - """Capture is why 8191 rows exist, so capture it -- and replay it again - after overwriting the operands in place, which a graph that baked in its - inputs would fail.""" - args = _capture_args(seed=9100, only_visible_blocks=True) - op(**args) # warm the Triton launch outside capture - torch.cuda.synchronize() + 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() - try: - with torch.cuda.graph(graph): - captured = op(**args) + 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() - _compare(captured, op.reference(**args)) - - fresh = _capture_args(seed=9101, only_visible_blocks=True) - for name in ("q", "index_cache"): - args[name].copy_(fresh[name]) - # A new logical->physical map for the same slot: the pages it maps, - # reshuffled in place. - table = args["block_table"] - mapped = table >= 0 - gen = torch.Generator(device=DEV).manual_seed(9102) - pages = table[mapped] - table[mapped] = pages[torch.randperm(pages.numel(), generator=gen, device=DEV)] - graph.replay() - torch.cuda.synchronize() - _compare(captured, op.reference(**args)) - finally: - del graph - torch.cuda.synchronize() - - -def _wrong_variants(args): - """Plausible mis-implementations in plain torch; none of them is the op.""" - q, index_cache, table = args["q"], args["index_cache"], args["block_table"].long() - requests = args["request_indices"].long() - tpb, ratio = args["tokens_per_block"], args["compress_ratio"] - rows, _, head_dim = q.shape - columns = _columns(args) - - def scores(*, offset=ratio - 1, relu=True, scale=head_dim**-0.5, mean=False, no_table=False): - anchors = torch.arange(columns, device=DEV) * ratio + offset - logical, in_page = anchors // tpb, anchors % tpb - physical = logical.expand(rows, columns) if no_table else table[requests][:, logical] - keys = index_cache[physical.clamp_min(0), in_page.expand(rows, columns)][:, :, 0, :] - per_head = torch.einsum("rhd,rbd->rbh", q.float(), keys.float()) - per_head = torch.relu(per_head) if relu else per_head - return (per_head.mean(-1) if mean else per_head.sum(-1)) * scale - - return { - "no_relu": scores(relu=False), - "mean_over_heads": scores(mean=True), - "key_from_the_first_group_member": scores(offset=0), - "unscaled": scores(scale=1.0), - "scaled_by_one_over_head_dim": scores(scale=1.0 / head_dim), - "block_table_ignored": scores(no_table=True), - } - - -@pytest.mark.parametrize("width", ["8pages", "260pages"]) -def test_wrong_variants_fail_the_gate(width) -> None: - """The default fp32 band is only a gate if a wrong answer fails it. - - Each variant is rejected by `_compare` and lands more than 1000x outside the - band the correct result clears; measured 3e5x to 4e6x. - """ - if width == "8pages": - args = _build(_case(64, 4), seed=1601) - else: - args = _build(_case(8, 8, pages=CAPTURE_PAGES), seed=9000) - gen = torch.Generator(device=DEV).manual_seed(9001) - args["query_positions"] = torch.randint(4096, 8192, (8,), generator=gen, device=DEV).to( - torch.int32 - ) - out = op(**args) - expected = op.reference(**args) - _compare(out, expected) - visible = _visible(args) - gate = 1e-5 + 1.3e-6 * expected[visible].abs().max() - assert (out[visible] - expected[visible]).abs().max() < gate - for name, variant in _wrong_variants(args).items(): - wrong = expected.clone() - wrong[visible] = variant[visible] - assert (wrong[visible] - out[visible]).abs().max() > 1000 * gate, name - with pytest.raises(AssertionError): - _compare(out, wrong) - + assert torch.equal(graph_out[visible], op(**args)[visible]), f"replay {trial}" -def test_rows_before_a_complete_group_see_nothing() -> None: - """Positions 0..7 at compress_ratio 4: 0,0,0,1,1,1,1,2 visible columns. - - A row that has not closed its first group comes back entirely `-inf`; the - caller's top-k has to know that row can happen. - """ - args = _build(_case(8, 2, layout="arange"), seed=1251) - out = op(**args) - _compare(out, op.reference(**args)) - assert torch.isfinite(out).sum(dim=1).tolist() == [0, 0, 0, 1, 1, 1, 1, 2] - -def test_empty_row_batch() -> None: - """Zero rows: the correctly shaped empty fp32 tensor, and no launch.""" - args = _build(_case(0, 2), seed=1301) - with _watch_launches() as watcher: - out = op(**args) - assert watcher.log == [], "an empty batch must not launch a kernel" - assert out.shape == (0, _columns(args)) and out.dtype == torch.float32 - - -def test_inputs_are_not_written() -> None: - """Nothing is in place, and the result is a fresh allocation.""" - args = _build(_case(64, 4), seed=1401) - snapshots = {k: v.clone() for k, v in args.items() if torch.is_tensor(v)} - out = op(**args) - torch.cuda.synchronize() - for name, snapshot in snapshots.items(): - assert torch.equal(args[name], snapshot), f"{name} was modified" - inputs = {args[k].untyped_storage().data_ptr() for k in snapshots} - assert out.untyped_storage().data_ptr() not in inputs - - -def test_float32_operands_use_a_reduced_precision_dot() -> None: - """fp32 q and cache are accepted and do NOT hold `_compare`'s band. - - `tl.dot` rounds fp32 operands to a tensor-core format. Measured: 8.1e-3 - max abs, 0.29 max relative, on scores to 7.8 -- where bf16 operands give - 9.5e-7. That is why no cell covers fp32. The loosened band below - (rtol 2e-3, atol 2e-2) is ~2.5x the measured error, and every wrong - variant still lands more than 25x outside it. - """ - args = _build(_case(64, 4, dtype=torch.float32), seed=1501) - out = op(**args) - expected = op.reference(**args) - visible = _visible(args) - with pytest.raises(AssertionError): - _compare(out, expected) - torch.testing.assert_close(out[visible], expected[visible], rtol=2e-3, atol=2e-2) - gate = (2e-2 + 2e-3 * expected[visible].abs()).max() - for name, variant in _wrong_variants(args).items(): - assert (variant[visible] - out[visible]).abs().max() > 25 * gate, name - - -def _mutated(args, **overrides): - return {**args, **overrides} - - -_REJECTIONS = { - "two_dimensional_q": ( - lambda a: _mutated(a, q=a["q"][:, 0]), - "QSA index queries must have shape", - ), - "three_dimensional_cache": ( - lambda a: _mutated(a, index_cache=a["index_cache"][:, :, 0]), - "QSA index cache must have shape", - ), - "two_kv_head_cache": ( - lambda a: _mutated(a, index_cache=a["index_cache"].expand(-1, -1, 2, -1)), - "QSA index cache must have shape", - ), - "int64_block_table": ( - lambda a: _mutated(a, block_table=a["block_table"].long()), - "two-dimensional int32 block table", - ), - "one_dimensional_block_table": ( - lambda a: _mutated(a, block_table=a["block_table"][0]), - "two-dimensional int32 block table", - ), - "empty_block_table": ( - lambda a: _mutated(a, block_table=a["block_table"][:, :0].contiguous()), - "requires a non-empty block table", - ), - "mismatched_dtypes": ( - lambda a: _mutated(a, index_cache=a["index_cache"].float()), - "dtypes must match", - ), - "mismatched_head_dims": ( - lambda a: _mutated(a, index_cache=a["index_cache"][..., :64].contiguous()), - "dimensions do not match", - ), - "compress_ratio_1": (lambda a: _mutated(a, compress_ratio=1), "compression ratio >= 2"), - "tokens_per_block_0": (lambda a: _mutated(a, tokens_per_block=0), "compression ratio >= 2"), - "float_metadata": ( - lambda a: _mutated(a, query_positions=a["query_positions"].float()), - "must use integer storage", - ), - "short_metadata": ( - lambda a: _mutated(a, query_positions=a["query_positions"][:4]), - "one entry per row", - ), - "host_q": (lambda a: _mutated(a, q=a["q"].cpu()), "requires CUDA tensors"), -} - - -@pytest.mark.parametrize("case", list(_REJECTIONS), ids=list(_REJECTIONS)) -def test_rejects(case) -> None: - """The callable's own refusals, messages as observed -- the reason this - entry has no `is_valid`.""" - mutate, message = _REJECTIONS[case] - args = mutate(_build(_case(16, 2), seed=1701)) - with pytest.raises(ValueError, match=message): +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_index_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_index_store.py index a8c3b5b8fd53..d6d7c5b03b47 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_index_store.py @@ -7,9 +7,7 @@ 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. Both planes are pre-filled with a sentinel so a stray write shows. - -Out-of-range pages, requests and slots are deliberately absent: each is a CUDA -device-side assert that poisons the process for every test after it. +A `hole` cell appends one row that resolves to no slot. """ import zlib @@ -45,7 +43,7 @@ def _compare(out, ref) -> None: DEV = "cuda" KEY_FILL = -7.0 POS_FILL = -7 -ENGINE_MAX_NUM_TOKENS = 8192 +ENGINE_MAX_NUM_TOKENS = 16384 def _bits(tensor: torch.Tensor) -> torch.Tensor: @@ -53,6 +51,35 @@ def _bits(tensor: torch.Tensor) -> torch.Tensor: 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 _shared_page_planes(spec, extras: dict): + """The NVFP4-KV pool's side planes: each physical page holds the `layers` + index-K pages and then the position page, so the index cache (the last + role) and the position cache are strided views of one byte allocation. + Every role holds `KEY_FILL` and the position page `POS_FILL`; the other + roles are recorded for the untouched check.""" + layers, tpb, hkv, head_dim, axes, pages = ( + spec[k] for k in ("layers", "tokens_per_block", "hkv", "head_dim", "axes", "pages") + ) + index_bytes = tpb * hkv * head_dim * 2 + page = torch.empty(pages, layers * index_bytes + tpb * axes * 4, dtype=torch.uint8, device=DEV) + page[:, : layers * index_bytes].view(torch.bfloat16).fill_(KEY_FILL) + page[:, layers * index_bytes :].view(torch.int32).fill_(POS_FILL) + index_cache = page[:, (layers - 1) * index_bytes : layers * index_bytes].view(torch.bfloat16) + index_cache = index_cache.view(pages, tpb, hkv, head_dim) + position_cache = page[:, layers * index_bytes :].view(torch.int32).view(pages, tpb, axes) + assert index_cache.stride() == (spec["page_stride"], hkv * head_dim, head_dim, 1) + assert position_cache.stride() == spec["position_strides"] + extras["roles"] = page[:, : (layers - 1) * index_bytes] + 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.""" gen = torch.Generator(device="cpu").manual_seed(seed) @@ -65,63 +92,53 @@ 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) - used = 0 + 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 += 1 table = table.to(DEV) - if spec["table_layout"] == "column_slice": - wide = torch.full((slots, width + 4), -1, dtype=torch.int32, device=DEV) - wide[:, 2 : 2 + width] = table - table = wide[:, 2 : 2 + width] - assert not table.is_contiguous() projection = torch.randn( rows, spec["index_heads"] + hkv, head_dim, generator=gen, dtype=torch.float32 ).to(device=DEV, dtype=spec["key_dtype"]) token_k = projection[:, spec["index_heads"] :] - if spec["key_layout"] == "contiguous": - token_k = token_k.contiguous() + 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] request_indices = torch.tensor(requests, dtype=spec["request_dtype"], device=DEV) logical_positions = torch.tensor(positions, dtype=spec["position_dtype"], device=DEV) - if spec["meta_layout"] == "strided": - wide_requests = torch.zeros(rows, 2, dtype=spec["request_dtype"], device=DEV) - wide_positions = torch.zeros(rows, 3, dtype=spec["position_dtype"], device=DEV) - wide_requests[:, 0] = request_indices - wide_positions[:, 1] = logical_positions - request_indices, logical_positions = wide_requests[:, 0], wide_positions[:, 1] - assert not logical_positions.is_contiguous() # Three genuinely different axes, as an image span has: t, t + 7, t + 14. coords = logical_positions.to(torch.int32)[:, None] + 7 * torch.arange( 3, dtype=torch.int32, device=DEV ) - if spec["coords_layout"] == "transposed": - # The engine's `[3, max_num_tokens]` position buffer, transposed. - buffer = torch.full((3, ENGINE_MAX_NUM_TOKENS), -1, dtype=torch.int32, device=DEV) - buffer[:, :rows] = coords.t() - coords = buffer[:, :rows].t() - elif spec["coords_layout"] == "stride0": - coords = logical_positions[:, None].expand(-1, 3) - assert coords.stride(1) == 0 - else: - coords = coords.contiguous() + # The engine's `[3, max_num_tokens]` position buffer, transposed. + buffer = torch.full((3, ENGINE_MAX_NUM_TOKENS), -1, dtype=torch.int32, device=DEV) + buffer[:, :rows] = coords.t() + coords = buffer[:, :rows].t() extras = {} - if spec["cache_layout"] == "role_view": - pool = torch.full((pages, 2, tpb, hkv, head_dim), KEY_FILL, device=DEV).to(torch.bfloat16) - index_cache = pool[:, 0] - extras["neighbour"] = (pool[:, 1], pool[:, 1].clone()) - assert not index_cache.is_contiguous() + position_cache = None + if spec["cache_layout"] == "shared_page": + index_cache, position_cache = _shared_page_planes(spec, extras) else: index_cache = torch.full((pages, tpb, hkv, head_dim), KEY_FILL, device=DEV).to( torch.bfloat16 ) - position_cache = torch.full((pages, tpb, axes), POS_FILL, dtype=torch.int32, device=DEV) + if position_cache is None: + position_cache = torch.full((pages, tpb, axes), POS_FILL, dtype=torch.int32, device=DEV) args = dict( token_k=token_k, position_coordinates=coords, @@ -139,117 +156,32 @@ 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())) - before = args["position_cache"].clone() ref = op.reference(**args) out = op(**args) assert out[0] is args["index_cache"] and out[1] is args["position_cache"] _compare(out, ref) - if "neighbour" in extras: - view, saved = extras["neighbour"] - assert torch.equal(_bits(view), _bits(saved)), "the neighbouring role moved" - if cell.spec["axes"] > 3: - assert torch.equal(out[1][..., 3:], before[..., 3:]), "axes past the coordinates moved" - if cell.spec["coords_layout"] == "stride0": - for axis in (1, 2): - assert torch.equal(out[1][..., axis], out[1][..., 0]) - - -def _small(seed: int, **overrides): - """Four decode rows at the checkpoint's call-site layouts.""" - spec = {**op.CELLS[0].spec, "segments": ((0, 1), (7, 1), (32, 1), (40, 1)), "pages": 8} - return _build({**spec, **overrides}, seed=seed)[0] - - -def test_only_the_named_slots_change() -> None: - """The exact set of (page, slot) pairs written, independent of the reference.""" - args = _small(seed=13) - before_index = args["index_cache"].clone() - before_position = args["position_cache"].clone() - op(**args) - table = args["block_table"] - written = {(int(table[r, p // 32]), p % 32) for r, p in enumerate((0, 7, 32, 40))} - changed_index = { - (int(p), int(s)) - for p, s in torch.nonzero( - (_bits(args["index_cache"]) != _bits(before_index)).flatten(2).any(-1) + if "roles" in extras: + fill = int(_bits(torch.tensor([KEY_FILL], dtype=torch.bfloat16))[0]) + assert bool((extras["roles"].view(torch.int16) == fill).all()), ( + "another index-K role of the shared page moved" ) - } - changed_position = { - (int(p), int(s)) - for p, s in torch.nonzero((args["position_cache"] != before_position).any(-1)) - } - assert changed_index == written, f"index plane touched {changed_index}, expected {written}" - assert changed_position == written, f"position plane touched {changed_position}" - - -def _one_row(seed: int): - """Row 0 of `_small` alone, so nothing else can land where it does.""" - args = _small(seed=seed) - for name in ("token_k", "position_coordinates", "request_indices", "logical_positions"): - args[name] = args[name][:1] - return args - - -def test_coordinates_wider_than_int32_wrap_rather_than_raise() -> None: - """int64 coordinates are truncated into the int32 plane with nothing raised.""" - args = _small(seed=39) - args["position_coordinates"] = torch.full((4, 3), 2**31 + 5, dtype=torch.int64, device=DEV) - op(**args) - page = int(args["block_table"][0, 0]) - assert int(args["position_cache"][page, 0, 0]) == -2147483643 - - -def test_unallocated_page_sentinel_writes_the_last_page() -> None: - """The `-1` sentinel is not skipped: it wraps to the pool's last page.""" - args = _one_row(seed=42) - pages = args["index_cache"].shape[0] - args["block_table"][0] = -1 - before = args["index_cache"].clone() - op(**args) - assert torch.equal( - _bits(args["index_cache"][pages - 1, 0]), _bits(args["token_k"][0].to(torch.bfloat16)) - ), "the -1 row was expected on the last page" - assert torch.equal(_bits(args["index_cache"][:-1]), _bits(before[:-1])) - - -def test_negative_logical_position_lands_on_the_last_table_column() -> None: - """`-1 // 32 == -1` and `-1 % 32 == 31`: the last column's page, last slot.""" - args = _one_row(seed=44) - args["logical_positions"][0] = -1 - args["block_table"][0, -1] = 5 - op(**args) - assert torch.equal( - _bits(args["index_cache"][5, 31]), _bits(args["token_k"][0].to(torch.bfloat16)) - ) -def test_two_rows_on_one_slot_store_one_of_them() -> None: - """Nothing orders duplicate destinations; the slot holds one of the two rows.""" - args = _small(seed=47) - args["logical_positions"][1] = args["logical_positions"][0] - args["request_indices"][1] = args["request_indices"][0] - args["token_k"] = torch.stack( - [torch.full((1, 128), float(v), dtype=torch.bfloat16, device=DEV) for v in (1, 2, 3, 4)] - ) - op(**args) - stored = float(args["index_cache"][int(args["block_table"][0, 0]), 0, 0, 0]) - assert stored in (1.0, 2.0), f"a duplicate slot held {stored}, neither row's value" +#: 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("decode MTP7: 32 requests")) +_MTP7_BATCH = {key: _MTP7_32[key] for key in ("segments", "pages", "slots")} -def test_repeats_are_stream_ordered_and_idempotent() -> None: - """A read enqueued after the store sees it; repeating the call repeats it.""" - args = _small(seed=49) - op(**args) - first = (args["index_cache"].clone(), args["position_cache"].clone()) - for _ in range(7): - op(**args) # no synchronize: the clone below is on the same stream - assert torch.equal(_bits(args["index_cache"]), _bits(first[0])) - assert torch.equal(args["position_cache"], first[1]) +def _small(seed: int, **overrides): + """Four decode rows at the checkpoint's call-site layouts.""" + spec = {**_MTP7_32, "segments": ((0, 1), (7, 1), (32, 1), (40, 1)), "pages": 8} + return _build({**spec, **overrides}, seed=seed)[0] def test_cuda_graph_capture_and_replay() -> None: - """Captured once, replayed against new positions written into the same buffer.""" - args = _small(seed=52) + """Captured once at the 32 MTP7 verify blocks (256 rows), replayed against new + positions written into the same buffer.""" + args = _small(seed=52, **_MTP7_BATCH) positions = args["logical_positions"] replayed = positions.clone() positions.zero_() @@ -272,116 +204,13 @@ def test_cuda_graph_capture_and_replay() -> None: _compare((args["index_cache"], args["position_cache"]), ref) -def test_the_stored_key_is_raw_not_normalized() -> None: - """Discrimination against a Gemma-normed key, the decode kernel's form. - - A key stored normed double-normalizes every anchor `qsa_prefill_compress` - builds. Old receipt at its seeds: 3.0625 max abs apart; the floor states a - bound rather than one draw. - """ - args = _small(seed=66) - out = op(**args) - gen = torch.Generator(device=DEV).manual_seed(68) - weight = torch.randn(128, generator=gen, device=DEV) * 0.5 - x = args["token_k"].float() - normed = (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * (1 + weight)).to( - torch.bfloat16 - ) - page = int(args["block_table"][0, 0]) - stored = out[0][page, 0] - assert torch.equal(_bits(stored), _bits(args["token_k"][0])) - distance = (stored.float() - normed[0].float()).abs().max().item() - assert distance > 0.25, f"a normed key differs from the raw one by only {distance}" - - -# ---------------------------------------------------------------- the guards -# -# Each case is shown first unguarded -- accepted, and silently wrong -- then -# refused under `validating`. - - -def _unguarded_then_refused(args, match: str) -> None: - op(**args) - with pytest.raises(AssertionError, match=match): - with validating(op): - op(**args) - - -def test_guard_refuses_two_dimensional_request_indices() -> None: - args = _small(seed=60) - args["request_indices"] = args["request_indices"][:, None] - _unguarded_then_refused(args, "one-dimensional") - - -def test_guard_refuses_metadata_of_unequal_length() -> None: - """A length-1 position vector collapses every row onto one slot.""" +def test_rejects_metadata_of_unequal_length() -> None: + """A shape the kernel would otherwise read past; the entry refuses it before + launching, so `is_valid` has nothing left to guard.""" args = _small(seed=61) - args["request_indices"] = torch.zeros(4, dtype=torch.int32, device=DEV) args["logical_positions"] = args["logical_positions"][:1] - before = args["index_cache"].clone() - op(**args) - changed = (_bits(args["index_cache"]) != _bits(before)).flatten(1).any(-1).sum() - assert int(changed) == 1, "four rows were expected to land on one page" - with pytest.raises(AssertionError, match="logical_positions has"): - with validating(op): - op(**args) - - -def test_guard_refuses_a_single_token_k_row() -> None: - """One key row broadcasts into every destination slot.""" - args = _small(seed=62) - args["token_k"] = args["token_k"][:1] - op(**args) - table = args["block_table"] - last = args["index_cache"][int(table[3, 1]), 8] - assert torch.equal(_bits(last), _bits(args["token_k"][0])), "row 0 should fill row 3's slot" - with pytest.raises(AssertionError, match="token_k has"): - with validating(op): - op(**args) - - -def test_guard_refuses_a_single_coordinate_row() -> None: - args = _small(seed=63) - args["position_coordinates"] = args["position_coordinates"][:1] - _unguarded_then_refused(args, "position_coordinates has") - - -def test_guard_refuses_a_page_size_that_is_not_the_cache_s() -> None: - """Position 40 at tokens_per_block 16 is page column 2, slot 8; its slot at - the cache's 32 is column 1, slot 8 -- which the call never writes.""" - args = _small(seed=64) - args["tokens_per_block"] = 16 - # Column 2 is unallocated for both rows reaching it; point it inside the pool - # at a page that is not the truth page. - args["block_table"][2:4, 2] = args["block_table"][0, 0] - truth = int(args["block_table"][3, 1]), 8 - before = args["index_cache"][truth].clone() - op(**args) - assert torch.equal(_bits(args["index_cache"][truth]), _bits(before)) - with pytest.raises(AssertionError, match="does not match index_cache"): - with validating(op): - op(**args) - - -def test_guard_refuses_a_position_plane_with_a_different_page_size() -> None: - """Refused outright: past 16 positions the unguarded write is a device-side - assert that would poison every later test, so it is not shown.""" - args = _small(seed=65) - args["position_cache"] = torch.full((8, 16, 3), POS_FILL, dtype=torch.int32, device=DEV) - with pytest.raises(AssertionError, match="does not match position_cache"): - with validating(op): - op(**args) - - -def test_guard_refuses_a_cache_with_more_heads_than_the_key() -> None: - """A one-head key broadcasts into both heads of a two-head cache.""" - args = _small(seed=67) - args["index_cache"] = torch.full((8, 32, 2, 128), KEY_FILL, device=DEV).to(torch.bfloat16) - op(**args) - page = int(args["block_table"][0, 0]) - assert torch.equal(_bits(args["index_cache"][page, 0, 1]), _bits(args["token_k"][0, 0])), ( - "the second head was expected to receive the broadcast" - ) - with pytest.raises(AssertionError, match="index-KV heads"): - with validating(op): - op(**args) + before = (args["index_cache"].clone(), args["position_cache"].clone()) + with pytest.raises(ValueError, match="equally long"): + op(**args) + assert torch.equal(_bits(args["index_cache"]), _bits(before[0])) + assert torch.equal(args["position_cache"], before[1]) 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 index 2a1fe4cc9440..caf83f14f1c7 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.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 @@ -50,25 +50,11 @@ def _bits(tensor: torch.Tensor) -> torch.Tensor: return tensor.contiguous().view(as_int) -def _hole(spec, table, live_requests, allocated_cols): - """One row that resolves to no slot, and the table edit it needs.""" - tpb, width, slots, pages = (spec[k] for k in ("tokens_per_block", "width", "slots", "pages")) - kind = spec["hole"] - if kind == "unallocated_page": - assert int(table[0, allocated_cols[0]]) == -1 - return 0, allocated_cols[0] * tpb - if kind == "page_past_pool": - table[live_requests] = -1 - table[live_requests, 0] = pages + 7 - return live_requests, 0 - if kind == "request_past_table": - return slots, 0 - if kind == "negative_request": - return -1, 0 - if kind == "negative_position": - return 0, -1 - assert kind == "column_past_table" - return 0, width * tpb +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): @@ -91,29 +77,21 @@ def _build(spec, seed: int): 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, len(segments), allocated_cols) + 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, head_dim, generator=gen, dtype=torch.float32 + 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 :] - if spec["k_layout"] == "head_major": - k = k.transpose(1, 2).contiguous().transpose(1, 2) - v = v.transpose(1, 2).contiguous().transpose(1, 2) - assert k.stride(2) != 1 - + v = qkv[:, spec["heads_q"] + hkv : spec["heads_q"] + 2 * hkv] table = table.to(DEV) - if spec["table_layout"] == "role_slice": - padded = torch.full((slots, 2, width), -1, dtype=torch.int32, device=DEV) - padded[:, 0] = table - table = padded[:, 0] - elif spec["table_layout"] == "transposed": - table = table.t().contiguous().t() - assert table.stride(1) != 1 pool = torch.full((pages, 2, hkv, tpb, head_dim), FILL, dtype=torch.bfloat16, device=DEV).to( spec["cache_dtype"] @@ -147,24 +125,20 @@ def test_certified_cells(cell, certifying): # noqa: F811 def _decode_args(seed: int, **overrides): - spec = op.CELLS[0].spec + """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) -def test_fp8_pool_saturates_an_out_of_range_row() -> None: - """1000.0 lands as e4m3's finite maximum, matching torch's own conversion.""" - args, _ = _decode_args(seed=4, cache_dtype=torch.float8_e4m3fn) - args["k"] = torch.full_like(args["k"], 1000.0) - out = op(**args) - page, slot = int(args["block_table"][0, 0]), 0 - stored = out[0][page, :, slot, :] - assert torch.equal(_bits(stored), _bits(args["k"][0].to(torch.float8_e4m3fn))) - assert float(stored[0, 0].float()) == 448.0 +#: 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, replayed against new positions written into the same buffer.""" - args, pool = _decode_args(seed=9) + """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_() @@ -187,37 +161,6 @@ def call(): _compare((args["k_cache"], args["v_cache"]), ref) -def test_op_rejects_operands_it_cannot_serve() -> None: - """The callable's own loud errors: why `is_valid` carries only the stride guard.""" - args, _ = _decode_args(seed=5) - row, cache = args["k"], args["k_cache"] - narrow = torch.zeros(row.shape[0], row.shape[1], 64, dtype=row.dtype, device=DEV) - cases = { - "QSA K/V inputs must use BF16 or FP16": dict(k=row.float(), v=row.float()), - "QSA K/V cache must use BF16, FP16, or FP8 E4M3": dict( - k_cache=cache.float(), v_cache=cache.float() - ), - "QSA K/V pairs must use matching dtypes": dict(v=args["v"].half()), - "QSA K/V cache token dimension does not match tokens_per_block": dict(tokens_per_block=16), - "QSA K/V inputs and caches have incompatible head geometry": dict(k=narrow, v=narrow), - "QSA K/V store request metadata must use integer storage": dict( - logical_positions=args["logical_positions"].float() - ), - "QSA K/V store metadata must have one entry per input row": dict( - request_indices=args["request_indices"][:1].contiguous() - ), - "QSA K/V store requires a two-dimensional int32 block table": dict( - block_table=args["block_table"].long() - ), - "QSA paged K/V store requires CUDA tensors": { - name: (t.cpu() if isinstance(t, torch.Tensor) else t) for name, t in args.items() - }, - } - for message, override in cases.items(): - with pytest.raises(ValueError, match=message): - op(**{**args, **override}) - - def test_guard_refuses_strided_row_metadata() -> None: """`is_valid`'s one guard, after the hazard it exists for. @@ -226,44 +169,22 @@ def test_guard_refuses_strided_row_metadata() -> None: buffer is read through its base pointer: accepted, and rows land elsewhere. """ args, pool = _decode_args(seed=10) - requests, positions = args["request_indices"], args["logical_positions"] + requests = args["request_indices"] rows = requests.numel() wide_requests = torch.full((rows, 2), rows - 1, dtype=requests.dtype, device=DEV) wide_requests[:, 0] = requests - wide_positions = torch.full((rows, 2), 31, dtype=positions.dtype, device=DEV) - wide_positions[:, 0] = positions ref = op.reference(**args) + strided = {**args, "request_indices": wide_requests[:, 0]} - for override in ( - dict(request_indices=wide_requests[:, 0]), - dict(logical_positions=wide_positions[:, 0]), - ): - pool.fill_(FILL) - out = op(**{**args, **override}) # 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(**{**args, **override}) - assert torch.equal(_bits(pool), _bits(before)), "a refused call still wrote" - - -def test_int64_position_past_int32_is_truncated_not_skipped() -> None: - """Hard-limit probe: the kernel casts each position to int32. - - `2^32 + 5` should resolve to a page column far past the table and be - skipped; truncated, it is position 5, a real slot, and it is written. - """ - args, pool = _decode_args(seed=11) - args["logical_positions"][0] = 2**32 + 5 - ref = op.reference(**args) - out = op(**args) - page = int(args["block_table"][0, 0]) - assert torch.equal(_bits(out[0][page, :, 5, :]), _bits(args["k"][0].to(out[0].dtype))), ( - "position 2^32 + 5 was expected to land on slot 5 of page column 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" ) - assert not torch.equal(_bits(out[0]), _bits(ref[0])), "the reference skips that row" + + 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_kv_store_nvfp4.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_nvfp4.py new file mode 100644 index 000000000000..67f8d87f6179 --- /dev/null +++ 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_nvfp4.py @@ -0,0 +1,369 @@ +# 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. + +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]`, +K and V the two role planes of each -- and the engine's `[slots, width]` table, a +row per request slot, more slots than requests, live rows' pages drawn from a +seeded permutation, `-1` past each allocation. Both pools are 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_nvfp4 import ( # noqa: E501 + qsa_paged_kv_store_nvfp4 as op, +) +from tensorrt_llm.bindings.internal import thop + +__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" + +DEV = "cuda" +DATA_FILL, SCALE_FILL = 0xA5, 0x5A +QUANT_MODE_NVFP4_KV = 8192 + + +def _compare(out, ref) -> None: + """Bitwise over all four planes: the store is a pure function of the row's + bits, and `reference` performs the same rounding steps.""" + names = ("k_cache", "v_cache", "k_scale_cache", "v_scale_cache") + for name, got, want in zip(names, out, ref): + assert got.dtype == want.dtype and got.shape == want.shape, ( + name, + got.shape, + want.shape, + ) + differ = (got != want).nonzero() + if differ.numel(): + first = tuple(differ[0].tolist()) + raise AssertionError( + f"{name}: {differ.shape[0]} bytes differ from the reference, first at " + f"{first}: {int(got[first]):#04x} vs {int(want[first]):#04x}" + ) + + +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 _fill_rows(spec, k, v, gen): + """Row values: random with per-block magnitudes over 2^-10..2^5, plus edge blocks + (zero, -0.0, an outlier) when the cell asks for them.""" + hkv, head_dim = spec["hkv"], spec["head_dim"] + rows = k.shape[0] + for t in (k, v): + base = torch.randn(rows, hkv, head_dim // 16, 16, generator=gen) + exponent = torch.randint(-10, 6, (rows, hkv, head_dim // 16, 1), generator=gen) + t.copy_((base * torch.pow(2.0, exponent.float())).reshape(t.shape).to(DEV)) + if spec["values"] == "edge": + k[0, 0, :16] = 0.0 + v[0, 1, 16:32] = -0.0 + k[1, 1, 32:48] = torch.tensor([-0.0] * 8 + [0.0] * 8) + k[2, 0, 5] = 900.0 + v[3, 0, 200] = -900.0 + k[4, 1, 64:80] = 1e-30 # far below the smallest e4m3 scale + else: + assert spec["values"] == "randn" + + +def _pool(spec): + """The allocation the cell's planes view: data and scales of `slot_pages` role + pages per slot, every byte its sentinel.""" + pages, slot_pages = spec["pages"], spec["slot_pages"] + hkv, tpb, head_dim = spec["hkv"], spec["tokens_per_block"], spec["head_dim"] + data = torch.full( + (pages, slot_pages, hkv, tpb, head_dim // 2), DATA_FILL, dtype=torch.uint8, device=DEV + ) + scales = torch.full( + (pages, slot_pages, hkv, tpb, head_dim // 16), SCALE_FILL, dtype=torch.uint8, device=DEV + ) + return data, scales + + +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 + for role in range(spec["slot_pages"]): + if role // 2 != spec["layer"]: + assert bool((data[:, role] == DATA_FILL).all()), f"data role {role} moved" + assert bool((scales[:, role] == SCALE_FILL).all()), f"scale role {role} moved" + + +def _build(spec, seed: int, pool=None): + """Real pool, table and rows for one cell; returns the op's kwargs, the pool, + and the packed QKV rows the K/V slices view. `pool` reuses an allocation + from `_pool(spec)`.""" + 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.zeros( + rows, spec["heads_q"] + 2 * hkv, head_dim, dtype=spec["row_dtype"], device=DEV + ) + k = qkv[:, spec["heads_q"] : spec["heads_q"] + hkv] + v = qkv[:, spec["heads_q"] + hkv :] + qkv[:, : spec["heads_q"]] = torch.randn(rows, spec["heads_q"], head_dim, generator=gen).to( + device=DEV, dtype=spec["row_dtype"] + ) + _fill_rows(spec, k, v, gen) + + layer = spec["layer"] + data, scales = _pool(spec) if pool is None else pool + args = dict( + k=k, + v=v, + k_cache=data[:, 2 * layer], + v_cache=data[:, 2 * layer + 1], + k_scale_cache=scales[:, 2 * layer], + v_scale_cache=scales[:, 2 * layer + 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.to(DEV), + tokens_per_block=tpb, + kv_scale_orig_quant=torch.tensor( + [1.0, spec["global_k"], spec["global_v"]], dtype=torch.float32, device=DEV + ), + ) + return args, (data, scales), qkv + + +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).""" + data, scales = pool + tpb = 32 + hq, hkv, d = 24, 2, 256 + ns = len(requests) + width = table.shape[1] + offsets = torch.zeros(1, ns, 2, width, dtype=torch.int32, device=DEV) + for row, r in enumerate(requests): + live = table[r].clamp_min(0) + offsets[0, row, 0] = 2 * live + offsets[0, row, 1] = 2 * live + 1 + pointers = torch.zeros(1, 2, 2, dtype=torch.int64) + pointers[0, 0, 0], pointers[0, 0, 1] = data.data_ptr(), scales.data_ptr() + kv = [c + n for c, n in zip(cached, new)] + ctx = [n if i < num_contexts else kv[i] for i, n in enumerate(new)] + scale_quant_orig = scale_orig_quant.reciprocal() + rows = qkv.reshape(qkv.shape[0], -1) + thop.attention( + q=rows, k=None, v=None, + output=torch.empty(rows.shape[0], hq * d, dtype=torch.float8_e4m3fn, device=DEV), + output_sf=None, workspace_=torch.empty(0, dtype=torch.int8, device=DEV), + sequence_length=torch.tensor(kv, dtype=torch.int32, device=DEV), + host_past_key_value_lengths=torch.tensor(kv, dtype=torch.int32), + host_total_kv_lens=torch.tensor( + [sum(kv[:num_contexts]), sum(kv[num_contexts:])], dtype=torch.int32 + ), + context_lengths=torch.tensor(ctx, dtype=torch.int32, device=DEV), + host_context_lengths=torch.tensor(ctx, dtype=torch.int32), + host_request_types=torch.tensor( + [0 if i < num_contexts else 1 for i in range(ns)], dtype=torch.int32 + ), + max_context_q_len_override=None, kv_cache_block_offsets=offsets, + host_kv_cache_pool_pointers=pointers, + host_kv_cache_pool_mapping=torch.zeros(1, 2, dtype=torch.int32), + cache_indirection=None, kv_scale_orig_quant=scale_orig_quant, kv_scale_quant_orig=scale_quant_orig, + out_scale=torch.ones(1, dtype=torch.float32, device=DEV), rotary_inv_freq=None, + rotary_cos_sin=None, latent_cache=None, q_pe=None, block_ids_per_seq=None, + attention_sinks=None, is_fused_qkv=True, update_kv_cache=True, + predicted_tokens_per_seq=p, local_layer_idx=0, num_heads=hq, num_kv_heads=hkv, + head_size=d, tokens_per_block=tpb, max_num_requests=ns, + max_context_length=width * tpb, max_seq_len=width * tpb, + attention_window_size=width * tpb, beam_width=1, mask_type=1, + quant_mode=QUANT_MODE_NVFP4_KV, q_scaling=1.0, position_embedding_type=0, 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, + use_paged_context_fmha=True, attention_input_type=0, is_mla_enable=False, + chunked_prefill_buffer_batch_size=1, 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, + mrope_rotary_cos_sin=None, mrope_position_deltas=None, helix_position_offsets=None, + helix_is_inactive_rank=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_for_cpp=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, num_contexts=num_contexts, + num_ctx_tokens=sum(new[:num_contexts]), + ) # fmt: skip + torch.cuda.synchronize() + + +def _thop_equality(spec, seed: int) -> None: + """Steps of a context, a chunked continuation, decode P=1 and MTP verify P=3 over + three requests, each written once by the attention op's append into one pool and + once by this op into another, at the cell's second-level scales: the two pools stay + bitwise equal, and this op's output equals its reference at every step.""" + base, ours, _ = _build(spec, seed=seed) + theirs = _snapshot(ours) + gen = torch.Generator(device="cpu").manual_seed(seed + 1) + table = base["block_table"] + cached = [0, 0, 0] + for new, nc, p in ( + ([70, 33, 5], 3, 1), + ([41, 64, 27], 3, 1), + ([1, 1, 1], 0, 1), + ([3] * 3, 0, 3), + ): + rows = sum(new) + qkv = torch.randn(rows, 28, 256, generator=gen) + qkv *= torch.pow(2.0, torch.randint(-6, 4, (rows, 28, 1), generator=gen).float()) + qkv = qkv.to(device=DEV, dtype=torch.bfloat16) + requests = [r for r, n in enumerate(new) for _ in range(n)] + positions = [c + i for c, n in zip(cached, new) for i in range(n)] + _thop_append(qkv, theirs, table, [0, 1, 2], new, cached, nc, p, base["kv_scale_orig_quant"]) + args = { + **base, + "k": qkv[:, 24:26], + "v": qkv[:, 26:28], + "request_indices": torch.tensor(requests, dtype=torch.int32, device=DEV), + "logical_positions": torch.tensor(positions, dtype=torch.int64, device=DEV), + } + ref = op.reference(**args) + out = op(**args) + _compare(out, ref) + for plane, (a, b) in enumerate(zip(ours, theirs)): + assert torch.equal(a, b), ( + f"step {new}: {('data', 'scales')[plane]} differ from the attention op's " + f"append in {int((a != b).sum())} bytes" + ) + 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 test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + seed = zlib.crc32(cell.why.encode()) + if cell.spec["thop_append"]: + _thop_equality(cell.spec, seed=seed) + return + args, pool, _ = _build(cell.spec, seed=seed) + before = _snapshot(pool) + 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"]), ( + "a cell with rows must write, an empty one must not" + ) + _assert_other_roles_untouched(cell.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.""" + prefill = next(c.spec for c in op.CELLS if c.why.startswith("prefill:")) + spec = {**prefill, "segments": tuple((29 + 61 * i, 8) for i in range(8)), "pages": 256} + _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.""" + args, _, _ = _decode_args(seed=5) + strided = torch.stack([args["request_indices"]] * 2, 1)[:, 0] + with pytest.raises(ValueError, match="request_indices must be contiguous"): + with validating(op): + op(**{**args, "request_indices": strided}) 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 index f7d69fa1df94..28ed2deb6c8a 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.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 @@ -44,7 +44,7 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: 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) are each more than 0.1 off and rejected. + 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) @@ -54,7 +54,6 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: DEV = "cuda" -ENGINE_MAX_SEQ_LEN = 8192 _KERNELS = { "_qsa_paged_sparse_gqa_kernel": "prefill", @@ -89,26 +88,6 @@ def launches(monkeypatch): # -------------------------------------------------------------------------- -def _segments(spec): - """`(first position, rows)` per request in row order, and a generation mask.""" - if "mixed" not in spec: - segments = list(spec["segments"]) - return segments, torch.zeros(sum(n for _, n in segments), dtype=torch.bool, device=DEV) - requests, context_lens = spec["mixed"] - # Chunk 0 starts at position 0, so one launch spans rows that see one token - # and rows at the full budget; later chunks are chunked-prefill continuations. - segments = [ - (0 if i == 0 else (907 * i) % (ENGINE_MAX_SEQ_LEN - n + 1), n) - for i, n in enumerate(context_lens) - ] - # MTP: two rows per generation request, positions across the budget both ways. - segments += [(100 + (499 * g) % 8000, 2) for g in range(requests - len(context_lens))] - generation = torch.cat( - [torch.full((n,), i >= len(context_lens), device=DEV) for i, (_, n) in enumerate(segments)] - ) - return segments, generation - - def _selected_tokens(positions, seed, compress_ratio, token_topk): """QSA's expansion layout, built in one pass. @@ -150,17 +129,14 @@ def _build(spec, seed): spec["head_dim"], spec["q_heads"], ) - segments, generation = _segments(spec) + 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)] - # `min_seq_len` maps pages past the last query row's position, so a probe - # can plant a token that is resident but not causally visible. - floor = spec.get("min_seq_len", 0) - mapped = [math.ceil(max(s + n, floor) / tpb) for s, n in 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) @@ -172,25 +148,21 @@ def _build(spec, seed): 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] - if spec["cache_layout"] == "contiguous": - k_cache, v_cache = k_cache.contiguous(), v_cache.contiguous() q_dtype = spec["q_dtype"] - if spec["q_layout"] == "qkv_slice": - qkv = torch.randn(rows, q_heads + 2 * kv, d, generator=gen, device=DEV).to(q_dtype) - q = qkv[:, :q_heads] - else: - q = torch.randn(rows, q_heads, d, generator=gen, device=DEV).to(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"]) - if "width" in spec: - selected = selected[:, : spec["width"]].contiguous() request_indices = slot_of.to(DEV)[req].to(spec["request_dtype"]) - if spec.get("holes"): - selected[1] = -1 # an all-padding row - request_indices[2] = slots + 67 # a request the table does not have - block_table[:, 0] = -1 # every request's first page unmapped + 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, @@ -204,7 +176,9 @@ def _build(spec, seed): ) if spec["prefill_pair"]: args.update(query_positions=positions, compress_ratio=spec["compress_ratio"]) - return args, generation + if gate is not None: + args.update(output_gate=gate) + return args def _bound(positions, spec): @@ -224,9 +198,9 @@ def _is_prefix_packed(table): def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811 certifying(op) spec = cell.spec - args, _ = _build(spec, seed=sum(map(ord, cell.why))) + args = _build(spec, seed=sum(map(ord, cell.why))) rows = args["q"].shape[0] - if spec["prefill_pair"] and not spec.get("holes"): + 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"] @@ -243,11 +217,7 @@ def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811 _compare(out, op.reference(**args)) -# -------------------------------------------------------------------------- -# the prefill kernel's bound -# -------------------------------------------------------------------------- - -#: The first cell is this checkpoint's own decode geometry; the probes below +#: 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) @@ -257,234 +227,24 @@ def _with(**overrides): return {**CKPT, **overrides} -def _first_unscanned_column(positions, tile=64): - """First column of each row the prefill kernel never reads (bound rounded - up to its 64-wide tile); none when that lands at or past the row's width.""" - final_topk = CKPT["token_topk"] + CKPT["compress_ratio"] - 1 - bound = _bound(positions, CKPT).clamp(max=final_topk) - return ((bound + tile - 1) // tile) * tile, final_topk - - -def _plant_past_the_bound(args): - """Put each row's own newest token in its first unscanned column.""" - positions = args["query_positions"] - first, final_topk = _first_unscanned_column(positions) - gap_rows = (first < final_topk).nonzero().flatten() - assert gap_rows.numel() > 0, "no row has an unscanned column to plant into" - planted = args["selected_tokens"].clone() - planted[gap_rows, first[gap_rows]] = positions[gap_rows].to(torch.int32) - return planted, gap_rows - - -def _sixty_four_token_rows(rows, seed): - """Rows seeing 64 tokens: the bound is exactly one tile and column 64 the - first unread one. It is planted with token 200 -- resident (each request - maps 256 tokens) and not otherwise selected.""" - spec = _with(segments=((63, 1),) * rows, slots=rows + 8, min_seq_len=256, prefill_pair=True) - args, _ = _build(spec, seed) - packed = args["selected_tokens"] - planted = packed.clone() - planted[:, 64] = 200 - return {**args, "selected_tokens": planted}, packed - - -def test_prefill_kernel_stops_at_the_causal_prefix(launches) -> None: - """An index past the bound is silently not attended to. - - The op matches the reference computed WITHOUT the planted index, and the - planted index moves the correct answer by more than 0.1, so the first - comparison is not vacuous. - """ - args, packed = _sixty_four_token_rows(300, seed=7000) - out = op(**args) - assert [name for name, _ in launches] == ["prefill"] - without_it = op.reference(**{**args, "selected_tokens": packed}) - _compare(out, without_it) - with_it = op.reference(**args) - assert (with_it.float() - without_it.float()).abs().max().item() > 0.1 - - -def test_split_k_reads_the_whole_row(launches) -> None: - """Below the program threshold the same operands attend to every column.""" - args, packed = _sixty_four_token_rows(100, seed=7100) - out = op(**args) - assert [name for name, _ in launches][0] == "splitk" - _compare(out, op.reference(**args)) - without_it = op.reference(**{**args, "selected_tokens": packed}) - assert (out.float() - without_it.float()).abs().max().item() > 0.1 - - -def test_prefill_bound_at_8191_rows() -> None: - """The bound is the same at 8191 rows, and reaches much further: a row - seeing 64 tokens is scanned for 64 of 2051 columns.""" - args, _ = _build(_with(segments=((0, 8191),), slots=8, prefill_pair=True), seed=11000) - planted, _ = _plant_past_the_bound(args) - out = op(**{**args, "selected_tokens": planted}) - without_it = op.reference(**args) - _compare(out, without_it) - with_it = op.reference(**{**args, "selected_tokens": planted}) - assert (with_it.float() - without_it.float()).abs().max().item() > 0.1 - - -def test_prefill_bound_reaches_the_generation_rows() -> None: - """A mixed launch's generation rows ride the prefill kernel, bound and all.""" - args, generation = _build( - _with(mixed=(32, (900, 200, 150, 28)), slots=32, prefill_pair=True), seed=22000 - ) - planted, gap_rows = _plant_past_the_bound(args) - assert bool(generation[gap_rows].any()), "no generation row has an unscanned column" - out = op(**{**args, "selected_tokens": planted}) - without_it = op.reference(**args) - _compare(out, without_it) - with_it = op.reference(**{**args, "selected_tokens": planted}) - generation_gaps = gap_rows[generation[gap_rows]] - moved = (with_it[generation_gaps].float() - without_it[generation_gaps].float()).abs().max() - assert moved.item() > 0.1, f"the planted indices must matter to the generation rows: {moved}" - - -# -------------------------------------------------------------------------- -# semantics a cell cannot state -# -------------------------------------------------------------------------- - - -def test_repeated_index_is_attended_twice() -> None: - """Selected indices are not deduplicated; each column is its own key.""" - args, _ = _build(_with(segments=((900, 1),), prefill_pair=False), seed=8700) - selected = torch.full((1, 64), -1, dtype=torch.int32, device=DEV) - selected[0, :3] = torch.tensor([5, 5, 40], dtype=torch.int32, device=DEV) - twice = op(**{**args, "selected_tokens": selected}) - _compare(twice, op.reference(**{**args, "selected_tokens": selected})) - once = selected.clone() - once[0, 1] = -1 - single = op.reference(**{**args, "selected_tokens": once}) - assert (twice.float() - single.float()).abs().max().item() > 1e-2 - - def test_calls_are_bitwise_deterministic() -> None: """Including a 32-split launch whose merge is a second kernel.""" - args, _ = _build(_with(segments=tuple((4000 + 37 * i, 2) for i in range(4))), seed=8800) + 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)) -def test_zero_rows_return_an_empty_output() -> None: - args, _ = _build(_with(segments=((900, 1),) * 4), seed=8400) - for name in ("q", "selected_tokens", "request_indices"): - args[name] = args[name][:0] - out = op(**args) - assert out.shape == (0, CKPT["q_heads"], CKPT["head_dim"]) and out.dtype == torch.bfloat16 - _compare(out, op.reference(**args)) - - -def test_gate_discriminates_wrong_variants() -> None: - """Six references, each wrong one way, all rejected by `_compare`. - - Built by feeding the entry's own reference a mutated input, so each differs - from the correct one in exactly the way it names. - """ - args, _ = _build(_with(segments=((2032, 272),), prefill_pair=True), seed=9000) - out = op(**args) - _compare(out, op.reference(**args)) - tpb, groups = CKPT["tokens_per_block"], CKPT["q_heads"] // CKPT["kv_heads"] - selected, table = args["selected_tokens"], args["block_table"] - - # Head h reads KV head h % Hkv instead of h // groups. - interleave = torch.arange(CKPT["q_heads"], device=DEV) - interleave = (interleave % groups) * CKPT["kv_heads"] + interleave // groups - identity = torch.arange(table.shape[1], dtype=torch.int32, device=DEV).expand_as(table) - tokens = selected.clamp_min(0) - transposed = (tokens % tpb) * tpb + (tokens // tpb).clamp(max=tpb - 1) - variants = { - "no_softmax_scale": dict(softmax_scale=1.0), - "interleaved_heads": dict(q=args["q"][:, interleave]), - "padding_attended": dict(selected_tokens=tokens), - "identity_pages": dict(block_table=identity.contiguous()), - "k_and_v_swapped": dict(k_cache=args["v_cache"], v_cache=args["k_cache"]), - "page_and_slot_transposed": dict( - selected_tokens=torch.where(selected >= 0, transposed, selected) - ), - } - for name, override in variants.items(): - wrong = op.reference(**{**args, **override}) - if name == "interleaved_heads": - wrong = wrong[:, torch.argsort(interleave)] - with pytest.raises(AssertionError): - _compare(out, wrong) - assert (out.float() - wrong.float()).abs().max().item() > 0.1, name - - # -------------------------------------------------------------------------- # what the op refuses # -------------------------------------------------------------------------- -_REJECTIONS = { - "float32_query": (lambda a: {**a, "q": a["q"].float()}, "QSA queries must use BF16 or FP16"), - "flat_query": ( - lambda a: {**a, "q": a["q"].reshape(a["q"].shape[0], -1)}, - "QSA queries must have shape [rows, query_heads, head_dim]", - ), - "int64_selection": ( - lambda a: {**a, "selected_tokens": a["selected_tokens"].long()}, - "QSA selected tokens must be a two-dimensional int32 tensor", - ), - "int64_block_table": ( - lambda a: {**a, "block_table": a["block_table"].long()}, - "QSA sparse GQA requires a two-dimensional int32 block table", - ), - "wrong_tokens_per_block": ( - lambda a: {**a, "tokens_per_block": a["tokens_per_block"] // 2}, - "QSA cache token dimension must match positive tokens_per_block", - ), - "negative_scale": ( - lambda a: {**a, "softmax_scale": -1.0}, - "QSA sparse GQA softmax scale must be finite and positive", - ), - "strided_request_indices": ( - lambda a: {**a, "request_indices": torch.zeros(8, dtype=torch.int32, device=DEV)[::2]}, - "QSA request indices must be contiguous", - ), - "float32_cache": ( - lambda a: {**a, "k_cache": a["k_cache"].float(), "v_cache": a["v_cache"].float()}, - "QSA K/V cache must use BF16, FP16, or FP8 E4M3", - ), - "positions_without_ratio": ( - lambda a: {**a, "query_positions": torch.zeros(4, dtype=torch.int64, device=DEV)}, - "QSA causal-prefix bounds require positions and compression ratio", - ), - "head_dim_192": ( - lambda a: { - **a, - "q": a["q"][..., :192].contiguous(), - "k_cache": a["k_cache"][..., :192], - "v_cache": a["v_cache"][..., :192], - }, - "QSA sparse GQA requires a power-of-two head dimension", - ), - "host_query": ( - lambda a: {**a, "q": a["q"].cpu()}, - "QSA paged sparse GQA requires CUDA tensors", - ), -} - -@pytest.mark.parametrize("case", list(_REJECTIONS), ids=list(_REJECTIONS)) -def test_rejects(case) -> None: - """The op's own loud refusals, verbatim -- the reason `is_valid` is empty. - 192 is a legal head size elsewhere in this engine and not here.""" - args, _ = _build(_with(segments=((600, 1),) * 4), seed=9100) +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() - mutate, message = _REJECTIONS[case] + 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(**mutate(args)) - - -def test_selection_width_is_checked_on_the_prefill_kernel_only() -> None: - """A width the compress ratio cannot produce: refused by the prefill - kernel, accepted (the pair is unused) by split-K.""" - args, _ = _build(_with(segments=((200, 272),), prefill_pair=True), seed=9300) - args["selected_tokens"] = args["selected_tokens"][:, :64].contiguous() - with pytest.raises(ValueError, match="incompatible with its compression ratio"): op(**args) - - args, _ = _build(_with(segments=((200, 100),), prefill_pair=True), seed=9400) - args["selected_tokens"] = args["selected_tokens"][:, :64].contiguous() - _compare(op(**args), 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_sparse_gqa_decode.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_decode.py new file mode 100644 index 000000000000..7d24c80e8471 --- /dev/null +++ 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_decode.py @@ -0,0 +1,118 @@ +# 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_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.""" + +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 + 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 + qsa_topk_expand_decode, +) + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 — path above; fixture by name + assert_within_ulp, + certifying, + validating, +) + +assert torch.cuda.is_available(), "qsa_paged_sparse_gqa_decode requires a CUDA device" + +DEV = "cuda" +Q_HEADS, KV_HEADS, D, TPB, CR, TOKEN_TOPK, PAGES, SLOTS = 24, 2, 256, 32, 4, 2048, 320, 6 + + +def _build(spec, seed): + """Paged K/V for one decode row per position, the selection QSA's expansion emits + (512 random complete groups, then the incomplete group, then -1), q and the gate as + column views of the target's 55-head fused projection row.""" + g = torch.Generator(device=DEV).manual_seed(seed) + pos = torch.tensor(spec["pos"], dtype=torch.int64, device=DEV) + rows = pos.numel() + pages_per = [int(p) // TPB + 1 for p in spec["pos"]] + pool_pages = sum(pages_per) + 5 + page_ids = torch.randperm(pool_pages, generator=g, device=DEV).to(torch.int32) + block_table = torch.full((SLOTS, PAGES), -1, dtype=torch.int32, device=DEV) + slot_of = list(range(1, rows + 1)) + cursor = 0 + for slot, count in zip(slot_of, pages_per): + block_table[slot, :count] = page_ids[cursor : cursor + count] + cursor += count + pool = torch.randn(pool_pages, 2, KV_HEADS, TPB, D, generator=g, device=DEV).to(torch.bfloat16) + k_cache, v_cache = pool[:, 0], pool[:, 1] + fused = torch.randn(rows, 55 * D, generator=g, device=DEV).to(torch.bfloat16) + q = fused.view(rows, 55, D)[:, :Q_HEADS] + gate = fused[:, 28 * D : (28 + Q_HEADS) * D] if spec["gated"] else None + block_topk = TOKEN_TOPK // CR + selected = torch.full((rows, TOKEN_TOPK + CR - 1), -1, dtype=torch.int32, device=DEV) + for r, p in enumerate(spec["pos"]): + visible = p + 1 + groups = visible // CR + chosen = torch.randperm(groups, generator=g, device=DEV)[: min(groups, block_topk)] + toks = (chosen[:, None] * CR + torch.arange(CR, device=DEV)).reshape(-1) + tail = torch.arange(groups * CR, visible, device=DEV) + row = torch.cat([toks, tail]).to(torch.int32) + selected[r, : row.numel()] = row + if spec.get("holes"): + selected[r, 1 : row.numel() : 2] = -1 + request_indices = torch.tensor(slot_of, dtype=torch.int32, device=DEV) + return [q, k_cache, v_cache, block_table, selected, request_indices, TPB, D**-0.5, 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 — fixture by name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31)) + out = op(*args) + ref = op.reference(*args) + assert out.dtype == torch.bfloat16 and out.shape == ref.shape and out.is_contiguous() + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) + + +def _replays_like_eager(step, refresh) -> None: + 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): + refresh() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(graph_out, step()), f"replay {trial}" + + +def test_reads_the_selection_after_the_wait_under_graph_replay() -> None: + """The target's top-k runs inside the graph right before the op: a PDL producer that + releases the op right after its own wait, before it writes the selection. A read of the + selection before the op's wait would see the previous replay's.""" + args = _build(dict(pos=(8700,), gated=True), seed=23) + pos = torch.tensor([8700], dtype=torch.int64, device=DEV) + ends = ((pos + 1) // CR).to(torch.int32) + seqlen = (pos + 1).to(torch.int32) + scores = torch.rand(1, 2304, device=DEV) + topk = torch.empty(1, TOKEN_TOPK // CR, dtype=torch.int32, device=DEV) + + def step(): + sel = qsa_topk_expand_decode(scores, ends, pos, seqlen, topk, CR, TOKEN_TOPK) + 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_paged_sparse_gqa_nvfp4.py new file mode 100644 index 000000000000..ab4313b61e21 --- /dev/null +++ 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_nvfp4.py @@ -0,0 +1,321 @@ +# 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. + +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, +D/16]` -- 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, packed at the front. + +Beside `_compare` against the reference, every cell runs QSA's bf16 sparse GQA +over a bf16 pool that holds the dequantized values and holds the two outputs to +the same band: the NVFP4 kernels are the bf16 ones with a dequantizing load. +""" + +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 + 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 + 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__ = ["../../../.."] +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" + +DEV = "cuda" +ENGINE_MAX_SEQ_LEN = 8192 +#: The e4m3fn NaN bit pattern. +E4M3_NAN = 0x7F + +_KERNELS = { + "_nvfp4_sparse_gqa_prefill_kernel": "prefill", + "_nvfp4_sparse_gqa_splitk_kernel": "splitk", + "_nvfp4_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(): + monkeypatch.setattr(_entry, name, _LaunchRecorder(getattr(_entry, name), tag, log)) + return log + + +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. + + After the dequantize the kernels are the bf16 sparse GQA kernels, whose band + this is (they round the online-softmax probabilities to the value dtype + before P.V); the dequantize itself is exact at unit scale and rounded once, + as `reference` rounds it, otherwise. + """ + 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) + + +def _selected_tokens(positions, seed, compress_ratio, token_topk): + """QSA's expansion layout: chosen complete groups, the trailing partial group, -1.""" + 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 _random_pool(pages, kv, tpb, d, gen): + """Pool bytes: any e2m1 nibble, and positive e4m3 scales from 2^-6 to 0.5.""" + data = torch.randint(0, 256, (pages, 2, kv, tpb, d // 2), generator=gen, device=DEV) + codes = torch.randint(0x08, 0x31, (pages, 2, kv, tpb, d // 16), generator=gen, device=DEV) + return data.to(torch.uint8), codes.to(torch.uint8) + + +def _build(spec, seed): + """Real paged state for one spec, keyed by the op's parameter names.""" + gen = torch.Generator(device=DEV).manual_seed(seed) + tpb, kv, d, q_heads = (spec[k] for k in ("tokens_per_block", "kv_heads", "head_dim", "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 not mapped or max(mapped) <= spec["pages"] + needed = sum(mapped) + 7 + slot_pages, layer = spec["slot_pages"], spec["layer"] + planes = slice(2 * layer, 2 * layer + 2) + # The requests sit on the pool's last `needed` pages, and every scale outside them is + # NaN, so a page address that strayed could not come back finite. + 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) + data = torch.zeros(pool_pages, slot_pages, kv, tpb, d // 2, dtype=torch.uint8, device=DEV) + scales = torch.full( + (pool_pages, slot_pages, kv, tpb, d // 16), E4M3_NAN, dtype=torch.uint8, device=DEV + ) + data[first:, planes], scales[first:, planes] = _random_pool(needed, kv, tpb, d, gen) + 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 + + assert spec["q_layout"] == "qkv_slice", spec["q_layout"] + 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"]) + request_indices = slot_of.to(DEV)[req].to(spec["request_dtype"]) + + args = dict( + q=q, + k_cache=data[:, 2 * layer], + v_cache=data[:, 2 * layer + 1], + k_scale_cache=scales[:, 2 * layer], + v_scale_cache=scales[:, 2 * layer + 1], + kv_scale_quant_orig=torch.tensor( + [1.0, spec["global_k"], spec["global_v"]], dtype=torch.float32, device=DEV + ), + 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"]) + return args + + +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) + 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] + 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") + } + twin.update(softmax_scale=args["softmax_scale"], 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) + + +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) + + +_OPERANDS = ( + "q", + "k_cache", + "v_cache", + "k_scale_cache", + "v_scale_cache", + "block_table", + "selected_tokens", + "request_indices", +) + + +def _capture_sweep(spec, seed, launches): + """The decode site under a CUDA graph, at each G for each rows-per-request: + capture, write a second state's operands in place, replay, and require the + eager call on that second state bitwise.""" + sweep = spec["capture"] + for per_request in sweep["rows_per_request"]: + for groups in sweep["groups"]: + _capture_once( + _decode_step(spec, groups, per_request), + seed + 100 * per_request + groups, + launches, + where=f"G={groups} x {per_request} rows", + ) + + +def _start(request: int, rows: int) -> int: + """A request's first position, across the budget.""" + return min(100 + (499 * request) % 8000, ENGINE_MAX_SEQ_LEN - rows) + + +def _decode_step(spec, groups, per_request): + """The decode site's step: `groups` requests x `per_request` rows.""" + segments = tuple((_start(g, per_request), per_request) for g in range(groups)) + return {**spec, "segments": segments} + + +def _capture_once(shape, seed, launches, where) -> None: + """One call of `shape` captured, a second state's operands written in place, + replayed: the replay must be bitwise the eager call on the second state.""" + where = f"{where}, [{shape['slots']}, {shape['pages']}] table" + args = _build(shape, seed) + fresh = _build(shape, seed + 50) + launches.clear() + op(**args) + assert [name for name, _ in launches] == list(shape["launches"]), where + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + try: + with torch.cuda.graph(graph): + captured = op(**args) + for name in _OPERANDS: + args[name].copy_(fresh[name]) + graph.replay() + torch.cuda.synchronize() + eager = op(**fresh) + assert torch.equal(captured, eager), f"replay differs from eager at {where}" + _compare(captured, op.reference(**fresh)) + finally: + del graph + torch.cuda.synchronize() + + +@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 + if "capture" in spec: + _capture_sweep(spec, seed=sum(map(ord, cell.why)), launches=launches) + return + args = _build(spec, seed=sum(map(ord, cell.why))) + rows = args["q"].shape[0] + if spec["prefill_pair"]: + 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)) + # QSA's bf16 kernels over the same values: inside the same band, not bitwise + # (the dequantizing load changes the compiled kernel). + _compare(out, _bf16_twin(args)) + + +CKPT = op.CELLS[0].spec + + +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_prefill_compress.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_prefill_compress.py index 05b1b5a03d10..cac393b9865c 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_prefill_compress.py @@ -28,24 +28,37 @@ assert torch.cuda.is_available(), "qsa_prefill_compress requires a CUDA device" -def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: - """Bitwise over the whole cache. - - Measured on sm_103: the kernel matches this fp32 reference on every - element of every cell, so there is no tolerance to derive. It has to - be bitwise -- a pool accumulated in bf16 instead of fp32 lands only - 0.03125 away, inside any bf16 band. Comparing every slot, not just the - anchors, is the untouched-slot lock. +def _compare(out: torch.Tensor, ref: torch.Tensor, flips: int = 0) -> None: + """Bitwise over the whole cache, except for `flips` elements one bf16 step off. + + Measured on sm_103: the kernel matches this fp32 reference on every element of every + cell but one, so there is no tolerance to derive. It has to be bitwise -- a pool + accumulated in bf16 instead of fp32 lands only 0.03125 away, inside any bf16 band. + Comparing every slot, not just the anchors, is the untouched-slot lock. + + The exception is the cell of 32 MTP7 verify blocks (64 anchors, `rounding_flips`): the + 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. """ assert out.dtype == ref.dtype, (out.dtype, ref.dtype) assert out.shape == ref.shape, (out.shape, ref.shape) - torch.testing.assert_close(out.view(torch.int16), ref.view(torch.int16), rtol=0, atol=0) + if flips == 0: + torch.testing.assert_close(out.view(torch.int16), ref.view(torch.int16), rtol=0, atol=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" + ) DEV = "cuda" ROTARY_DIM = 64 THETA = 1.0e7 MAX_POSITIONS = 16384 +SENTINEL = 0x3C3C def _cos_sin() -> torch.Tensor: @@ -77,37 +90,36 @@ def _build(spec, seed: int): table[r, :cols] = order[used : used + cols] used += cols p = torch.arange(cols * tpb).view(cols, tpb) - if spec["distinct_axes"]: - coords = (p, (3 * p + 5 * r) % 500, (5 * p + 11 * r) % 500) - else: - coords = (p, p, p) + 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" 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"]) - if spec["duplicate"]: - logical, request = torch.cat([logical, logical]), torch.cat([request, request]) table, index_cache, position_cache = table.to(DEV), index_cache.to(DEV), position_cache.to(DEV) extras = {} - if spec["layout"] == "strided": - wide_positions = torch.randint(0, 500, (pages, tpb, 6), dtype=torch.int32, device=DEV) - wide_positions[..., ::2] = position_cache - position_cache = wide_positions[..., ::2] - wide_table = torch.randint(0, pages, (slots, 2 * width), dtype=torch.int32, device=DEV) - wide_table[:, ::2] = table - table = wide_table[:, ::2] - pool = torch.randn(2 * pages, tpb, 1, head_dim, device=DEV).to(torch.bfloat16) - pool[pages:] = index_cache - index_cache = pool[pages:] - extras["neighbour"] = (pool[:pages], pool[:pages].clone()) - elif spec["layout"] == "head_dim_strided": - wide = torch.randn(pages, tpb, 1, 2 * head_dim, device=DEV).to(torch.bfloat16) - wide[..., ::2] = index_cache - index_cache = wide[..., ::2] - extras["neighbour"] = (wide[..., 1::2], wide[..., 1::2].clone()) + 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. + roles = spec["roles"] + index_bytes = tpb * head_dim * 2 + page = torch.empty(pages, roles * index_bytes + tpb * 3 * 4, dtype=torch.uint8, device=DEV) + others = page[:, : (roles - 1) * index_bytes].view(torch.int16) + others.fill_(SENTINEL) + shared_index = page[:, (roles - 1) * index_bytes : roles * index_bytes] + shared_index = shared_index.view(torch.bfloat16).view(pages, tpb, 1, head_dim) + shared_position = page[:, roles * index_bytes :].view(torch.int32).view(pages, tpb, 3) + shared_index.copy_(index_cache) + shared_position.copy_(position_cache) + assert shared_index.stride(0) == roles * tpb * head_dim + tpb * 3 * 2 + assert shared_position.stride() == (roles * tpb * head_dim // 2 + tpb * 3, 3, 1) + index_cache, position_cache = shared_index, shared_position + extras["neighbour"] = (others, others.clone()) + extras["position_page"] = (position_cache, position_cache.clone()) + else: + assert spec["layout"] == "plain", spec["layout"] args = dict( logical_positions=logical.to(DEV), @@ -140,41 +152,33 @@ def test_certified_cells(cell, certifying): # noqa: F811 ref = op.reference(**args) out = op(**args) assert out is args["index_cache"] - _compare(out, ref) + _compare(out, ref, cell.spec.get("rounding_flips", 0)) changed = (out.view(torch.int16) != before.view(torch.int16)).any(-1).sum() assert int(changed) == _anchors(args), "exactly the anchor slots are written" if "neighbour" in extras: view, saved = extras["neighbour"] - assert torch.equal(view, saved), "wrote outside the strided view" + assert torch.equal(view, saved), "wrote into another index role of the page" + 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 under mRoPE, at the checkpoint geometry.""" + """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] -def test_text_rows_reduce_mrope_to_plain_rope() -> None: - """With t == h == w, the mRoPE selection is an identity, bit for bit.""" - args = _scene(seed=7, distinct_axes=False) - plain = dict(args, index_cache=args["index_cache"].clone(), mrope_section=None) - op(**args) - op(**plain) - assert torch.equal(args["index_cache"], plain["index_cache"]) - - -def test_a_second_call_over_the_same_rows_is_not_idempotent() -> None: - """Recorded, not tolerated: the anchor slot is one of its own inputs.""" - args = _scene(seed=14) - once = op(**args).clone() - twice = op(**args) - changed = (twice.view(torch.int16) != once.view(torch.int16)).any(-1).sum() - assert int(changed) == _anchors(args) +#: The cell of 32 MTP7 verify blocks (256 rows, two anchors each): the segments, pool and table. +_MTP7_32 = next(c.spec for c in op.CELLS if c.why.startswith("decode MTP7: 32 requests")) +_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.""" - args = _scene(seed=16) + """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) @@ -194,65 +198,7 @@ def test_cuda_graph_capture_and_replay() -> None: cache.copy_(before) graph.replay() torch.cuda.synchronize() - _compare(cache, ref) - - -def test_stream_ordering_needs_no_explicit_synchronization() -> None: - """The launch is on the current stream; a read enqueued after it sees it.""" - args = _scene(seed=17) - before = args["index_cache"].clone() - ref = op.reference(**args) - for _ in range(20): - args["index_cache"].copy_(before) - _compare(op(**args), ref) - - -def test_callable_rejects_its_documented_domain_errors() -> None: - """The callable's own loud errors: why `is_valid` carries only stride guards.""" - args = _scene(seed=13) - cache, logical, request = ( - args["index_cache"], - args["logical_positions"], - args["request_indices"], - ) - before = cache.clone() - cases = [ - ("request indices must match logical positions", dict(request_indices=request[:-1])), - ( - "request metadata must be one-dimensional", - dict(logical_positions=logical.view(2, -1), request_indices=request.view(2, -1)), - ), - ("request metadata must use integer storage", dict(logical_positions=logical.float())), - ( - "requires a two-dimensional int32 block table", - dict(block_table=args["block_table"].long()), - ), - ("must have one KV head", dict(index_cache=cache.expand(-1, -1, 2, -1).contiguous())), - ( - "position cache must contain three axes", - dict(position_cache=args["position_cache"][..., :2].contiguous()), - ), - ( - "position cache must use int32 storage", - dict(position_cache=args["position_cache"].long()), - ), - ("index cache must use BF16", dict(index_cache=cache.half())), - ("pages do not match tokens_per_block", dict(tokens_per_block=16)), - ("packed 2D cos/sin tensor", dict(cos_sin=args["cos_sin"][:, :-1].contiguous())), - ("half-width rotary embedding", dict(cos_sin=args["cos_sin"][:, :32].contiguous())), - ("power-of-two head and compression widths", dict(compress_ratio=3)), - ( - "norm weight does not match the index head dimension", - dict(k_norm_weight=args["k_norm_weight"][:-1].contiguous()), - ), - ("requires CUDA tensors", dict(block_table=args["block_table"].cpu())), - ("MRoPE sections do not cover the rotary width", dict(mrope_section=(1, 1, 1))), - ("MRoPE sections do not cover the rotary width", dict(mrope_section=(16, 16))), - ] - for message, override in cases: - with pytest.raises(ValueError, match=message): - op(**{**args, **override}) - assert torch.equal(cache, before), "a rejected call still wrote" + _compare(cache, ref, batch["rounding_flips"]) def _strided(tensor: torch.Tensor, dim: int) -> torch.Tensor: @@ -266,109 +212,26 @@ def _strided(tensor: torch.Tensor, dim: int) -> torch.Tensor: return wide[tuple(index)] -@pytest.mark.parametrize( - "name, dim, match", - [ - ("logical_positions", 0, "logical_positions must be contiguous"), - ("request_indices", 0, "request_indices must be contiguous"), - ("k_norm_weight", 0, "k_norm_weight must be contiguous"), - ("cos_sin", 1, "unit stride"), - ], -) -def test_guard_refuses_the_strides_the_callable_takes_silently(name, dim, match) -> None: - """Shown unguarded first -- accepted and wrong -- then refused. - - Old receipt's max abs error against the right answer: 4.72, 6.20, 0.80 - and 7.16 respectively; the floor below states a bound, not one draw. +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, **{name: _strided(args[name], dim)}) - assert not bad[name].is_contiguous() + 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 {name} was expected to be read wrongly; off by {error}" + 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=match): + 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" - - -def _variant(args, defect): - """The compression with one deliberate defect, in plain torch.""" - cache, cos_sin = args["index_cache"], args["cos_sin"] - cr, tpb = args["compress_ratio"], args["tokens_per_block"] - out = cache.clone() - logical = args["logical_positions"].long() - anchored = (logical + 1) % cr == 0 - anchor, request = logical[anchored], args["request_indices"].long()[anchored] - members = anchor[:, None] - (cr - 1) + torch.arange(cr, device=DEV) - pages = args["block_table"].long()[request[:, None], members // tpb] - slots = members % tpb - raw = cache[pages, slots, 0] - if defect == "pool_in_bfloat16": - acc = raw[:, 0] - for i in range(1, cr): - acc = (acc + raw[:, i]).to(torch.bfloat16) - pooled = (acc / cr).to(torch.bfloat16).float() - else: - pooled = (raw.float().sum(dim=1) / cr).to(torch.bfloat16).float() - scale = torch.rsqrt(pooled.mul(pooled).sum(-1, keepdim=True) / cache.shape[3] + args["eps"]) - weight = args["k_norm_weight"].float() - gamma = weight if defect == "plain_weight" else weight + 1.0 - normed = (pooled * scale * gamma).to(cos_sin.dtype) - source = -1 if defect == "rope_at_last_member" else 0 - t, h, w = args["position_cache"][pages[:, source], slots[:, source], :3].long().unbind(-1) - n = cos_sin.shape[1] // 2 - q = torch.arange(n, device=DEV) - section = None if defect == "ignore_mrope" else args["mrope_section"] - if section is None: - pos = t[:, None].expand(-1, n) - else: - pos = torch.where( - (q % 3 == 1) & (q < 3 * section[1]), - h[:, None], - torch.where((q % 3 == 2) & (q < 3 * section[2]), w[:, None], t[:, None]), - ) - cos, sin = cos_sin[pos, q], cos_sin[pos, n + q] - rot = normed[:, : 2 * n] - if defect == "gptj_pair_split": - first, second = rot[:, ::2], rot[:, 1::2] - head = torch.stack([first * cos - second * sin, second * cos + first * sin], -1).flatten(1) - else: - first, second = rot[:, :n], rot[:, n:] - head = torch.cat([first * cos - second * sin, second * cos + first * sin], -1) - out[pages[:, -1], slots[:, -1], 0] = torch.cat([head, normed[:, 2 * n :]], -1).to( - torch.bfloat16 - ) - return out - - -@pytest.mark.parametrize( - "defect, floor", - [ - ("rope_at_last_member", 1.0), - ("plain_weight", 1.0), - ("ignore_mrope", 1.0), - ("gptj_pair_split", 1.0), - ("pool_in_bfloat16", 0.0), - ], -) -def test_wrong_variants_land_outside_the_gate(defect: str, floor: float) -> None: - """Every defect differs in at least one bit; four of them by more than 1.0. - - Pooling in bf16 moves the result by only ~0.03 (old receipt: 2124 elements, - max 0.03125) -- inside any bf16 band, which is why the gate is bitwise. - """ - args = _scene(seed=15) - ref = op.reference(**args) - assert torch.equal(_variant(args, None).view(torch.int16), ref.view(torch.int16)) - wrong = _variant(args, defect) - out = op(**args) - _compare(out, ref) - assert int((out.view(torch.int16) != wrong.view(torch.int16)).sum()) > 0, defect - assert (out.float() - wrong.float()).abs().max().item() > floor 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_sparse_gqa_rowtile.py new file mode 100644 index 000000000000..fdbdc68c264d --- /dev/null +++ b/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 @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_sparse_gqa_rowtile entry. + +Same paged state as `qsa_paged_sparse_gqa`'s test: one `[pages, 2, kv, tpb, D]` +HND pool, the engine's slot table (more slots than requests, a seeded +permutation, `-1` past each sequence), and selection rows 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 that entry's: this op replaces it on context launches. +""" + +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 + 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 + qsa_sparse_gqa_rowtile as op, +) + +__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.""" + 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, 2048): + stop = min(start + 2048, 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 _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] + assert spec["cache_layout"] == "pool_planes" and 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) + + +@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] + # 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) + 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"]) + assert launches[3][1] == (groups, 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)) + + +CKPT = op.CELLS[0].spec +assert (CKPT["q_heads"], CKPT["kv_heads"], CKPT["head_dim"]) == (24, 2, 256) + + +def _with(segments, **overrides): + return {**CKPT, "segments": tuple(segments), **overrides} + + +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) 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 new file mode 100644 index 000000000000..c8e63aaa8343 --- /dev/null +++ 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 @@ -0,0 +1,103 @@ +# 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) + 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_qsa_topk_select.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_select.py index 11421ec14395..e0ee54787a3a 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qsa_topk_select.py @@ -67,29 +67,18 @@ def _logits(kind: str, rows: int, cols: int, gen: torch.Generator) -> torch.Tens def _build(spec: dict, seed: int): gen = torch.Generator(device=DEV).manual_seed(seed) rows, cols, index_topk = spec["rows"], spec["cols"], spec["index_topk"] - if spec.get("strided"): - backing = _logits(spec["logits"], rows, 2 * cols, gen).view(rows, cols, 2) - logits = backing[:, :, 0] - assert logits.stride() == (2 * cols, 2) - else: - logits = _logits(spec["logits"], rows, cols, gen) - - kind = spec["windows"][0] - if kind in ("positions", "chunk"): - if kind == "positions": - positions = torch.tensor(spec["windows"][1], device=DEV) - else: - positions = spec["windows"][1] + torch.arange(rows, device=DEV) - assert positions.numel() == rows - # QSA's per-row causal bound: complete compressed blocks at the position. - row_ends = ((positions + 1) // COMPRESS_RATIO).clamp_max(cols).to(torch.int32) - row_starts = torch.zeros(rows, dtype=torch.int32, device=DEV) - elif kind == "explicit": - row_starts = torch.tensor(spec["windows"][1], dtype=torch.int32, device=DEV) - row_ends = torch.tensor(spec["windows"][2], dtype=torch.int32, device=DEV) + logits = _logits(spec["logits"], rows, cols, gen) + + kind, start = spec["windows"] + if kind == "positions": + positions = torch.tensor(start, device=DEV) else: - row_starts = torch.zeros(rows, dtype=torch.int32, device=DEV) - row_ends = torch.randint(0, cols + 1, (rows,), generator=gen, device=DEV).to(torch.int32) + assert kind == "chunk", kind + positions = start + torch.arange(rows, device=DEV) + assert positions.numel() == rows + # QSA's per-row causal bound: complete compressed blocks at the position. + row_ends = ((positions + 1) // COMPRESS_RATIO).clamp_max(cols).to(torch.int32) + row_starts = torch.zeros(rows, dtype=torch.int32, device=DEV) if spec["tail"] == "nan": columns = torch.arange(cols, device=DEV)[None, :] @@ -108,42 +97,12 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 _compare(out, op.reference(*args)) -def test_selection_honours_row_ends_and_is_not_a_prefix() -> None: - """The two wrong answers that would still look plausible. - - The global maxima sit outside every window, and the in-window winners are - the *last* visible columns: a selector that ignored `row_ends` would take - the former, one that returned the first k columns would miss the latter. - """ - gen = torch.Generator(device=DEV).manual_seed(2) - rows, cols, top_k = 6, 1024, 64 - logits = torch.rand(rows, cols, generator=gen, device=DEV) - logits[:, 256:] += 100.0 - logits[:, 192:256] += 10.0 - row_starts = torch.zeros(rows, dtype=torch.int32, device=DEV) - row_ends = torch.full((rows,), 256, dtype=torch.int32, device=DEV) - indices = torch.full((rows, top_k), SENTINEL, dtype=torch.int32, device=DEV) - out = op(logits, row_starts, row_ends, indices, top_k) - _compare(out, op.reference(logits, row_starts, row_ends, indices, top_k)) - assert bool((out < 256).all()), "selected a column beyond row_ends" - assert bool((out >= 192).all()), "did not find the planted in-window maxima" - - -def test_nan_inside_the_window_is_selected() -> None: - """Why the note makes finiteness a precondition: NaN outranks every value.""" - gen = torch.Generator(device=DEV).manual_seed(11) - logits = torch.randn(1, 32, generator=gen, device=DEV) - logits[0, :4] = float("nan") - row_starts = torch.zeros(1, dtype=torch.int32, device=DEV) - row_ends = torch.full((1,), 32, dtype=torch.int32, device=DEV) - out = op(logits, row_starts, row_ends, torch.empty(1, 8, dtype=torch.int32, device=DEV), 8) - assert {0, 1, 2, 3} <= set(out[0].tolist()), "NaN is no longer selected; update the note" - - def test_cuda_graph_capture_and_replay() -> None: - """Captured, then replayed against new scores and a new causal bound.""" + """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 gen = torch.Generator(device=DEV).manual_seed(8) - rows, cols = 8, 2048 + 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) @@ -168,21 +127,6 @@ def test_cuda_graph_capture_and_replay() -> None: _compare(indices, op.reference(logits, row_starts, row_ends, indices, BLOCK_TOPK)) -def test_rejects_out_of_contract() -> None: - """The op's own loud refusals -- the reason `is_valid` does not restate them.""" - 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) - out = torch.empty(rows, top_k, dtype=torch.int32, device=DEV) - with pytest.raises(RuntimeError, match="expected scalar type Float"): - op(logits.to(torch.bfloat16), row_starts, row_ends, out, top_k) - with pytest.raises(RuntimeError, match="expected scalar type Int"): - op(logits, row_starts, row_ends, out.long(), top_k) - with pytest.raises(RuntimeError, match="one entry per row"): - op(logits, row_starts[:2], row_ends[:2], out, top_k) - - def _layout_case(): torch.manual_seed(10) rows, cols, top_k = 4, 128, 8 @@ -192,25 +136,6 @@ def _layout_case(): return logits, row_starts, row_ends, top_k -def test_guard_refuses_a_wider_output_buffer() -> None: - """Accepted by the op, which packs the rows at a stride of `index_topk` - regardless: row 1's answer starts inside row 0's slots and the last row is - never written. Shown unguarded first, then refused.""" - logits, row_starts, row_ends, top_k = _layout_case() - rows = logits.shape[0] - wide = torch.full((rows, top_k + 4), -99, dtype=torch.int32, device=DEV) - op(logits, row_starts, row_ends, wide, top_k) - torch.cuda.synchronize() - packed = wide.flatten()[: rows * top_k].view(rows, top_k) - _compare(packed, op.reference(logits, row_starts, row_ends, packed, top_k)) - assert not bool((wide[:, top_k:] == -99).all()), "the hazard stopped reproducing" - assert bool((wide[-1] == -99).all()), "the hazard stopped reproducing" - - with pytest.raises(AssertionError, match=r"indices must be \[num_rows, index_topk\]"): - with validating(op): - op(logits, row_starts, row_ends, wide, 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 @@ -220,14 +145,3 @@ def test_guard_refuses_the_token_budget_default() -> None: with pytest.raises(AssertionError, match="index_topk=2048"): with validating(op): op(logits, row_starts, row_ends, indices) - - -def test_guard_refuses_a_non_contiguous_output() -> None: - """A row slice of a wider arena: the same flat addressing, the same fault.""" - logits, row_starts, row_ends, top_k = _layout_case() - arena = torch.full((logits.shape[0], top_k + 4), -99, dtype=torch.int32, device=DEV) - view = arena[:, :top_k] - assert not view.is_contiguous() - with pytest.raises(AssertionError, match="contiguous"): - with validating(op): - op(logits, row_starts, row_ends, view, top_k) 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_attn_decode_prologue.py new file mode 100644 index 000000000000..24b1b58b3521 --- /dev/null +++ b/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 @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_attn_decode_prologue entry: every cell bitwise the +composition it replaces (`qwen4_exp_qk_norm_rope`, `qsa_paged_kv_store`, +`qsa_paged_index_store`, `qsa_prefill_compress`) on every buffer either writes -- the +projection, the index query, the K/V pool, the index plane, the position cache; a +stale-read check 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.qwen4_exp_attn_decode_prologue import ( # noqa: E501 + qwen4_exp_attn_decode_prologue as op, +) + +__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" + +DEV = "cuda" +HQ, HK, D, IH, ID, IQ_COL, W = 24, 2, 256, 4, 128, 13312, 14080 +TPB, CR, ROT, TABLE, SLOTS, WIDTH = 32, 4, 64, 16384, 8, 300 +GEOM = dict( + 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, +) # fmt: skip + + +def _build(spec, seed): + g = torch.Generator(device=DEV).manual_seed(seed) + pos = list(spec["pos"]) + rows = len(pos) + pages = SLOTS * WIDTH // 2 + 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 r, p in enumerate(pos): + n = p // TPB + 1 + block_table[r + 1, :n] = perm[used : used + n] + used += n + kv = torch.randn(pages, 2, HK, TPB, D, generator=g, device=DEV).to(torch.bfloat16) + index_cache = torch.randn(pages, TPB, 1, ID, generator=g, device=DEV).to(torch.bfloat16) + position_cache = torch.randint( + 0, TABLE, (pages, TPB, 3), generator=g, device=DEV, dtype=torch.int32 + ) + request = torch.arange(1, rows + 1, 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) + angles = torch.rand(TABLE, 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, W, generator=g, device=DEV) * 2.0).to(torch.bfloat16) + index_q = torch.zeros(rows, IH + 2, ID, dtype=torch.bfloat16, device=DEV) + weights = [ + (0.3 * torch.randn(n, generator=g, device=DEV)).to(torch.bfloat16) for n in (D, D, ID, ID) + ] + return dict( + proj=proj, index_q=index_q[:, :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=kv[:, 0], v_cache=kv[:, 1], index_cache=index_cache, position_cache=position_cache, + request_indices=request, logical_positions=logical, block_table=block_table, **GEOM, + ) # fmt: skip + + +_WRITTEN = ("proj", "index_q", "k_cache", "v_cache", "index_cache", "position_cache") + + +def _clone(args): + out = dict(args) + kv = args["k_cache"]._base if args["k_cache"]._base is not None else None + if kv is not None: + kv2 = kv.clone() + out["k_cache"], out["v_cache"] = kv2[:, 0], kv2[:, 1] + iq = args["index_q"]._base if args["index_q"]._base is not None else args["index_q"] + iq2 = iq.clone() + out["index_q"] = iq2[:, : args["index_q"].shape[1]] + for name in ("proj", "index_cache", "position_cache"): + out[name] = args[name].clone() + return out + + +@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 = _clone(args) + op.reference(**want) + got = _clone(args) + op(**got) + torch.cuda.synchronize() + for name in _WRITTEN: + assert torch.equal(got[name], want[name]), name + + +def test_reads_after_the_wait_under_graph_replay() -> 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.""" + args = _build(dict(pos=(8699, 8700)), seed=23) + src = args["proj"].clone() + work = _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() + for trial in range(3): + src.copy_((torch.randn(src.shape, device=DEV) * 2.0).to(src.dtype)) + before = _clone(work) + graph.replay() + torch.cuda.synchronize() + want = _clone(before) + want["proj"].copy_(src) + op.reference(**want) + for name in _WRITTEN: + assert torch.equal(work[name], want[name]), f"replay {trial}: {name}" 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/attention/test_modeling_v2_qwen4_exp_qk_norm_rope.py new file mode 100644 index 000000000000..d5404666e2c6 --- /dev/null +++ b/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 @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_qk_norm_rope entry.""" + +import functools +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.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 + qwen4_exp_qk_norm_rope as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 — path above; fixture by name + assert_within_ulp, + certifying, + validating, +) + +assert torch.cuda.is_available(), "qwen4_exp_qk_norm_rope requires a CUDA device" + +THETA = 1e7 +EPS = 1e-6 +TABLE_ROWS = 262144 # the checkpoint's max_position_embeddings +#: Spec keys that shape the test's inputs rather than the op call. +_NOT_OP_ARGS = ("tokens", "max_position", "proj_width", "index_scratch_heads") + + +def _seed(text: str) -> int: + return zlib.crc32(text.encode()) + + +@functools.lru_cache(maxsize=1) +def _table() -> torch.Tensor: + """The target's rope table, built the way `post_load_weights` builds it.""" + rot = CKPT["rotary_dim"] + exponent = torch.arange(0, rot, 2, dtype=torch.float64, device="cuda") + inv_freq = (1.0 / (THETA ** (exponent / rot))).to(torch.float32) + positions = torch.arange(TABLE_ROWS, dtype=torch.float32, device="cuda") + angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0) + return torch.cat([torch.cos(angles), torch.sin(angles)], dim=1).contiguous() + + +def _distinct_rows(tokens: int, max_position: int, g: torch.Generator) -> torch.Tensor: + """A contiguous int32 `[3, T]` whose three rows differ at every token.""" + row0 = torch.randint(0, max_position, (tokens,), generator=g, device="cuda") + d1 = torch.randint(1, max_position, (tokens,), generator=g, device="cuda") + d2 = torch.randint(1, max_position - 1, (tokens,), generator=g, device="cuda") + d2 = d2 + (d2 >= d1).long() + rows = torch.stack([row0, (row0 + d1) % max_position, (row0 + d2) % max_position]) + return rows.to(torch.int32).contiguous() + + +def _build(spec, seed): + """`(proj, scratch, kwargs)`: the fused projection row, the index scratch + whose first `index_heads` heads are the `index_q` view, and the op's + keyword arguments. Norm weights near 1 put q/k at their largest magnitude + under the Gemma `1 + w`.""" + g = torch.Generator(device="cuda").manual_seed(seed) + t = spec["tokens"] + proj = torch.randn(t, spec["proj_width"], generator=g, device="cuda").to(torch.bfloat16) + scratch = torch.randn( + t, spec["index_scratch_heads"], spec["index_head_dim"], generator=g, device="cuda" + ).to(torch.bfloat16) + + def weight(n): + return (torch.randn(n, generator=g, device="cuda") * 0.1 + 1.0).to(torch.bfloat16) + + positions = _distinct_rows(t, spec.get("max_position", 32768), g) + kwargs = {k: v for k, v in spec.items() if k not in _NOT_OP_ARGS} + kwargs.update( + index_q=scratch[:, : spec["index_heads"]], + position_ids=positions, + cos_sin=_table(), + q_weight=weight(spec["head_dim"]), + k_weight=weight(spec["head_dim"]), + index_q_weight=weight(spec["index_head_dim"]), + eps=EPS, + ) + return proj, scratch, kwargs + + +def _deviation(out: torch.Tensor, ref: torch.Tensor) -> str: + ulp = 2.0**-8 + o, r = out.float(), ref.float() + row = r.abs().amax(dim=-1, keepdim=True).clamp_min(1e-9) + elt = ((o - r).abs() / row).max().item() / ulp + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ulp + return f"max abs {(o - r).abs().max().item():.6f}, elt {elt:.2f} ulp, rms {rms:.3f} ulp" + + +def _main_cols(spec) -> int: + return (spec["num_heads_q"] + spec["num_heads_k"]) * spec["head_dim"] + + +@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 + proj, scratch, kwargs = _build(spec, seed=_seed(cell.why)) + before, scratch_before = proj.clone(), scratch.clone() + ref_proj, ref_iq = op.reference(proj, **kwargs) + assert torch.equal(proj, before), "the reference mutated its input" + out_proj, out_iq = op(proj, **kwargs) + assert out_proj is proj and out_iq is kwargs["index_q"] + main = _main_cols(spec) + t = spec["tokens"] + flat_ref = ref_iq.reshape(t, -1) + flat_out = out_iq.reshape(t, -1) + print(f"\n {cell.why[:50]}: main {_deviation(proj[:, :main], ref_proj[:, :main])}") + print(f" {cell.why[:50]}: index {_deviation(flat_out, flat_ref)}") + # fp32 throughout and one rounding, against an fp64 reference over the same + # table: the two can differ by one output rounding step, which is up to 2 + # ulp of the row scale (an element just above a power of two in a row whose + # largest magnitude sits just below the next one). Measured: max 1.48 ulp + # element-wise, relative RMS <= 0.011 ulp, so the RMS gate is tight. + assert_within_ulp(proj[:, :main], ref_proj[:, :main], element_ulp=2.0, rms_ulp=0.25) + assert_within_ulp(flat_out, flat_ref, element_ulp=2.0, rms_ulp=0.25) + assert torch.equal(proj[:, main:], before[:, main:]), "a column past the main k heads moved" + heads = spec["index_heads"] + assert torch.equal(scratch[:, heads:], scratch_before[:, heads:]), ( + "index_q wrote past its heads" + ) + + +def test_cuda_graph_replay_reads_fresh_inputs() -> None: + """Captured once at M=4, replayed after the inputs are rewritten in place: + the replay equals an eager call on the new inputs, bit for bit.""" + spec = dict(tokens=4, **CKPT) + proj, scratch, kwargs = _build(spec, seed=_seed("graph")) + proj_new, _, kwargs_new = _build(spec, seed=_seed("graph-new")) + eager_proj, eager_scratch = proj_new.clone(), scratch.clone() + eager_kwargs = dict(kwargs_new, index_q=eager_scratch[:, : CKPT["index_heads"]]) + op(eager_proj, **eager_kwargs) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + op(proj.clone(), **dict(kwargs, index_q=scratch.clone()[:, : CKPT["index_heads"]])) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(proj, **kwargs) + proj.copy_(proj_new) + for name in ("position_ids", "q_weight", "k_weight", "index_q_weight"): + kwargs[name].copy_(kwargs_new[name]) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(proj, eager_proj) + assert torch.equal(kwargs["index_q"], eager_kwargs["index_q"]) + + +def test_guard_refuses_a_position_past_the_table() -> None: + """An out-of-range position would read past the rope table.""" + proj, _, kwargs = _build(dict(tokens=3, **CKPT), seed=_seed("guard")) + kwargs["position_ids"][0, 0] = TABLE_ROWS + with validating(op), pytest.raises(AssertionError, match="past the rope table"): + op(proj, **kwargs) 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 index 8843c837443b..2a83d03b3fc4 100644 --- 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 @@ -4,12 +4,11 @@ `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. Four ways of getting a 2-D vision rope subtly wrong are run too, so -the gate reads against the distance to a wrong answer. +multiply. """ import zlib -from typing import List, Tuple +from typing import Tuple import pytest import torch @@ -98,146 +97,8 @@ def test_certified_cells(cell, certifying): # noqa: F811 _compare((q_out, k_out), (q_ref, k_ref)) -# --- test-local positions and angles, for the structural checks and the wrong -# operators below ------------------------------------------------------------ - - -def _positions(grid, merge) -> List[Tuple[int, int]]: - positions: List[Tuple[int, int]] = [] - for frames, height, width in grid: - frame = [] - for block_row in range(height // merge): - for block_col in range(width // merge): - for intra_row in range(merge): - for intra_col in range(merge): - frame.append((block_row * merge + intra_row, block_col * merge + intra_col)) - positions.extend(frame * frames) - return positions - - -def _angles(positions, rope_dim, theta, inv_scale_dim=None) -> torch.Tensor: - """`[tokens, rope_dim]` fp32 on cuda; height half first, then width.""" - axis_dim = rope_dim // 2 - denom = rope_dim if inv_scale_dim is None else inv_scale_dim - inverse = torch.tensor( - [theta ** (-(2.0 * j) / denom) for j in range(axis_dim)], dtype=torch.float64 - ) - pos = torch.tensor(positions, dtype=torch.float64).reshape(-1, 2) - return (pos[:, :, None] * inverse).flatten(1).to(device="cuda", dtype=torch.float32) - - -def _rotate_by_pairs(x, angles) -> torch.Tensor: - rope_dim = angles.shape[-1] - fp32 = x.float() - pair = torch.complex(fp32[..., :rope_dim], fp32[..., rope_dim:]) - rotated = pair * torch.polar(torch.ones_like(angles), angles).unsqueeze(-2) - return torch.cat((rotated.real, rotated.imag), dim=-1).to(x.dtype) - - -def _wrong_interleaved(q, grid): - """GPT-J pairing: dims (2j, 2j+1) instead of (j, j + rope_dim).""" - rope_dim = q.shape[-1] // 2 - angles = _angles(_positions(grid, MERGE), rope_dim, THETA) - pair = torch.view_as_complex(q.float().unflatten(-1, (rope_dim, 2)).contiguous()) - rotor = torch.polar(torch.ones_like(angles), angles).unsqueeze(-2) - return torch.view_as_real(pair * rotor).flatten(-2).to(q.dtype) - - -def _wrong_axes_swapped(q, grid): - rope_dim = q.shape[-1] // 2 - axis_dim = rope_dim // 2 - angles = _angles(_positions(grid, MERGE), rope_dim, THETA) - return _rotate_by_pairs(q, torch.cat((angles[:, axis_dim:], angles[:, :axis_dim]), dim=-1)) - - -def _wrong_row_major(q, grid): - positions = [] - for frames, height, width in grid: - positions.extend([(r, c) for r in range(height) for c in range(width)] * frames) - return _rotate_by_pairs(q, _angles(positions, q.shape[-1] // 2, THETA)) - - -def _wrong_full_head_inv_freq(q, grid): - head_dim = q.shape[-1] - angles = _angles(_positions(grid, MERGE), head_dim // 2, THETA, inv_scale_dim=head_dim) - return _rotate_by_pairs(q, angles) - - -WRONG_VARIANTS = { - "interleaved_pairing": _wrong_interleaved, - "axes_swapped": _wrong_axes_swapped, - "row_major_positions": _wrong_row_major, - "full_head_inv_freq": _wrong_full_head_inv_freq, -} - - -@pytest.mark.parametrize("variant", sorted(WRONG_VARIANTS)) -def test_discrimination_against_wrong_operators(variant: str) -> None: - """At 1024 bf16 tokens each wrong operator lands > 0.5 away (measured - 8.5-9.5) and > 100x the correct deviation; the entry's gate refuses it.""" - grid = [(1, 32, 32)] - q, k = _operands(_num_tokens(grid), HEADS, HEADS, HEAD_DIM, torch.bfloat16, seed=23) - grid_thw = torch.tensor(grid, dtype=torch.int32) - - q_out, k_out = op(q, k, grid_thw, MERGE, THETA) - q_ref, k_ref = op.reference(q, k, grid_thw, MERGE, THETA) - correct = _max_deviation(q_out, q_ref) - wrong_q = WRONG_VARIANTS[variant](q, grid) - wrong = _max_deviation(wrong_q, q_ref) - print(f"[discrimination/{variant}] correct={correct:.3e} wrong={wrong:.3e}") - assert wrong > 0.5, f"{variant} is only {wrong} away -- not a discriminating case" - assert wrong > 100.0 * max(correct, 1e-6) - with pytest.raises(AssertionError): - _compare((wrong_q, k_out), (q_ref, k_ref)) - - -def test_origin_token_is_the_identity() -> None: - """Patch (0, 0) has every angle zero: bitwise unchanged. No other is.""" - grid = [(1, 4, 4)] - q, k = _operands(16, HEADS, HEADS, HEAD_DIM, torch.bfloat16, seed=3) - q_out, k_out = op(q, k, torch.tensor(grid, dtype=torch.int32), MERGE, THETA) - assert torch.equal(q_out[0], q[0]) and torch.equal(k_out[0], k[0]) - assert not torch.equal(q_out[1:], q[1:]) - - -def test_rotation_preserves_the_pair_norms() -> None: - grid = [(1, 16, 24)] - q, k = _operands(_num_tokens(grid), HEADS, HEADS, HEAD_DIM, torch.float32, seed=11) - rope_dim = HEAD_DIM // 2 - q_out, k_out = op(q, k, torch.tensor(grid, dtype=torch.int32), MERGE, THETA) - for before, after in ((q, q_out), (k, k_out)): - pair_before = before[..., :rope_dim] ** 2 + before[..., rope_dim:] ** 2 - pair_after = after[..., :rope_dim] ** 2 + after[..., rope_dim:] ** 2 - torch.testing.assert_close(pair_after, pair_before, rtol=1e-5, atol=1e-5) - - -def test_height_and_width_occupy_separate_halves_of_the_angle_table() -> None: - """A column-0 token moves only in the height lanes of both rotation halves, - a row-0 token only in the width lanes.""" - grid = [(1, 4, 4)] - q, k = _operands(16, HEADS, HEADS, HEAD_DIM, torch.float32, seed=17) - rope_dim = HEAD_DIM // 2 - axis_dim = rope_dim // 2 - q_out, _ = op(q, k, torch.tensor(grid, dtype=torch.int32), MERGE, THETA) - positions = _positions(grid, MERGE) - row_only = [i for i, (r, c) in enumerate(positions) if r != 0 and c == 0] - col_only = [i for i, (r, c) in enumerate(positions) if r == 0 and c != 0] - assert row_only and col_only - for t in row_only: - torch.testing.assert_close(q_out[t, :, axis_dim:rope_dim], q[t, :, axis_dim:rope_dim]) - torch.testing.assert_close( - q_out[t, :, rope_dim + axis_dim :], q[t, :, rope_dim + axis_dim :] - ) - assert not torch.equal(q_out[t, :, :axis_dim], q[t, :, :axis_dim]) - for t in col_only: - torch.testing.assert_close(q_out[t, :, :axis_dim], q[t, :, :axis_dim]) - torch.testing.assert_close( - q_out[t, :, rope_dim : rope_dim + axis_dim], q[t, :, rope_dim : rope_dim + axis_dim] - ) - assert not torch.equal(q_out[t, :, axis_dim:rope_dim], q[t, :, axis_dim:rope_dim]) - - 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) @@ -247,63 +108,6 @@ def test_grid_placement_and_integer_dtype_do_not_change_the_result() -> None: assert torch.equal(q_host, q_other) and torch.equal(k_host, k_other) -def test_theta_changes_the_result() -> None: - grid = [(1, 4, 4)] - q, k = _operands(16, HEADS, HEADS, HEAD_DIM, torch.float32, seed=31) - grid_thw = torch.tensor(grid, dtype=torch.int32) - default, _ = op(q, k, grid_thw, MERGE, THETA) - other, _ = op(q, k, grid_thw, MERGE, 1.0e6) - assert not torch.equal(default, other) - - -def test_cuda_graph_replay_matches_eager() -> None: - """Capturable; the grid is baked in at capture.""" - grid_thw = torch.tensor([(1, 16, 16)], dtype=torch.int32) - q, k = _operands(256, HEADS, HEADS, HEAD_DIM, torch.bfloat16, seed=37) - eager_q, eager_k = op(q, k, grid_thw, MERGE, THETA) - - warmup = torch.cuda.Stream() - warmup.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup): - for _ in range(3): - op(q, k, grid_thw, MERGE, THETA) - torch.cuda.current_stream().wait_stream(warmup) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured_q, captured_k = op(q, k, grid_thw, MERGE, THETA) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(captured_q, eager_q) and torch.equal(captured_k, eager_k) - - -# --- guards ------------------------------------------------------------------ -# -# The indivisible-side and short-grid hazards are not driven unguarded here: -# the unwritten position rows are uninitialized memory, and an out-of-range -# value there trips a device-side assert in the angle gather that poisons the -# CUDA context for every later test in the process. - - -def test_guard_refuses_a_grid_side_not_divisible_by_the_merge() -> None: - grid = [(1, 4, 5)] - q, k = _operands(20, HEADS, HEADS, HEAD_DIM, torch.bfloat16, seed=2) - with pytest.raises(AssertionError, match="divisible by spatial_merge_size"): - with validating(op): - op(q, k, torch.tensor(grid, dtype=torch.int32), MERGE, THETA) - - -def test_guard_refuses_a_token_count_that_does_not_match_the_grid() -> None: - q, k = _operands(17, HEADS, HEADS, HEAD_DIM, torch.bfloat16, seed=2) - with pytest.raises(AssertionError, match="describes 16 tokens but q/k carry 17"): - with validating(op): - op(q, k, torch.tensor([(1, 4, 4)], dtype=torch.int32), MERGE, THETA) - q, k = _operands(64, HEADS, HEADS, HEAD_DIM, torch.bfloat16, seed=2) - with pytest.raises(AssertionError, match="describes 1 tokens but q/k carry 64"): - with validating(op): - op(q, k, torch.tensor([(1, 1, 1)], dtype=torch.int32), 1, THETA) - - 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) @@ -314,13 +118,3 @@ def test_guard_refuses_a_one_token_k() -> None: with pytest.raises(AssertionError, match="must agree on tokens and head_dim"): with validating(op): op(q, k[:1], grid_thw, MERGE, THETA) - - -def test_head_dim_not_divisible_by_4_is_loud() -> None: - """Not guarded: the angle table comes out wider than the head and the - broadcast raises.""" - grid_thw = torch.tensor([(1, 4, 4)], dtype=torch.int32) - for head_dim in (70, 71): - q, k = _operands(16, HEADS, HEADS, head_dim, torch.bfloat16, seed=7) - with pytest.raises(RuntimeError): - op(q, k, grid_thw, MERGE, THETA) 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/embedding/test_modeling_v2_ngram_hash.py index e88fb04f3fb5..a402400fc12d 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/embedding/test_modeling_v2_ngram_hash.py @@ -3,9 +3,8 @@ """GPU test for the ngram_hash entry. Exact integer arithmetic, gated bitwise against the entry's scalar reference. -A wrong hash is still a valid row index, so beyond the cells: six near-miss -derivations must each differ, the per-head bands, the EOS reset semantics, -row independence, the tail slice and chunked prefill, and the four guards. +Beyond the cells: the EOS reset that `ngram_token_history`'s EOS seed relies +on, the narrow-multiplier guard and graph capture. """ import zlib @@ -41,7 +40,6 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: DEV = "cuda" EOS = CKPT_EOS_TOKEN_ID -NGRAM_SIZE = len(CKPT_MULTIPLIERS) NGRAM_HEADS = len(CKPT_HEAD_VOCAB_SIZES) PADDED_TABLE_ROWS = 320001536 # sum(head_vocab_sizes) rounded up to a multiple of 128 @@ -81,7 +79,7 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res eos_positions=spec.get("eos_positions", ()), lead_eos=spec.get("lead_eos", True), eos_density=spec.get("eos_density", 0.0), - ).to(spec.get("ids_dtype", torch.int64)) + ) before = ids.clone() args = (ids, *_constants(), EOS, spec["num_query_tokens"]) out = op(*args) @@ -91,74 +89,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res _compare(out, op.reference(*args)) -def _variant(rows, num_query_tokens, variant): - """The scalar formula with one specific mistake in it.""" - multipliers = CKPT_MULTIPLIERS[::-1] if variant == "multipliers_reversed" else CKPT_MULTIPLIERS - heads_per_ngram = NGRAM_HEADS // (NGRAM_SIZE - 1) - out = [] - for row in rows: - segment_start, last_eos = [], -1 - for position, token in enumerate(row): - if variant == "inclusive_segment" and token == EOS: - last_eos = position - segment_start.append(last_eos + 1) - if variant != "inclusive_segment" and token == EOS: - last_eos = position - row_out = [] - for position in range(len(row) - num_query_tokens, len(row)): - context = [] - for shift in range(NGRAM_SIZE): - source = position - shift - (1 if variant == "extra_shift" and shift else 0) - if variant == "no_eos_reset": - inside = source >= 0 - else: - inside = position - segment_start[position] >= shift and source >= 0 - context.append(row[source] if inside else EOS) - heads = [] - for ngram in range(2, NGRAM_SIZE + 1): - mixed = context[0] * multipliers[0] - for index in range(1, ngram): - mixed ^= context[index] * multipliers[index] - block = ngram - 2 - if variant == "blocks_swapped": - block = NGRAM_SIZE - 2 - block - low = block * heads_per_ngram - for head in range(heads_per_ngram): - offset = 0 if variant == "no_offsets" else CKPT_HEAD_OFFSETS[low + head] - heads.append(mixed % CKPT_HEAD_VOCAB_SIZES[low + head] + offset) - row_out.append(heads) - out.append(row_out) - return torch.tensor(out, dtype=torch.long, device=DEV) - - -def test_wrong_variants_are_separated() -> None: - """Six near misses: each in range, inside the table, and different.""" - ids = _tokens(2, CKPT_CONTEXT_LEN + 512, seed=5, eos_positions=(30, 31, 200, 400)) - got = _hash(ids, 512) - _compare(got, op.reference(ids, *_constants(), EOS, 512)) - for variant in ( - "extra_shift", - "no_eos_reset", - "inclusive_segment", - "multipliers_reversed", - "no_offsets", - "blocks_swapped", - ): - wrong = _variant(ids.tolist(), 512, variant) - differing = int((got != wrong).sum()) - print(f" {variant:<22} {differing} indices differ") - assert differing > 0, f"{variant} is indistinguishable from the correct result" - assert 0 <= int(wrong.min()) and int(wrong.max()) < PADDED_TABLE_ROWS - - -def test_every_head_lands_in_its_own_band() -> None: - got = _hash(_tokens(2, CKPT_CONTEXT_LEN + 512, seed=9, eos_positions=(100,)), 512) - for head in range(NGRAM_HEADS): - column = got[:, :, head] - assert int(column.min()) >= CKPT_HEAD_OFFSETS[head] - assert int(column.max()) < CKPT_HEAD_OFFSETS[head] + CKPT_HEAD_VOCAB_SIZES[head] - - def test_eos_restarts_the_window() -> None: """One past an EOS hashes as a sequence's first token, and only there.""" token = 4242 @@ -169,74 +99,6 @@ def test_eos_restarts_the_window() -> None: assert not torch.equal(_hash(fresh, 1), _hash(no_eos, 1)) -def test_eos_token_itself_ends_its_segment() -> None: - """An EOS hashes with the segment it closes, not as a fresh start.""" - closing = torch.tensor([[11, 22, 33, EOS]], device=DEV) - opening = torch.tensor([[EOS] * (CKPT_CONTEXT_LEN + 1)], device=DEV) - assert not torch.equal(_hash(closing, 1), _hash(opening, 1)) - - -def test_rows_are_independent() -> None: - ids = _tokens(4, CKPT_CONTEXT_LEN + 64, seed=13, eos_positions=(20,)) - batched = _hash(ids, 64) - for row in range(ids.shape[0]): - assert torch.equal(batched[row : row + 1], _hash(ids[row : row + 1], 64)) - - -def test_tail_slice_is_exact() -> None: - """Asking for Q rows equals asking for all and keeping the last Q.""" - ids = _tokens(2, 200, seed=17, eos_positions=(50, 51)) - whole = _hash(ids, 200) - for num_query_tokens in (1, 3, 64, 199): - assert torch.equal(_hash(ids, num_query_tokens), whole[:, -num_query_tokens:]) - - -def test_chunked_prefill_agrees_with_one_shot() -> None: - """Two chunks, the second carrying the `ngram_size - 1` cached ids.""" - ids = _tokens(2, CKPT_CONTEXT_LEN + 96, seed=19, eos_positions=(30, 60)) - split = CKPT_CONTEXT_LEN + 48 - first = _hash(ids[:, :split], 48) - second = _hash(ids[:, split - CKPT_CONTEXT_LEN :], 48) - assert torch.equal(torch.cat([first, second], dim=1), _hash(ids, 96)) - - -# The four guards: each shows the unguarded op answering plausibly, then the -# refusal under `validating`. - - -def test_guard_refuses_num_query_tokens_beyond_the_history() -> None: - ids = _tokens(1, CKPT_CONTEXT_LEN, seed=31) - unguarded = op(ids, *_constants(), EOS, CKPT_CONTEXT_LEN + 1) - assert unguarded.shape == (1, CKPT_CONTEXT_LEN - 1, NGRAM_HEADS), ( - "the hazard this guard covers is gone: it used to return fewer rows than asked" - ) - for bad in (CKPT_CONTEXT_LEN + 1, -1): - with pytest.raises(AssertionError, match="num_query_tokens"): - with validating(op): - op(ids, *_constants(), EOS, bad) - - -def test_guard_refuses_a_ragged_head_count() -> None: - ids = _tokens(1, CKPT_CONTEXT_LEN + 4, seed=37) - multipliers, sizes, offsets = _constants() - unguarded = op(ids, multipliers, sizes[:15], offsets[:15], EOS, 4) - assert unguarded.shape[-1] == 14, "the hazard this guard covers is gone" - with pytest.raises(AssertionError, match="multiple of"): - with validating(op): - op(ids, multipliers, sizes[:15], offsets[:15], EOS, 4) - - -def test_guard_refuses_an_over_long_offsets_tensor() -> None: - ids = _tokens(1, CKPT_CONTEXT_LEN + 4, seed=41) - multipliers, sizes, offsets = _constants() - long_offsets = torch.cat([offsets, offsets]) - unguarded = op(ids, multipliers, sizes, long_offsets, EOS, 4) - assert torch.equal(unguarded, _hash(ids, 4)), "the hazard this guard covers is gone" - with pytest.raises(AssertionError, match="same length"): - with validating(op): - op(ids, multipliers, sizes, long_offsets, EOS, 4) - - def test_guard_refuses_narrow_multipliers() -> None: ids = _tokens(1, CKPT_CONTEXT_LEN + 4, seed=47) multipliers, sizes, offsets = _constants() @@ -249,13 +111,6 @@ def test_guard_refuses_narrow_multipliers() -> None: op(ids, narrow, sizes, offsets, EOS, 4) -def test_single_multiplier_raises_on_its_own() -> None: - ids = _tokens(1, CKPT_CONTEXT_LEN + 4, seed=43) - multipliers, sizes, offsets = _constants() - with pytest.raises(ZeroDivisionError): - op(ids, multipliers[:1], 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() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_bf16_gemv_pdl.py new file mode 100644 index 000000000000..ba0d3a05c9df --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/gemm/test_modeling_v2_bf16_gemv_pdl.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the bf16_gemv_pdl entry. + +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. +""" + +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.gemm.bf16_gemv_pdl import ( # noqa: E501 + bf16_gemv_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(), "bf16_gemv_pdl requires a CUDA device" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Accumulation order only: 4 bf16 ulp element-wise against the row scale, + 2 ulp of relative RMS, as for `low_m_bf16_mm` and `cublas_mm`.""" + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0, ulp=ULP[torch.bfloat16]) + + +def _build(m, n, k, seed): + """`mat_b` is a `[N, K]` weight's `.t()`, as the target stores its projections.""" + g = torch.Generator(device="cuda").manual_seed(seed) + mat_a = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + weight = (torch.randn(n, k, generator=g, device="cuda") * k**-0.5).to(torch.bfloat16) + return mat_a, weight.t() + + +@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 + mat_a, mat_b = _build(s["m"], s["n"], s["k"], seed=zlib.crc32(cell.why.encode()) % 2**31) + out = op(mat_a, mat_b, s["pdl"]) + assert out.shape == (s["m"], s["n"]) and out.dtype == torch.bfloat16 and out.is_contiguous() + _compare(out, op.reference(mat_a, mat_b, s["pdl"])) + + +def test_dependent_chain_under_graph_replay() -> 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) + + def chain(): + h = x + outs = [] + for w in weights: + h = op(h, w, pdl) + outs.append(h) + return outs + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(2): + chain() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outs = chain() + for trial in range(3): + g = torch.Generator(device="cuda").manual_seed(100 + trial) + x.copy_(torch.randn(x.shape, generator=g, device="cuda").to(x.dtype)) + graph.replay() + torch.cuda.synchronize() + h = x + for w, out in zip(weights, outs): + # 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_low_m_bf16_mm.py b/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 new file mode 100644 index 000000000000..d3b3a17e22a0 --- /dev/null +++ b/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 @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the low_m_bf16_mm entry. + +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 + +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 + 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 + plan, +) + +__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(), "low_m_bf16_mm requires a CUDA device" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Accumulation order only: 4 bf16 ulp element-wise against the row scale, + 2 ulp of relative RMS -- the band `cublas_mm` uses, for the same reason. + Every branch sums the same bf16 products in fp32 and rounds once; what is + left between kernel and fp32 reference is the final bf16 rounding (half an + ulp) plus the summation order.""" + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0, ulp=ULP[torch.bfloat16]) + + +def _build(m, n, k, seed): + """`mat_b` is a `[N, K]` weight's `.t()`, as a target stores its projections.""" + g = torch.Generator(device="cuda").manual_seed(seed) + mat_a = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + weight = (torch.randn(n, k, generator=g, device="cuda") * k**-0.5).to(torch.bfloat16) + return mat_a, weight.t() + + +@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 + assert plan(s["m"], s["n"], s["k"]).kind == s["kind"], plan(s["m"], s["n"], s["k"]) + mat_a, mat_b = _build(s["m"], s["n"], s["k"], seed=zlib.crc32(cell.why.encode()) % 2**31) + out = op(mat_a, mat_b, s["pdl"]) + assert out.shape == (s["m"], s["n"]) and out.dtype == torch.bfloat16 and out.is_contiguous() + _compare(out, op.reference(mat_a, mat_b, s["pdl"])) + + +#: `(M, N, K, branch)` of the table's shapes and M buckets the cells do not drive: +#: what `plan` picks for each, held without a launch. +_PLAN_PINS = ( + (1, 10240, 2560, "direct"), + (1, 48, 2560, "direct"), + (32, 48, 2560, "direct"), + (8, 2560, 6144, "splitk"), + (2, 16480, 2560, "cublas"), + (1, 2560, 6144, "direct"), + (3, 12288, 2560, "splitk"), + (8, 12288, 2560, "splitk"), + (16, 12288, 2560, "cublas"), + (16, 512, 2560, "direct"), + (24, 512, 2560, "splitk"), + (16, 640, 2560, "splitk"), + (4, 1280, 2560, "direct"), + (8, 1296, 2560, "splitk"), + (1, 1, 2560, "direct"), + (1, 16480, 2560, "direct"), + (2, 14080, 2560, "direct"), + (3, 14080, 2560, "splitk"), + (8, 14080, 2560, "splitk"), + (16, 14080, 2560, "cublas"), + (1, 2560, 2560, "direct"), + (4, 2560, 2560, "direct"), + (8, 2560, 2560, "splitk"), + (1, 2560, 640, "cublas"), + (8192, 2560, 6144, "cublas"), + (1, 1024, 2048, "cublas"), +) + + +def test_plan_pins_the_branch_of_the_other_table_shapes() -> None: + picked = [(m, n, k, want, plan(m, n, k).kind) for m, n, k, want in _PLAN_PINS] + wrong = [p for p in picked if p[3] != p[4]] + 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/hc/test_modeling_v2_qwen4_exp_gated_residual.py b/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 index 4fa6444fc13d..9d2a95df0a97 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_gated_residual.py @@ -4,8 +4,9 @@ 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 --- and every mis-reading of it keeps the shapes. So beyond the cells, one test -measures how far outside the entry's band each wrong variant lands. +-- and every mis-reading of it keeps the shapes; the cells' band sits 32x..192x +inside the nearest wrong formula (measured). Beyond the cells: graph replay and +one loud misuse. """ import zlib @@ -13,7 +14,6 @@ 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.hc.qwen4_exp_gated_residual import ( # noqa: E501 gate_ratio, @@ -127,104 +127,11 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res _compare((mixed, inject), op.reference(*args)) -def test_stack_final_mixer_matches_the_injecting_call() -> None: - """The injection branch is independent of `mixed`, bitwise.""" - _, (normed, down, up, inject_w, c) = _build((256,), seed=1) - mixed_only, none_weights = op(normed, down, up, None, c) - assert none_weights is None - with_inject, _ = op(normed, down, up, inject_w, c) - assert torch.equal(mixed_only, with_inject) - - -def test_flattened_call_matches_batched() -> None: - _, (normed, down, up, inject_w, c) = _build((3, 17), seed=2) - batched_mixed, batched_inject = op(normed, down, up, inject_w, c) - flat_mixed, flat_inject = op(normed.reshape(-1, WIDTH), down, up, inject_w, c) - assert torch.equal(flat_mixed, batched_mixed.reshape(-1, HIDDEN)) - assert torch.equal(flat_inject, batched_inject.reshape(-1, HC_COUNT)) - - -def test_far_above_the_checkpoint_range_the_outputs_stay_bounded() -> None: - """Finite, `|mixed| <= max|normed|`, `inject` in [0, 2], to 4096x the - stream. The float64 comparison is not the question up there: the sigmoid - saturates and bf16 rounding of its logit flips an all-but-binary gate.""" - _, (base, down, up, inject_w, c) = _build((256,), seed=6) - assert base.float().abs().max().item() > 20.0, "operands are not at the checkpoint's range" - for scale in (1.0, 8.0, 64.0, 512.0, 4096.0): - # Scale the normed stream: an RMS norm would undo a scaled input. - normed = (base.float() * scale).to(torch.bfloat16) - mixed, inject = op(normed, down, up, inject_w, c) - assert torch.isfinite(mixed.float()).all(), scale - assert torch.isfinite(inject.float()).all(), scale - bound = normed.float().abs().max().item() * (1.0 + torch.finfo(torch.bfloat16).eps) - assert mixed.float().abs().max().item() <= bound, scale - assert (inject.float() >= 0.0).all() and (inject.float() <= 2.0).all(), scale - - -def test_wrong_variants_land_outside_the_band() -> None: - """Each mis-reading keeps the shapes; each must land > 5x outside the band - the correct result sits inside.""" - raw, args = _build((512,), seed=7) - normed, down, up, inject_w, c = args - want_mixed, want_inject = op.reference(*args) - got_mixed, got_inject = op(*args) - print(f"\n correct mixed {gate_ratio(got_mixed, want_mixed):.4g}x the band") - print(f" correct inject {gate_ratio(got_inject, want_inject):.4g}x the band") - - x = normed.double() - streams = x.unflatten(-1, (c, HIDDEN)) - logits = x @ down.double().t() - gate = torch.sigmoid(F.silu(logits / c) @ up.double().t()).unflatten(-1, (c, HIDDEN)) - inject_logits = x @ inject_w.double().t() - wrong_mixed = { - "sum instead of mean": lambda: (gate * streams).sum(dim=-2), - "no /hc_count before silu": lambda: ( - torch.sigmoid(F.silu(logits) @ up.double().t()).unflatten(-1, (c, HIDDEN)) * streams - ).mean(dim=-2), - "silu and sigmoid swapped": lambda: ( - F.silu(torch.sigmoid(logits / c) @ up.double().t()).unflatten(-1, (c, HIDDEN)) * streams - ).mean(dim=-2), - "stream axis taken last": lambda: ( - gate.transpose(-1, -2).reshape(-1, c, HIDDEN) * streams - ).mean(dim=-2), - "gates the raw stream": lambda: (gate * raw.double().unflatten(-1, (c, HIDDEN))).mean( - dim=-2 - ), - "mean over hidden": lambda: (gate * streams) - .mean(dim=-1) - .repeat_interleave(HIDDEN // c, dim=-1), - } - wrong_inject = { - "missing the factor of 2": lambda: torch.sigmoid(inject_logits / c), - "no /hc_count on the logit": lambda: 2.0 * torch.sigmoid(inject_logits), - "tanh instead of 2*sigmoid": lambda: torch.tanh(inject_logits / c), - "silu instead of 2*sigmoid": lambda: F.silu(inject_logits / c), - } - for want, variants in ((want_mixed, wrong_mixed), (want_inject, wrong_inject)): - for name, build in variants.items(): - ratio = gate_ratio(build().to(torch.bfloat16), want) - print(f" {name:<28} {ratio:10.4g}x the band") - assert ratio > 5.0, f"the band does not separate {name!r} ({ratio:.4g}x)" - with pytest.raises(AssertionError): - _compare((build().to(torch.bfloat16), None), (want, None)) - - def test_malformed_operands_raise_rather_than_return() -> None: - """Why there is no `is_valid`: each of these is loud on its own.""" + """Why there is no `is_valid`: a width the weights do not take is loud on its own.""" _, (normed, down, up, inject_w, c) = _build((32,), seed=9) with pytest.raises(RuntimeError): op(normed[:, : WIDTH - 8], down, up, inject_w, c) - with pytest.raises(RuntimeError): - op(normed, down.float(), up, inject_w, c) - odd = WIDTH + 1 - with pytest.raises(RuntimeError): - op( - torch.randn(32, odd, dtype=torch.bfloat16, device=DEV), - torch.randn(LOWRANK, odd, dtype=torch.bfloat16, device=DEV), - torch.randn(odd, LOWRANK, dtype=torch.bfloat16, device=DEV), - torch.randn(3, odd, dtype=torch.bfloat16, device=DEV), - 3, - ) def test_cuda_graph_capture_and_replay() -> None: 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_first_cute.py new file mode 100644 index 000000000000..2cd74f1c0b3d --- /dev/null +++ b/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 @@ -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 qwen4_exp_hc_decode_first_cute entry. + +Cells against the entry's float64 reference: `hc_out` within one ulp (plus the +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. +""" + +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_first_cute import ( # noqa: E501 + qwen4_exp_hc_decode_first_cute as op, +) + +__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" + +DEV = "cuda" +HC_COUNT, HIDDEN, RANK = 4, 2560, 320 +DOWN_RMS, INJECT_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0296, 0.0107, 0.477, 6.625 +PART_EPS = 2.0**-16 + + +def _band(got, want, eps) -> float: + w = want.double() + return ((got.double() - w).abs() / (eps * w.abs().max().item() + eps * w.abs())).max().item() + + +def _ulp(got, want) -> float: + eps, tiny = torch.finfo(want.dtype).eps, torch.finfo(want.dtype).tiny + w = want.double() + return ((got.double() - w).abs() / (eps * w.abs() + tiny)).max().item() + + +def _compare(out, ref) -> None: + (hc, part, ss), (ref_hc, ref_part, ref_ss) = out, ref + assert hc.dtype == ref_hc.dtype and hc.shape == ref_hc.shape + assert part.dtype == torch.float32 and part.shape == ref_part.shape and part.stride(2) == 1 + assert ss.dtype == torch.float32 and ss.shape == ref_ss.shape + assert _ulp(hc, ref_hc) <= 1.0, f"hc_out {_ulp(hc, ref_hc):.4g} ulp" + assert _band(part, ref_part, PART_EPS) <= 1.0, f"part {_band(part, ref_part, PART_EPS):.4g}x" + assert torch.allclose(ss, ref_ss, rtol=1e-5, atol=0), "ss" + + +def _routed(block_out, topk_weights, topk_rows): + if topk_rows is None: + return block_out + valid = (topk_rows >= 0).unsqueeze(-1) + rows = block_out.double()[topk_rows.clamp(min=0).long()] + return ( + torch.where(valid, rows * topk_weights.double().unsqueeze(-1), 0.0) + .sum(1) + .to(block_out.dtype) + ) + + +def _hc_band_ratio( + hc, ref_hc, block_out, inject, shared, gate, topk_weights=None, topk_rows=None +) -> float: + """The `hc_out` band at the MoE sites: one ulp of `hc_out`, plus one + ulp of the combine (and with the fold, of the routed row) times `|inject|`.""" + feps, tiny = torch.finfo(ref_hc.dtype).eps, torch.finfo(ref_hc.dtype).tiny + routed = _routed(block_out, topk_weights, topk_rows) + s = torch.sigmoid(gate.double().reshape(-1, 1)) + b = (routed.double() + s * shared.double()).to(routed.dtype).double().abs() + if topk_rows is not None: + b = b + routed.double().abs() + band = feps * ref_hc.double().abs() + ( + feps * b.unsqueeze(-2) * inject.double().abs().unsqueeze(-1) + ).flatten(-2) + return ((hc.double() - ref_hc.double()).abs() / (band + tiny)).max().item() + + +def _compare_moe(args, out, ref) -> None: + hyper_input, block_out, inject, gamma, first, c, n, shared, gate = args[:9] + topk_weights, topk_rows = args[9:] if len(args) > 9 else (None, None) + (hc, part, ss), (ref_hc, _, _) = out, ref + ratio = _hc_band_ratio(hc, ref_hc, block_out, inject, shared, gate, topk_weights, topk_rows) + assert ratio <= 1.0, f"hc_out {ratio:.4g}x the band" + _compare((hc, part, ss), op.reference(hc, None, None, gamma, first, c, n)) + + +def _build(spec, seed): + """The test operands at one row.""" + rows, hc_count, hidden, dtype = 1, HC_COUNT, HIDDEN, torch.bfloat16 + width = hc_count * hidden + n_cols = spec.get("n_cols", RANK + hc_count) + g = torch.Generator(device=DEV).manual_seed(seed) + + def normal(*shape, scale): + return torch.randn(*shape, generator=g, device=DEV) * scale + + hyper_input = normal(rows, width, scale=3.0).to(dtype) + gamma = normal(width, scale=GAMMA_RMS) + gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX + gamma = gamma.to(dtype) + n_rows = -(-n_cols // 128) * 128 + weight = torch.zeros(n_rows, width, device=DEV) + weight[:n_cols] = normal(n_cols, width, scale=DOWN_RMS) + weight[n_cols - hc_count : n_cols] *= INJECT_RMS / DOWN_RMS + first = weight.to(dtype).t() + block_out = inject = shared = gate_logits = None + if spec.get("inject", True): + block_out = normal(rows, hidden, scale=2.0).to(dtype) + inject = (2.0 * torch.sigmoid(normal(rows, hc_count, scale=1.0))).to(dtype) + if spec.get("moe", False): + shared = normal(rows, hidden, scale=2.0).to(dtype) + gate_logits = normal(rows, 1, scale=2.0).to(dtype) + top_k = spec.get("fold", 0) + if not top_k: + return hyper_input, block_out, inject, gamma, first, hc_count, n_cols, shared, gate_logits + n_perm = rows * top_k + 37 + block_out = normal(n_perm, hidden, scale=2.0).to(dtype) + order = torch.randperm(n_perm, generator=g, device=DEV)[: rows * top_k] + topk_rows = order.reshape(rows, top_k).to(torch.int32) + if spec.get("drop"): + topk_rows[:, -3:] = -1 + topk_weights = torch.softmax(normal(rows, top_k, scale=1.0), -1).to(dtype) + return ( + hyper_input, + block_out, + inject, + gamma, + first, + hc_count, + n_cols, + shared, + gate_logits, + topk_weights, + topk_rows, + ) + + +@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)) + before = [a.clone() if isinstance(a, torch.Tensor) else a for a in args] + out = op(*args) + for a, b in zip(args, before): + assert not isinstance(a, torch.Tensor) or torch.equal(a, b), "every input is read-only" + if args[1] is None: + assert out[0] is args[0], "norm-only returns hyper_input itself" + if args[7] is None: + _compare(out, op.reference(*args)) + else: + _compare_moe(args, out, op.reference(*args)) + + +def test_reads_after_the_wait_under_graph_replay() -> None: + """A producer kernel rewrites every activation inside the graph right before this op. + With PDL the op may start before the producer finishes; reading any activation before + `griddepcontrol.wait` would return the previous replay's value.""" + args = list(_build(dict(moe=True, fold=10), seed=8)) + sources = [a.clone() for a in (args[0], args[1], args[2], args[7], args[8], args[9])] + targets = (args[0], args[1], args[2], args[7], args[8], args[9]) + + def step(): + for dst, src in zip(targets, sources): + dst.copy_(src * 1.0) # an elementwise producer kernel per operand + return op(*args) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(2): + 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(dict(moe=True, fold=10), seed=100 + trial) + for src, new in zip(sources, (fresh[0], fresh[1], fresh[2], fresh[7], fresh[8], fresh[9])): + src.copy_(new) + graph.replay() + torch.cuda.synchronize() + 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_decode_mix_fp4_cute.py new file mode 100644 index 000000000000..5958e069eb6a --- /dev/null +++ b/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 @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the two entries of qwen4_exp_hc_decode_mix_fp4_cute: the NVFP4 build +(`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.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + 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 + qwen4_exp_hc_decode_mix_cute, + qwen4_exp_hc_decode_mix_fp4_cute, +) + +__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" + +DEV = "cuda" +HC_COUNT, HIDDEN, RANK, EPS = 4, 2560, 320, 1e-6 +UP_RMS, GAMMA_RMS, GAMMA_ABSMAX = 0.0681, 0.477, 6.625 + +ENTRIES = {"fp4": qwen4_exp_hc_decode_mix_fp4_cute, "plain": qwen4_exp_hc_decode_mix_cute} +CASES = [(name, cell) for name, entry in ENTRIES.items() for cell in entry.CELLS] + + +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 _band(got, want) -> float: + eps = torch.finfo(want.dtype).eps + w = want.double() + return ((got.double() - w).abs() / (eps * w.abs().max().item() + eps * w.abs())).max().item() + + +def _compare(out, ref) -> None: + (mixed, inj), (ref_mixed, ref_inj) = out[:2], ref[:2] + assert (inj is None) == (ref_inj is None) + for got, want in [(mixed, ref_mixed)] + ([] if inj is None else [(inj, ref_inj)]): + assert got.dtype == want.dtype and got.shape == want.shape and got.is_contiguous() + assert _band(got, want) <= 1.0, f"{_band(got, want):.4g}x the band" + + +def _unpack(xq): + return torch.stack([xq & 0xF, xq >> 4], -1).flatten(-2) + + +def _check_quant(mixed, xq, xsf, gs) -> None: + """The NVFP4 outputs exact on the returned `mixed`: scales bitwise, e2m1 codes + bitwise outside the window where the approximate reciprocals may pick either side.""" + want_q, want_sf = nvfp4_reference(mixed, gs) + 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" + stray = (_unpack(xq) != _unpack(want_q)) & ~nvfp4_near_tie(mixed, gs) + assert not stray.any(), f"{stray.sum().item()} e2m1 codes differ outside the tie window" + + +def _build(spec, seed, fp4=True): + """The entry's arguments; `fp4` appends the NVFP4 build's global scale.""" + rows, hc_count, hidden, rank, dtype = 1, HC_COUNT, HIDDEN, RANK, torch.bfloat16 + with_inject = spec.get("with_inject", True) + width = hc_count * hidden + g = torch.Generator(device=DEV).manual_seed(seed) + + def normal(*shape, scale): + return torch.randn(*shape, generator=g, device=DEV) * scale + + hc_out = normal(rows, width, scale=3.0).to(dtype) + if spec.get("layout") == "row_slice": + buf = normal(2 * rows, width, scale=1.0).to(dtype) + buf[::2] = hc_out + hc_out = buf[::2] + ss = hc_out.double().unflatten(-1, (hc_count, hidden)).pow(2).sum(-1).t().float().contiguous() + gamma = normal(width, scale=GAMMA_RMS) + gamma[0], gamma[width // 2] = GAMMA_ABSMAX, -GAMMA_ABSMAX + gamma = gamma.to(dtype) + n_pad = -(-(rank + hc_count) // 4) * 4 + part = normal(hc_count, rows, n_pad + 4, scale=60.0)[:, :, :n_pad] + up = normal(width, rank, scale=UP_RMS).to(dtype).t() + args = [part, ss, up, hc_out, gamma, EPS, hc_count, with_inject] + if fp4: + # A calibrated activation global scale: 448 * 6 / amax of the mix (~|8|). + args.append(torch.tensor([448.0 * 6.0 / 8.0], device=DEV, dtype=torch.float32)) + return args + + +@pytest.mark.parametrize( + "name,cell", CASES, ids=[f"{name}-{cell.why[:44]}" for name, cell in CASES] +) +def test_certified_cells(name, cell, certifying) -> None: # noqa: F811 — fixture by name + op = ENTRIES[name] + certifying(op) + fp4 = name == "fp4" + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % (2**31), fp4=fp4) + out = op(*args) + assert out[0].shape == (1, HIDDEN) + _compare(out, op.reference(*args)) + if fp4: + _check_quant(out[0], out[2], out[3], args[8]) + + +def test_reads_after_the_wait_under_graph_replay() -> None: + """Producer kernels rewrite `part`, `ss` and `hc_out` inside the graph right before + the op; with PDL a read before `griddepcontrol.wait` would see the previous replay's. + The NVFP4 build: the MoE sites' variant, with the most outputs.""" + op = qwen4_exp_hc_decode_mix_fp4_cute + args = _build(dict(), seed=8) + targets = (args[0], args[1], args[3]) + sources = [t.clone() for t in targets] + + def step(): + for dst, src in zip(targets, sources): + dst.copy_(src * 1.0) + return op(*args) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(2): + 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(dict(), seed=100 + trial) + for src, new in zip(sources, (fresh[0], fresh[1], fresh[3])): + src.copy_(new) + graph.replay() + 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_tc_first.py b/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 new file mode 100644 index 000000000000..4ec79d79cbe3 --- /dev/null +++ b/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 @@ -0,0 +1,92 @@ +# 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.""" + +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 + qwen4_exp_hc_decode_tc_first as op, +) + +__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_tc_first requires a CUDA device" + +DEV = "cuda" +RANK = 320 + + +def _ratio(out, ref) -> float: + """Worst element's distance from `ref` as a fraction of the band: one ulp of + the element (`eps * |ref|`) plus `eps / 16` of the tensor's largest value. + An fp32-accumulated 10240-term projection rounded once can sit one ulp off + the float64 value at a rounding midpoint; the absolute term covers SiLU's + cancellation near zero, where an fp32 logit error of ~1e-6 relative to the + row's scale is many ulps of a tiny output.""" + eps = torch.finfo(ref.dtype).eps + r = ref.double() + return ((out.double() - r).abs() / (eps * r.abs() + eps / 16 * r.abs().max())).max().item() + + +def _compare(out, ref) -> None: + 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 and got.is_contiguous() + assert _ratio(got, want) <= 1.0, f"{_ratio(got, want):.4g}x the band" + + +def _build(spec, seed): + hc, hidden = 4, 2560 + rows, width, inject = spec["rows"], hc * hidden, spec.get("inject", True) + g = torch.Generator(device=DEV).manual_seed(seed) + # This checkpoint's layer statistics: down ~N(0, 0.0296), block_inject ~N(0, 0.0107), + # the normed stream unit-RMS per stream times (1 + gamma), gamma ~N(0, 0.477). + down = torch.randn(RANK, width, generator=g, device=DEV) * 0.0296 + parts = [down] + ([torch.randn(hc, width, generator=g, device=DEV) * 0.0107] if inject else []) + w = torch.cat(parts) + w = torch.cat([w, torch.zeros(384 - w.shape[0], width, device=DEV)]) + gamma = torch.randn(width, generator=g, device=DEV) * 0.477 + x = torch.randn(rows, hc, hidden, generator=g, device=DEV) + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) + normed = (x.flatten(-2) * (1.0 + gamma)).to(torch.bfloat16) + return normed, w.to(torch.bfloat16).t(), RANK, hc, inject + + +@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)) + low, inject = op(*args) + assert low.shape == (cell.spec["rows"], RANK) + _compare((low, inject), op.reference(*args)) + + +def test_guard_refuses_more_than_one_tile_of_rows() -> None: + normed, first, rank, c, wi = _build(dict(rows=129), seed=4) + with validating(op), pytest.raises(AssertionError): + op(normed, first, rank, c, wi) + + +def test_cuda_graph_capture_and_replay() -> None: + normed, first, rank, c, wi = _build(dict(rows=8), 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=8), 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_inject_norm.py b/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 new file mode 100644 index 000000000000..b69e82f5c120 --- /dev/null +++ b/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 @@ -0,0 +1,202 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_hc_inject_norm 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. +""" + +import sys +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_inject_norm import ( # noqa: E501 + qwen4_exp_hc_inject_norm as op, +) + +__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" + +#: The op's module (its kernels and dispatch guards). +M = sys.modules[type(op).__module__] + +DEV = "cuda" +HC_COUNT = 4 +HIDDEN = 2560 +EPS = 1e-6 +# This checkpoint's grouped-norm weight statistics (layer-0 hyper-connection norm). +GAMMA_RMS = 0.477 +GAMMA_ABSMAX = 6.625 + + +def _ulp_ratio(out: torch.Tensor, ref: torch.Tensor) -> float: + """Worst element's distance from `ref` in units of `eps(dtype) * |ref|` + (one ulp of the reference value, or a little more). 1 means one ulp.""" + eps = torch.finfo(ref.dtype).eps + ref64 = ref.double() + tiny = torch.finfo(ref.dtype).tiny + return ((out.double() - ref64).abs() / (eps * ref64.abs() + tiny)).max().item() + + +def _compare(out, ref) -> None: + """Both outputs within one ulp of the float64 result rounded once. + + The op rounds each output once from an fp32 value whose own error is + ~1e-6 relative (a 16-bit product is exact in fp32; the RMS is a 2560-term + fp32 sum), so it can differ from the once-rounded float64 value only where + the two straddle a rounding midpoint: by exactly one ulp.""" + (hc, normed), (ref_hc, ref_normed) = out, ref + for got, want in ((hc, ref_hc), (normed, ref_normed)): + assert got.dtype == want.dtype and got.shape == want.shape + if got.numel(): + ratio = _ulp_ratio(got, want) + assert ratio <= 1.0, f"{ratio:.4g} ulp" + + +def _combine(block_out, shared, gate_logits) -> torch.Tensor: + """The MoE block output in float64, rounded once to `block_out`'s dtype.""" + s = torch.sigmoid(gate_logits.double().reshape(-1, 1)) + return (block_out.double() + s * shared.double()).to(block_out.dtype) + + +def _routed(block_out, topk_weights, topk_rows): + """The routed rows: `block_out` itself, or with the fold the float64 finalize of + the permuted rows rounded once.""" + if topk_rows is None: + return block_out + valid = (topk_rows >= 0).unsqueeze(-1) + picked = block_out.double()[topk_rows.clamp(min=0).long()] + w = topk_weights.double().unsqueeze(-1) + return torch.where(valid, picked * w, 0.0).sum(1).to(block_out.dtype) + + +def _compare_moe(args, out, ref) -> None: + """MoE-combine cells. + + `hc_out`: the op's combine is an fp32 value rounded once, the reference's a + float64 value rounded once; where the two straddle a rounding midpoint they + differ by one ulp of the block output, and the accumulate carries that + scaled by `inject`. So the band is one ulp of `hc_out` plus one ulp of the + combine times `|inject|`, element by element -- still well under one ulp + of `hc_out` wherever the block term is small against the residual. With the + finalize fold the routed row has its own rounding (fp32 FMAs against the + float64 sum), one more ulp of the routed row times `|inject|`. + + `normed`: one ulp of the float64 norm of the residual the op stored, + `reference(hc_out, None, None, ...)` -- the dataflow the op states.""" + hyper_input, block_out, inject, gamma, eps, c, shared, gate = args[:8] + topk_weights, topk_rows = args[8:] if len(args) > 8 else (None, None) + (hc, normed), (ref_hc, _) = out, ref + assert hc.dtype == ref_hc.dtype and hc.shape == ref_hc.shape + if not hc.numel(): + return + feps, tiny = torch.finfo(ref_hc.dtype).eps, torch.finfo(ref_hc.dtype).tiny + routed = _routed(block_out, topk_weights, topk_rows) + b = _combine(routed, shared, gate).double().abs() + if topk_rows is not None: + b = b + routed.double().abs() + band = feps * ref_hc.double().abs() + ( + feps * b.unsqueeze(-2) * inject.double().abs().unsqueeze(-1) + ).flatten(-2) + ratio = ((hc.double() - ref_hc.double()).abs() / (band + tiny)).max().item() + assert ratio <= 1.0, f"hc_out {ratio:.4g}x the band" + own = op.reference(hc, None, None, gamma, eps, c)[1] + assert _ulp_ratio(normed, own) <= 1.0, f"normed {_ulp_ratio(normed, own):.4g} ulp" + + +def _build(spec, seed): + rows, hc_count, hidden, dtype = spec["rows"], HC_COUNT, HIDDEN, torch.bfloat16 + 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 + gamma = gamma.to(dtype) + block_out = inject = shared = gate_logits = None + if spec.get("inject", True): + block_out = (torch.randn(rows, hidden, generator=g, device=DEV) * 2.0).to(dtype) + inject = (2.0 * torch.sigmoid(torch.randn(rows, hc_count, generator=g, device=DEV))).to( + dtype + ) + if spec.get("moe", False): + shared = (torch.randn(rows, hidden, generator=g, device=DEV) * 2.0).to(dtype) + gate_logits = (torch.randn(rows, 1, generator=g, device=DEV) * 2.0).to(dtype) + top_k = spec.get("fold", 0) + if not top_k: + return hyper_input, block_out, inject, gamma, EPS, hc_count, shared, gate_logits + # The runner's unfinalized output: permuted FC2 rows (more rows than slots, as + # its padded buffer has), each slot's row and its routing weight. + n_perm = rows * top_k + 37 + block_out = (torch.randn(n_perm, hidden, generator=g, device=DEV) * 2.0).to(dtype) + order = torch.randperm(n_perm, generator=g, device=DEV)[: rows * top_k] + topk_rows = order.reshape(rows, top_k).to(torch.int32) + if spec.get("drop"): + topk_rows[:, -3:] = -1 + topk_weights = torch.softmax(torch.randn(rows, top_k, generator=g, device=DEV), -1).to(dtype) + return ( + hyper_input, + block_out, + inject, + gamma, + EPS, + hc_count, + shared, + gate_logits, + topk_weights, + topk_rows, + ) + + +@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())) + before = [a.clone() if isinstance(a, torch.Tensor) else a for a in args] + hc, normed = op(*args) + for a, b in zip(args, before): + assert not isinstance(a, torch.Tensor) or torch.equal(a, b), "every input is read-only" + if args[1] is None: + assert hc is args[0], "norm-only returns hyper_input itself" + else: + assert hc.is_contiguous() + assert args[0].numel() == 0 or hc.data_ptr() != args[0].data_ptr(), "hc_out is fresh" + assert normed.is_contiguous() + if args[6] is None: + _compare((hc, normed), op.reference(*args)) + else: + _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 + CuTe guard) at the last row count before the switch.""" + rows = M._FEW_WARPS_MIN_ROWS - 1 + args = _build(dict(rows=rows, moe=True, fold=10), seed=500 + rows) + calls = [] + cute_fold = M._fold_cute + + def recording(*a): + calls.append(a[0].shape[0]) + cute_fold(*a) + + monkeypatch.setattr(M, "_fold_cute", recording) + hc, normed = op(*args) + assert calls == [rows], "the decode MoE site runs the CuTe kernel" + monkeypatch.setattr(M, "_fold_cute_ok", lambda *a: False) + want_hc, want_normed = op(*args) + assert torch.equal(hc, want_hc) + assert torch.equal(normed, want_normed) 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_lowrank_act.py new file mode 100644 index 000000000000..8641b2b8b7b5 --- /dev/null +++ b/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 @@ -0,0 +1,79 @@ +# 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.""" + +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 + qwen4_exp_hc_lowrank_act as op, +) + +__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" + +DEV = "cuda" +RANK = 320 + + +def _ulp_ratio(out, ref) -> float: + eps = torch.finfo(ref.dtype).eps + tiny = torch.finfo(ref.dtype).tiny + return ((out.double() - ref.double()).abs() / (eps * ref.double().abs() + tiny)).max().item() + + +def _compare(out, ref) -> None: + """One ulp: each output is one fp32 activation of an fp32 input, rounded + once, so it can differ from the once-rounded float64 value only at a + rounding midpoint.""" + 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 and got.is_contiguous() + assert _ulp_ratio(got, want) <= 1.0, f"{_ulp_ratio(got, want):.4g} ulp" + + +def _build(spec, seed): + hc_count = spec.get("hc_count", 4) + g = torch.Generator(device=DEV).manual_seed(seed) + # The first projection's logits on this checkpoint reach |67| (down) / |48| (inject). + logits = torch.randn(spec["rows"], spec.get("logits_width", 384), generator=g, device=DEV) * 20 + return logits, RANK, hc_count, spec.get("inject", True), spec.get("dtype", torch.bfloat16) + + +@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)) + 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean.py new file mode 100644 index 000000000000..3e392c8e6f75 --- /dev/null +++ b/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 @@ -0,0 +1,78 @@ +# 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 entry: cells in the band +`qwen4_exp_gated_residual`'s test derives for `mixed`, the guard, 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.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 import ( # noqa: E501 + qwen4_exp_hc_up_gate_mean as op, +) + +__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" + +DEV = "cuda" + + +def _compare(out, ref) -> None: + """`rtol = eps`, `atol = eps * max|ref|` (`gate_ratio <= 1`), the band + `qwen4_exp_gated_residual` derives for `mixed`.""" + assert out.dtype == ref.dtype and out.shape == ref.shape and out.is_contiguous() + assert gate_ratio(out, ref) <= 1.0, f"{gate_ratio(out, ref):.4g}x the band" + + +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) + rows, width = spec["rows"], hc * hidden + # This checkpoint's statistics: low = silu(.) of the bottleneck (|x| < ~3), the up + # weight ~N(0, 0.068) (gate logits within about +-8), the normed stream ~+-30. + low = torch.nn.functional.silu(torch.randn(rows, rank, generator=g, device=DEV) * 1.5).to( + torch.bfloat16 + ) + up_w = (torch.randn(width, rank, generator=g, device=DEV) * 0.068).to(torch.bfloat16) + normed = (torch.randn(rows, width, generator=g, device=DEV) * 6).to(torch.bfloat16) + return low, up_w.t(), normed, hc + + +@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)) + mixed = op(*args) + assert mixed.shape == (cell.spec["rows"], cell.spec.get("hidden", 2560)) + _compare(mixed, op.reference(*args)) + + +def test_guard_refuses_a_contiguous_up() -> None: + low, up, normed, c = _build(dict(rows=4), seed=4) + with validating(op), pytest.raises(AssertionError): + op(low, up.contiguous(), normed, c) + + +def test_cuda_graph_capture_and_replay() -> None: + low, up, normed, c = _build(dict(rows=8), seed=5) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(low, up, normed, c) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + g_mixed = op(low, up, normed, c) + fresh = _build(dict(rows=8), seed=6) + low.copy_(fresh[0]) + normed.copy_(fresh[2]) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(g_mixed, op(fresh[0], up, fresh[2], c)) 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__tp1/hc/test_modeling_v2_qwen4_exp_hc_up_gate_mean_fp4.py new file mode 100644 index 000000000000..eecc9d836bef --- /dev/null +++ b/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 @@ -0,0 +1,125 @@ +# 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.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + 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.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.hc.qwen4_exp_hc_up_gate_mean_fp4 import ( # noqa: E501 + qwen4_exp_hc_up_gate_mean_fp4 as op, +) + +__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_fp4 requires a CUDA device" + +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 _unpack(xq: torch.Tensor) -> torch.Tensor: + return torch.stack([xq & 0xF, xq >> 4], -1).flatten(-2) + + +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) + 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) + 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) + rows, width = spec["rows"], hc * hidden + # This checkpoint's statistics: low = silu(.) of the bottleneck (|x| < ~3), the up + # weight ~N(0, 0.068) (gate logits within about +-8), the normed stream ~+-30. + low = torch.nn.functional.silu(torch.randn(rows, rank, generator=g, device=DEV) * 1.5).to( + torch.bfloat16 + ) + up_w = (torch.randn(width, rank, generator=g, device=DEV) * 0.068).to(torch.bfloat16) + normed = (torch.randn(rows, width, generator=g, device=DEV) * 6).to(torch.bfloat16) + # A calibrated activation global scale: 448 * 6 / amax over a typical batch. + gs = torch.tensor([448.0 * 6.0 / 24.0], device=DEV, dtype=torch.float32) + return low, up_w.t(), normed, hc, gs + + +def _seed(cell) -> int: + return zlib.crc32(cell.why.encode()) % (2**31) + + +@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=_seed(cell)) + out = op(*args) + assert out[0].shape == (cell.spec["rows"], cell.spec.get("hidden", 2560)) + _compare(out, op.reference(*args), args[4]) + + +def test_guards() -> None: + """The kernel loads `global_scale[0]` only; a per-row scale is refused.""" + low, up, normed, c, gs = _build(dict(rows=4), seed=4) + with validating(op), pytest.raises(AssertionError): + op(low, up, normed, c, torch.cat([gs, gs])) + + +def test_cuda_graph_capture_and_replay() -> None: + low, up, normed, c, gs = _build(dict(rows=8), seed=5) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + op(low, up, normed, c, gs) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + g_out = op(low, up, normed, c, gs) + fresh = _build(dict(rows=8), seed=6) + low.copy_(fresh[0]) + normed.copy_(fresh[2]) + gs.mul_(0.5) + graph.replay() + torch.cuda.synchronize() + want = op(fresh[0], up, fresh[2], c, gs) + 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_post_inject.py index 65853efa4a69..e609e8632440 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/hc/test_modeling_v2_qwen4_exp_post_inject.py @@ -2,9 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the qwen4_exp_post_inject entry. -The cells are gated bitwise (16-bit) by `_compare`. Beyond them: -the separation from the plausible wrong operators, the bitwise identity and -one-hot claims, the leading-dim guard, and graph capture. +The cells are gated bitwise (16-bit) by `_compare`. Beyond them: the +leading-dim guard and graph capture. """ import zlib @@ -23,24 +22,18 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: - """Bitwise for 16-bit operands: the product of two bf16 (or fp16) - values is exact in fp32, so op and reference round the same fp32 sum - once -- measured bit-exact at every token count. fp32 operands round the - product, and whether it is fused into the add is the kernel's choice: - torch's default band there.""" + """Bitwise for 16-bit operands: the product of two bf16 values is exact in + fp32, so op and reference round the same fp32 sum once -- measured + bit-exact at every token count.""" 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 in (torch.bfloat16, torch.float16): - assert torch.equal(out, ref), (out.float() - ref.float()).abs().max().item() - else: - torch.testing.assert_close(out, ref) + assert torch.equal(out, ref), (out.float() - ref.float()).abs().max().item() DEV = "cuda" HC_COUNT = 4 HIDDEN = 2560 -WIDTH = HC_COUNT * HIDDEN def _operands(leading, hc_count=HC_COUNT, hidden=HIDDEN, dtype=torch.bfloat16, seed=0): @@ -60,11 +53,7 @@ def _build(spec, seed): hyper_input, block_out, inject = _operands(spec["leading"], hc_count, hidden, dtype, seed) layout = spec.get("layout", "contiguous") width = hc_count * hidden - if layout == "row_slice": - buf = torch.randn(2 * hyper_input.shape[0], width, device=DEV).to(dtype) - buf[::2] = hyper_input - hyper_input = buf[::2] - elif layout == "col_slice": + if layout == "col_slice": buf = torch.randn(hyper_input.shape[0], 2 * width, device=DEV).to(dtype) buf[:, :width] = hyper_input hyper_input = buf[:, :width] @@ -85,104 +74,16 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res _compare(out, op.reference(hyper_input, block_out, inject)) -def test_rounds_once_in_fp32_not_twice_in_bf16() -> None: - """The bitwise gate is not the bf16-multiply-then-bf16-add form (0.0625 apart).""" - for num_tokens in (64, 2048): - hyper_input, block_out, inject = _operands((num_tokens,), seed=100 + num_tokens) - out = op(hyper_input, block_out, inject) - double_round = ( - hyper_input.unflatten(-1, (HC_COUNT, HIDDEN)) - + block_out.unsqueeze(-2) * inject.unsqueeze(-1) - ).flatten(-2) - assert torch.equal(out, op.reference(hyper_input, block_out, inject)) - assert (out.float() - double_round.float()).abs().max().item() > 0.0, ( - "the two rounding orders did not separate" - ) - - -def test_wrong_operators_are_separated() -> None: - """Stream-minor packing, no gate, gate on the residual: each > 1.0 away.""" - hyper_input, block_out, inject = _operands((64,), seed=3) - out = op(hyper_input, block_out, inject).float() - hi32, bo32, in32 = hyper_input.float(), block_out.float(), inject.float() - wrong = { - "stream_minor": ( - hi32.unflatten(-1, (HIDDEN, HC_COUNT)) + bo32.unsqueeze(-1) * in32.unsqueeze(-2) - ).flatten(-2), - "no_gate": (hi32.unflatten(-1, (HC_COUNT, HIDDEN)) + bo32.unsqueeze(-2)).flatten(-2), - "gate_on_residual": ( - hi32.unflatten(-1, (HC_COUNT, HIDDEN)) * in32.unsqueeze(-1) + bo32.unsqueeze(-2) - ).flatten(-2), - } - for name, variant in wrong.items(): - dev = (out - variant).abs().max().item() - print(f" {name:<18} max deviation {dev:.6g}") - assert dev > 1.0, (name, dev) - - -def test_zero_gate_is_the_identity() -> None: - hyper_input, block_out, _ = _operands((33,), seed=4) - inject = torch.zeros(33, HC_COUNT, dtype=torch.bfloat16, device=DEV) - assert torch.equal(op(hyper_input, block_out, inject), hyper_input) - - -def test_one_hot_gate_touches_exactly_one_stream() -> None: - """The property a token-broadcast bug destroys; asserted bitwise.""" - hyper_input, block_out, _ = _operands((16,), seed=5) - streams = hyper_input.unflatten(-1, (HC_COUNT, HIDDEN)) - for live in range(HC_COUNT): - inject = torch.zeros(16, HC_COUNT, dtype=torch.bfloat16, device=DEV) - inject[:, live] = 1.0 - out = op(hyper_input, block_out, inject).unflatten(-1, (HC_COUNT, HIDDEN)) - for s in range(HC_COUNT): - if s == live: - expected = (streams[:, s].float() + block_out.float()).to(torch.bfloat16) - assert torch.equal(out[:, s], expected), (live, s) - else: - assert torch.equal(out[:, s], streams[:, s]), (live, s) - - def test_guard_refuses_unequal_leading_dims() -> None: - """Unguarded, a 1-D `inject` or `block_out` broadcasts one row over every token.""" + """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)) - unguarded = op(hyper_input, block_out[0], inject) - assert unguarded.shape == hyper_input.shape, "the hazard this guard covers is gone" - for args in ( - (hyper_input, block_out, inject[0]), - (hyper_input, block_out[0], inject), - (hyper_input, block_out[:4], inject), - (hyper_input, block_out, inject[:4]), - ): - with pytest.raises(AssertionError, match="leading dims must match"): - with validating(op): - op(*args) - - -def test_loud_refusals_are_not_guarded() -> None: - """A width that does not factor, and a host operand, raise on their own.""" - _, block_out, inject = _operands((8,), seed=11) - wrong = torch.randn(8, WIDTH + 16, dtype=torch.bfloat16, device=DEV) - with pytest.raises(RuntimeError, match="don't multiply up to the size of dim"): - op(wrong, block_out, inject) - hyper_input, _, inject = _operands((8,), seed=12) - with pytest.raises(RuntimeError, match="Expected all tensors to be on the same device"): - op(hyper_input, torch.randn(8, HIDDEN, dtype=torch.bfloat16), inject) - - -def test_mixed_dtypes_promote_the_output() -> None: - """Recorded, not guarded: an fp32 operand silently widens the residual stream.""" - hyper_input, block_out, _ = _operands((8,), seed=13) - fp32_inject = torch.rand(8, HC_COUNT, dtype=torch.float32, device=DEV) * 2.0 - out = op(hyper_input, block_out, fp32_inject) - assert out.dtype == torch.float32 - _compare(out, op.reference(hyper_input, block_out, fp32_inject)) - fp32_block_out = torch.randn(8, HIDDEN, dtype=torch.float32, device=DEV) - _, _, inject = _operands((8,), seed=13) - assert op(hyper_input, fp32_block_out, inject).dtype == torch.float32 + 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: 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_gather_grouped_gemm_act_fusion.py new file mode 100644 index 000000000000..f23d2789ee81 --- /dev/null +++ b/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 @@ -0,0 +1,399 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the CuteDSL NVFP4 FC1 entry of the Qwen3.8-Flash-Next target. + +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. +""" + +import functools +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.moe.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion import ( # noqa: E501 + E2M1, + E2M1_MAX, + GATE_UP_GROUP, + Fc1Expected, + e2m1_codes, + 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 + 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__ = ["../../../.."] +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" + +DEV = "cuda" +E, K, H, INTER, VEC = 512, 10, 2560, 640, 16 +SWIGLU = int(ActivationType.Swiglu) +#: The activation's input scale: |x| = 5 maps to the e4m3 maximum, as the +#: checkpoint's per-layer `input_scale` does for its own calibration maximum. +ISC1 = 5.0 / (448 * 6) +#: The SwiGLU output maximum `global_sf` is calibrated to (see `_ckpt` in the entry). +ACT_AMAX = 48.0 +GSF = { + "calibrated": 448 * 6 / ACT_AMAX, + "subnormal": 2.0**-12 * 448 * 6 / ACT_AMAX, +} + +#: Tie windows of `_compare`, derived in its docstring. +SF_TIE = 2.0**-20 +CODE_TIE = 2.0**-17 +#: At most this fraction of block scales / codes may take the other side of a tie +#: (4 of them when that is more). +FLIP_BUDGET = 1e-5 + + +# ── the gate ────────────────────────────────────────────────────────────── + + +def _e2m1_rank(packed: torch.Tensor) -> torch.Tensor: + """Packed codes -> signed magnitude rank in -7..7 per element; -0 and +0 are both 0.""" + lo, hi = (packed & 0xF).long(), (packed >> 4).long() + codes = torch.stack([lo, hi], dim=-1).reshape(packed.shape[0], -1) + return torch.where(codes >= 8, -(codes & 7), codes & 7) + + +def _e2m1_rank_of(v: torch.Tensor) -> torch.Tensor: + codes = e2m1_codes(v).long() + return torch.where(codes >= 8, -(codes & 7), codes & 7) + + +def _rank_value(rank: torch.Tensor) -> torch.Tensor: + lut = torch.tensor(E2M1, dtype=torch.float64, device=rank.device) + return lut[rank.abs()] * torch.sign(rank).double() + + +def _deviation(out: tuple[torch.Tensor, torch.Tensor], ref: Fc1Expected) -> dict: + """Every disagreement between the op and the reference on the rows it writes. + + Block scales: equal, or the two neighbours of an e4m3 midpoint that the exact + pre-rounding scale lies within `SF_TIE` (relative) of -- a `sf_flip`; anything + else is `sf_outside`. Codes are re-derived under the op's OWN block scale (so + a flipped scale does not cascade into its 16 codes): equal, or the two + neighbours of an e2m1 midpoint that the exact `scaled / sf` lies within + `CODE_TIE` of -- a `code_flip`; anything else is `code_outside`. Blocks whose + scale is zero dequantize to zero whatever their codes hold and are not + code-checked. `*_far` counts the disagreements that are not between + neighbours at all. + """ + c, c_sf = out + m, half = c.shape + inter = half * 2 + nb = c_sf.numel() // m + vec = inter // nb + rows = ref.rows + s_op = unswizzle(c_sf.view(torch.uint8), m, nb)[rows].long() + s_ref = unswizzle(ref.c_sf.view(torch.uint8), m, nb)[rows].long() + pre = ref.scaled.view(-1, nb, vec).abs().amax(-1) / E2M1_MAX + + v_op, v_ref = e4m3_value(s_op), e4m3_value(s_ref) + sf_diff = s_op != s_ref + finite = (s_op < 0x7F) & (s_ref < 0x7F) + sf_adj = sf_diff & finite & ((s_op - s_ref).abs() == 1) + sf_dist = ((pre - (v_op + v_ref) / 2).abs() / pre.clamp_min(1e-300)).masked_fill(~sf_adj, 0.0) + sf_flip = sf_adj & (sf_dist <= SF_TIE) + sf_outside = sf_diff & ~sf_flip + + rank_op = _e2m1_rank(c.view(torch.uint8)[rows]) + scale = v_op.repeat_interleave(vec, dim=1) + live = scale > 0 + q = torch.where(live, ref.scaled / scale, torch.zeros_like(scale)) + rank_exp = _e2m1_rank_of(q) + code_diff = live & (rank_op != rank_exp) + code_adj = code_diff & ((rank_op - rank_exp).abs() == 1) + mid = (_rank_value(rank_op) + _rank_value(rank_exp)) / 2 + code_dist = (q - mid).abs().masked_fill(~code_adj, 0.0) + code_flip = code_adj & (code_dist <= CODE_TIE) + code_outside = code_diff & ~code_flip + + return dict( + rows=int(rows.sum()), + blocks=s_op.numel(), + sf_flips=int(sf_flip.sum()), + sf_outside=int(sf_outside.sum()), + sf_far=int((sf_diff & ~sf_adj).sum()), + elements=int(live.sum()), + code_flips=int(code_flip.sum()), + code_outside=int(code_outside.sum()), + code_far=int((code_diff & ~code_adj).sum()), + ) + + +def _compare(out: tuple[torch.Tensor, torch.Tensor], ref: Fc1Expected) -> None: + """Bitwise on every written row, except at rounding ties of bounded width. + + The op accumulates in fp32 on the tensor cores, evaluates SwiGLU with + `ex2.approx` / `rcp.approx`, and divides by the block scale through + `rcp.approx`; the reference is exact. The op's pre-rounding value is + therefore off by a handful of fp32 roundings (2^-24 relative each), and + only where the exact value lies that close to a rounding midpoint may the + op land on its other side. Everywhere else the two agree bit for bit. + + - Block scale: a flip only within `SF_TIE` = 2^-20 (16 fp32 roundings) + relative of an e4m3 midpoint. + - Code: a flip only within `CODE_TIE` = 2^-17 of an e2m1 midpoint, in units + of the block scale -- 16 fp32 roundings of a block maximum of at most 6, + rounded up to a power of two. + - Flips: at most `FLIP_BUDGET` = 1e-5 of the scales and of the codes (4 when + that is fewer), so the op may not take the far side of every near-tie. + + Every wrongly prepared operand puts most of the block scales outside the gate. + """ + c, c_sf = out + assert c.dtype == torch.float4_e2m1fn_x2, c.dtype + assert c_sf.dtype == torch.uint8, c_sf.dtype + assert c.shape == ref.c.shape, (tuple(c.shape), tuple(ref.c.shape)) + assert c_sf.shape == ref.c_sf.shape, (tuple(c_sf.shape), tuple(ref.c_sf.shape)) + d = _deviation(out, ref) + assert d["rows"] > 0, d + assert d["sf_outside"] == 0, f"block scales off by more than a tie: {d}" + assert d["code_outside"] == 0, f"codes off by more than a tie: {d}" + assert d["sf_flips"] <= max(4, FLIP_BUDGET * d["blocks"]), f"too many scale ties flipped: {d}" + assert d["code_flips"] <= max(4, FLIP_BUDGET * d["elements"]), ( + f"too many code ties flipped: {d}" + ) + + +# ── operands ────────────────────────────────────────────────────────────── + + +@functools.lru_cache(maxsize=None) +def _experts(num_experts: int = E): + """Checkpoint-format routed experts: per expert `[up; gate]` NVFP4 rows with + LINEAR e4m3 block scales and one `weight_scale_2` shared by both halves, as the + checkpoint stores them. Expert magnitudes spread over 2^-1..2^1 so a wrong + per-expert alpha shows.""" + g = torch.Generator(device=DEV).manual_seed(20261013) + spread = torch.exp2(torch.rand(num_experts, generator=g, device=DEV) * 2 - 1) + w = torch.empty(num_experts, 2 * INTER, H // 2, dtype=torch.uint8, device=DEV) + s = torch.empty(num_experts, 2 * INTER, H // VEC, dtype=torch.uint8, device=DEV) + ws2 = torch.empty(num_experts, dtype=torch.float32, device=DEV) + for e in range(num_experts): + dense = torch.randn(2 * INTER, H, generator=g, device=DEV) * (float(spread[e]) / H**0.5) + dense = dense.to(torch.bfloat16) + gs = (448 * 6 / dense.abs().amax().float()).reshape(1) + q, sf = torch.ops.trtllm.fp4_quantize(dense, gs, VEC, False, False) + w[e] = q.view(2 * INTER, H // 2) + s[e] = sf.view(2 * INTER, H // VEC) + ws2[e] = 1.0 / gs + return w, s, ws2 + + +def _row_perm(n: int) -> torch.Tensor: + """Destination row -> source row of the `[up; gate]` stack, in `GATE_UP_GROUP`-row + groups.""" + idx = torch.arange(n, device=DEV) + return idx.view(2, n // (2 * GATE_UP_GROUP), GATE_UP_GROUP).transpose(0, 1).reshape(-1) + + +@functools.lru_cache(maxsize=None) +def _weights(): + """`(weight, weight_scale, alpha)` in the layout CuteDSL reads -- what the + target's load-time transform produces.""" + w, s, ws2 = _experts() + perm = _row_perm(2 * INTER) + w_k = w[:, perm].contiguous() + s_k = torch.ops.trtllm.block_scale_interleave(s[:, perm].contiguous()).view(s.shape) + alpha = (ISC1 * ws2).contiguous() + return w_k.view(torch.float4_e2m1fn_x2), s_k, alpha + + +def _activation(tokens: int, seed: int, zero_tokens: int = 0): + """The target's FC1 input: bf16 hidden rows through `fp4_quantize(x, 1/isc1, + 16, False, False)`, scales viewed `[T, H/16]`.""" + g = torch.Generator(device=DEV).manual_seed(seed) + x = torch.randn(tokens, H, generator=g, device=DEV).to(torch.bfloat16) + x[:zero_tokens] = 0 + gs = torch.tensor([1.0 / ISC1], dtype=torch.float32, device=DEV) + q, sf = torch.ops.trtllm.fp4_quantize(x, gs, VEC, False, False) + return q.view(torch.float4_e2m1fn_x2), sf.view(tokens, H // VEC) + + +def _ids(dist: str, tokens: int, seed: int) -> torch.Tensor: + """Expert ids `[tokens, K]` int32: each row's top-K of uniform scores.""" + assert dist == "router", dist + g = torch.Generator(device=DEV).manual_seed(seed) + scores = torch.rand(tokens, E, device=DEV, generator=g) + return scores.topk(K, dim=-1).indices.to(torch.int32) + + +def _sort(ids: torch.Tensor, tile: int): + """moe_sort as the CuteDSL MoE calls it; the routing weights do not reach FC1.""" + scales = torch.full(ids.shape, 0.1, dtype=torch.float32, device=DEV) + return torch.ops.trtllm.moe_sort( + token_selected_experts=ids, + token_final_scales=scales, + num_experts=E, + top_k=K, + local_expert_offset=0, + local_num_experts=E, + tile_tokens_dim=tile, + ) + + +def _args(x, x_sf, sort, tile, gsf): + """The keyword arguments the target passes, by the op's own names.""" + weight, weight_scale, alpha = _weights() + t2e, t2mn, _e2p, p2e, _total, n_tiles = sort + return dict( + input=x, + weight=weight, + input_scale=x_sf, + weight_scale=weight_scale, + alpha=alpha, + tile_idx_to_group_idx=t2e, + tile_idx_to_mn_limit=t2mn, + permuted_idx_to_expanded_idx=p2e, + num_non_exiting_tiles=n_tiles, + global_sf=torch.tensor([gsf], dtype=torch.float32, device=DEV), + num_experts=E, + top_k=K, + num_local_experts=E, + local_expert_offset=0, + tile_size=tile, + scaling_vector_size=VEC, + activation_type=SWIGLU, + ) + + +def _problem(tokens, tile, seed, dist="router", gsf="calibrated", zero_tokens=0): + x, x_sf = _activation(tokens, seed, zero_tokens) + sort = _sort(_ids(dist, tokens, seed + 1), tile) + return _args(x, x_sf, sort, tile, GSF[gsf]), sort + + +def _call(args, mma_n): + """One op call with the autotuner tactic whose MMA tile is `(tile_size, mma_n)`, + selected through the tuner's own capture/replay so the call still goes through + the op.""" + tuner = AutoTuner.get() + with tuner.capture() as cap: + op(**args) + for config in cap: + ((_runner, tactic),) = config + if isinstance(tactic, tuple) and tactic[0] == (args["tile_size"], mma_n): + with tuner.replay(config): + return op(**args) + raise AssertionError(f"no tactic with MMA tile ({args['tile_size']}, {mma_n})") + + +def _pairs(out, sort, tokens): + """Each (token, slot) pair's output row and its block-scale row, in expanded + order -- what FC2 and the combine see of it, whatever row the sort gave it.""" + c, c_sf = out + e2p = sort[2].reshape(-1).long()[: tokens * K] + sf = unswizzle(c_sf, c.shape[0], c_sf.numel() // c.shape[0]) + return c.view(torch.uint8)[e2p], sf[e2p] + + +def _clone(v): + """A copy that does not depend on copy kernels for the fp4 dtype.""" + if not isinstance(v, torch.Tensor): + return v + return v.view(torch.uint8).clone().view(v.dtype) if v.element_size() == 1 else v.clone() + + +def _assert_pairs_equal(a, b, what): + assert torch.equal(a[0], b[0]), f"{what}: codes differ" + assert torch.equal(a[1], b[1]), f"{what}: block scales differ" + + +def _cell_id(cell): + s = cell.spec + return ( + f"t{s['tile_size']}-n{s['mma_n']}-T{s['tokens']}-{s['dist']}-{s['gsf']}-z{s['zero_tokens']}" + ) + + +# ── the certified cells ─────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[_cell_id(c) for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + s = cell.spec + assert (s["num_experts"], s["top_k"], s["hidden"], s["intermediate"]) == (E, K, H, INTER) + seed = zlib.crc32(cell.why.encode()) + args, sort = _problem(s["tokens"], s["tile_size"], seed, s["dist"], s["gsf"], s["zero_tokens"]) + before = {k: _clone(v) for k, v in args.items() if isinstance(v, torch.Tensor)} + out = _call(args, s["mma_n"]) + _compare(out, op.reference(**args)) + for k, v in before.items(): + assert torch.equal(args[k].view(torch.uint8), v.view(torch.uint8)), f"{k} was written" + if s["zero_tokens"]: + zero = sort[2][: s["zero_tokens"]].reshape(-1).long() + c, c_sf = out + sf = unswizzle(c_sf, c.shape[0], INTER // VEC) + assert bool((sf[zero] == 0).all()), "an all-zero token's block scales are not zero" + assert bool((c.view(torch.uint8)[zero] & 0x77 == 0).all()), "a zero token's codes are not 0" + + +# ── invariance: a pair's output depends on the pair alone ───────────────── + + +def test_tiles_and_tactics_agree_per_pair(certifying): # noqa: F811 + """Tile 128 / 256, MMA N 128 / 256 and the tuner's fallback all produce every + pair's codes and block scales bit for bit the same at the 8192-row prefill, so a + cell certifies the tactic the autotuner picks in serving as well as the one it + ran.""" + certifying(op) + tokens = 8192 + x, x_sf = _activation(tokens, seed=tokens) + ids = _ids("router", tokens, seed=tokens + 1) + seen = {} + for tile in (128, 256): + sort = _sort(ids, tile) + args = _args(x, x_sf, sort, tile, GSF["calibrated"]) + for mma_n in (128, 256): + seen[(tile, mma_n)] = _pairs(_call(args, mma_n), sort, tokens) + seen[(tile, "fallback")] = _pairs(op(**args), sort, tokens) + if tile == 128: + _compare(op(**args), op.reference(**args)) + 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_cute_dsl_nvfp4_grouped_gemm.py new file mode 100644 index 000000000000..f744f13a4402 --- /dev/null +++ b/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 @@ -0,0 +1,375 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the CuteDSL NVFP4 FC2 entry of the Qwen3.8-Flash-Next target. + +The entry owns the contract. This file builds the checkpoint-shaped operands -- +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. +""" + +import functools +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.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 + cute_dsl_nvfp4_grouped_gemm as op, +) +from tensorrt_llm._torch.autotuner import AutoTuner + +__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" + +DEV = "cuda" +E, K, H, INTER, VEC = 512, 10, 2560, 640, 16 +NB = INTER // VEC +#: The SwiGLU output maximum FC1's `global_sf` is calibrated to (see `_ckpt` in the entry). +ACT_AMAX = 48.0 +GSF = { + "calibrated": 448 * 6 / ACT_AMAX, + "subnormal": 2.0**-12 * 448 * 6 / ACT_AMAX, +} + +#: Tie window of `_compare`, in units of the element's accumulation magnitude +#: (derived in its docstring). +TIE = 2.0**-20 +#: At most this fraction of the gated elements with a nonzero magnitude may take +#: the other side of a tie (4 of them when that is more). +FLIP_BUDGET = 1e-5 + + +# ── the gate ────────────────────────────────────────────────────────────── + + +def _bf16_order(x: torch.Tensor) -> torch.Tensor: + """bf16 -> an integer that is monotonic in its value, consecutive between + neighbouring bf16 values, and equal for +0 and -0.""" + bits = x.view(torch.int16).int() & 0xFFFF + mag = bits & 0x7FFF + return torch.where(bits >= 0x8000, -mag, mag) + + +def _deviation(out: torch.Tensor, ref: Fc2Expected) -> dict: + """Every disagreement between the op and the reference on the rows it writes. + + An element is `equal` when its value matches the reference's `out` (+0 and + -0 alike). Otherwise it is a `flip` when the two are neighbouring bf16 + values and the exact value lies within `TIE * magnitude` of the midpoint + between them, and `outside` when it is anything else. `far` counts the + disagreements that are not between neighbours at all. + """ + span = ref.exact.shape[0] + idx = torch.nonzero(ref.rows[:span]).squeeze(1) + o = out[idx] + r = ref.out[idx] + v = ref.exact[idx] + s = ref.magnitude[idx] + of, rf = o.double(), r.double() + finite = torch.isfinite(of) + diff = (of != rf) | ~finite + adj = diff & finite & ((_bf16_order(o) - _bf16_order(r)).abs() == 1) + dist = ((v - (of + rf) / 2).abs() / s.clamp_min(1e-300)).masked_fill(~adj, 0.0) + flip = adj & (dist <= TIE) + outside = diff & ~flip + return dict( + rows=int(idx.numel()), + elements=o.numel(), + nonzero=int((s > 0).sum()), + nonfinite=int((~finite).sum()), + flips=int(flip.sum()), + outside=int(outside.sum()), + far=int((diff & ~adj).sum()), + ) + + +def _compare(out: torch.Tensor, ref: Fc2Expected) -> None: + """Bitwise on every written row, except at bf16 rounding ties of derived width. + + The op accumulates the 640 products of a row in fp32 on the tensor cores -- + ten MMA k-steps of 64 -- multiplies by alpha in fp32 and rounds to bf16. Each + step rounds a partial sum no larger than the element's `magnitude` (`|alpha| * + sum |a_k w_k|`), so the op's fp32 value is within 11 roundings of `2^-24 * + magnitude` of the exact one; only where the exact value lies that close to a + bf16 midpoint may the op land on the other neighbour. Everywhere else the two + agree bit for bit. + + - Flip: only to the neighbouring bf16 value, and only within `TIE` = 2^-20 + (16 fp32 roundings) of `magnitude` from the midpoint. + - Flips: at most `FLIP_BUDGET` = 1e-5 of the gated elements whose magnitude is + nonzero (4 when that is fewer) -- padding rows do not widen the budget. + - No NaN or infinity. + + Every wrongly prepared operand puts most of the pair rows' elements outside the + gate, most of them not even a bf16 neighbour away. + """ + assert out.dtype == torch.bfloat16, out.dtype + assert out.shape == ref.out.shape, (tuple(out.shape), tuple(ref.out.shape)) + d = _deviation(out, ref) + assert d["rows"] > 0, d + assert d["nonfinite"] == 0, f"non-finite output on written rows: {d}" + assert d["outside"] == 0, f"elements off by more than a tie: {d}" + assert d["flips"] <= max(4, FLIP_BUDGET * d["nonzero"]), f"too many ties flipped: {d}" + + +# ── operands ────────────────────────────────────────────────────────────── + + +def _swizzle(sf: torch.Tensor) -> torch.Tensor: + """`[rows, cols]` scale bytes -> the flat 128x4-swizzled buffer; inverse of `unswizzle`.""" + rows, cols = sf.shape + return sf.reshape(rows // 128, 4, 32, cols // 4, 4).permute(0, 3, 2, 1, 4).reshape(-1) + + +def _quantized_experts(num_experts: int, rows: int, cols: int, seed: int): + """Per expert, a Gaussian `[rows, cols]` matrix quantized the way the checkpoint + stores it: NVFP4 codes, LINEAR e4m3 block scales, one `weight_scale_2`. Expert + magnitudes spread over 2^-1..2^1 so a wrong per-expert alpha shows.""" + g = torch.Generator(device=DEV).manual_seed(seed) + spread = torch.exp2(torch.rand(num_experts, generator=g, device=DEV) * 2 - 1) + w = torch.empty(num_experts, rows, cols // 2, dtype=torch.uint8, device=DEV) + s = torch.empty(num_experts, rows, cols // VEC, dtype=torch.uint8, device=DEV) + ws2 = torch.empty(num_experts, dtype=torch.float32, device=DEV) + for e in range(num_experts): + dense = torch.randn(rows, cols, generator=g, device=DEV) * (float(spread[e]) / cols**0.5) + dense = dense.to(torch.bfloat16) + gs = (448 * 6 / dense.abs().amax().float()).reshape(1) + q, sf = torch.ops.trtllm.fp4_quantize(dense, gs, VEC, False, False) + w[e] = q.view(rows, cols // 2) + s[e] = sf.view(rows, cols // VEC) + ws2[e] = 1.0 / gs + return w, s, ws2 + + +@functools.lru_cache(maxsize=None) +def _down_proj(): + """The checkpoint's `down_proj` experts: `[E, H, I/2]` codes, `[E, H, I/16]` + linear scales, `weight_scale_2` per expert.""" + return _quantized_experts(E, H, INTER, seed=20261013) + + +@functools.lru_cache(maxsize=None) +def _weights(gsf: str = "calibrated"): + """`(weight, weight_scale, alpha)` in the layout CuteDSL reads -- what the target's + load-time transform produces. `alpha = down_proj.input_scale * weight_scale_2`, + with the input scale the inverse of FC1's `global_sf`.""" + w, s, ws2 = _down_proj() + s_k = torch.ops.trtllm.block_scale_interleave(s.contiguous()).view(s.shape) + alpha = (ws2 / GSF[gsf]).contiguous() + return w.contiguous().view(torch.float4_e2m1fn_x2), s_k, alpha + + +def _pair_rows(pairs: int, seed: int, zero_pairs: int = 0) -> torch.Tensor: + """SwiGLU outputs of FC1, one bf16 `[I]` row per (token, slot) pair, in expanded + order: `up * silu(gate)` of Gaussian halves.""" + g = torch.Generator(device=DEV).manual_seed(seed) + up = torch.randn(pairs, INTER, generator=g, device=DEV) * 2 + gate = torch.randn(pairs, INTER, generator=g, device=DEV) * 2 + y = (up * gate * torch.sigmoid(gate)).to(torch.bfloat16) + y[:zero_pairs] = 0 + return y + + +def _fc2_input(y: torch.Tensor, sort, gsf: str, garbage_seed: int = 0): + """This op's input as the CuteDSL FC1 returns it: each pair's row NVFP4-quantized + with `global_sf` at its permuted row, block scales 128x4-swizzled, the padding + rows of live tiles zero, and the rows past the live tiles holding stale bytes.""" + t2e, _t2mn, e2p, p2e, _total, n_tiles = sort + m = p2e.numel() + tile = m // t2e.numel() + span = int(n_tiles) * tile + x = torch.zeros(m, INTER, dtype=torch.bfloat16, device=DEV) + x[e2p.reshape(-1).long()[: y.shape[0]]] = y + gs = torch.tensor([GSF[gsf]], dtype=torch.float32, device=DEV) + q, sf = torch.ops.trtllm.fp4_quantize(x, gs, VEC, False, False) + q = q.view(torch.uint8).reshape(m, INTER // 2) + sf = sf.view(torch.uint8).reshape(m, NB) + g = torch.Generator(device=DEV).manual_seed(garbage_seed) + q[span:] = torch.randint(0, 256, q[span:].shape, generator=g, device=DEV, dtype=torch.uint8) + sf[span:] = torch.randint(0, 256, sf[span:].shape, generator=g, device=DEV, dtype=torch.uint8) + return q.view(torch.float4_e2m1fn_x2), _swizzle(sf).contiguous() + + +def _ids(dist: str, tokens: int, seed: int) -> torch.Tensor: + """Expert ids `[tokens, K]` int32: each row's top-K of uniform scores.""" + assert dist == "router", dist + g = torch.Generator(device=DEV).manual_seed(seed) + scores = torch.rand(tokens, E, device=DEV, generator=g) + return scores.topk(K, dim=-1).indices.to(torch.int32) + + +def _sort(ids: torch.Tensor, tile: int): + """moe_sort as the CuteDSL MoE calls it; the routing weights do not reach FC2.""" + scales = torch.full(ids.shape, 0.1, dtype=torch.float32, device=DEV) + return torch.ops.trtllm.moe_sort( + token_selected_experts=ids, + token_final_scales=scales, + num_experts=E, + top_k=K, + local_expert_offset=0, + local_num_experts=E, + tile_tokens_dim=tile, + ) + + +def _args(x, x_sf, sort, tile, gsf="calibrated"): + """The keyword arguments the target passes, by the op's own names.""" + weight, weight_scale, alpha = _weights(gsf) + return dict( + input=x, + weight=weight, + input_scale=x_sf, + weight_scale=weight_scale, + alpha=alpha, + tile_idx_to_group_idx=sort[0], + num_non_exiting_tiles=sort[5], + num_experts=E, + top_k=K, + num_local_experts=E, + local_expert_offset=0, + tile_size=tile, + output_dtype=torch.bfloat16, + scaling_vector_size=VEC, + ) + + +def _problem(tokens, tile, seed, dist="router", gsf="calibrated", zero_tokens=0): + sort = _sort(_ids(dist, tokens, seed + 1), tile) + y = _pair_rows(tokens * K, seed, zero_pairs=zero_tokens * K) + x, x_sf = _fc2_input(y, sort, gsf, garbage_seed=seed + 2) + return _args(x, x_sf, sort, tile, gsf), sort + + +def _call(args, mma_n, cluster_n): + """One op call with the autotuner tactic whose MMA tile is `(tile_size, mma_n)` + and cluster `(tile_size // 128, cluster_n)`, selected through the tuner's own + capture/replay so the call still goes through the op.""" + tuner = AutoTuner.get() + with tuner.capture() as cap: + op(**args) + tile = args["tile_size"] + want = ((tile, mma_n), (tile // 128, cluster_n)) + for config in cap: + ((_runner, tactic),) = config + if tactic == want: + with tuner.replay(config): + return op(**args) + raise AssertionError(f"no tactic {want}") + + +def _pairs(out, sort, tokens): + """Each (token, slot) pair's output row, in expanded order -- what the combine + reads, whatever row the sort gave it.""" + return out[sort[2].reshape(-1).long()[: tokens * K]] + + +def _clone(v): + """A copy that does not depend on copy kernels for the fp4 dtype.""" + if not isinstance(v, torch.Tensor): + return v + return v.view(torch.uint8).clone().view(v.dtype) if v.element_size() == 1 else v.clone() + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.view(torch.int16) + + +def _assert_pairs_equal(a, b, what): + assert torch.equal(_bits(a), _bits(b)), f"{what}: rows differ" + + +def _cell_id(cell): + s = cell.spec + return ( + f"t{s['tile_size']}-n{s['mma_n']}-c{s['cluster_n']}-T{s['tokens']}-{s['dist']}" + f"-{s['gsf']}-z{s['zero_tokens']}" + ) + + +# ── the certified cells ─────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[_cell_id(c) for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + s = cell.spec + assert (s["num_experts"], s["top_k"], s["hidden"], s["intermediate"]) == (E, K, H, INTER) + seed = zlib.crc32(cell.why.encode()) + args, sort = _problem(s["tokens"], s["tile_size"], seed, s["dist"], s["gsf"], s["zero_tokens"]) + before = {k: _clone(v) for k, v in args.items() if isinstance(v, torch.Tensor)} + out = _call(args, s["mma_n"], s["cluster_n"]) + _compare(out, op.reference(**args)) + for k, v in before.items(): + assert torch.equal(_u8(args[k]), _u8(v)), f"{k} was written" + if s["zero_tokens"]: + zero = sort[2][: s["zero_tokens"]].reshape(-1).long() + assert bool((out[zero] == 0).all()), "an all-zero token's rows are not zero" + + +# ── invariance: a pair's row depends on the pair alone ──────────────────── + + +def test_tiles_and_tactics_agree_per_pair(certifying): # noqa: F811 + """Tile 128 / 256, MMA N 128 / 256, cluster N 1 / 2 and the tuner's fallback + all produce every pair's row bit for bit the same at the 8192-row prefill, so a + cell certifies the tactic the autotuner picks in serving as well as the one it + ran.""" + certifying(op) + tokens = 8192 + ids = _ids("router", tokens, seed=tokens + 1) + y = _pair_rows(tokens * K, seed=tokens) + seen = {} + for tile in (128, 256): + sort = _sort(ids, tile) + x, x_sf = _fc2_input(y, sort, "calibrated", garbage_seed=tile) + args = _args(x, x_sf, sort, tile) + for mma_n in (128, 256): + for cluster_n in (1, 2): + seen[(tile, mma_n, cluster_n)] = _pairs(_call(args, mma_n, cluster_n), sort, tokens) + fallback = op(**args) + seen[(tile, "fallback")] = _pairs(fallback, sort, tokens) + _compare(fallback, op.reference(**args)) + first = seen[(128, 128, 1)] + for key, got in seen.items(): + _assert_pairs_equal(got, first, f"T={tokens} {key} vs (128, 128, 1)") + + +# ── the guard ───────────────────────────────────────────────────────────── + + +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_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 new file mode 100644 index 000000000000..ddef4b35ea36 --- /dev/null +++ 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 @@ -0,0 +1,367 @@ +# 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`. +BACKEND_TOLERANCE = 3e-3 + + +# ── 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_fp4_block_scale_moe_renorm_permuted.py b/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 new file mode 100644 index 000000000000..fbdc5c576c43 --- /dev/null +++ b/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 @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for ops/moe/fp4_block_scale_moe_renorm_permuted. + +Per cell: the emulated finalize of the three outputs is **bitwise** the +finalized entry's output on identical operands (so the fold a consumer makes +with the same arithmetic is bitwise what the runner's finalize stored), the +outputs have the documented shapes, dtypes and index map, and the finalized +rows sit inside the pre-routed entry's band against the fp32 reference. The +operand construction is the finalized entry's test's, loaded by path. +""" + +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.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 + 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 + fp4_block_scale_moe_renorm_routed as finalized, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "fp4_block_scale_moe_renorm_permuted requires a CUDA device" + + +def _routed_test(): + path = Path(__file__).resolve().parent / "test_modeling_v2_fp4_block_scale_moe_renorm_routed.py" + spec = importlib.util.spec_from_file_location("_fp4_moe_renorm_routed_test", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +R = _routed_test() +H = R.H + + +def _args(args, data, sf, logits, scalars, spec, top_k=None): + return ( + logits, + data, + sf.view(torch.float8_e4m3fn), + args["gemm1_weights"], + args["gemm1_weights_scale"], + args["gemm2_weights"], + args["gemm2_weights_scale"], + *scalars, + spec["num_experts"], + spec["top_k"] if top_k is None else top_k, + spec["intermediate"], + ) + + +@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, k, hidden = spec["tokens"], spec["top_k"], spec["hidden"] + seed = zlib.crc32(cell.why.encode()) % 2**31 + args, ref, data, sf, x, logits, ids, wts, g2, scalars = R._operands(spec, seed) + call = _args(args, data, sf, logits, scalars, spec) + + rows, topk_w, topk_rows = op(*call) + assert rows.dtype == torch.bfloat16 and rows.dim() == 2 and rows.shape[1] == hidden + assert topk_w.shape == (t, k) and topk_w.dtype == torch.bfloat16 + assert topk_rows.shape == (t, k) and topk_rows.dtype == torch.int32 + # Every slot maps to its own row (no expert parallelism: no -1 slots). + assert int(topk_rows.min()) >= 0 and int(topk_rows.max()) < rows.shape[0] + assert topk_rows.flatten().unique().numel() == t * k, "two slots share a permuted row" + # The weights are the routing's (as a set per token; the slot order is the runner's). + assert torch.equal(topk_w.sort(-1).values, wts.sort(-1).values), "routing weights differ" + + # 1. Bitwise the finalized entry once finalized with the runner's arithmetic. + y = finalize(rows, topk_w, topk_rows) + y_fin = finalized(*call) + differing = (y.view(torch.int16) != y_fin.view(torch.int16)).any(dim=1).sum().item() + assert differing == 0, f"{differing} of {t} rows differ from the finalized entry" + + # 2. The fp32 reference, through the pre-routed entry's gate. + 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(): + spec = op.CELLS[0].spec + args, _, data, sf, _, logits, _, _, _, scalars = R._operands(spec, 13) + 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(): + spec = next(c.spec for c in op.CELLS if c.spec["tokens"] == 1) + args, _, data, sf, _, logits, _, _, _, scalars = R._operands(spec, 43) + 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_block_scale_moe_renorm_routed.py new file mode 100644 index 000000000000..c0ead50efd89 --- /dev/null +++ b/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 @@ -0,0 +1,127 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for ops/moe/fp4_block_scale_moe_renorm_routed. + +Two claims per cell: the logits-routed call is **bitwise** the composition it +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) is the pre-routed entry's test's, loaded by path so the two tests +cannot drift apart on what a valid operand is. Beyond the cells: the refusal +of an uncertified (num_experts, top_k) pair. +""" + +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.moe.fp4_block_scale_moe_renorm_routed import ( # noqa: E501 + CERTIFIED_PAIRS, +) +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 + 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 + renorm_moe_routing_op, +) + +__extra_import_path__ = ["../../../.."] +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[4] + / "moe" + / "test_modeling_v2_fp4_block_scale_moe_runner.py" + ) + spec = importlib.util.spec_from_file_location("_fp4_moe_runner_test_helpers", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +H = _helpers() + + +def _logits(spec, gen): + t, e = spec["tokens"], spec["num_experts"] + if spec["dist"] == "coarse": + v = torch.randint(-1, 2, (t, e), generator=gen, device=DEV).float() + else: + v = torch.randn(t, e, generator=gen, device=DEV) + return v.to(spec["logits_dtype"]).contiguous() + + +def _operands(spec, seed): + gen = torch.Generator(device=DEV).manual_seed(seed) + args, ref = H._many_experts() + data, sf, x = H._rand_activation(spec["tokens"], spec["hidden"], G1, gen) + logits = _logits(spec, gen) + ids, wts = renorm_moe_routing_op(logits, spec["top_k"], torch.bfloat16) + g2 = H._calibrate_g2(x, ids, ref) + scalars = H._scalars(spec["num_experts"], G1, g2) + return args, ref, data, sf, x, logits, ids, wts, g2, scalars + + +def _routed(args, data, sf, logits, scalars, spec, top_k=None): + return op( + logits, + data, + sf.view(torch.float8_e4m3fn), + args["gemm1_weights"], + args["gemm1_weights_scale"], + args["gemm2_weights"], + args["gemm2_weights_scale"], + *scalars, + spec["num_experts"], + spec["top_k"] if top_k is None else top_k, + spec["intermediate"], + ) + + +@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 + seed = zlib.crc32(cell.why.encode()) % 2**31 + args, ref, data, sf, x, logits, ids, wts, g2, (s1, sg, s2) = _operands(spec, seed) + + y = _routed(args, data, sf, logits, (s1, sg, s2), 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, + args, + spec["num_experts"], + spec["top_k"], + topk_ids=ids, + topk_weights=wts, + output1_scale_scalar=s1, + output1_scale_gate_scalar=sg, + output2_scale_scalar=s2, + )[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. + exp = op.reference(logits, x, spec["top_k"], ref["up"], ref["gate"], ref["down"], g2) + H._compare(y, exp.to(torch.bfloat16)) + + +def test_uncertified_pair_refused(): + spec = op.CELLS[0].spec + assert (spec["num_experts"], 9) not in CERTIFIED_PAIRS + args, _, data, sf, _, logits, _, _, _, scalars = _operands(spec, 13) + with validating(op), pytest.raises(AssertionError, match="certified Renormalize"): + _routed(args, data, sf, logits, scalars, spec, top_k=9) 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp4_moe_renorm_cute_dsl.py new file mode 100644 index 000000000000..5f2aeabd5cf1 --- /dev/null +++ b/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 @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for ops/moe/fp4_moe_renorm_cute_dsl. + +The experts are the pre-routed entry's test's (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, +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. +""" + +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 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 + 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 + renorm_moe_routing_op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "fp4_moe_renorm_cute_dsl 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[4] + / "moe" + / "test_modeling_v2_fp4_block_scale_moe_runner.py" + ) + spec = importlib.util.spec_from_file_location("_fp4_moe_runner_test_helpers", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +H = _helpers() +_CD = None + + +def _cute_dsl_experts(ref): + """This entry's operands from the raw expert tensors, through `weights.py`.""" + global _CD + if _CD is None: + up_p, gt_p, up_s, gt_s, dn_p, dn_s = ref["raw"] + e = up_p.shape[0] + fc1_w = torch.stack([W._expert_fc1_w_cd((up_p[i], gt_p[i]), None) for i in range(e)]) + fc1_s = torch.stack([W._expert_fc1_s_cd((up_s[i], gt_s[i]), None) for i in range(e)]) + fc2_w = dn_p.contiguous() + fc2_s = torch.stack([W._expert_fc2_s_cd(dn_s[i], None) for i in range(e)]) + _CD = (fc1_w, fc1_s, fc2_w, fc2_s) + return _CD + + +def _operands(spec, seed): + gen = torch.Generator(device=DEV).manual_seed(seed) + args, ref = H._many_experts() + data, sf, x = H._rand_activation(spec["tokens"], spec["hidden"], G1, gen) + logits = torch.randn(spec["tokens"], spec["num_experts"], generator=gen, device=DEV).to( + spec["logits_dtype"] + ) + ids, _ = renorm_moe_routing_op(logits, spec["top_k"], torch.bfloat16) + g2 = H._calibrate_g2(x, ids, ref) + return args, ref, data, sf, x, logits, g2 + + +def _cute_dsl(ref, data, sf, logits, g2, spec): + fc1_w, fc1_s, fc2_w, fc2_s = _cute_dsl_experts(ref) + s1, sg, s2 = H._scalars(spec["num_experts"], G1, g2) + gsf = torch.tensor([g2], dtype=torch.float32, device=DEV) + return op( + logits, + data, + sf, + fc1_w, + fc1_s, + sg, + fc2_w, + fc2_s, + s2, + gsf, + spec["num_experts"], + spec["top_k"], + ) + + +@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 + seed = zlib.crc32(cell.why.encode()) % 2**31 + args, ref, data, sf, x, logits, g2 = _operands(spec, seed) + y = _cute_dsl(ref, data, sf, logits, g2, spec) + assert y.shape == (spec["tokens"], spec["hidden"]) and y.dtype == torch.bfloat16 + 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(): + spec = {**op.CELLS[0].spec, "tokens": 1024} + args, ref, data, sf, x, logits, g2 = _operands(spec, 9) + fc1_w, fc1_s, fc2_w, fc2_s = _cute_dsl_experts(ref) + 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp8_block_scale_moe_runner.py index dc92b8dd9c6b..b2d77dcf4336 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fp8_block_scale_moe_runner.py @@ -3,9 +3,9 @@ """GPU test for the fp8_block_scale_moe_runner catalog entry. The entry owns the contract: `CELLS`, `reference`, `is_valid`. This -file turns a cell into operands and holds the claims no cell carries: the -intermediate requantization pinned bit-exactly, how far each wrong reference -lands, the expert-window bookkeeping, and the guards. +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. """ import zlib @@ -13,24 +13,15 @@ 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.moe.fp8_block_scale_moe_runner import ( # noqa: E501 - ACT_AMAX_FLOOR, BLOCK, - E4M3_MAX, - dequantize_block_weight, - requantize_1x128, ) 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 as op, ) __extra_import_path__ = ["../../../.."] -from _validating import ( # noqa: E402,F401 — path above; fixture by name - assert_within_ulp, - certifying, - validating, -) +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" @@ -53,8 +44,7 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: 0.342. 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, which is why the intermediate requantization is also - pinned bit-exactly by the test. + -- the closest. """ assert out.dtype == ref.dtype == torch.bfloat16, (out.dtype, ref.dtype) assert out.shape == ref.shape, (out.shape, ref.shape) @@ -101,26 +91,11 @@ 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, - gemm1_clamp_limit=None, - foreign_ids=False, - seed=0, + 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, for one configuration. - - Weights are drawn for the whole routing space and the window sliced out - of them, so a window and the full stack hold the same experts. - """ - local_num_experts = num_experts if local_num_experts is None else local_num_experts - stack = max(num_experts, local_expert_offset + local_num_experts) - w1, w1s, w2, w2s = _weights(hidden, inter, stack, seed=hidden + inter + stack) - window = slice(local_expert_offset, local_expert_offset + local_num_experts) + """The op's arguments by name, pre-routed, every expert local.""" + assert local_expert_offset == 0 and local_num_experts in (None, num_experts) + w1, w1s, w2, w2s = _weights(hidden, inter, num_experts, seed=hidden + inter + num_experts) 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. @@ -128,30 +103,28 @@ def _args( logits = torch.randn(tokens, num_experts, device=DEV, generator=g) weights, ids = torch.topk(torch.softmax(logits, dim=-1), top_k, dim=-1) ids = ids.to(torch.int32) - if foreign_ids: - ids[0, 0], ids[1, 0], ids[2, 0] = -1, num_experts, 40 return dict( routing_logits=None, routing_bias=None, hidden_states=hs, hidden_states_scale=hs_scale, - gemm1_weights=w1[window].contiguous(), - gemm1_weights_scale=w1s[window].contiguous(), - gemm2_weights=w2[window].contiguous(), - gemm2_weights_scale=w2s[window].contiguous(), + gemm1_weights=w1, + gemm1_weights_scale=w1s, + gemm2_weights=w2, + gemm2_weights_scale=w2s, num_experts=num_experts, top_k=top_k, num_fused_shared_experts=None, n_group=None, topk_group=None, intermediate_size=inter, - local_expert_offset=local_expert_offset, - local_num_experts=local_num_experts, + local_expert_offset=0, + local_num_experts=num_experts, routed_scaling_factor=None, routing_method_type=1, topk_weights=weights.to(torch.bfloat16).contiguous(), topk_ids=ids.contiguous(), - gemm1_clamp_limit=gemm1_clamp_limit, + gemm1_clamp_limit=None, ) @@ -162,48 +135,6 @@ def _ulp(out: torch.Tensor, ref: torch.Tensor): return d.abs().max().item(), d.pow(2).mean().sqrt().item() -def _variant( - a, swap=False, quant_fc1_out=True, quant_act=True, count_repeats=True, requant=requantize_1x128 -): - """`reference` with one step changed -- the wrong readings it must refute.""" - tokens, hidden = a["hidden_states"].shape - inter, nb = a["intermediate_size"], hidden // BLOCK - act_in = ( - a["hidden_states"].float().reshape(tokens, nb, BLOCK) - * a["hidden_states_scale"].t().reshape(tokens, nb, 1) - ).reshape(tokens, hidden) - ids = a["topk_ids"].long() - local = ids - a["local_expert_offset"] - live = (local >= 0) & (local < a["local_num_experts"]) - if not count_repeats: - first = torch.ones_like(live) - for j in range(1, ids.shape[1]): - first[:, j] = (ids[:, j : j + 1] != ids[:, :j]).all(dim=1) - live &= first - out = torch.zeros(tokens, hidden, device=DEV) - with true_fp32_matmul(): - for e in torch.unique(local[live]).tolist(): - tok, slot = ((local == e) & live).nonzero(as_tuple=True) - y1 = ( - act_in[tok] - @ dequantize_block_weight(a["gemm1_weights"][e], a["gemm1_weights_scale"][e]).t() - ) - if quant_fc1_out: - y1 = requant(y1) - up, gate = y1[:, :inter], y1[:, inter:] - if swap: - up, gate = gate, up - act = torch.nn.functional.silu(gate) * up - if quant_act: - act = requant(act) - y = ( - act - @ dequantize_block_weight(a["gemm2_weights"][e], a["gemm2_weights_scale"][e]).t() - ) - out.index_add_(0, tok, y * a["topk_weights"][tok, slot].float().unsqueeze(1)) - return out.to(torch.bfloat16) - - def _seed(cell) -> int: return zlib.crc32(cell.why.encode()) @@ -223,162 +154,7 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th _compare(out, ref) -# ── the two internal requantizations, read out bit-exactly ──────────────── - - -def _pow2_requantize(x: torch.Tensor) -> torch.Tensor: - """The MX-style round-up power-of-two block scale -- a wrong recipe here.""" - shp = x.shape - xb = x.reshape(*shp[:-1], shp[-1] // BLOCK, BLOCK) - amax = xb.abs().amax(dim=-1, keepdim=True).clamp(min=ACT_AMAX_FLOOR) - scale = torch.exp2(torch.ceil(torch.log2(amax / E4M3_MAX))) - coded = (xb / scale).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn) - return (coded.float() * scale).reshape(shp) - - -@pytest.mark.parametrize("tokens,width", [(16, 256), (64, 640), (256, 128)]) -def test_intermediate_is_fp8_block_quantized(tokens, width): - """Identity down projection, one slot of weight 1.0: the output is the - requantized intermediate, one bf16 round away. The recipe reproduces - >=99.9% of it bit-exactly (measured 99.95-100%); every neighbouring recipe - reproduces under 10%. The misses are FC1 sums on an e4m3 rounding boundary. - """ - num_experts = 8 - a = _args(tokens, width, width, num_experts, 1, seed=width + tokens) - eye = torch.eye(width, device=DEV).to(torch.float8_e4m3fn) - a["gemm2_weights"] = eye.expand(num_experts, width, width).contiguous() - a["gemm2_weights_scale"] = torch.ones( - num_experts, width // BLOCK, width // BLOCK, device=DEV, dtype=torch.float32 - ) - a["topk_weights"] = torch.ones(tokens, 1, device=DEV, dtype=torch.bfloat16) - out = op(**a) - exact = (op.reference(**a) == out).float().mean().item() - print(f" T={tokens} width={width}: recipe reproduces {100 * exact:.2f}%") - assert exact >= 0.999, f"modelled recipe reproduced only {100 * exact:.2f}%" - - assert torch.equal(_variant(a), op.reference(**a)), "the variant model drifted from reference" - wrong = { - "power-of-two scale": _variant(a, requant=_pow2_requantize), - "no activation requant": _variant(a, quant_act=False), - "no FC1-output requant": _variant(a, quant_fc1_out=False), - "gate/up swapped": _variant(a, swap=True), - } - for label, pred in wrong.items(): - frac = (pred == out).float().mean().item() - print(f" {label}: {100 * frac:.2f}%") - assert frac <= 0.10, f"{label} reproduced {100 * frac:.2f}% -- not discriminated" - - -# ── what the gate separates ─────────────────────────────────────────────── - - -def test_wrong_variants_land_outside_the_gate(): - """Every operand-level mistake, measured against one call at the checkpoint.""" - a = _args(1024, **CKPT, seed=77) - out = op(**a) - _compare(out, op.reference(**a)) - good = _ulp(out, op.reference(**a))[1] - for label, kw in ( - ("gate/up swapped", dict(swap=True)), - ("no FC1-output requant", dict(quant_fc1_out=False)), - ("no activation requant", dict(quant_act=False)), - ("neither requant", dict(quant_fc1_out=False, quant_act=False)), - ): - elt, rms = _ulp(out, _variant(a, **kw)) - print(f" {label}: {elt:.2f} ulp element-wise, {rms:.3f} ulp RMS") - # The RMS gate separates all four; the element gate does not separate - # a missing activation requant, hence the bit-exact pin above. - assert rms > 0.75, f"{label} is inside the RMS gate" - assert rms >= 2.0 * good, f"{label} is within 2x of the correct reference" - - -def test_transposed_activation_scale_is_a_silent_wrong_answer(): - """At `T == H/128` both of the op's dimension checks pass on the transpose.""" - a = _args(20, **CKPT, seed=99) - assert a["hidden_states_scale"].shape == (20, 20) - _compare(op(**a), op.reference(**a)) - swapped = op(**dict(a, hidden_states_scale=a["hidden_states_scale"].t().contiguous())) - elt, rms = _ulp(swapped, op.reference(**a)) - print(f" transposed: {elt:.2f} ulp element-wise, {rms:.3f} ulp RMS") - assert rms > 10.0, "the transposed layout must be visibly wrong" - - -# ── expert parallelism and the window ───────────────────────────────────── - - -def test_expert_parallel_windows_sum_to_the_full_stack(): - """Four 128-wide windows of 512: each matches its own reference, a token - with no slot in the window is zero, and the four sum to the full call.""" - full = _args(64, **CKPT, seed=7) - out = op(**full) - _compare(out, op.reference(**full)) - ids = full["topk_ids"].long() - summed = torch.zeros(out.shape, device=DEV) - for offset in (0, 128, 256, 384): - a = _args(64, **CKPT, local_expert_offset=offset, local_num_experts=128, seed=7) - assert torch.equal(a["topk_ids"], full["topk_ids"]) - got = op(**a) - _compare(got, op.reference(**a)) - dead = (got == 0).all(dim=-1) - expected_dead = ~((ids >= offset) & (ids < offset + 128)).any(dim=-1) - assert torch.equal(dead, expected_dead), "a token with no slot in the window must be zero" - summed += got.float() - # Each window rounds its partial to bf16 first: inside the gate, not bitwise. - _compare(summed.to(torch.bfloat16), op.reference(**full)) - - -def test_window_past_the_routing_space_is_all_zero(): - a = _args(64, 256, 128, 8, 2, local_expert_offset=8, local_num_experts=4, seed=22) - assert bool((op(**a) == 0).all()) - - -def test_repeated_expert_ids_are_counted_once_or_twice_by_token_count(): - """Counted once or twice depending on which routing kernel the token count - selects. Records the verdict per token count; asserts only that each call - matches one of the two readings, not that both readings occur.""" - seen = set() - for tokens in (16, 256): - a = _args(tokens, 256, 128, 8, 2, seed=tokens * 3 + 256) - a["topk_ids"][:, 1] = a["topk_ids"][:, 0] - out = op(**a) - _, twice = _ulp(out, _variant(a, count_repeats=True)) - _, once = _ulp(out, _variant(a, count_repeats=False)) - verdict = "once" if once < twice else "twice" - print(f" T={tokens}: counted {verdict} (once={once:.3f} twice={twice:.3f} ulp RMS)") - assert min(once, twice) <= 0.75, "neither reading of a repeat matches" - seen.add(verdict) - print(f" verdicts across token counts: {sorted(seen)}") - - -# ── arguments ───────────────────────────────────────────────────────────── - - -def test_gemm1_clamp_limit_is_live(): - a = _args(16, **CKPT, gemm1_clamp_limit=1.0, seed=16) - clamped = op(**a) - _compare(clamped, op.reference(**a)) - _, rms = _ulp(clamped, op.reference(**dict(a, gemm1_clamp_limit=None))) - print(f" clamp=1.0 vs the unclamped reference: {rms:.3f} ulp RMS") - assert rms > 0.75, "the clamp must change the result" - b = dict(a) - del b["gemm1_clamp_limit"] - assert torch.equal(op(**dict(a, gemm1_clamp_limit=None)), op(**b)) - - -def test_inert_arguments_are_bitwise_identical(): - a = _args(16, **CKPT, seed=16) - base = op(**a) - for label, kw in ( - ("routed_scaling_factor=1.0", dict(routed_scaling_factor=1.0)), - ("routed_scaling_factor=2.5", dict(routed_scaling_factor=2.5)), - ("routing_method_type=0", dict(routing_method_type=0)), - ("routing_method_type=2", dict(routing_method_type=2)), - ("routing_method_type=4", dict(routing_method_type=4)), - ("n_group=1, topk_group=1", dict(n_group=1, topk_group=1)), - ("tune_max_num_tokens=128", dict(tune_max_num_tokens=128)), - ("use_dp=True", dict(use_dp=True)), - ): - assert torch.equal(op(**dict(a, **kw)), base), f"{label} changed the result" +# ── what a cell cannot carry ────────────────────────────────────────────── def test_autotuned_tactic_matches_the_fallback_tactic(): @@ -401,68 +177,6 @@ def test_autotuned_tactic_matches_the_fallback_tactic(): AutoTuner.get().clear_cache() -def test_registered_fake_returns_the_wrong_hidden_size(): - """An upstream defect, recorded so nothing relies on it: the fake assumes - packed fp4 hidden states and doubles the width.""" - a = _args(16, 256, 128, 8, 2, seed=3) - real = op(**a) - with torch._subclasses.FakeTensorMode(): - fake_args = { - k: (torch.empty(v.shape, dtype=v.dtype, device=DEV) if torch.is_tensor(v) else v) - for k, v in a.items() - } - fake_shape = tuple(op(**fake_args).shape) - print(f" fake says {fake_shape}, the kernel returns {tuple(real.shape)}") - assert fake_shape == (real.shape[0], 2 * real.shape[1]) - - -def test_determinism_and_purity(): - a = _args(16, 512, 256, 16, 4, seed=25) - snaps = {k: v.clone() for k, v in a.items() if torch.is_tensor(v)} - assert torch.equal(op(**a), op(**a)) - for k, v in snaps.items(): - assert torch.equal(a[k], v), f"{k} was mutated" - - -def test_output_buffer_semantics(): - """`output=` is written in place and the call returns an empty `[0]` tensor; - rows of a taller buffer past `T` are untouched.""" - a = _args(16, 256, 128, 8, 2, seed=3) - fresh = op(**a) - buf = torch.full((32, 256), 7.0, device=DEV, dtype=torch.bfloat16) - ret = op(**dict(a, output=buf[:16])) - assert ret.shape == (0,) and ret.dtype == torch.bfloat16 - assert torch.equal(buf[:16], fresh) - assert bool((buf[16:] == 7.0).all()), "rows past T must be untouched" - - -def test_rejections(): - """The op's own loud checks; none of these needs a guard.""" - a = _args(16, 256, 128, 8, 2, seed=3) - - def refuses(label, **kw): - with pytest.raises(RuntimeError) as excinfo: - op(**dict(a, **kw)) - torch.cuda.synchronize() - print(f" {label}: {str(excinfo.value).splitlines()[0][:100]}") - - refuses("bf16 hidden_states", hidden_states=a["hidden_states"].to(torch.bfloat16)) - refuses("bf16 hidden_states_scale", hidden_states_scale=a["hidden_states_scale"].bfloat16()) - refuses("1-D hidden_states_scale", hidden_states_scale=a["hidden_states_scale"].flatten()) - refuses("fp16 weight scale", gemm1_weights_scale=a["gemm1_weights_scale"].half()) - refuses("int64 topk_ids", topk_ids=a["topk_ids"].long()) - refuses("fp32 topk_weights", topk_weights=a["topk_weights"].float()) - refuses("no routing given", topk_ids=None, topk_weights=None) - refuses("intermediate_size mismatch", intermediate_size=256) - refuses("num_experts == top_k", num_experts=2) - refuses("fused shared experts pre-routed", num_fused_shared_experts=1) - refuses("output of the wrong dtype", output=torch.empty(16, 256, device=DEV).half()) - refuses( - "output of the wrong shape", - output=torch.empty(16, 128, device=DEV, dtype=torch.bfloat16), - ) - - # ── the guards ──────────────────────────────────────────────────────────── @@ -474,39 +188,8 @@ def _strided(t: torch.Tensor) -> torch.Tensor: return view -def test_unguarded_hazards_are_silent(): - """What `is_valid` exists for: the raw op takes each of these and answers.""" - a = _args(16, 256, 128, 8, 2, seed=11) - good = op.reference(**a) - _, rms = _ulp(op(**dict(a, hidden_states=_strided(a["hidden_states"]))), good) - print(f" strided hidden_states accepted: {rms:.3f} ulp RMS") - assert rms > 10.0, "a strided operand must be visibly wrong, or the guard is pointless" - base = op(**a) - assert torch.equal(op(**dict(a, act_type=1)), base), "act_type is no longer ignored" - bias = torch.tensor([1e3, -1e3] * 4, device=DEV, dtype=torch.bfloat16) - assert torch.equal(op(**dict(a, routing_bias=bias)), base), ( - "routing_bias is no longer ignored on the pre-routed path" - ) - - -def test_guards_refuse(): +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): - ok = op(**a) - for name in ( - "hidden_states", - "hidden_states_scale", - "gemm1_weights", - "gemm1_weights_scale", - "gemm2_weights", - "gemm2_weights_scale", - "topk_weights", - "topk_ids", - ): - with pytest.raises(AssertionError, match="must be contiguous"): - op(**dict(a, **{name: _strided(a[name])})) - with pytest.raises(AssertionError, match="silently ignored"): - op(**dict(a, act_type=1)) - with pytest.raises(AssertionError, match="routing_bias is ignored"): - op(**dict(a, routing_bias=torch.zeros(8, device=DEV, dtype=torch.bfloat16))) - assert torch.equal(op(**a), ok) + 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py index 5aaa3b5107d6..035ca9c348ce 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_fused_sigmoid_gate_mul_add.py @@ -3,9 +3,8 @@ """GPU test for the fused_sigmoid_gate_mul_add catalog entry. 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: gate layouts -bitwise interchangeable, aliasing, the empty batch, the wrong-variant -discrimination and the loud rejections. +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. """ import zlib @@ -24,38 +23,27 @@ DEV = "cuda" HIDDEN = 2560 -PAD = 64 def _operands(spec, seed): - """`(final, gate, shared, output, buffers)`; `buffers` holds the padded - allocations a `dest="padded"` cell writes a window of.""" + """`(final, gate, shared)`: the in-place combine's operands.""" t, h, dt = spec["tokens"], spec["hidden"], spec["dtype"] + assert spec["dest"] == "inplace", spec["dest"] g = torch.Generator(device=DEV).manual_seed(seed) def rand(*shape, dtype=dt): return torch.randn(*shape, device=DEV, generator=g).to(dtype) - buffers = {} - if spec["dest"] == "padded": - buffers = {k: rand(t + 2, h + PAD) for k in ("final", "shared", "output")} - final, shared, output = (buffers[k][:t, :h] for k in ("final", "shared", "output")) - assert output.stride() == (h + PAD, 1) and not output.is_contiguous() - else: - final, shared = rand(t, h), rand(t, h) - output = torch.empty_like(final) if spec["dest"] == "output" else None + final, shared = rand(t, h), rand(t, h) gate = rand(t, 1, dtype=spec["gate_dtype"]) - layout = spec["gate_layout"] - if layout == "flat": - gate = gate.reshape(-1) - elif layout == "row": - gate = gate.reshape(1, t) - elif layout == "strided": + 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() - return final, gate, shared, output, buffers + else: + assert spec["gate_layout"] == "column", spec["gate_layout"] + return final, gate, shared # ── the certified cells ─────────────────────────────────────────────────── @@ -64,68 +52,21 @@ def rand(*shape, dtype=dt): @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) - final, gate, shared, output, buffers = _operands(cell.spec, zlib.crc32(cell.why.encode())) + final, gate, shared = _operands(cell.spec, zlib.crc32(cell.why.encode())) before = {k: v.clone() for k, v in dict(final=final, gate=gate, shared=shared).items()} - snapshots = {k: v.clone() for k, v in buffers.items()} - ref = op.reference(before["final"], gate, shared, output) + ref = op.reference(before["final"], gate, shared) - out = op(final, gate, shared, output) + out = op(final, gate, shared) - destination = final if output is None else output - assert out.data_ptr() == destination.data_ptr(), "the op returns what it wrote" + assert out.data_ptr() == final.data_ptr(), "the op returns what it wrote" assert out.dtype == final.dtype and out.shape == final.shape torch.testing.assert_close(out, ref) assert torch.equal(gate, before["gate"]) and torch.equal(shared, before["shared"]) - if output is not None: - assert torch.equal(final, before["final"]), "an output buffer leaves the input alone" - if buffers: - t, h = cell.spec["tokens"], cell.spec["hidden"] - o, o0 = buffers["output"], snapshots["output"] - assert torch.equal(o[:t, h:], o0[:t, h:]), "wrote into the row padding" - assert torch.equal(o[t:], o0[t:]), "wrote past num_tokens" - assert torch.equal(buffers["final"], snapshots["final"]) - assert torch.equal(buffers["shared"], snapshots["shared"]) # ── what a cell cannot carry ────────────────────────────────────────────── -def test_gate_layouts_are_the_same_call_bit_for_bit(): - spec = dict( - tokens=33, - hidden=HIDDEN, - dtype=torch.bfloat16, - gate_dtype=torch.bfloat16, - gate_layout="column", - dest="output", - ) - final, gate, shared, _, _ = _operands(spec, 3) - wide = torch.zeros(33, 4, dtype=torch.bfloat16, device=DEV) - wide[:, 2:3] = gate - results = [ - op(final.clone(), g, shared) - for g in (gate, gate.reshape(-1), gate.reshape(1, 33), wide[:, 2:3]) - ] - for r in results[1:]: - assert torch.equal(r, results[0]) - - -def test_the_output_buffer_may_alias_the_shared_expert_output(): - spec = dict( - tokens=17, - hidden=HIDDEN, - dtype=torch.bfloat16, - gate_dtype=torch.bfloat16, - gate_layout="column", - dest="inplace", - ) - final, gate, shared, _, _ = _operands(spec, 8) - expected = op.reference(final, gate, shared) - returned = op(final, gate, shared, output=shared) - assert returned.data_ptr() == shared.data_ptr() - torch.testing.assert_close(shared, expected) - - 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) @@ -134,88 +75,9 @@ def test_zero_tokens_is_a_no_op(): assert returned.shape == (0, HIDDEN) and returned.data_ptr() == final.data_ptr() -def test_wrong_variants_land_far_outside_the_default_gate(): - """The default bf16 band (atol 1e-5, rtol 1.6e-2) separates the correct - combine from every plausible wrong one by more than 100x.""" - spec = dict( - tokens=17, - hidden=HIDDEN, - dtype=torch.bfloat16, - gate_dtype=torch.bfloat16, - gate_layout="column", - dest="inplace", - ) - final, gate, shared, _, _ = _operands(spec, 9) - correct = op.reference(final, gate, shared) - actual = op(final.clone(), gate, shared) - width = 1e-5 + 1.6e-2 * correct.float().abs() - s = torch.sigmoid(gate.reshape(-1, 1).float()) - f, sh = final.float(), shared.float() - - def ratio(wrong: torch.Tensor) -> float: - return ((wrong.float() - correct.float()).abs() / width).max().item() - - assert ratio(actual) <= 1.0 - for name, wrong in { - "no sigmoid, raw gate": f + gate.float().reshape(-1, 1) * sh, - "gate applied to the routed half": s * f + sh, - "one gate for the whole batch": f + s.reshape(-1)[0] * sh, - "multiply instead of add": f * s * sh, - "gate dropped, plain add": f + sh, - }.items(): - r = ratio(wrong.to(torch.bfloat16)) - print(f" {name:36s} {r:10.1f}x the gate") - assert r > 100.0, f"{name} is indistinguishable from the correct result" - - -@pytest.mark.parametrize( - "case", - [ - "final_3d", - "shared_shape_mismatch", - "shared_dtype_mismatch", - "gate_numel_too_small", - "gate_numel_too_large", - "output_shape_mismatch", - "output_dtype_mismatch", - "output_last_stride", - "final_last_stride_without_an_output", - ], -) -def test_metadata_misuse_raises(case: str) -> None: - """Every metadata violation is loud, which is why the entry has no guard.""" +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, gate = ( - torch.randn(4, 128, **bf), - torch.randn(4, 128, **bf), - torch.randn(4, 1, **bf), - ) - output = None - if case == "final_3d": - final = torch.randn(2, 2, 128, **bf) - elif case == "shared_shape_mismatch": - shared = torch.randn(4, 64, **bf) - elif case == "shared_dtype_mismatch": - shared = torch.randn(4, 128, dtype=torch.float16, device=DEV) - elif case == "gate_numel_too_small": - gate = torch.randn(3, 1, **bf) - elif case == "gate_numel_too_large": - gate = torch.randn(5, 1, **bf) - elif case == "output_shape_mismatch": - output = torch.empty(4, 64, **bf) - elif case == "output_dtype_mismatch": - output = torch.empty(4, 128, dtype=torch.float16, device=DEV) - elif case == "output_last_stride": - output = torch.empty(4, 256, **bf)[:, ::2] - elif case == "final_last_stride_without_an_output": - final = torch.randn(4, 256, **bf)[:, ::2] + final, shared = torch.randn(4, 128, **bf), torch.randn(4, 128, **bf) with pytest.raises(AssertionError): - op(final, gate, shared, output=output) - - -def test_cpu_operands_are_rejected_by_triton() -> None: - final = torch.randn(4, 128, dtype=torch.bfloat16) - shared = torch.randn(4, 128, dtype=torch.bfloat16) - gate = torch.randn(4, 1, dtype=torch.bfloat16) - with pytest.raises(ValueError, match="Pointer argument cannot be accessed from Triton"): - op(final, gate, shared) + 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py new file mode 100644 index 000000000000..8d97d4b0dd4f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_sort.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""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. +""" + +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.moe.moe_sort import ( # noqa: E501 + moe_sort as op, +) + +__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" + +DEV = "cuda" +E, K = 512, 10 + +_NAMES = ( + "tile_idx_to_expert_idx", + "tile_idx_to_mn_limit", + "expanded_idx_to_permuted_idx", + "permuted_idx_to_expanded_idx", + "total_num_padded_tokens", + "num_non_exiting_tiles", +) + + +def _compare(out: list[torch.Tensor], ref: list[torch.Tensor]) -> None: + """Integer outputs, so exact -- up to the one freedom the op has. + + Counts and the tile arrays' written prefix are bitwise. The order of an + expert's rows is the kernel's atomic order (it varies between identical calls + from T = 5), so `permuted_idx_to_expanded_idx` is gated per expert run as a + set: each real row's value sorted within its run must equal the reference's + run. `expanded_idx_to_permuted_idx` must then be the exact inverse: -1 + exactly where the expert is not local, and otherwise a real row of the right + run that maps back. Padding rows and tiles past `num_non_exiting_tiles` are + not written by the op and are not compared. + """ + assert len(out) == len(ref) == len(_NAMES), (len(out), len(ref)) + for name, o, r in zip(_NAMES, out, ref): + assert o.dtype == r.dtype == torch.int32, (name, o.dtype, r.dtype) + assert o.shape == r.shape, (name, tuple(o.shape), tuple(r.shape)) + t2e, t2mn, e2p, p2e, total, n_tiles = out + r_t2e, r_t2mn, r_e2p, r_p2e, r_total, r_n_tiles = ref + assert torch.equal(n_tiles.cpu(), r_n_tiles.cpu()), (n_tiles, r_n_tiles) + assert torch.equal(total.cpu(), r_total.cpu()), (total, r_total) + n = int(r_n_tiles) + assert torch.equal(t2e[:n], r_t2e[:n]), "tile_idx_to_expert_idx differs" + assert torch.equal(t2mn[:n], r_t2mn[:n]), "tile_idx_to_mn_limit differs" + + tile = p2e.numel() // t2e.numel() + num_expanded = e2p.numel() + real = r_p2e >= 0 + rows = torch.nonzero(real).squeeze(1) + run = r_t2e[rows // tile].long() + got = p2e[real].long() + assert bool(((got >= 0) & (got < num_expanded)).all()), ( + "permuted_idx_to_expanded_idx holds an out-of-range value on a real row" + ) + assert torch.equal( + torch.sort(run * num_expanded + got).values, run * num_expanded + r_p2e[real].long() + ), "an expert's run of permuted_idx_to_expanded_idx holds the wrong pairs" + + flat, r_flat = e2p.reshape(-1).long(), r_e2p.reshape(-1).long() + not_local = r_flat < 0 + assert bool((flat[not_local] == -1).all()), "a non-local pair is not -1" + p = flat[~not_local] + assert bool(((p >= 0) & (p < p2e.numel())).all()), "a local pair points outside the rows" + assert bool(real[p].all()), "a local pair points at a padding row" + assert torch.equal(p2e[p].long(), torch.arange(num_expanded, device=p.device)[~not_local]), ( + "expanded_idx_to_permuted_idx is not the inverse of permuted_idx_to_expanded_idx" + ) + + +def _ids(dist, tokens, tile, seed, num_experts=E, top_k=K): + """Expert ids, `[tokens, top_k]` int32, every row naming distinct experts.""" + g = torch.Generator(device=DEV).manual_seed(seed) + if dist == "router": + scores = torch.rand(tokens, num_experts, device=DEV, generator=g) + return scores.topk(top_k, dim=-1).indices.to(torch.int32) + if dist == "tile_edges": + edge = (7, 300, 511) + scores = torch.rand(tokens, num_experts, device=DEV, generator=g) + scores[:, list(edge)] = -1.0 + ids = scores.topk(top_k, dim=-1).indices.to(torch.int32) + for slot, (expert, rows) in enumerate(zip(edge, (tile - 1, tile, tile + 1))): + ids[:rows, slot] = expert + return ids + raise ValueError(dist) + + +def _args(ids, tile, seed=0, scales_dtype=torch.float32): + """The CuteDSL call site's keyword arguments at the checkpoint's values.""" + g = torch.Generator(device=DEV).manual_seed(seed + 1) + scales = torch.rand(ids.shape, device=DEV, generator=g).to(scales_dtype) + return dict( + token_selected_experts=ids, + token_final_scales=scales, + num_experts=E, + top_k=K, + local_expert_offset=0, + local_num_experts=E, + tile_tokens_dim=tile, + ) + + +def _check(args): + """Drive the op and gate it through the entry's own reference and `_compare`.""" + out = op(**args) + ref = op.reference(**args) + _compare(out, ref) + return out, ref + + +def _cell_id(cell): + s = cell.spec + return f"t{s['tile_tokens_dim']}-T{s['tokens']}-{s['dist']}" + + +# ── the certified cells ─────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[_cell_id(c) for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + s = cell.spec + seed = zlib.crc32(cell.why.encode()) + ids = _ids(s["dist"], s["tokens"], s["tile_tokens_dim"], seed) + args = _args(ids, s["tile_tokens_dim"], seed, s["scales_dtype"]) + assert (s["num_experts"], s["top_k"]) == (args["num_experts"], args["top_k"]) + assert (s["local_expert_offset"], s["local_num_experts"]) == (0, E) + ids_before = ids.clone() + scales_before = args["token_final_scales"].clone() + out, ref = _check(args) + assert torch.equal(ids, ids_before), "token_selected_experts was written" + assert torch.equal(args["token_final_scales"], scales_before), "token_final_scales was written" + if s["dist"] == "tile_edges": + tile = s["tile_tokens_dim"] + n = int(ref[5]) + limits = out[1][:n].long() + starts = torch.arange(n, device=DEV) * tile + fill = (limits - starts).tolist() + experts = out[0][:n].tolist() + per_expert = {e: [] for e in (7, 300, 511)} + for e, f in zip(experts, fill): + if e in per_expert: + per_expert[e].append(f) + assert per_expert == {7: [tile - 1], 300: [tile], 511: [tile, 1]}, per_expert + + +# ── CUDA graphs: captured row counts, on new ids ────────────────────────── + +#: The largest decode capture: 32 requests at MTP verify width 4, cluster kernel. +_CAPTURED_ROWS = 128 + + +def _capture(args): + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(2): + op(**args) + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_out = op(**args) + return graph, static_out + + +def test_captured_rows_replay_on_new_ids(certifying): # noqa: F811 + """The graph is replayed on two fresh id sets and gated against the reference, + and against an eager call on the same ids: counts and tile arrays bitwise. At + the target's tile, 128.""" + certifying(op) + tile, rows = 128, _CAPTURED_ROWS + args = _args(_ids("router", rows, tile, seed=rows), tile, seed=rows) + graph, static_out = _capture(args) + for trial in (1, 2): + fresh = _ids("router", rows, tile, seed=1000 * trial + rows) + args["token_selected_experts"].copy_(fresh) + graph.replay() + eager = op(**dict(args, token_selected_experts=fresh.clone())) + ref = op.reference(**args) + _compare(static_out, ref) + _compare(eager, ref) + n = int(ref[5]) + for i in (0, 1): + assert torch.equal(static_out[i][:n], eager[i][:n]), (trial, i) + for i in (4, 5): + 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_unpermute_inplace.py new file mode 100644 index 000000000000..1ff8171c9b9e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_moe_unpermute_inplace.py @@ -0,0 +1,268 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the moe_unpermute_inplace entry of the Qwen3.8-Flash-Next target. + +The entry owns the contract. This file lays a cell's FC2 rows out the way the +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. +""" + +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.moe.moe_unpermute_inplace import ( # noqa: E501 + moe_unpermute_inplace as op, +) + +__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" + +DEV = "cuda" +E, K, H = 512, 10, 2560 +NAN = float("nan") + +#: The largest decode capture: 32 requests at MTP verify width 4. +_CAPTURED_ROWS = 128 + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Bitwise. + + `reference` is the computation the op performs, in its order and with its + roundings, so no band is needed and none is given. The gate sees the + accumulation order: the same sum taken in reverse slot order, with each + product rounded before its add, as the float64 sum rounded once, or + accumulated in bf16, differs in many elements and is rejected. A failure in a + few isolated one-ulp elements points at the kernel's rounding (a lost + multiply-add contraction, say); a wrong row or scale moves whole rows. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (tuple(out.shape), tuple(ref.shape)) + bits = {2: torch.int16, 4: torch.int32}[out.element_size()] + differ = out.view(bits) != ref.view(bits) + n = int(differ.sum()) + if n: + first = torch.nonzero(differ)[0].tolist() + raise AssertionError( + f"{n} of {out.numel()} elements differ bitwise from the slot-order fp32 " + f"sum; first at {first}: {out[tuple(first)].item()} vs " + f"{ref[tuple(first)].item()}" + ) + + +def _spec(tokens, tile_tokens_dim=128, values="normal"): + """The fields `_problem` reads; the same names as a cell's spec.""" + return dict(tokens=tokens, tile_tokens_dim=tile_tokens_dim, values=values) + + +def _ids(tokens: int, seed: int) -> torch.Tensor: + """Expert ids, `[tokens, K]` int32: each row's top-10 of uniform scores.""" + g = torch.Generator(device=DEV).manual_seed(seed) + scores = torch.rand(tokens, E, device=DEV, generator=g) + return scores.topk(K, dim=-1).indices.to(torch.int32) + + +def _scales(tokens: int, seed: int) -> torch.Tensor: + """fp32 routing weights, a softmax over each row's ten slots.""" + g = torch.Generator(device=DEV).manual_seed(seed + 1) + return torch.softmax(torch.randn(tokens, K, device=DEV, generator=g), dim=-1) + + +def _rows(tokens: int, seed: int, values: str = "normal") -> torch.Tensor: + """The FC2 output row of each (token, slot) pair, `[tokens, K, H]` bf16.""" + g = torch.Generator(device=DEV).manual_seed(seed + 2) + v = torch.randn(tokens, K, H, device=DEV, generator=g) + if values == "wide": + v = v * torch.exp2(torch.randint(-30, 31, v.shape, device=DEV, generator=g).float()) + elif values != "normal": + raise ValueError(values) + return v.to(torch.bfloat16) + + +def _sort(ids: torch.Tensor, scales: torch.Tensor, tile: int): + """moe_sort as the CuteDSL MoE calls it: (expanded_idx_to_permuted_idx, padded rows).""" + out = torch.ops.trtllm.moe_sort( + token_selected_experts=ids, + token_final_scales=scales, + num_experts=E, + top_k=K, + local_expert_offset=0, + local_num_experts=E, + tile_tokens_dim=tile, + ) + return out[2], out[3].numel() + + +def _place(rows: torch.Tensor, e2p: torch.Tensor, num_rows: int) -> torch.Tensor: + """The FC2 output as the op receives it: each live pair's row at its permuted + index, NaN in every row the op has no business reading.""" + perm = torch.full((num_rows, H), NAN, dtype=rows.dtype, device=DEV) + flat = e2p.reshape(-1) + live = flat >= 0 + perm[flat[live].long()] = rows.reshape(-1, H)[live] + return perm + + +def _output(tokens: int) -> torch.Tensor: + return torch.full((tokens, H), NAN, dtype=torch.bfloat16, device=DEV) + + +def _problem(spec: dict, seed: int): + """A cell's arguments, plus the per-pair rows and ids they were built from.""" + tokens = spec["tokens"] + ids = _ids(tokens, seed) + scales = _scales(tokens, seed) + e2p, num_rows = _sort(ids, scales, spec["tile_tokens_dim"]) + rows = _rows(tokens, seed, spec["values"]) + args = dict( + permuted_input=_place(rows, e2p, num_rows), + output=_output(tokens), + expanded_idx_to_permuted_idx=e2p, + topk_scales=scales, + ) + return args + + +def _bits(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.uint8) + + +def _same_bits(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.shape == b.shape and a.dtype == b.dtype and torch.equal(_bits(a), _bits(b)) + + +def _cell_id(cell): + s = cell.spec + return f"t{s['tile_tokens_dim']}-T{s['tokens']}-{s['values']}" + + +# ── the certified cells ─────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[_cell_id(c) for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + s = cell.spec + assert (s["num_experts"], s["top_k"], s["hidden_size"]) == (E, K, H) + assert (s["dtype"], s["scales_dtype"]) == (torch.bfloat16, torch.float32) + args = _problem(s, zlib.crc32(cell.why.encode())) + before = { + n: args[n].clone() + for n in ("permuted_input", "expanded_idx_to_permuted_idx", "topk_scales") + } + out = op(**args) + assert out is args["output"], "the op must return the tensor it wrote" + _compare(out, op.reference(**args)) + for name, snapshot in before.items(): + assert _same_bits(args[name], snapshot), f"{name} was written" + + +# ── CUDA graphs: captured row counts, on new data ───────────────────────── + + +def _capture(fn): + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(2): + fn() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static = fn() + return graph, static + + +def test_captured_rows_replay_on_new_data(certifying): # noqa: F811 + """The sort, the FC2 rows' placement and this op are captured together, as the + target runs them; the graph is replayed on two fresh sets of ids, scales and + rows. The permuted buffer is not cleared between replays, so the rows the new + sort does not use hold the previous placement's rows. Each replay is gated + against the reference, and is bitwise equal to an eager call on the same data + through its own sort. At the target's tile, 128. + """ + certifying(op) + tile = 128 + stale_rows = 0 + tokens = _CAPTURED_ROWS + ids = _ids(tokens, seed=tokens) + scales = _scales(tokens, seed=tokens) + rows = _rows(tokens, seed=tokens) + num_rows = _sort(ids, scales, tile)[1] + perm = torch.full((num_rows, H), NAN, dtype=torch.bfloat16, device=DEV) + out = _output(tokens) + + def step(ids=ids, scales=scales, rows=rows, perm=perm, out=out, tile=tile): + e2p = _sort(ids, scales, tile)[0] + perm.index_copy_(0, e2p.reshape(-1).long(), rows.reshape(-1, H)) + op( + permuted_input=perm, + output=out, + expanded_idx_to_permuted_idx=e2p, + topk_scales=scales, + ) + return e2p + + graph, e2p = _capture(step) + for trial in (1, 2): + seed = 1000 * trial + tokens + ids.copy_(_ids(tokens, seed)) + scales.copy_(_scales(tokens, seed)) + rows.copy_(_rows(tokens, seed)) + out.fill_(NAN) + graph.replay() + args = dict( + permuted_input=perm, + output=out, + expanded_idx_to_permuted_idx=e2p, + topk_scales=scales, + ) + _compare(out, op.reference(**args)) + used = torch.zeros(num_rows, dtype=torch.bool, device=DEV) + used[e2p.reshape(-1).long()] = True + stale_rows += int((~used & ~perm.isnan().any(1)).sum()) + + e2p_eager, n = _sort(ids.clone(), scales.clone(), tile) + eager = op( + permuted_input=_place(rows.clone(), e2p_eager, n), + output=_output(tokens), + expanded_idx_to_permuted_idx=e2p_eager, + topk_scales=scales.clone(), + ) + assert _same_bits(out, eager), trial + 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_renorm_moe_routing_op.py index fcf05d43e29e..5c7fe656dd28 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_renorm_moe_routing_op.py @@ -3,9 +3,8 @@ """GPU test for the renorm_moe_routing_op catalog entry. The entry owns the contract. This file turns a cell into logits and holds the -claims no cell carries: the return order, the tie-break against `torch.topk`, -the sibling op one flag away, the sweep of the old certified surface, the edge -behaviour, and the guards. +claims no cell carries: equivalence to HF's `norm_topk_prob` routing, and the +contiguity guard. """ import zlib @@ -95,211 +94,32 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th ) -def test_old_certified_surface(): - """The rest of the surface the entry was certified over, through the entry's - own reference and gate: the token column in both output dtypes, every input - dtype against every output dtype, and the num_experts / topk sweep across - the kernel's power-of-two instantiation boundaries.""" - for tokens in (1, 2, 7, 8, 64, 512, 2048, 4096, 8192): - logits = _logits(tokens, 512, torch.bfloat16, seed=tokens) - for out_dtype in (torch.float32, torch.bfloat16): - _check(logits, 10, out_dtype) - base = _logits(256, 512, torch.float32, magnitude=2.0, seed=8) - for in_dtype in (torch.bfloat16, torch.float16, torch.float32): - for out_dtype in (torch.float32, torch.bfloat16, None): - _check(base.to(in_dtype), 10, out_dtype) - for num_experts in (1, 2, 8, 31, 32, 33, 64, 96, 128, 129, 256, 384, 511, 512): - _check(_logits(64, num_experts, torch.bfloat16, seed=num_experts), min(10, num_experts)) - logits = _logits(256, 512, torch.bfloat16, seed=9) - for topk in (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 16): # topk 0: test_degenerate_shapes - _check(logits, topk, torch.bfloat16) - - -# ── the traps ───────────────────────────────────────────────────────────── - - -def test_return_order_is_ids_then_weights(): - logits = _logits(256, 512, torch.bfloat16, seed=1) - first, second = op(logits, 10, torch.bfloat16) - assert first.dtype == torch.int32 and second.dtype == torch.bfloat16 - top1 = torch.sort(logits, dim=-1, descending=True, stable=True).indices[:, 0] - assert torch.equal(first[:, 0], top1.to(torch.int32)) +# ── what a cell cannot carry ────────────────────────────────────────────── def test_equivalent_to_hf_norm_topk_prob(): """Softmax over all experts, top-k, divide by the selected sum: the same operator up to rounding (5.96e-08 measured); skipping the renormalization lands three orders of magnitude outside that.""" - for dtype in (torch.bfloat16, torch.float32): - logits = _logits(1024, 512, dtype, magnitude=2.0, seed=2) - _, weights = op(logits, 10, torch.float32) - hf_values, _ = torch.topk(torch.softmax(logits.float(), dim=-1), 10, dim=-1) - hf_weights = hf_values / hf_values.sum(-1, keepdim=True) - # ids can differ on tie rows; the sorted weight multiset cannot - torch.testing.assert_close( - weights.sort(-1, descending=True).values, - hf_weights.sort(-1, descending=True).values, - rtol=1e-5, - atol=1e-7, - ) - assert (weights - hf_values).abs().max().item() > 1e-2 - - -def test_distinct_from_default_moe_routing_op(): - """The sibling one flag away: same ids, weights that do not sum to 1.""" - logits = _logits(512, 512, torch.bfloat16, seed=3) - ids, weights = op(logits, 10, torch.float32) - other_ids, other_weights = torch.ops.trtllm.default_moe_routing_op(logits, 10, torch.float32) - assert torch.equal(ids, other_ids) - row_sums = other_weights.sum(-1) - assert float(row_sums.max()) < 0.99 - assert (weights - other_weights).abs().max().item() > 1e-2 + logits = _logits(1024, 512, torch.bfloat16, magnitude=2.0, seed=2) + _, weights = op(logits, 10, torch.float32) + hf_values, _ = torch.topk(torch.softmax(logits.float(), dim=-1), 10, dim=-1) + hf_weights = hf_values / hf_values.sum(-1, keepdim=True) + # ids can differ on tie rows; the sorted weight multiset cannot torch.testing.assert_close( - other_weights / row_sums.unsqueeze(-1), weights, rtol=1e-6, atol=1e-7 + weights.sort(-1, descending=True).values, + hf_weights.sort(-1, descending=True).values, + rtol=1e-5, + atol=1e-7, ) + assert (weights - hf_values).abs().max().item() > 1e-2 -def test_ties_break_toward_lower_expert_index(): - """Opposite to torch.topk, and common at 512 bf16 experts.""" - logits = torch.zeros(4, 64, dtype=torch.bfloat16, device=DEV) - logits[0, :] = 1.0 - logits[1, 3] = logits[1, 7] = 5.0 - logits[2, ::2] = 2.0 - logits[3, 0] = logits[3, 63] = 1.0 - ids, _ = _check(logits, 4, torch.float32) - expected = [[0, 1, 2, 3], [3, 7, 0, 1], [0, 2, 4, 6], [0, 63, 1, 2]] - assert ids.tolist() == expected - - dense = _logits(4096, 512, torch.bfloat16, seed=4) - ids, _ = _check(dense, 10, torch.float32) - torch_topk_ids = torch.topk(dense, 10, dim=-1).indices.to(torch.int32) - disagreeing = int((ids != torch_topk_ids).any(-1).sum()) - selected = torch.gather(dense, 1, ids.long()) - tied = int((selected[:, :-1] == selected[:, 1:]).any(-1).sum()) - print(f" {tied} of 4096 rows tie inside the top-10; torch.topk disagrees on {disagreeing}") - assert disagreeing > 0 and tied > dense.shape[0] // 4 - for tokens, num_experts, topk in ((512, 128, 8), (256, 32, 4)): - _check(_logits(tokens, num_experts, torch.bfloat16, "coarse", seed=5), topk, torch.float32) - - -def test_slot_order_and_weight_monotonicity(): - logits = _logits(512, 512, torch.bfloat16, seed=6) - ids, weights = _check(logits, 10, torch.float32) - selected = torch.gather(logits.float(), 1, ids.long()) - assert bool((selected[:, :-1] >= selected[:, 1:]).all()) - assert bool((weights[:, :-1] >= weights[:, 1:]).all()) - - -def test_output_dtype_defaults_to_fp32_and_casts_after_the_softmax(): - logits = _logits(1024, 512, torch.float32, magnitude=3.0, seed=7) - default_ids, default_weights = op(logits, 10) - assert default_weights.dtype == torch.float32 - explicit_ids, explicit_weights = op(logits, 10, torch.float32) - assert torch.equal(default_ids, explicit_ids) and torch.equal(default_weights, explicit_weights) - _, bf16_weights = _check(logits, 10, torch.bfloat16) - order = torch.sort(logits, dim=-1, descending=True, stable=True).indices[:, :10] - narrowed = torch.gather(logits, 1, order).to(torch.bfloat16).float() - e = torch.exp(narrowed - narrowed.max(-1, keepdim=True).values) - cast_first = (e / e.sum(-1, keepdim=True)).to(torch.bfloat16) - differing = int((bf16_weights != cast_first).sum()) - print(f" cast-before-softmax disagrees on {differing} of {bf16_weights.numel()}") - assert differing > bf16_weights.numel() // 10 - - -# ── edges ───────────────────────────────────────────────────────────────── - - -def test_degenerate_shapes(): - ids, weights = op(_logits(8, 512, torch.bfloat16), 0, torch.float32) - assert ids.shape == (8, 0) and weights.shape == (8, 0) - ids, weights = op(torch.empty(0, 512, dtype=torch.bfloat16, device=DEV), 10, torch.float32) - assert ids.shape == (0, 10) and weights.shape == (0, 10) - assert ids.dtype == torch.int32 and weights.dtype == torch.float32 - - -def test_input_not_mutated_and_deterministic(): - logits = _logits(4096, 512, torch.bfloat16, seed=10) - snapshot = logits.clone() - ids_a, weights_a = op(logits, 10, torch.bfloat16) - ids_b, weights_b = op(logits, 10, torch.bfloat16) - assert torch.equal(logits, snapshot) - assert torch.equal(ids_a, ids_b) and torch.equal(weights_a, weights_b) - assert weights_a.data_ptr() not in (weights_b.data_ptr(), logits.data_ptr()) - stream = torch.cuda.Stream() - with torch.cuda.stream(stream): - ids_c, weights_c = op(logits, 10, torch.bfloat16) - stream.synchronize() - assert torch.equal(ids_a, ids_c) and torch.equal(weights_a, weights_c) - - -def test_non_finite_logits_poison_the_whole_row(): - logits = _logits(3, 512, torch.float32) - logits[0, 5], logits[1, 7] = float("inf"), float("nan") - logits[2, :] = float("-inf") - ids, weights = op(logits, 10, torch.float32) - assert bool(torch.isnan(weights).all()) - assert int(ids[0, 0]) == 5 and int(ids[1, 0]) == 7 - # -inf is a soft mask: excluded while topk fits in the finite experts... - masked = _logits(4, 512, torch.float32, seed=1) - masked[:, :8] = float("-inf") - masked_ids, masked_weights = op(masked, 10, torch.float32) - assert bool(torch.isfinite(masked_weights).all()) and int(masked_ids.min()) >= 8 - # ...and at weight exactly 0 once it does not - starved = _logits(4, 512, torch.float32, seed=2) - starved[:, 7:] = float("-inf") - starved_ids, starved_weights = op(starved, 10, torch.float32) - assert bool(torch.isfinite(starved_weights).all()) - assert bool((starved_weights[:, 7:] == 0).all()) and bool((starved_weights[:, :7] > 0).all()) - assert int(starved_ids.max()) < 512 - - -def test_op_rejects_unsupported_domains(): - """Loud on the op's own; none of these needs a guard.""" - logits = _logits(8, 512, torch.bfloat16) - cases = [ - ("1D router_logits", lambda: op(logits[0], 6)), - ("3D router_logits", lambda: op(logits.unsqueeze(0), 6)), - ("fp64 router_logits", lambda: op(logits.double(), 6)), - ("int32 router_logits", lambda: op(logits.int(), 6)), - ("fp16 output_dtype", lambda: op(logits, 6, torch.float16)), - ("fp16 output_dtype, fp16 logits", lambda: op(logits.half(), 6, torch.float16)), - ("fp64 output_dtype", lambda: op(logits, 6, torch.float64)), - ("int32 output_dtype", lambda: op(logits, 6, torch.int32)), - ("topk 17", lambda: op(logits, 17)), - ("negative topk", lambda: op(logits, -1)), - ("num_experts 513", lambda: op(_logits(8, 513, torch.bfloat16), 6)), - ("cpu router_logits", lambda: op(logits.cpu(), 6)), - ] - for tag, call in cases: - with pytest.raises((RuntimeError, ValueError, NotImplementedError)): - call() - - -# ── the guards ──────────────────────────────────────────────────────────── - - -def test_unguarded_hazards_are_silent(): - """What `is_valid` exists for: the raw op answers both of these.""" - wide = _logits(8, 1024, torch.bfloat16, seed=12) - sliced = wide[:, :512] - dense = wide.reshape(-1)[: 8 * 512].reshape(8, 512).contiguous() - strided_ids, _ = op(sliced, 10, torch.float32) - assert torch.equal(strided_ids, op(dense, 10, torch.float32)[0]), ( - "the strided read is no longer a dense reinterpretation" - ) - assert not torch.equal(strided_ids, op.reference(sliced, 10, torch.float32)[0]) - narrow = _logits(4, 8, torch.bfloat16, seed=13) - padded_ids, padded_weights = op(narrow, 12, torch.float32) - assert bool((padded_ids[:, 8:] == 0).all()) and bool((padded_weights[:, 8:] == 0).all()) +# ── the guard ───────────────────────────────────────────────────────────── -def test_guards_refuse(): +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): - with pytest.raises(AssertionError, match="contiguous"): - op(wide[:, :512], 10) - with pytest.raises(AssertionError, match="contiguous"): - op(wide.t(), 6) - with pytest.raises(AssertionError, match="exceeds num_experts"): - op(_logits(4, 8, torch.bfloat16), 12) - _check(wide[:, :512].contiguous(), 10, torch.float32) + 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py index a6491f715bd3..3d86b01b2f55 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_flashinfer_gemma_rmsnorm.py @@ -34,14 +34,9 @@ def randn(*s): rows, off = shape[0], spec["col_offset"] buffer = randn(rows, spec["row_width"]) x = buffer[:, off : off + width] - elif "row_step" in spec: - buffer = randn(shape[0] * spec["row_step"], width) - x = buffer[:: spec["row_step"]] else: buffer = x = randn(*shape) - weight = (torch.randn(width, generator=g, device="cuda") * spec.get("weight_scale", 1.0)).to( - dtype - ) + weight = torch.randn(width, generator=g, device="cuda").to(dtype) return buffer, x, weight, spec["eps"] @@ -62,65 +57,14 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res torch.testing.assert_close(out, op.reference(x, weight, eps)) -def test_weight_bias_is_one_not_zero() -> None: - """Discrimination: the scale is (1 + w), not w. - - Measured: 0 away from the (1 + w) reference, ~3.98 from the plain-w one. - """ - _, x, weight, eps = _build(dict(shape=(4, 5120), dtype=torch.bfloat16, eps=1e-6), seed=0) - out = op(x, weight, eps) - xf = x.float() - plain = (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) * weight.float()).to(x.dtype) - gemma = op.reference(x, weight, eps) - - torch.testing.assert_close(out, gemma) - gemma_dev = (out.float() - gemma.float()).abs().max().item() - plain_dev = (out.float() - plain.float()).abs().max().item() - assert plain_dev > 100 * max(gemma_dev, 2.0**-8), (gemma_dev, plain_dev) - - def test_guard_refuses_fp8_before_the_kernel_answers() -> None: - """Both fp8 dtypes are in the DSL dtype table: accepted, and a tensor comes back.""" - for dtype in (torch.float8_e4m3fn, torch.float8_e5m2): - x = torch.randn(4, 2560, device="cuda").to(dtype) - w = torch.randn(2560, device="cuda").to(dtype) - unguarded = torch.ops.trtllm.flashinfer_gemma_rmsnorm(x, w, 1e-6) - assert unguarded.dtype is 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) - - -def test_unsupported_dtypes_raise() -> None: - for dtype in (torch.float64, torch.int8, torch.uint8): - x = torch.ones(4, 2560, dtype=dtype, device="cuda") - w = torch.ones(2560, dtype=dtype, device="cuda") - with pytest.raises(KeyError): + """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) - - -def test_shape_dtype_and_device_mismatches_raise() -> None: - """All loud, which is why `is_valid` restates none of them.""" - w = torch.randn(2560, dtype=torch.bfloat16, device="cuda") - x = torch.randn(4, 2560, dtype=torch.bfloat16, device="cuda") - with pytest.raises(ValueError): - op(x, w.float(), 1e-6) - with pytest.raises(ValueError): - op(x, torch.randn(2048, dtype=torch.bfloat16, device="cuda"), 1e-6) - with pytest.raises(ValueError): - op(x, w.cpu(), 1e-6) - with pytest.raises(ValueError): - op(x.cpu(), w.cpu(), 1e-6) - with pytest.raises(ValueError): - op(torch.randn(2560, dtype=torch.bfloat16, device="cuda"), w, 1e-6) - # a strided normalized dim: "Mismatched mX.strides[1] ... expected to be 1" - with pytest.raises(ValueError, match="strides"): - op(torch.randn(8, 2560, 2, dtype=torch.bfloat16, device="cuda")[..., 0], w, 1e-6) - with pytest.raises(ValueError): - op( - torch.randn(2, 4, 8, 128, dtype=torch.bfloat16, device="cuda"), - torch.randn(128, dtype=torch.bfloat16, device="cuda"), - 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_layer_norm.py index e297e7674808..0e07d7975518 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_layer_norm.py @@ -29,13 +29,6 @@ WEIGHT_LO, WEIGHT_HI = -1.210938, 17.5 BIAS_LO, BIAS_HI = -25.625, 14.875 -#: torch.testing's untouched defaults, restated for the discrimination test. -DEFAULT_TOL = { - torch.bfloat16: (1.6e-2, 1e-5), - torch.float16: (1e-3, 1e-5), - torch.float32: (1.3e-6, 1e-5), -} - def _gen(seed: int) -> torch.Generator: return torch.Generator(device=DEVICE).manual_seed(seed) @@ -50,186 +43,25 @@ def _affine(gen, width=WIDTH, dtype=torch.bfloat16): def _build(spec, seed): gen = _gen(seed) shape, dtype = spec["shape"], spec["dtype"] - if spec.get("layout", "contiguous") == "column_slice": - rows, width = shape - wide = torch.randn(rows, 2 * width, generator=gen, device=DEVICE).to(dtype) - x = wide[:, width:] - assert not x.is_contiguous() and x.stride(-1) == 1 - else: - x = torch.randn(shape, generator=gen, device=DEVICE).to(dtype) - weight, bias = _affine(gen, width=shape[-1], dtype=spec.get("param_dtype", dtype)) - return x, weight, (bias if spec.get("bias", True) else None), EPS + x = torch.randn(shape, generator=gen, device=DEVICE).to(dtype) + weight, bias = _affine(gen, width=shape[-1], dtype=dtype) + return x, weight, bias, EPS @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) x, weight, bias, eps = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) - before = [t.clone() for t in (x, weight, bias) if t is not None] + before = [t.clone() for t in (x, weight, bias)] out = op(x, weight, bias, eps) assert out.shape == x.shape and out.dtype == x.dtype assert out.is_contiguous() and out.data_ptr() != x.data_ptr() - after = [t for t in (x, weight, bias) if t is not None] + after = [x, weight, bias] assert all(torch.equal(a, b) for a, b in zip(before, after)), "inputs must not be mutated" - if x.dim() == 3: - assert torch.equal(out.reshape(-1, WIDTH), op(x.reshape(-1, WIDTH), weight, bias, eps)) - if not x.is_contiguous(): - assert torch.equal(out, op(x.contiguous(), weight, bias, eps)) 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 _fp64(x, weight, bias, eps): - xd = x.double() - mean = xd.mean(dim=-1, keepdim=True) - var = (xd - mean).square().mean(dim=-1, keepdim=True) - out = (xd - mean) / torch.sqrt(var + eps) * weight.double() - return out + bias.double() if bias is not None else out - - -def _worst_ratio(got, expected_fp64, dtype): - """max |got - expected| / (atol + rtol * |expected|) at `dtype`'s default gate.""" - rtol, atol = DEFAULT_TOL[dtype] - return ((got.double() - expected_fp64).abs() / (atol + rtol * expected_fp64.abs())).max().item() - - -def test_bf16_result_is_the_correctly_rounded_reference(): - """The bf16 gate is set by the output dtype, not by the kernel. - - Where the two roundings differ, the bound is against the summed terms' - magnitude: `x_hat * weight` and `bias` cancel outright on some elements at - this checkpoint's affine, and a bound relative to that residue bounds - nothing. fp32 evaluation costs <= ~2^-24 of the term magnitude and each - bf16 store <= 2^-8 of the result, so 2^-7 covers both twice over. - """ - gen = _gen(99) - x = torch.randn(4096, WIDTH, generator=gen, device=DEVICE).to(torch.bfloat16) - weight, bias = _affine(gen) - - out = op(x, weight, bias, EPS) - rounded = op.reference(x, weight, bias, EPS) - agree = (out == rounded).double().mean().item() - assert agree > 0.999, f"only {agree:.6%} of elements are the correctly rounded value" - - xd = x.double() - mean = xd.mean(dim=-1, keepdim=True) - var = (xd - mean).square().mean(dim=-1, keepdim=True) - terms = ((xd - mean) / torch.sqrt(var + EPS) * weight.double()).abs() + bias.double().abs() - deviation = (out.double() - rounded.double()).abs() - assert (deviation <= 2.0**-7 * terms).all(), (deviation / terms).max().item() - - -def test_wrong_variants_land_outside_the_same_gate(): - """The default bf16 gate separates this computation from its neighbours.""" - gen = _gen(3) - x = torch.randn(4096, WIDTH, generator=gen, device=DEVICE).to(torch.bfloat16) - weight, bias = _affine(gen) - - expected = _fp64(x, weight, bias, EPS) - xd, wd, bd = x.double(), weight.double(), bias.double() - mean = xd.mean(dim=-1, keepdim=True) - var = (xd - mean).square().mean(dim=-1, keepdim=True) - wrong = { - "rmsnorm": xd / torch.sqrt(xd.square().mean(dim=-1, keepdim=True) + EPS) * wd + bd, - "bias_dropped": _fp64(x, weight, None, EPS), - "weight_as_delta": (xd - mean) / torch.sqrt(var + EPS) * (wd + 1.0) + bd, - "unbiased_variance": (xd - mean) - / torch.sqrt((xd - mean).square().sum(dim=-1, keepdim=True) / (WIDTH - 1) + EPS) - * wd - + bd, - } - - out = op(x, weight, bias, EPS) - correct_ratio = _worst_ratio(out, expected, torch.bfloat16) - print(f"\n correct error/allowance {correct_ratio:.3e}") - assert correct_ratio < 1.0, correct_ratio - for name, value in wrong.items(): - ratio = _worst_ratio(value.to(torch.bfloat16), expected, torch.bfloat16) - print(f" {name:18s} error/allowance {ratio:.3e}") - assert ratio > 1.0, f"{name} is inside the gate at error/allowance = {ratio:.6e}" - with pytest.raises(AssertionError): - torch.testing.assert_close(out, value.to(torch.bfloat16)) - - -def test_eps_is_inside_the_square_root(): - """Only observable where `var` is comparable to `eps`: a near-constant fp32 - row. The kernel must sit orders of magnitude nearer `1/sqrt(var + eps)` - than `1/(sqrt(var) + eps)`.""" - gen = _gen(17) - x = torch.full((4, WIDTH), 1.0, device=DEVICE, dtype=torch.float32) - x[:, ::2] += 1e-3 - weight, bias = _affine(gen, dtype=torch.float32) - - out = op(x, weight, bias, EPS).double() - xd, wd, bd = x.double(), weight.double(), bias.double() - mean = xd.mean(dim=-1, keepdim=True) - var = (xd - mean).square().mean(dim=-1, keepdim=True) - inside = (xd - mean) / torch.sqrt(var + EPS) * wd + bd - outside = (xd - mean) / (torch.sqrt(var) + EPS) * wd + bd - to_inside = (out - inside).abs().max().item() - to_outside = (out - outside).abs().max().item() - assert to_outside > 100.0 * to_inside, (to_inside, to_outside) - - -def test_zero_variance_row_returns_the_bias_exactly(): - gen = _gen(31) - x = torch.full((4, WIDTH), 2.5, device=DEVICE, dtype=torch.bfloat16) - weight, bias = _affine(gen) - out = op(x, weight, bias, EPS) - assert torch.isfinite(out.float()).all() - assert torch.equal(out, bias.expand_as(out).contiguous()) - - -def test_runs_unchanged_under_inference_mode(): - gen = _gen(55) - x = torch.randn(256, WIDTH, generator=gen, device=DEVICE).to(torch.bfloat16) - weight, bias = _affine(gen) - eager = op(x, weight, bias, EPS) - with torch.inference_mode(): - inferred = op(x, weight, bias, EPS).clone() - assert torch.equal(inferred, eager) - - -def test_exactly_one_kernel_is_launched(): - """The atomicity claim: one `_layer_norm_fwd_1pass_kernel`, no repacking copy.""" - from torch.profiler import ProfilerActivity, profile - - gen = _gen(66) - x = torch.randn(1024, WIDTH, generator=gen, device=DEVICE).to(torch.bfloat16) - weight, bias = _affine(gen) - op(x, weight, bias, EPS) # warm the Triton JIT - torch.cuda.synchronize() - - with profile(activities=[ProfilerActivity.CUDA]) as prof: - op(x, weight, bias, EPS) - torch.cuda.synchronize() - - launched = [ - event.name - for event in prof.events() - if str(event.device_type) == "DeviceType.CUDA" and event.device_time_total > 0 - ] - assert launched == ["_layer_norm_fwd_1pass_kernel"], launched - - -def test_capturable_in_a_cuda_graph(): - gen = _gen(77) - x = torch.randn(512, WIDTH, generator=gen, device=DEVICE).to(torch.bfloat16) - weight, bias = _affine(gen) - static_x = x.clone() - op(static_x, weight, bias, EPS) # compile before capture - torch.cuda.synchronize() - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - static_out = op(static_x, weight, bias, EPS) - - static_x.copy_(x) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(static_out, op(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) @@ -237,32 +69,3 @@ def test_rejects_a_weight_whose_length_is_not_the_normalized_width(): weight, bias = _affine(gen) with pytest.raises(AssertionError): op(x, weight[:-1].contiguous(), bias, EPS) - with pytest.raises(AssertionError): - op(x, weight, bias[:-1].contiguous(), EPS) - - -def test_rejects_a_feature_dim_that_does_not_fit_the_fused_kernel(): - width = 65536 - x = torch.randn(2, width, device=DEVICE).to(torch.bfloat16) - weight = torch.ones(width, device=DEVICE, dtype=torch.bfloat16) - bias = torch.zeros(width, device=DEVICE, dtype=torch.bfloat16) - with pytest.raises(RuntimeError, match="doesn't support feature dim"): - op(x, weight, bias, EPS) - - -def test_accepts_an_fp8_activation_and_returns_fp8(): - """Measured, not certified: `empty_like(x)` makes the output e4m3, the - arithmetic is still fp32, and the store lands within one e4m3 ulp (2^-3 - relative, 2^-9 smallest subnormal) of the exact value.""" - gen = _gen(111) - x = torch.randn(64, WIDTH, generator=gen, device=DEVICE).to(torch.bfloat16) - weight, bias = _affine(gen) - x8 = x.to(torch.float8_e4m3fn) - - out = op(x8, weight, bias, EPS) - - assert out.dtype == torch.float8_e4m3fn - assert torch.isfinite(out.float()).all() - expected = _fp64(x8.float(), weight, bias, EPS) - deviation = (out.float().double() - expected).abs() - assert (deviation <= 2.0**-3 * expected.abs() + 2.0**-9).all(), deviation.max().item() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py index 381ef675e1fe..21569aeef4e0 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_rms_norm_gated_token_major.py @@ -17,20 +17,12 @@ assert torch.cuda.is_available(), "rms_norm_gated_token_major requires a CUDA device" -def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: - """fp8 output: byte-exact, as measured. Otherwise torch's default band.""" - assert out.dtype == ref.dtype, (out.dtype, ref.dtype) - assert out.shape == ref.shape, (out.shape, ref.shape) - if out.dtype == torch.float8_e4m3fn: - assert torch.equal(out.view(torch.uint8), ref.view(torch.uint8)), ( - "fp8 output is byte-identical to the unfused reference" - ) - return - torch.testing.assert_close(out, ref) +#: 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 GDN in_proj row the old column-slice measurement used: [q k v | z] = 10240 | 6144. -_PROJ_WIDTH, _Z_OFFSET = 16384, 10240 +#: 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 _seed(text: str) -> int: @@ -45,129 +37,44 @@ def _build(spec, seed): def randn(*shape): return torch.randn(*shape, generator=g, device="cuda").to(dtype) - if spec.get("x") == "strided_rows": - x = randn(t * heads, 2 * n)[:, :n] - else: - x = randn(t * heads, n) - if "x_scale" in spec: - x = x * spec["x_scale"] + x = randn(t * heads, n) layout = spec.get("z", "packed") - if layout == "column_slice": - z = randn(t, _PROJ_WIDTH)[:, _Z_OFFSET : _Z_OFFSET + heads * n].view(t, heads, n) - elif layout == "head_strided": - z = randn(t, heads, 2 * n)[:, :, :n] + if layout == "in_proj": + z = randn(t, _IN_PROJ_WIDTH)[:, _Z_OFFSET : _Z_OFFSET + heads * n].view(t, heads, n) else: + assert layout == "packed", layout z = randn(t, heads, n) weight = randn(n) - fp8_scale = None - if "fp8_scale" in spec: - fp8_scale = torch.tensor([spec["fp8_scale"]], dtype=torch.float32, device="cuda") - return x, z, weight, spec["eps"], fp8_scale, spec["gate_activation"] + return x, z, weight, spec["eps"], None, spec["gate_activation"] @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)) - x, z, weight, _, fp8_scale, _ = args - if cell.spec.get("z") == "column_slice": - assert z.stride()[1:] == (cell.spec["width"], 1) - if cell.spec["tokens"] > 1: - assert z.stride(0) == _PROJ_WIDTH and not z.is_contiguous() - if cell.spec.get("x") == "strided_rows": - assert x.stride(0) == 2 * cell.spec["width"] and not x.is_contiguous() + x, z, weight = args[:3] + assert z.stride() == (_IN_PROJ_WIDTH, cell.spec["width"], 1) and not z.is_contiguous() before = [x.clone(), z.clone(), weight.clone()] out = op(*args) - assert out.shape == x.shape and out.is_contiguous() - assert out.dtype == (torch.float8_e4m3fn if fp8_scale is not None else x.dtype) + assert out.shape == x.shape and out.is_contiguous() and out.dtype == x.dtype for t, b in zip((x, z, weight), before): assert torch.equal(t, b), "inputs are read-only" - if fp8_scale is not None: - # ... and no byte is the e4m3 NaN pattern, so the byte gate compares numbers. - assert torch.isfinite(out.float()).all() - _compare(out, op.reference(*args)) - - -@pytest.mark.parametrize("scale_value", [1.0, 0.5, 0.0234375]) -def test_fp8_byte_gate_discriminates(scale_value: float) -> None: - """Each wrong variant of the reference misses the byte gate on most elements.""" - spec = dict(op.CELLS[-1].spec, fp8_scale=scale_value) - x, z, weight, eps, fp8_scale, act = _build(spec, seed=8) - out = op(x, z, weight, eps, fp8_scale, act) - _compare(out, op.reference(x, z, weight, eps, fp8_scale, act)) - - gated_x = x * torch.sigmoid(z.reshape(x.shape).float()).to(x.dtype) - wrong = { - "sigmoid gate instead of silu": op.reference(x, z, weight, eps, fp8_scale, "sigmoid"), - "gate applied before the norm": op.reference(gated_x, z, weight, eps, fp8_scale, act), - "weight omitted": op.reference(x, z, torch.ones_like(weight), eps, fp8_scale, act), - } - if scale_value != 1.0: # at a scale of 1 the two directions coincide - wrong["scale multiplied instead of divided"] = ( - op.reference(x, z, weight, eps, None, act).float() * fp8_scale - ).to(torch.float8_e4m3fn) - for name, other in wrong.items(): - matched = (out.view(torch.uint8) == other.view(torch.uint8)).float().mean().item() - assert matched < 0.5, f"{name}: byte-identical on {matched:.4f} of elements" - - -def test_the_two_gates_differ() -> None: - """A kernel ignoring `gate_activation` would fail one of the gate cells, not both.""" - x, z, weight, eps, _, _ = _build(op.CELLS[1].spec, seed=3) - assert not torch.equal( - op(x, z, weight, eps, None, "silu"), op(x, z, weight, eps, None, "sigmoid") - ) + 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(dict(op.CELLS[1].spec, tokens=4), seed=12) + 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" - for bad in (long_weight, weight[:64]): - with pytest.raises(AssertionError, match="weight must be 1-D of length"): - with validating(op): - op(x, z, bad, eps, None, act) - - -def test_guard_refuses_a_malformed_fp8_scale() -> None: - """Unguarded, a two-element scale quantizes against element 0 and a bf16 - scale is accepted -- neither raises.""" - x, z, weight, eps, _, act = _build(dict(op.CELLS[1].spec, tokens=4), seed=12) - one = torch.tensor([0.5], dtype=torch.float32, device="cuda") - two = torch.tensor([0.5, 0.25], dtype=torch.float32, device="cuda") - half = torch.tensor([0.5], dtype=torch.bfloat16, device="cuda") - by_two = torch.ops.trtllm.rms_norm_gated_token_major(x, z, weight, eps, two, act) - by_one = torch.ops.trtllm.rms_norm_gated_token_major(x, z, weight, eps, one, act) - assert torch.equal(by_two.view(torch.uint8), by_one.view(torch.uint8)) - by_half = torch.ops.trtllm.rms_norm_gated_token_major(x, z, weight, eps, half, act) - assert by_half.dtype == torch.float8_e4m3fn - - with pytest.raises(AssertionError, match="fp8_scale must hold exactly one element"): - with validating(op): - op(x, z, weight, eps, two, act) - with pytest.raises(AssertionError, match="fp8_scale must be float32"): + with pytest.raises(AssertionError, match="weight must be 1-D of length"): with validating(op): - op(x, z, weight, eps, half, act) - - -def test_rejects_unsupported_gate_activation() -> None: - x, z, weight, eps, _, _ = _build(dict(op.CELLS[1].spec, tokens=4), seed=10) - with pytest.raises(ValueError, match="gate_activation must be 'silu' or 'sigmoid'"): - op(x, z, weight, eps, None, "gelu") - - -def test_rejects_z_that_does_not_cover_x() -> None: - x, z, weight, eps, _, act = _build(dict(op.CELLS[1].spec, tokens=4, heads=8), seed=11) - with pytest.raises(AssertionError, match="does not match x shape"): - op(x, z[:, :7, :].contiguous(), weight, eps, None, act) - with pytest.raises(ValueError, match="not enough values to unpack"): - op(x, z.reshape(32, 128), weight, eps, None, act) + 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_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 index e08dc99a1dfe..f93fdf147cd6 100644 --- 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 @@ -1,13 +1,8 @@ # 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. +"""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.""" -Every mis-reading of this formula keeps the shapes and stays bounded by -`|value|`, so beyond the cells one test measures how far outside the entry's -band each wrong reading lands, and two pin the floor and the zero-dot gate. -""" - -import math import zlib import pytest @@ -112,12 +107,10 @@ def _build(spec, seed): dots = spec["dots"] key, query, value = _single_channel_operands(dots, spec["value_dtype"]) return key, query, value, len(dots), spec["clamp_min"] - hc_count = spec.get("hc_count", HC_COUNT) - dtype = spec.get("dtype", torch.bfloat16) key, query, value = _make_operands( - seed, dtype, spec["leading"], hc_count, spec.get("value_dtype") + seed, torch.bfloat16, spec["leading"], value_dtype=spec.get("value_dtype") ) - return key, query, value, hc_count, spec.get("clamp_min", CLAMP_MIN) + 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]) @@ -134,119 +127,6 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res _compare(out, op.reference(*args)) -def test_flattened_call_matches_batched() -> None: - key, query, value = _make_operands(11, torch.bfloat16, (3, 17)) - batched = op(key, query, value, HC_COUNT) - flat = op(key.reshape(-1, WIDTH), query.reshape(-1, WIDTH), value.reshape(-1, HIDDEN), HC_COUNT) - assert torch.equal(flat, batched.reshape(-1, WIDTH)) - - -def test_output_is_bounded_by_the_value_it_gates() -> None: - """sigmoid is in (0, 1): no stream amplifies its value, to 4096x the operands.""" - key, query, value = _make_operands(40, torch.bfloat16, (256,)) - for scale in (1.0, 8.0, 64.0, 512.0, 4096.0): - got = op( - (key.float() * scale).to(torch.bfloat16), - (query.float() * scale).to(torch.bfloat16), - value, - HC_COUNT, - ) - assert torch.isfinite(got.float()).all(), scale - bound = value.float().abs().max().item() * (1.0 + torch.finfo(torch.bfloat16).eps) - assert got.float().abs().max().item() <= bound, scale - - -def test_clamp_floor_moves_fp32_and_not_bf16() -> None: - """Sub-floor dots move by > 1e-4 relative in fp32; bf16 is bitwise unmoved.""" - dots = (0.0, 5e-8, -5e-8, 0.05) - key, query, value32 = _single_channel_operands(dots, torch.float32) - floored = op(key, query, value32, len(dots), CLAMP_MIN) - unfloored = op(key, query, value32, len(dots), 0.0) - apart = (floored - unfloored).abs().unflatten(-1, (len(dots), HIDDEN)) - relative = (apart / floored.abs().unflatten(-1, (len(dots), HIDDEN))).amax(dim=(0, 2)) - print(f"\n fp32 relative movement per stream: {relative.tolist()}") - assert relative[0].item() == 0.0 and relative[3].item() == 0.0 - assert relative[1].item() > 1e-4 and relative[2].item() > 1e-4 - - _, _, value16 = _single_channel_operands(dots, torch.bfloat16) - assert torch.equal( - op(key, query, value16, len(dots), CLAMP_MIN), op(key, query, value16, len(dots), 0.0) - ) - - -def test_exact_zero_dot_gives_exactly_half_the_value() -> None: - key, query, value = _single_channel_operands((0.0,) * HC_COUNT, torch.float32) - assert torch.equal(op(key, query, value, HC_COUNT), value.repeat(1, HC_COUNT) * 0.5) - - -def _signed_sqrt(x): - return x.abs().clamp_min(CLAMP_MIN).sqrt() * x.sign() - - -def test_wrong_variants_land_outside_the_band() -> None: - """Each reading stays shaped and bounded; each must land > 5x outside.""" - key, query, value = _make_operands(50, torch.bfloat16, (2048,)) - want = op.reference(key, query, value, HC_COUNT) - correct = gate_ratio(op(key, query, value, HC_COUNT), want) - print(f"\n correct {correct:.4g}x the band") - assert correct <= 1.0 - - keys = key.double().unflatten(-1, (HC_COUNT, HIDDEN)) - queries = query.double().unflatten(-1, (HC_COUNT, HIDDEN)) - values = value.double().unsqueeze(-2) - dot = (keys * queries).sum(dim=-1, keepdim=True) - scaled = dot / math.sqrt(HIDDEN) - - def interleaved(): - k = key.double().unflatten(-1, (HIDDEN, HC_COUNT)).transpose(-1, -2) - q = query.double().unflatten(-1, (HIDDEN, HC_COUNT)).transpose(-1, -2) - g = (k * q).sum(dim=-1, keepdim=True) / math.sqrt(HIDDEN) - return torch.sigmoid(_signed_sqrt(g)) * values - - variants = ( - ("no square root", lambda: torch.sigmoid(scaled) * values), - ( - "unsigned square root", - lambda: torch.sigmoid(scaled.abs().clamp_min(CLAMP_MIN).sqrt()) * values, - ), - ("no 1/sqrt(D) scale", lambda: torch.sigmoid(_signed_sqrt(dot)) * values), - ("divided by D", lambda: torch.sigmoid(_signed_sqrt(dot / HIDDEN)) * values), - ( - "scale after the root", - lambda: torch.sigmoid(_signed_sqrt(dot) / math.sqrt(HIDDEN)) * values, - ), - ("sigmoid before the root", lambda: _signed_sqrt(torch.sigmoid(scaled)) * values), - ("tanh instead of sigmoid", lambda: torch.tanh(_signed_sqrt(scaled)) * values), - ("streams read interleaved", interleaved), - ) - for name, build in variants: - wrong = build().flatten(-2).to(torch.bfloat16) - ratio = gate_ratio(wrong, want) - print(f" {name:<26} {ratio:10.4g}x the band") - assert ratio > 5.0, f"the band does not separate {name!r} ({ratio:.4g}x)" - with pytest.raises(AssertionError): - _compare(wrong, want) - - -def test_fp32_chain_differs_from_the_bf16_chain() -> None: - """The literal bf16 chain falls outside the band; the entry does not.""" - key, query, value = _make_operands(90, torch.bfloat16, (2048,)) - want = op.reference(key, query, value, HC_COUNT) - got = op(key, query, value, HC_COUNT) - native = (key.unflatten(-1, (HC_COUNT, HIDDEN)) * query.unflatten(-1, (HC_COUNT, HIDDEN))).sum( - dim=-1, keepdim=True - ) / math.sqrt(HIDDEN) - native = (torch.sigmoid(_signed_sqrt(native)) * value.unsqueeze(-2)).flatten(-2) - differing = (got != native).float().mean().item() - print( - f"\n fp32 chain {gate_ratio(got, want):.4g}x, bf16 chain {gate_ratio(native, want):.4g}x" - ) - print(f" output elements whose bits differ: {differing:.4%}") - assert gate_ratio(got, want) <= 1.0 - assert gate_ratio(native, want) > 1.0 - assert differing > 0.0 - - 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,)) @@ -263,15 +143,6 @@ def test_guard_refuses_unequal_leading_dims() -> None: op(key, query, value[:1], HC_COUNT) -def test_malformed_widths_raise_rather_than_return() -> None: - """Why the width is not guarded: both mistakes are loud.""" - key, query, value = _make_operands(70, torch.bfloat16, (32,)) - with pytest.raises(RuntimeError): - op(key[:, : WIDTH - 8], query[:, : WIDTH - 8], value, HC_COUNT) - with pytest.raises(RuntimeError): - op(key, query, value, 3) - - def test_cuda_graph_capture_and_replay() -> None: key, query, value = _make_operands(100, torch.bfloat16, (8,)) warmup = torch.cuda.Stream() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ngram_token_history.py index b6069fe13f5b..0e16ce5bdedf 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ngram_token_history.py @@ -5,9 +5,9 @@ Integer bookkeeping, gated bitwise -- window and the whole published pool -- 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. -A wrong window is a window of valid ids, so beyond the cells: six near-miss -assemblies must each differ, hand-computed windows pin the lifecycle, chunked -prefill is checked step by step, and the two guards. +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. """ import zlib @@ -69,39 +69,27 @@ def _qsl(lengths, dtype, qsl_from="literal"): def _build(spec, seed): - """`(args, backing_pool)`: the op's positional arguments, and the tensor - the pool argument views (itself, unless `layout == "pool_view"`).""" + """The op's positional arguments for one cell.""" g = torch.Generator().manual_seed(seed) - ngram_size = spec.get("ngram_size", NGRAM) lengths = spec["lengths"] total = sum(lengths) ids = torch.randint(0, CKPT_VOCAB_SIZE, (total,), generator=g, dtype=torch.long) density = spec.get("eos_density", 0.0) if density > 0.0: ids[torch.rand(total, generator=g) < density] = EOS - layout = spec.get("layout", "contiguous") - pool_rows = spec["pool_slots"] * (2 if layout == "pool_view" else 1) - backing = torch.randint( - 0, CKPT_VOCAB_SIZE, (pool_rows, ngram_size - 1), generator=g, dtype=torch.long + pool = torch.randint( + 0, CKPT_VOCAB_SIZE, (spec["pool_slots"], CONTEXT_LEN), generator=g, dtype=torch.long ).to(DEV) - pool = backing[::2] if layout == "pool_view" else backing - input_ids = ids.to(DEV, spec.get("ids_dtype", torch.int64)) - if layout == "strided_ids": - wide = torch.zeros(2 * total, dtype=input_ids.dtype, device=DEV) - wide[::2] = input_ids - input_ids = wide[::2] - assert input_ids.stride(0) == 2 index_dtype = spec.get("index_dtype", torch.int32) - args = ( - input_ids, + return ( + ids.to(DEV, spec.get("ids_dtype", torch.int64)), pool, torch.tensor(spec["slots"], dtype=index_dtype, device=DEV), torch.tensor(spec["carries"], dtype=torch.bool, device=DEV), _qsl(lengths, spec.get("qsl_dtype", torch.int32), spec.get("qsl_from", "literal")), EOS, - ngram_size, + NGRAM, ) - return args, backing def _expect(args): @@ -112,9 +100,8 @@ def _expect(args): @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, backing = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) read_only = [a.clone() for a in (args[0], args[2], args[3], args[4])] - backing_before = backing.clone() pool_before = args[1].clone() expected = _expect(args) window = op(*args) @@ -124,35 +111,11 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res # The whole-pool comparison locks the unnamed slots; this keeps it from # being vacuous -- every cell has a sequence that publishes new ids. assert not torch.equal(args[1], pool_before), "no slot moved" - if cell.spec.get("layout") == "pool_view": - assert torch.equal(backing[1::2], backing_before[1::2]), "the view wrote outside itself" # ----------------------------------------------------------- multi-step -def test_chunked_prefill_carries_context_across_steps() -> None: - """One slot over three forward steps (96, 24, 7 tokens), each against the - reference from the pool the previous step left.""" - pool = torch.randint(0, CKPT_VOCAB_SIZE, (5, CONTEXT_LEN), device=DEV) - g = torch.Generator().manual_seed(4001) - for step, (length, carry) in enumerate(((96, False), (24, True), (7, True))): - ids = torch.randint(0, CKPT_VOCAB_SIZE, (length,), generator=g).to(DEV) - if step == 0: - ids[torch.rand(length, generator=g) < 0.04] = EOS - args = ( - ids, - pool, - torch.tensor([2], dtype=torch.int32, device=DEV), - torch.tensor([carry], device=DEV), - torch.tensor([0, length], dtype=torch.int32, device=DEV), - EOS, - NGRAM, - ) - expected = _expect(args) - _compare((op(*args), pool), expected) - - def test_chunked_prefill_matches_one_shot() -> None: """96 tokens split at 41 over two calls: window and pool bitwise one-shot.""" g = torch.Generator().manual_seed(5001) @@ -173,120 +136,6 @@ def call(pool, chunk, carry): assert torch.equal(pool, one_shot_pool) -# ----------------------------------------------------------- hand-computed - - -def _literal(ids, pool, slots, carries, starts): - window = op( - torch.tensor(ids, device=DEV), - pool, - torch.tensor(slots, dtype=torch.int32, device=DEV), - torch.tensor(carries, device=DEV), - torch.tensor(starts, dtype=torch.int32, device=DEV), - EOS, - NGRAM, - ) - return window.tolist() - - -def test_fresh_sequence_seeds_with_eos_not_the_pool() -> None: - pool = torch.full((4, CONTEXT_LEN), 1234, device=DEV) - window = _literal([11, 22, 33], pool, [2], [False], [0, 3]) - assert window == [[EOS, EOS, 11], [EOS, 11, 22], [11, 22, 33]] - assert pool[2].tolist() == [22, 33] - - -def test_carry_row_is_oldest_first_and_publishes_its_tail() -> None: - pool = torch.tensor([[70, 80]], device=DEV) - assert _literal([90, 100], pool, [0], [True], [0, 2]) == [[70, 80, 90], [80, 90, 100]] - assert pool[0].tolist() == [90, 100] - pool = torch.tensor([[70, 80]], device=DEV) - assert _literal([90], pool, [0], [True], [0, 1]) == [[70, 80, 90]] - assert pool[0].tolist() == [80, 90], "part carry, part new" - - -def test_zero_token_rows_keep_or_reseed() -> None: - pool = torch.tensor([[70, 80], [11, 12], [13, 14]], device=DEV) - window = _literal([90, 91], pool, [0, 1, 2], [True, True, False], [0, 2, 2, 2]) - assert window == [[70, 80, 90], [80, 90, 91]] - assert pool.tolist() == [[90, 91], [11, 12], [EOS, EOS]] - - -def test_sequence_boundary_is_not_crossed() -> None: - pool = torch.tensor([[70, 80], [170, 180]], device=DEV) - window = _literal([90, 91, 92, 190, 191], pool, [0, 1], [True, True], [0, 3, 5]) - assert window[3:] == [[170, 180, 190], [180, 190, 191]] - - -def test_in_stream_eos_is_passed_through_raw() -> None: - pool = torch.tensor([[70, 80]], device=DEV) - window = _literal([EOS, 91, 92], pool, [0], [True], [0, 3]) - assert window == [[70, 80, EOS], [80, EOS, 91], [EOS, 91, 92]] - - -# ----------------------------------------------------------- discrimination - - -def test_wrong_variants_are_separated() -> None: - """Six mis-assemblies, each in-vocabulary; each must differ in window or pool.""" - spec = dict( - lengths=(17, 1, 5, 1), - slots=(3, 0, 2, 1), - carries=(False, True, True, True), - pool_slots=6, - eos_density=0.08, - ) - args, _ = _build(spec, seed=9001) - ids, pool, slots, carries, qsl = args[:5] - start = pool.clone() - window = op(*args) - published = pool - - def ref(**override): - kw = dict( - input_ids=ids, - context_pool=start.clone(), - state_indices=slots, - has_initial_states=carries, - query_start_loc=qsl, - eos_token_id=EOS, - ngram_size=NGRAM, - ) - kw.update(override) - return op.reference(**kw) - - right_window, right_pool = ref() - - def publishes_the_head(): - head_pool = start.clone() - starts = qsl.tolist() - for s, slot in enumerate(slots.tolist()): - if starts[s + 1] > starts[s]: - head_pool[slot] = right_window[starts[s + 1] - 1, :-1] - return right_window, head_pool - - variants = { - "newest_first": lambda: (right_window.flip(1), right_pool), - "zero_seed": lambda: ref(eos_token_id=0), - "ignores_carry_flags": lambda: ref(has_initial_states=torch.ones_like(carries)), - # One flat stream: each sequence's first tokens read its neighbour's tail. - "packs_across_boundaries": lambda: ( - ref( - state_indices=slots[:1], - has_initial_states=carries[:1], - query_start_loc=qsl[[0, -1]], - )[0], - right_pool, - ), - "rotated_pool_read": lambda: ref(context_pool=start.roll(1, dims=1)), - "publishes_the_head": publishes_the_head, - } - for name, build in variants.items(): - wrong_window, wrong_pool = build() - same = torch.equal(window, wrong_window) and torch.equal(published, wrong_pool) - assert not same, f"variant {name} is indistinguishable" - - # ----------------------------------------------------------- guards @@ -314,18 +163,6 @@ def _misses(args, bad_index, bad_value): return not (torch.equal(window, expected[0]) and torch.equal(bad[1], expected[1])) -def test_guard_refuses_carry_flags_of_the_wrong_length() -> None: - """Unguarded, one flag broadcasts: sequence 1 reads its stale slot, not the seed.""" - args = _two_sequences() - one_flag = torch.ones(1, dtype=torch.bool, device=DEV) - assert _misses(args, 3, one_flag), "the hazard this guard covers is gone" - bad = list(args) - bad[3] = one_flag - with pytest.raises(AssertionError, match="one flag per sequence"): - with validating(op): - op(*bad) - - def test_guard_refuses_a_short_query_start_loc() -> None: """Unguarded, a 2-entry vector for 2 sequences misattributes the tail.""" args = _two_sequences() @@ -338,37 +175,6 @@ def test_guard_refuses_a_short_query_start_loc() -> None: op(*bad) -@pytest.mark.parametrize( - "index,value,match", - [ - (1, lambda a: torch.zeros(4, CONTEXT_LEN + 1, dtype=torch.long, device=DEV), "must match"), - (1, lambda a: a[1].to(torch.int32), "Index put requires the source and destination"), - (3, lambda a: a[3].to(torch.int32), "where expected condition to be a boolean"), - (2, lambda a: a[2].reshape(-1, 1), "Index is supposed to be an empty tensor or a vector"), - (6, lambda a: 1, "must match the size of tensor"), - ], - ids=["pool_too_wide", "int32_pool", "int_flags", "2d_state_indices", "ngram_size_1"], -) -def test_loud_refusals_are_not_guarded(index, value, match) -> None: - args = list(_two_sequences()) - args[index] = value(args) - with pytest.raises(RuntimeError, match=match): - op(*args) - - -def test_two_dimensional_input_ids_is_refused() -> None: - with pytest.raises(RuntimeError, match="must match the size of tensor"): - op( - torch.zeros(1, 4, dtype=torch.long, device=DEV), - torch.zeros(2, CONTEXT_LEN, dtype=torch.long, device=DEV), - torch.zeros(1, dtype=torch.int32, device=DEV), - torch.ones(1, dtype=torch.bool, device=DEV), - torch.tensor([0, 1], dtype=torch.int32, device=DEV), - EOS, - NGRAM, - ) - - # ----------------------------------------------------------- capture @@ -382,7 +188,7 @@ def test_cuda_graph_capture_and_replay() -> None: ids_dtype=torch.int32, qsl_dtype=torch.int64, ) - args, _ = _build(spec, seed=11001) + args = _build(spec, seed=11001) start = args[1].clone() side = torch.cuda.Stream() 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_gate_conv_inject_norm.py new file mode 100644 index 000000000000..b2729a8e1ee6 --- /dev/null +++ b/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 @@ -0,0 +1,181 @@ +# 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. + +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. +""" + +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.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 + ple_gate_conv_inject_norm as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402,F401 — path above; fixture by name + assert_within_ulp, + certifying, + validating, +) + +assert torch.cuda.is_available(), "ple_gate_conv_inject_norm requires a CUDA device" + +DEV = "cuda" +EPS = 1e-6 +PAD = -1 +GAMMA_RMS = 0.477 + + +def _ulp_ratio(out: torch.Tensor, ref: torch.Tensor) -> torch.Tensor: + """Per-element distance from `ref` in units of one ulp of `|ref|`.""" + eps = torch.finfo(ref.dtype).eps + tiny = torch.finfo(ref.dtype).tiny + ref64 = ref.double() + return (out.double() - ref64).abs() / (eps * ref64.abs() + tiny) + + +def _stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + ratio = _ulp_ratio(out, ref) + return ratio.max().item(), (out != ref).double().mean().item() + + +#: The band. Every rounding point of the composition is kept, so the op and the +#: float64 reference differ only where an fp32 reduction (three sums of squares +#: and the gate's dot per stream) or fp32 `exp` lands a value across a bf16 +#: midpoint -- measured on 2e-5..2.3e-4 of the elements. A flipped upstream +#: element then propagates, and through `hc + ple` it meets cancellation +#: (`hc ~ -ple`), so a per-element ulp is unbounded there: 73 ulp of a +#: 0.003-sized element in a row whose largest is ~10. The stream and its norm +#: are therefore gated in units of their row's scale (`assert_within_ulp`): +#: 2 ulp of the row maximum per element and 0.1 ulp relative RMS. The +#: composition this replaces sits at 0.87 / 0.33 of a row-max ulp against the +#: same reference (measured: normed / hc_out at 128 rows), the op at 0.72 / +#: 0.38. The published conv state is one rounding of the conv-input norm, +#: whose only perturbation is a one-ulp flip of a gated element: 2 ulp per +#: element (measured worst 1.78, the old composition 1.56). +ROW_ULP, RMS_ULP, STATE_ULP = 2.0, 0.1, 2.0 +#: The per-element row-max gate is a tail statistic, and a context-shaped cell +#: draws ~100x more elements than a 128-row one. Against the same float64 +#: 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. +CONTEXT_ROWS, CONTEXT_ROW_ULP = 1024, 4.0 + + +def _compare(out, ref) -> None: + row_ulp = CONTEXT_ROW_ULP if out[0].shape[0] > CONTEXT_ROWS else ROW_ULP + for name, got, want in zip(("hc_out", "normed", "conv_states"), out, ref): + assert got.dtype == want.dtype and got.shape == want.shape, name + if not got.numel(): + continue + if name == "conv_states": + worst = _ulp_ratio(got, want).max().item() + assert worst <= STATE_ULP, f"{name}: {worst:.4g} ulp" + else: + assert_within_ulp(got, want, row_ulp, RMS_ULP) + + +def _build(spec, seed): + g = torch.Generator(device=DEV).manual_seed(seed) + gc = torch.Generator().manual_seed(seed) + lengths = spec["lengths"] + c, d, dtype = CKPT_HC_COUNT, CKPT_HIDDEN, torch.bfloat16 + w = c * d + rows = sum(lengths) + num_seqs = len(lengths) + pool_slots = spec.get("pool_slots", num_seqs + 2) + + def rand(*shape, scale=1.0): + return (torch.randn(*shape, generator=g, device=DEV) * scale).to(dtype) + + # Streams at different scales, as the trained residual's are. + scales = torch.tensor([0.5 * 2.0 ** (i % 4) for i in range(c)], device=DEV) + hc = (rand(rows, w, scale=3.0).float() * scales.repeat_interleave(d)).to(dtype) + key = rand(rows, w) + value = rand(rows, d, scale=0.5) + norms = [rand(w, scale=GAMMA_RMS) for _ in range(4)] + conv_weight = rand(w, 1, 4, scale=0.3) + conv_states = rand(pool_slots, w, 3 * CKPT_DILATION) + order = (torch.randperm(pool_slots - 2, generator=gc) + 1).tolist() + slots = ([0, pool_slots - 1] + order)[:num_seqs] if pool_slots >= 64 else list(range(num_seqs)) + slots = torch.tensor(slots, dtype=torch.int32, device=DEV) + starts = [0] + for length in lengths: + starts.append(starts[-1] + length) + qsl = torch.tensor(starts, dtype=torch.int64, device=DEV) + max_rows = spec.get("max_rows", max(lengths)) + carries = spec.get("carries") + carries = None if carries is None else torch.tensor(carries, dtype=torch.bool, device=DEV) + return [ + hc, key, value, norms[0], norms[1], norms[2], conv_weight, conv_states, slots, qsl, + norms[3], EPS, c, CKPT_CLAMP_MIN, CKPT_DILATION, PAD, max_rows, carries, + ] # fmt: skip + + +def _run(args): + """`(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) + return hc_out, normed, args[7] + + +@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) + assert torch.equal(args[0], hc_before), "hc is read-only" + assert out[0].is_contiguous() and out[1].is_contiguous() + _compare(out, ref) + for name, got, want in zip(("hc_out", "normed", "conv_states"), out, ref): + worst, frac = _stats(got, want) + print( + f"\n {cell.why[:40]:<40} {name:<12} worst {worst:.3g} ulp, differ {frac:.2e}", end="" + ) + + +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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_ngram_embed.py new file mode 100644 index 000000000000..c46ea883dc07 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ple/test_modeling_v2_ple_ngram_embed.py @@ -0,0 +1,180 @@ +# 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. + +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. +""" + +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.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 + 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 + CKPT_EOS_TOKEN_ID, + CKPT_HEAD_DIM, + CKPT_SCALE, + CKPT_TABLE_ROWS, + ENGINE_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 + 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, +) + +__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" + +DEV = "cuda" +EOS = CKPT_EOS_TOKEN_ID +VOCAB = 248320 +FP8 = torch.float8_e4m3fn +#: Keeps the hosts alive for as long as the tables they back. +_HOSTS = [] + + +def _scattered(n: int, pool_slots: int, g) -> list: + """`n` distinct slots over the whole pool, both ends included.""" + middle = (torch.randperm(pool_slots - 2, generator=g) + 1).tolist() + slots = ([0, pool_slots - 1] + middle)[:n] + assert len(set(slots)) == n and max(slots) < pool_slots + return slots + + +def _geometry(spec, g): + """`(multipliers, vocab, offsets)` for the cell: the checkpoint's own over + the full-height table, the checkpoint's multipliers over small stand-in head + vocabularies otherwise.""" + if spec.get("table") == "ckpt": + return CKPT_MULTIPLIERS, CKPT_HEAD_VOCAB_SIZES, CKPT_HEAD_OFFSETS + mult = CKPT_MULTIPLIERS + vocab = tuple( + 1000 + 7 * h + int(x) + for h, x in enumerate(torch.randint(0, 50, (len(CKPT_HEAD_VOCAB_SIZES),), generator=g)) + ) + offsets, acc = [], 0 + for v in vocab: + offsets.append(acc) + acc += v + return mult, vocab, tuple(offsets) + + +def _table(kind: str, rows: int, touched: torch.Tensor, g) -> torch.Tensor: + if kind == "ckpt": + table = torch.empty(CKPT_TABLE_ROWS, CKPT_HEAD_DIM, dtype=FP8, device=DEV) + table[touched] = torch.randn(touched.numel(), CKPT_HEAD_DIM, generator=g).to(FP8).to(DEV) + 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 src.to(DEV) + + +def _with_eos(x: torch.Tensor, density: float, g) -> torch.Tensor: + if density > 0.0: + x = torch.where(torch.rand(x.shape, generator=g) < density, torch.full_like(x, EOS), x) + return x + + +def _build(spec, seed): + g = torch.Generator().manual_seed(seed) + lengths = spec["lengths"] + n = len(CKPT_MULTIPLIERS) + num_seqs = len(lengths) + pool_slots = spec.get("pool_slots", num_seqs + 2) + density = spec.get("eos_density", 0.0) + ids = _with_eos(torch.randint(0, VOCAB, (sum(lengths),), generator=g), density, g) + ids = ids.to(torch.int32).to(DEV) + pool = _with_eos(torch.randint(0, VOCAB, (pool_slots, n - 1), generator=g), density, g).to(DEV) + slots = _scattered(num_seqs, pool_slots, g) if pool_slots >= 64 else list(range(num_seqs)) + slots = torch.tensor(slots, dtype=torch.int32, device=DEV) + carries = spec.get("carries") + carries = None if carries is None else torch.tensor(carries, dtype=torch.bool, device=DEV) + starts = [0] + for length in lengths: + starts.append(starts[-1] + length) + qsl = torch.tensor(starts, dtype=torch.int64, device=DEV) + mult, vocab, offsets = _geometry(spec, g) + mult = torch.tensor(mult, dtype=torch.int64, device=DEV) + vocab = torch.tensor(vocab, dtype=torch.int64, device=DEV) + offsets = torch.tensor(offsets, dtype=torch.int64, device=DEV) + # The rows this batch touches, from the two scalar references, so that the + # full-height table needs bytes only where they are read. + seeded = torch.ones(num_seqs, dtype=torch.bool, device=DEV) if carries is None else carries + window, _ = ngram_token_history.reference(ids, pool, slots, seeded, qsl, EOS, n) + touched = ngram_hash.reference(window, mult, vocab, offsets, EOS, 1).reshape(-1) + rows = int(offsets[-1] + vocab[-1]) + table = _table(spec.get("table", "small"), rows, touched, g) + scale = torch.tensor([CKPT_SCALE], dtype=torch.bfloat16, device=DEV).reshape(1, 1, 1) + max_rows = spec.get("max_rows", max(lengths)) + return [ids, pool, slots, carries, qsl, table, scale, mult, vocab, offsets, EOS, max_rows] + + +def _compare(out, ref) -> None: + (emb, pool), (ref_emb, ref_pool) = out, ref + assert emb.dtype == ref_emb.dtype and emb.shape == ref_emb.shape + assert emb.is_contiguous() + assert torch.equal(emb, ref_emb), ( + f"embedding differs at {(emb != ref_emb).any(-1).nonzero().flatten()[:8].tolist()}" + ) + assert torch.equal(pool, ref_pool), "published pool differs" + + +@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())) + ref = op.reference(*args) + pool_before = args[1].clone() + out = op(*args) + _compare((out, args[1]), ref) + if cell.spec.get("table") == "host": + assert args[1].data_ptr() != 0 and not torch.equal(args[1], pool_before) + + +def test_guard_refuses_short_query_start_loc() -> None: + args = _build(dict(lengths=(1, 1, 1)), seed=3) + args[4] = args[4][:-1] + with validating(op), pytest.raises(AssertionError): + 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) + 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): + graph_out = op(*args) + fresh = _build(dict(lengths=(1,) * 16, pool_slots=ENGINE_SLOTS, eos_density=0.1), seed=12) + args[0].copy_(fresh[0]) + pool_before = args[1].clone() + want, want_pool = op.reference(*args) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(graph_out, want) + assert torch.equal(args[1], want_pool) + assert not torch.equal(args[1], pool_before) 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_bf16_to_fp32.py new file mode 100644 index 000000000000..3bf21aa13b26 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_bf16_to_fp32.py @@ -0,0 +1,67 @@ +# 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.""" + +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 + bf16_to_fp32 as op, +) + +__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" + + +def _build(spec, seed): + shape = spec["shape"] + if spec.get("source") == "all_codes": + return ( + torch.arange(-32768, 32768, dtype=torch.int32, device="cuda") + .to(torch.int16) + .view(torch.bfloat16) + .reshape(shape) + ) + g = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(shape, generator=g, device="cuda") * 8).to(torch.bfloat16) + + +def _bitwise(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.shape == b.shape and torch.equal(a.view(torch.int32), b.view(torch.int32)) + + +@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) + x = _build(cell.spec, zlib.crc32(cell.why.encode())) + out = op(x) + 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_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 index a58ee2bdde8b..ef52eab0be66 100644 --- 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 @@ -3,9 +3,8 @@ """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: the direction, the round trip through the -quantizer, degenerate and overflowing scales, graph capture, and the loud -rejections. +the claims no cell carries: graph capture with the scale read at replay, and +the loud rejection of an empty batch. """ import zlib @@ -52,11 +51,6 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: DEV = "cuda" -E4M3_MAX = 448.0 -# e4m3: a normal value's relative half-step and the smallest subnormal's -# absolute half-step, in units of the scale. -E4M3_REL_HALF_STEP = 2.0**-4 -E4M3_SUBNORMAL_HALF_STEP = 2.0**-10 _TABLE = {} @@ -79,28 +73,13 @@ def _input(spec, seed: int) -> torch.Tensor: _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] - elif source == "all_codes": - x = torch.arange(256, dtype=torch.uint8, device=DEV).view(shape).view(torch.float8_e4m3fn) - elif source == "unaligned": - n = shape[0] * shape[1] - base = _fp8(torch.randn(n + 16, device=DEV, generator=g) * 8.0) - x = base[1 : 1 + n].view(shape) - assert x.data_ptr() % 16 != 0 else: - x = _fp8(torch.randn(*shape, device=DEV, generator=g) * 8.0) + 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 _scales(spec, value: float, rank: int) -> torch.Tensor: - if not spec.get("scale_slice"): - return _scale(value, rank, spec["dtype"]) - buf = torch.tensor([[9.0, value, 5.0]], dtype=spec["dtype"], device=DEV) - sliced = buf[:, 1:2] - assert sliced.shape == (1, 1) and sliced.storage_offset() == 1 - return sliced - - 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]) @@ -114,7 +93,7 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th certifying(op) x = _input(cell.spec, zlib.crc32(cell.why.encode())) for value in cell.spec["scales"]: - scales = _scales(cell.spec, value, x.dim()) + 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)) @@ -126,55 +105,6 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th # ── what a cell cannot carry ────────────────────────────────────────────── -def test_multiplies_rather_than_divides(): - x = _fp8(torch.tensor([[1.0, 2.0, -3.0, 0.5]], device=DEV)) - for value in (0.5, PLE_SCALE): - out = op(x, _scale(value, 2, torch.float32)) - assert torch.equal(_bits(out), _bits(x.float() * value)) - assert not torch.allclose(out, x.float() / value) - - -def test_top_of_the_range_at_the_checkpoint_scale(): - top = op(_fp8(torch.tensor([[E4M3_MAX, -E4M3_MAX]], device=DEV)), _scale(PLE_SCALE, 2)) - # 448 * PLE_SCALE = 0.08929443359375 exactly; bf16's nearest is this - assert top.float().flatten().tolist() == [0.08935546875, -0.08935546875] - - -def test_round_trip_through_the_static_quantizer(): - """The error floor is e4m3 itself; a reciprocated scale lands >100x out.""" - g = torch.Generator(device=DEV).manual_seed(15) - x = (torch.randn(256, 5120, device=DEV, generator=g) * 2.0).to(torch.bfloat16) - value = 0.05 - scales = _scale(value, 2, torch.float32) - q, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(x, scales) - back = op(q, scales) - _compare(back, op.reference(q, scales)) - bound = E4M3_REL_HALF_STEP * x.float().abs() + E4M3_SUBNORMAL_HALF_STEP * value - assert bool(((back - x.float()).abs() <= bound).all()) - wrong = op(q, _scale(1.0 / value, 2, torch.float32)) - assert float(((wrong - x.float()).abs() / bound)[x.float().abs() > 0.1].min()) > 100.0 - - -def test_degenerate_scale_values_are_accepted_and_silent(): - y = _fp8(torch.tensor([[1.0, -1.0, 0.0, E4M3_MAX]], device=DEV)) - out = op(y, _scale(0.0, 2)) - assert out.float().flatten().tolist() == [0.0, -0.0, 0.0, 0.0] - assert not bool(torch.isnan(out.float()).any()) - assert op(y, _scale(-2.0, 2)).float().flatten().tolist() == [-2.0, 2.0, -0.0, -896.0] - out = op(y, _scale(float("inf"), 2)).float().flatten() - assert out[0].item() == float("inf") and out[1].item() == float("-inf") - assert torch.isnan(out[2]) - assert bool(torch.isnan(op(y, _scale(float("nan"), 2)).float()).all()) - - -def test_overflow_goes_to_infinity_not_to_a_clamp(): - y = _fp8(torch.tensor([[1.0, E4M3_MAX]], device=DEV)) - out = op(y, _scale(1000.0, 2, torch.float16)).float().flatten() - assert out[0].item() == 1000.0 and out[1].item() == float("inf") - out = op(y, _scale(3.0e38, 2, torch.bfloat16)).float().flatten() - assert torch.isfinite(out[0]) and out[1].item() == float("inf") - - 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) @@ -204,48 +134,7 @@ def test_cuda_graph_capture_and_replay(): _compare(static_out, op.reference(fresh, _scale(0.5, 3))) -def test_deterministic_and_operands_untouched(): - g = torch.Generator(device=DEV).manual_seed(17) - x = _fp8(torch.randn(64, 16, 160, device=DEV, generator=g) * 8.0) - scales = _scale(PLE_SCALE, 3) - x_before, s_before = x.clone(), scales.clone() - out1, out2 = op(x, scales), op(x, scales) - assert torch.equal(_bits(out1), _bits(out2)) - assert torch.equal(x.view(torch.uint8), x_before.view(torch.uint8)) - assert torch.equal(_bits(scales), _bits(s_before)) - assert out1.data_ptr() not in (x.data_ptr(), scales.data_ptr()) - - -def test_rejects_unsupported(): - """Every metadata violation is loud, which is why the entry has no guard.""" - x = _fp8(torch.randn(4, 128, device=DEV)) - scales = _scale(PLE_SCALE, 2) - bf = dict(dtype=torch.bfloat16, device=DEV) - cases = { - "1-D input": lambda: op(x.reshape(-1), _scale(0.5, 1)), - "0-D input": lambda: op(_fp8(torch.tensor(1.0, device=DEV)), _scale(0.5, 0)), - "zero rows": lambda: op(_fp8(torch.randn(0, 128, device=DEV)), scales), - "zero cols": lambda: op(_fp8(torch.randn(4, 0, device=DEV)), scales), - "column slice": lambda: op(_fp8(torch.randn(8, 256, device=DEV))[:, :128], scales), - "transpose": lambda: op(_fp8(torch.randn(128, 8, device=DEV)).t(), scales), - "bfloat16 input": lambda: op(torch.randn(4, 128, **bf), scales), - "float8_e5m2 input": lambda: op( - torch.randn(4, 128, device=DEV).to(torch.float8_e5m2), scales - ), - "int8 input": lambda: op( - torch.randint(-8, 8, (4, 128), dtype=torch.int8, device=DEV), scales - ), - "float64 scales": lambda: op(x, _scale(0.5, 2, torch.float64)), - "cpu scales": lambda: op(x, scales.cpu()), - "cpu input": lambda: op(x.cpu(), scales), - "cpu input and scales": lambda: op(x.cpu(), scales.cpu()), - "scales rank below": lambda: op(x, _scale(0.5, 1)), - "scales rank above": lambda: op(x, _scale(0.5, 3)), - "0-D scales": lambda: op(x, _scale(0.5, 0)), - "scales dim not 1": lambda: op(x, torch.full((1, 2), 0.5, **bf)), - "scales leading dim not 1": lambda: op(x, torch.full((4, 1), 0.5, **bf)), - "empty scales": lambda: op(x, torch.empty((1, 0), **bf)), - } - for name, fn in cases.items(): - with pytest.raises((RuntimeError, NotImplementedError)): - fn() +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/quantization/test_modeling_v2_fp8_quantize_1x128.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_fp8_quantize_1x128.py index 2a65ca954ee7..caa087518151 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/quantization/test_modeling_v2_fp8_quantize_1x128.py @@ -2,10 +2,10 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the fp8_quantize_1x128 catalog entry. -The entry owns the contract. This file turns a cell into an input and holds -what a cell cannot carry: that every code disagreement is a rounding tie, the -`use_ue8m0` bit-exactness, the wrong-recipe discrimination, and the edge -behaviour the entry's `note` records. +The entry owns the contract. This file turns a cell into an input, holds +every code disagreement of a cell to a rounding tie, and adds what a cell +cannot carry: CUDA-graph capture and replay, and the loud rejection of a +strided input. """ import zlib @@ -14,8 +14,6 @@ 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 - AMAX_FLOOR, - E4M3_MAX, SCALE_ULP_GATE, block_scales, ) @@ -79,10 +77,6 @@ def _compare( DEV = "cuda" -# e4m3: 3 mantissa bits, so a normal value carries at most a 2^-4 relative -# half-step; the smallest subnormal 2^-9 has a 2^-10 half-step, in scale units. -E4M3_REL_HALF_STEP = 2.0**-4 -E4M3_SUBNORMAL_HALF_STEP = 2.0**-10 # A code can only disagree when the exact product sits within the first # division's 2 ulp of an e4m3 rounding midpoint. TIE_REL_GATE = 2.0 * 2.0**-23 @@ -98,16 +92,6 @@ def _input(tokens: int, hidden: int, seed: int, magnitude: float = 3.0) -> torch return (torch.randn(tokens, hidden, device=DEV, generator=g) * magnitude).to(torch.bfloat16) -def _amax(x: torch.Tensor) -> torch.Tensor: - """Exact per-block amax `[T, H/128]` fp32, floored as the kernel floors it.""" - t, n = x.shape - return x.abs().reshape(t, n // 128, 128).amax(dim=2).float().clamp(min=AMAX_FLOOR) - - -def _ulp_f32(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - return a.contiguous().view(torch.int32).long() - b.contiguous().view(torch.int32).long() - - def _assert_misses_are_ties(x: torch.Tensor, data: torch.Tensor, ref_data: torch.Tensor) -> float: """Every code disagreement is a rounding tie: the exact product sits within the scale's 2-ulp band of the midpoint between the two codes.""" @@ -141,205 +125,7 @@ def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves th print(f" {cell.why[:44]}: {100 * exact:.4f}% bit-exact, worst tie {rel:.3e}") -# ── the flag, the scale and the direction ───────────────────────────────── - - -@pytest.mark.parametrize("tokens,hidden", [(1, 2560), (17, 2560), (64, 512), (1024, 2560)]) -def test_use_ue8m0_is_bit_exact_and_a_different_quantization(tokens, hidden): - """Once the dequant scale is a power of two its reciprocal is exact, so the - approximate division cannot reach the result.""" - x = _input(tokens, hidden, seed=5 + tokens) - data, sf = op(x, True) - ref_data, ref_sf = op.reference(x, True) - assert torch.equal(sf.view(torch.int32), ref_sf.view(torch.int32)) - assert torch.equal(data.view(torch.uint8), ref_data.view(torch.uint8)) - assert bool((torch.frexp(sf)[0] == 0.5).all()) - d0, sf0 = op(x, False) - sf_changed = (sf != sf0).float().mean().item() - d_changed = (data.view(torch.uint8) != d0.view(torch.uint8)).float().mean().item() - print(f" ue8m0 T={tokens} H={hidden}: scales changed {sf_changed:.5f}, codes {d_changed:.5f}") - assert sf_changed > 0.9 and d_changed > 0.5 - assert bool((sf >= sf0 * (1 - 2**-20)).all()) and bool((sf <= sf0 * 2 * (1 + 2**-20)).all()) - - -def test_scale_is_amax_over_448_either_spelling(): - """Within 4 ulp of both `amax / 448` and the kernel's `1 / (448 / amax)`.""" - x = _input(1024, 2560, seed=21) - _, sf = op(x, False) - amax = _amax(x) - direct = (amax / E4M3_MAX).t() - two_step = (1.0 / (E4M3_MAX / amax)).t() - u_direct = _ulp_f32(sf, direct).abs().max().item() - u_two_step = _ulp_f32(sf, two_step).abs().max().item() - exact = int((sf.view(torch.int32) == two_step.contiguous().view(torch.int32)).sum()) - print( - f" max|ulp| vs amax/448 {u_direct}, vs 1/(448/amax) {u_two_step}; exact {exact}/{sf.numel()}" - ) - assert u_direct <= SCALE_ULP_GATE and u_two_step <= SCALE_ULP_GATE - # the amax itself is recoverable from the scale to within those ulps - rec = sf.t() * E4M3_MAX - assert bool(((rec - amax).abs() <= amax * (SCALE_ULP_GATE + 1) * 2**-23).all()) - - -def test_scale_multiplies_the_code(): - """Round trip bounded by e4m3's own constants; the reciprocal reading is - four orders of magnitude off.""" - x = _input(64, 2560, seed=2, magnitude=2.0) - data, sf = op(x, False) - cols = _expand(sf.t().contiguous(), 2560) - deq = data.float() * cols - bound = (E4M3_REL_HALF_STEP + (SCALE_ULP_GATE + 1) * 2**-23) * x.float().abs() - bound = bound + E4M3_SUBNORMAL_HALF_STEP * cols - assert bool(((deq - x.float()).abs() <= bound).all()) - - def rel_rms(a): - return ((a - x.float()).pow(2).mean().sqrt() / x.float().pow(2).mean().sqrt()).item() - - mul, div = rel_rms(deq), rel_rms(data.float() / cols) - print(f" multiply relRMS={mul:.5f} divide relRMS={div:.5g}") - assert mul < 0.05 and div > 100.0 - assert data.float().abs().max().item() == E4M3_MAX - - -def test_discrimination_against_wrong_recipes(): - """Each wrong reading of the amax map, scored against the same call.""" - t, n = 256, 2560 - x = _input(t, n, seed=99) - data, sf = op(x, False) - amax = _amax(x) - - def score(a_cols, amax_max=E4M3_MAX): - qs_cols = amax_max / a_cols - d = (x.float() * qs_cols).to(torch.float8_e4m3fn) - agree = (d.view(torch.uint8) == data.view(torch.uint8)).float().mean().item() - rel = (((1.0 / qs_cols)[:, ::128].t() - sf).abs() / sf).max().item() - return agree, rel - - mant, exp = torch.frexp(amax / E4M3_MAX) - pow2 = torch.ldexp(torch.ones_like(amax), torch.where(mant == 0.5, exp - 1, exp)) - variants = { - "correct": (_expand(amax, n), E4M3_MAX), - "ue8m0 power-of-two scale": (_expand(pow2 * E4M3_MAX, n), E4M3_MAX), - "one scale per tensor": (_expand(amax.max().expand_as(amax), n), E4M3_MAX), - "one scale per token row": ( - _expand(amax.amax(1, keepdim=True).expand_as(amax), n), - E4M3_MAX, - ), - "blocks down the token axis": ( - _amax(x.t().contiguous()).t().repeat_interleave(128, dim=0)[:t].contiguous(), - E4M3_MAX, - ), - "misaligned block boundaries": ( - torch.roll(_expand(_amax(torch.roll(x, -64, dims=1)), n), 64, dims=1).contiguous(), - E4M3_MAX, - ), - "e4m3 max taken as 240": (_expand(amax, n), 240.0), - "signed max instead of amax": ( - _expand(x.float().reshape(t, n // 128, 128).amax(dim=2).clamp(min=AMAX_FLOOR), n), - E4M3_MAX, - ), - } - for name, (a_cols, mx) in variants.items(): - agree, rel = score(a_cols, mx) - print(f" {name:34s} {agree:.5f} {rel:.4e}") - if name == "correct": - assert agree >= 0.995 and rel <= (SCALE_ULP_GATE + 1) * 2**-23 - else: - assert agree < 0.9 and rel > 1e-3, name - - -def test_transposed_scale_buffer_is_undetectable_at_the_square_shape(): - x = _input(20, 2560, seed=4, magnitude=1.0) - _, sf = op(x, False) - tr = sf.t().contiguous() - assert sf.shape == tr.shape == (20, 20) and tr.dtype == sf.dtype and tr.is_contiguous() - spread = ((tr - sf).abs() / sf).max().item() - print(f" worst relative difference between the buffer and its transpose: {spread:.4f}") - assert spread > 0.1 - - -def test_width_not_a_multiple_of_128_takes_its_own_tail_amax(): - x = torch.zeros(1, 144, dtype=torch.bfloat16, device=DEV) - x[0, 0], x[0, 130] = 4.0, 0.25 - _, sf = op(x, False) - assert abs(sf[0, 0].item() - 4.0 / E4M3_MAX) <= 4.0 / E4M3_MAX * 2**-20 - assert abs(sf[1, 0].item() - 0.25 / E4M3_MAX) <= 0.25 / E4M3_MAX * 2**-20 - - -# ── edges ───────────────────────────────────────────────────────────────── - - -def test_zero_block_takes_the_amax_floor(): - x = torch.zeros(4, 256, dtype=torch.bfloat16, device=DEV) - x[0, 0] = 1.0 - data, sf = op(x, False) - floor = AMAX_FLOOR / E4M3_MAX - assert abs(sf[0, 0].item() - 1.0 / E4M3_MAX) <= 1.0 / E4M3_MAX * 2**-20 - for b in range(2): - for t in range(4): - if (b, t) != (0, 0): - assert abs(sf[b, t].item() - floor) <= floor * 2**-20, (b, t, sf[b, t].item()) - assert bool((data[1:].float() == 0).all()) and bool((data[0, 1:].float() == 0).all()) - - -def test_nonfinite_input(): - """inf poisons its whole 128-block; NaN is confined to its own element.""" - x = torch.zeros(1, 128, dtype=torch.bfloat16, device=DEV) - x[0, 0], x[0, 1] = float("inf"), 1.0 - data, sf = op(x, False) - assert sf.flatten().tolist() == [float("inf")] - assert torch.isnan(data[0, 0].float()) and data[0, 1].float().item() == 0.0 - - x = torch.zeros(1, 128, dtype=torch.bfloat16, device=DEV) - x[0, 0], x[0, 1] = float("nan"), 1.0 - data, sf = op(x, False) - assert abs(sf[0, 0].item() - 1.0 / E4M3_MAX) <= 1.0 / E4M3_MAX * 2**-20 - assert torch.isnan(data[0, 0].float()) and data[0, 1].float().item() == E4M3_MAX - assert bool((data[0, 2:].float() == 0).all()) - - -def test_saturation_never_exceeds_the_block_max(): - x = torch.zeros(1, 256, dtype=torch.bfloat16, device=DEV) - x[0, 0], x[0, 1], x[0, 3], x[0, 128] = 1.0e4, 1.0, -1.0e4, 2.0**-12 - data, sf = op(x, False) - assert data[0, 0].float().item() == E4M3_MAX and data[0, 3].float().item() == -E4M3_MAX - assert data[0, 128].float().item() == E4M3_MAX - assert bool(torch.isfinite(data.float()).all()) - # bf16 rounds 1e4 to 9984, and that is the magnitude the scale carries - assert abs(sf[0, 0].item() - 9984.0 / E4M3_MAX) <= 9984.0 / E4M3_MAX * 2**-20 - - -def test_storage_padding_and_density(): - """`data` views a 4-row-padded allocation; both outputs are dense.""" - for t in (1, 2, 3, 5, 20, 64): - data, sf = op(_input(t, 2560, seed=t), False) - assert data.untyped_storage().size() == (t + 3) // 4 * 4 * 2560, t - assert data.is_contiguous() and data.stride() == (2560, 1) - assert sf.is_contiguous() and sf.untyped_storage().size() >= sf.numel() * 4 - assert sf.stride() == (t, 1) or sf.shape[0] == 1 - _, sf = op(_input(1, 128, seed=0), False) - assert sf.shape == (1, 1) and sf.is_contiguous() - - -def test_empty_input(): - data, sf = op(torch.empty(0, 2560, dtype=torch.bfloat16, device=DEV), False) - assert data.shape == (0, 2560) and data.dtype == torch.float8_e4m3fn - assert sf.shape == (20, 0) and sf.dtype == torch.float32 - - -def test_registered_fake_matches_eager_shapes(): - from torch._subclasses.fake_tensor import FakeTensorMode - - cells = [(1, 2560), (5, 2560), (20, 2560), (8192, 2560), (7, 144)] - eager = {} - for t, n in cells: - d, sf = op(_input(t, n, seed=0)) - eager[(t, n)] = ((tuple(d.shape), d.dtype), (tuple(sf.shape), sf.dtype)) - with FakeTensorMode(): - for t, n in cells: - d, sf = op(torch.empty(t, n, dtype=torch.bfloat16, device=DEV)) - got = ((tuple(d.shape), d.dtype), (tuple(sf.shape), sf.dtype)) - assert got == eager[(t, n)], (t, n, got, eager[(t, n)]) +# ── what a cell cannot carry ────────────────────────────────────────────── def test_cuda_graph_capture_and_replay(): @@ -368,31 +154,8 @@ def test_cuda_graph_capture_and_replay(): assert torch.equal(a_sf.view(torch.int32), e_sf.view(torch.int32)) -def test_input_untouched_and_deterministic(): - x = _input(128, 2560, seed=3) - before = x.clone() - d1, s1 = op(x, False) - d2, s2 = op(x, False) - assert torch.equal(x, before) - assert torch.equal(d1.view(torch.uint8), d2.view(torch.uint8)) - assert torch.equal(s1.view(torch.int32), s2.view(torch.int32)) - - -def test_rejects_unsupported(): - """Every precondition is the op's own and loud, which is why there is no guard.""" - bf = dict(dtype=torch.bfloat16, device=DEV) - cases = { - "float16 input": lambda: op(torch.randn(4, 128, dtype=torch.float16, device=DEV)), - "float32 input": lambda: op(torch.randn(4, 128, device=DEV)), - "float8 input": lambda: op(torch.randn(4, 128, **bf).to(torch.float8_e4m3fn)), - "cpu input": lambda: op(torch.randn(4, 128, dtype=torch.bfloat16)), - "1-D input": lambda: op(torch.randn(128, **bf)), - "3-D input": lambda: op(torch.randn(2, 4, 128, **bf)), - "width 120": lambda: op(torch.randn(4, 120, **bf)), - "width 2568": lambda: op(torch.randn(4, 2568, **bf)), - "transposed": lambda: op(torch.randn(128, 8, **bf).t()), - "column slice": lambda: op(torch.randn(8, 5120, **bf)[:, :2560]), - } - for name, fn in cases.items(): - with pytest.raises((RuntimeError, NotImplementedError)): - fn() +def test_rejects_a_strided_input(): + """The op's own precondition, loud, which is why there is no guard.""" + x = torch.randn(8, 5120, dtype=torch.bfloat16, device=DEV)[:, :2560] + with pytest.raises((RuntimeError, NotImplementedError)): + op(x) 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_fwd.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_fwd.py index 2e691b0f38fe..5ca8d32dfed5 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_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_causal_conv1d_fwd.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the causal_conv1d_fwd entry.""" -import os import zlib from typing import Optional @@ -57,51 +56,29 @@ def _seed(text: str) -> int: def _build(spec, seed): """`__call__`'s positional arguments for one spec, freshly allocated.""" g = torch.Generator(device="cuda").manual_seed(seed) - dim, width = spec["dim"], spec.get("width", 4) - dtype, layout = spec.get("dtype", torch.bfloat16), spec.get("layout", "in_place") + dim, width, dtype = spec["dim"], 4, torch.bfloat16 def randn(*shape): return torch.randn(*shape, generator=g, device="cuda").to(dtype) - out = None - if "batch3d" in spec: - batch, seqlen = spec["batch3d"] - x = randn(batch, dim, seqlen) - qsl = None - else: - lengths = spec["lengths"] - total = sum(lengths) - starts = [0] - for length in lengths: - starts.append(starts[-1] + length) - qsl = torch.tensor(starts, dtype=torch.int32, device="cuda") - token_major = randn(total, dim) - if layout == "channel_last": - x = token_major.t() - assert x.stride() == (1, dim) - out = torch.zeros(total, dim, dtype=dtype, device="cuda").t() - else: - x = token_major.t().contiguous() - if layout == "out": - out = torch.zeros_like(x) + 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) - bias = randn(dim) if spec.get("bias", False) else None - pool = None if spec["pool_slots"] is None else randn(spec["pool_slots"], dim, width - 1) - slots = spec["slots"] - cache_indices = None if slots is None else torch.tensor(slots, dtype=torch.int32, device="cuda") - initial = spec["initial"] - his = None if initial is None else torch.tensor(initial, dtype=torch.bool, device="cuda") - silu = spec.get("silu", True) - pad = spec.get("pad_slot_id", PAD_SLOT_ID) - return [x, weight, bias, pool, qsl, cache_indices, his, silu, pad, out] + 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, out); strides preserved.""" + """The mutable operands cloned (x, conv_states).""" got = list(args) - for i in (0, 3, 9): - if got[i] is not None: - got[i] = got[i].clone() + for i in (0, 3): + got[i] = got[i].clone() return got @@ -118,9 +95,7 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res args = _build(cell.spec, seed=_seed(cell.why)) ref = op.reference(*args) returned, pool, got = _run(args) - assert returned is (got[0] if got[9] is None else got[9]), "returns the tensor it wrote" - if got[9] is not None: - assert torch.equal(got[0], args[0]), "x is read-only when out is given" + 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; @@ -129,135 +104,28 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res assert not torch.equal(pool[slot], args[3][slot]), slot -def test_in_place_and_out_of_place_are_bitwise_identical() -> None: - spec = dict(dim=10240, lengths=(64, 9), pool_slots=4, slots=(1, 2), initial=(True, True)) - in_place, pool_a, _ = _run(_build(spec, seed=1)) - out_of_place, pool_b, _ = _run(_build(dict(spec, layout="out"), seed=1)) - assert torch.equal(in_place, out_of_place) and torch.equal(pool_a, pool_b) - - -def test_channel_last_and_channel_major_are_bitwise_identical() -> None: - """The two kernels agree bitwise under the default SiLU form only.""" - spec = dict( - dim=10240, lengths=(37, 1, 128), pool_slots=8, slots=(3, 0, 5), initial=(True, False, True) - ) - major, pool_a, _ = _run(_build(spec, seed=11)) - last, pool_b, _ = _run(_build(dict(spec, layout="channel_last"), seed=11)) - assert torch.equal(pool_a, pool_b) - if os.environ.get("TRTLLM_CAUSAL_CONV1D_SILU_MODE", "") in ("", "0", "1"): - assert torch.equal(major, last) - else: - torch.testing.assert_close(major, last) - - -def test_the_published_state_is_the_raw_input_tail() -> None: - """Discrimination: the state is the pre-convolution input, not the activation.""" - spec = dict(dim=512, lengths=(48,), pool_slots=1, slots=(0,), initial=(True,), layout="out") - args = _build(spec, seed=12) - returned, pool, _ = _run(args) - x = args[0] - assert torch.equal(pool[0], x[:, -3:]) - assert (pool[0].float() - returned[:, -3:].float()).abs().max().item() > 0.1 - - -def test_the_initial_state_is_read() -> None: - """Discrimination: flipping the carry moves the first W-1 columns, nothing else.""" - spec = dict(dim=512, lengths=(64,), pool_slots=1, slots=(0,)) - with_carry, _, _ = _run(_build(dict(spec, initial=(True,)), seed=13)) - without, _, _ = _run(_build(dict(spec, initial=(False,)), seed=13)) - assert (with_carry[:, :3].float() - without[:, :3].float()).abs().max().item() > 0.1 - assert torch.equal(with_carry[:, 3:], without[:, 3:]) - - -_REJECTED = ( - "channel_last_in_place", - "channel_last_with_channel_major_out", - "channel_major_with_channel_last_out", - "out_dtype", - "out_shape", - "bias_dtype", - "x_float64", - "query_start_loc_int64", - "has_initial_state_length", - "conv_states_dtype", - "x_rank", -) - - -@pytest.mark.parametrize("case", _REJECTED) -def test_op_rejects_loudly_what_is_valid_does_not_restate(case: str) -> None: - """Everything the op refuses itself -- the reason `is_valid` stays narrow.""" - dim, tokens = 512, 64 - base = dict(dim=dim, lengths=(30, 34), pool_slots=2, slots=(0, 1), initial=(True, True)) - args = _build(dict(base, pool_slots=None, slots=None, initial=None), seed=14) - token_major = torch.randn(tokens, dim, dtype=torch.bfloat16, device="cuda") - pool = torch.randn(2, dim, 3, dtype=torch.bfloat16, device="cuda") - slots = torch.tensor([0, 1], dtype=torch.int32, device="cuda") - if case == "channel_last_in_place": - args[0] = token_major.t() - elif case == "channel_last_with_channel_major_out": - args[0] = token_major.t() - args[9] = torch.empty(dim, tokens, dtype=torch.bfloat16, device="cuda") - elif case == "channel_major_with_channel_last_out": - args[9] = torch.empty(tokens, dim, dtype=torch.bfloat16, device="cuda").t() - elif case == "out_dtype": - args[9] = torch.empty(dim, tokens, dtype=torch.float16, device="cuda") - elif case == "out_shape": - args[9] = torch.empty(dim, tokens + 1, dtype=torch.bfloat16, device="cuda") - elif case == "bias_dtype": - args[2] = torch.randn(dim, dtype=torch.float32, device="cuda") - elif case == "x_float64": - args[0] = args[0].double() - args[1] = args[1].double() - elif case == "query_start_loc_int64": - args[4] = args[4].long() - elif case == "has_initial_state_length": - args[3], args[5] = pool, slots - args[6] = torch.ones(3, dtype=torch.bool, device="cuda") - elif case == "conv_states_dtype": - args[3], args[5] = pool.float(), slots - args[6] = torch.ones(2, dtype=torch.bool, device="cuda") - elif case == "x_rank": - args[0] = torch.randn(2, dim, 32, dtype=torch.bfloat16, device="cuda") - with pytest.raises(RuntimeError): - with validating(op): - op(*args) - - -def _guard_case(case: str): - """A fresh argument list exhibiting one silent violation.""" +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,)) - if case in ("width_one", "width_five"): - spec.update( - width=1 if case == "width_one" else 5, pool_slots=None, slots=None, initial=None - ) args = _build(spec, seed=15) g = torch.Generator(device="cuda").manual_seed(16) - if case == "weight_dtype": - args[1] = args[1].float() - elif case == "token_stride": - 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 - elif case == "state_stride": - wide = torch.randn(1, dim, 6, generator=g, device="cuda").to(torch.bfloat16) - args[3] = wide[:, :, ::2] + 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 -@pytest.mark.parametrize( - "case", ["width_one", "width_five", "weight_dtype", "token_stride", "state_stride"] -) -def test_guard_refuses_what_the_op_answers_silently_wrong(case: str) -> None: +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(*_guard_case(case)) - unguarded = _guard_case(case) # strided views must not be cloned dense + 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(*_guard_case(case)) + 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_causal_conv1d_triton_update.py b/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 index 2d465304f39b..7f8a993b6785 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_triton_update.py @@ -8,7 +8,6 @@ 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 - CKPT_DIM, ENGINE_SLOTS, ) 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 @@ -49,61 +48,43 @@ def _compare(out: tuple, ref: tuple) -> None: def _build(spec, seed): """`__call__`'s positional arguments for one spec, freshly allocated.""" g = torch.Generator(device="cuda").manual_seed(seed) - dim, batch = spec["dim"], spec["batch"] - width, state_len = spec.get("width", 4), spec.get("state_len", 3) - seqlen, dtype = spec.get("seqlen", 2), spec.get("dtype", torch.bfloat16) - x_form = spec.get("x_form", "token_major") + dim, batch, seqlen, width = spec["dim"], spec["batch"], spec["seqlen"], 4 def randn(*shape): - return torch.randn(*shape, generator=g, device="cuda").to(dtype) - - if x_form == "2d": - x = randn(batch, dim) - elif x_form == "contiguous": - x = randn(batch, dim, seqlen) - else: # the call site: an [n, p, C] block seen as [n, C, p] - x = randn(batch, seqlen, dim).transpose(1, 2) - conv_state = randn(spec["pool_slots"], dim, state_len) + 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) + assert x.stride() == (seqlen * wide, 1, wide) + conv_state = randn(spec["pool_slots"], dim, width - 1) weight = randn(dim, width) - bias = randn(dim) if spec.get("bias", False) else None - slots = spec["slots"] - if slots == "random": - indices = torch.randperm(spec["pool_slots"], generator=g, device="cuda")[:batch] - indices = indices.to(torch.int32) - elif slots is None: - indices = None - else: - indices = torch.tensor(slots, dtype=torch.int32, device="cuda") - accepted = spec.get("accepted") - if accepted is not None: - accepted = torch.tensor(accepted, dtype=torch.int32, device="cuda") - rows = spec.get("snapshots", 0) - window = rows_idx = None - if rows: - steps = 1 if x_form == "2d" else seqlen - window = torch.zeros(rows, steps, dim, width - 1, dtype=dtype, device="cuda") - rows_idx = torch.arange(batch, dtype=torch.int32, device="cuda") + 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" + ) + rows_idx = torch.arange(batch, dtype=torch.int32, device="cuda") return [ x, conv_state, weight, - bias, - spec.get("activation", "silu"), + None, + "silu", indices, - spec.get("pad_slot_id", PAD_SLOT_ID), - accepted, + PAD_SLOT_ID, + None, window, rows_idx, - spec.get("launch_dependent_kernels", False), + False, ] def _run(args): """Drive the op on copies of the two buffers; return the triple compare takes.""" got = list(args) - got[1] = args[1].clone() - if args[8] is not None: - got[8] = args[8].clone() + got[1], got[8] = args[1].clone(), args[8].clone() returned = op(*got) return returned, got[1], got[8] @@ -123,103 +104,13 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res assert not torch.equal(got[1][slot], args[1][slot]), slot -def test_the_last_snapshot_is_what_the_pool_holds() -> None: - """Independent of the reference: snapshot `seqlen - 1` equals the advanced slot.""" - spec = dict(dim=512, batch=6, seqlen=4, pool_slots=16, slots="random", snapshots=6) - args = _build(spec, seed=1) - _, pool, snapshots = _run(args) - for b, slot in enumerate(args[5].tolist()): - assert torch.equal(snapshots[b, 3], pool[slot]), b - - -def test_fp32_products_are_the_wrong_reference() -> None: - """Discrimination: the gate is default-tolerance tight, and the natural - wrong guess -- multiplying in fp32 -- lands outside it. - - The rival is `reference` on fp32 copies of the operands, rounded once at - the end. Measured at 10240 channels: the faithful reference differs on - 2.4e-5 of elements (one-ulp `exp` cases); the rival on 0.37 of them. - """ - spec = dict(dim=CKPT_DIM, batch=4, pool_slots=128, slots="random") - args = _build(spec, seed=0) - got = _run(args) - faithful = op.reference(*args) - _compare(got, faithful) - fp32_args = list(args) - for i in (0, 1, 2): - fp32_args[i] = args[i].float() - rival_out = op.reference(*fp32_args)[0].to(args[0].dtype) - rival = (rival_out, faithful[1], faithful[2]) - with pytest.raises(AssertionError): - _compare(got, rival) - faithful_diff = (got[0] != faithful[0]).float().mean().item() - rival_diff = (got[0] != rival_out).float().mean().item() - print(f"\nelements differing: faithful {faithful_diff:.3e}, rival {rival_diff:.3e}") - assert faithful_diff < 1e-3 and rival_diff > 0.3, (faithful_diff, rival_diff) - - -def test_token_major_and_contiguous_x_are_bitwise_identical() -> None: - args = _build(dict(dim=512, batch=3, seqlen=4, pool_slots=16, slots=(2, 5, 1)), seed=11) - assert not args[0].is_contiguous() - strided = _run(args) - packed = _run([args[0].contiguous()] + args[1:]) - assert torch.equal(strided[0], packed[0]) and torch.equal(strided[1], packed[1]) - - -def test_launch_dependent_kernels_is_bitwise_inert() -> None: - spec = dict(dim=512, batch=4, pool_slots=16, slots="random", snapshots=4) - off = _run(_build(spec, seed=12)) - on = _run(_build(dict(spec, launch_dependent_kernels=True), seed=12)) - assert all(torch.equal(a, b) for a, b in zip(off, on)) - - -def test_width_five_is_silently_wrong_and_refused() -> None: - """Unguarded, width 5 runs and misses the reference; the guard refuses it.""" - spec = dict(dim=512, batch=2, width=5, state_len=4, pool_slots=8, slots=(1, 3)) - args = _build(spec, seed=13) - ref = op.reference(*args) - with pytest.raises(AssertionError): - _compare(_run(args), ref) - with pytest.raises(AssertionError, match="width must be 2, 3 or 4"): - with validating(op): - op(*_build(spec, seed=13)) - - -def test_a_short_pool_writes_into_the_next_slot_and_is_refused() -> None: - """Unguarded, a 2-column pool at W 4 writes past its slot into an unnamed - neighbour; the guard refuses it.""" - spec = dict(dim=512, batch=2, state_len=2, pool_slots=8, slots=(1, 4)) - args = _build(spec, seed=14) - _, pool, _ = _run(args) - unnamed = [s for s in range(8) if s not in (1, 4)] - assert not torch.equal(pool[unnamed], args[1][unnamed]), "the hazard this guard covers is gone" - with pytest.raises(AssertionError, match="must hold at least 3 columns"): - with validating(op): - op(*_build(spec, seed=14)) - - -@pytest.mark.parametrize( - "case,overrides,match", - [ - ("narrow_pool", None, "must cover all 512 channels"), - ("short_weight", None, "must cover all 512 channels"), - ("short_snapshot_buffer", None, "intermediate_conv_window must be at least"), - # The accepted-token form needs W-1 + (seqlen-1) = 4 columns at seqlen 2. - ("accepted_needs_a_longer_pool", dict(accepted=(1, 1)), "must hold at least 4 columns"), - ], -) -def test_guard_refuses_undersized_buffers(case, overrides, match) -> None: - """Accepted by the kernel and read or written out of bounds (measured); - refused by the guard before the launch.""" - spec = dict(dim=512, batch=2, pool_slots=8, slots=(0, 2), snapshots=2) - spec.update(overrides or {}) +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) - if case == "narrow_pool": - args[1] = torch.randn(8, 256, 3, dtype=torch.bfloat16, device="cuda") - elif case == "short_weight": - args[2] = torch.randn(256, 4, dtype=torch.bfloat16, device="cuda") - elif case == "short_snapshot_buffer": - args[8] = torch.zeros(2, 1, 512, 3, dtype=torch.bfloat16, device="cuda") - with pytest.raises(AssertionError, match=match): + 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) 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_update.py index 4b61e62ec676..bbd3b86e9575 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_causal_conv1d_update.py @@ -39,72 +39,48 @@ def _compare( PAD_SLOT_ID = -1 -#: Rows of prefill ahead of the decode rows in the mixed-batch cell. +#: 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 some layouts view.""" + """`__call__`'s arguments for one spec, plus the buffers x and the pool view.""" g = torch.Generator(device="cuda").manual_seed(seed) - dim, batch = spec["dim"], spec["batch"] - width, state_len, seqlen = spec.get("width", 4), spec.get("state_len", 3), spec.get("seqlen", 1) - dtype, x_form = spec.get("dtype", torch.bfloat16), spec.get("x_form", "3d") + dim, batch, layers = spec["dim"], spec["batch"], spec.get("pool_layers", 1) def randn(*shape): - return torch.randn(*shape, generator=g, device="cuda").to(dtype) - - extras = {} - if x_form == "2d": - x = randn(batch, dim) - elif x_form == "2d_row_slice": - extras["buf"] = randn(_PREFILL_ROWS + batch, dim) - x = extras["buf"][_PREFILL_ROWS:] - elif x_form == "token_major": - x = randn(batch, seqlen, dim).transpose(1, 2) - assert x.stride() == (seqlen * dim, 1, dim) - else: - x = randn(batch, dim, seqlen) - if spec.get("strided_state", False): - extras["wide"] = randn(spec["pool_slots"], dim, 2 * state_len) - conv_state = extras["wide"][:, :, ::2] - else: - conv_state = randn(spec["pool_slots"], dim, state_len) - weight = randn(dim, width) - bias = randn(dim) if spec.get("bias", False) else None + return torch.randn(*shape, generator=g, device="cuda").to(torch.bfloat16) + + buf = randn(_PREFILL_ROWS + batch, dim + _IN_PROJ_TAIL) + x = buf[_PREFILL_ROWS:, :dim] + assert x.stride() == (dim + _IN_PROJ_TAIL, 1) + pool_buf = randn(spec["pool_slots"], layers, dim, 3) + conv_state = pool_buf[:, layers - 1] + weight = randn(dim, 4) slots = spec["slots"] if slots == "random": indices = torch.randperm(spec["pool_slots"], generator=g, device="cuda")[:batch] indices = indices.to(torch.int32) - elif slots is None: - indices = None else: indices = torch.tensor(slots, dtype=torch.int32, device="cuda") - cache_seqlens = spec.get("cache_seqlens") - if cache_seqlens is not None: - cache_seqlens = torch.tensor(cache_seqlens, dtype=torch.int32, device="cuda") - silu, pad = spec.get("silu", True), spec.get("pad_slot_id", PAD_SLOT_ID) - return [x, conv_state, weight, bias, silu, cache_seqlens, indices, pad], extras + args = [x, conv_state, weight, None, True, None, indices, PAD_SLOT_ID] + return args, dict(buf=buf, pool_buf=pool_buf) def _run(args, extras): """Drive the op on copies of x and the pool, through the same views.""" got = list(args) - if "buf" in extras: - got_buf = extras["buf"].clone() - got[0] = got_buf[_PREFILL_ROWS:] - else: - got[0] = args[0].clone() - if "wide" in extras: - got_wide = extras["wide"].clone() - got[1] = got_wide[:, :, ::2] - else: - got[1] = args[1].clone() + got_buf, got_pool = extras["buf"].clone(), extras["pool_buf"].clone() + got[0] = got_buf[_PREFILL_ROWS:, : args[0].shape[1]] + got[1] = got_pool[:, got_pool.shape[1] - 1] returned = op(*got) assert returned is got[0], "returns the x it was handed" - if "buf" in extras: - assert torch.equal(got_buf[:_PREFILL_ROWS], extras["buf"][:_PREFILL_ROWS]) - if "wide" in extras: - assert torch.equal(got_wide[:, :, 1::2], extras["wide"][:, :, 1::2]) + assert torch.equal(got_buf[:_PREFILL_ROWS], extras["buf"][:_PREFILL_ROWS]) + dim = args[0].shape[1] + assert torch.equal(got_buf[:, dim:], extras["buf"][:, dim:]), "columns past x written" + assert torch.equal(got_pool[:, :-1], extras["pool_buf"][:, :-1]), "another layer written" return returned, got[1] @@ -112,70 +88,24 @@ def _run(args, extras): def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name certifying(op) args, extras = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) - weight, bias = args[2].clone(), None if args[3] is None else args[3].clone() + weight = args[2].clone() ref = op.reference(*args) returned, state = _run(args, extras) assert torch.equal(args[2], weight), "weight is read-only" - assert bias is None or torch.equal(args[3], bias), "bias 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"]: - assert not torch.equal(state[slot], args[1][slot]), slot + if slot != PAD_SLOT_ID: + assert not torch.equal(state[slot], args[1][slot]), slot -def _small(**overrides): +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)) - spec.update(overrides) - return _build(spec, seed=14)[0] - - -_REJECTIONS = { - "width_1": ( - "causal_conv1d only supports width between 2 and 4", - dict(dim=128, width=1, state_len=1), - ), - "width_5": ( - "causal_conv1d only supports width between 2 and 4", - dict(dim=128, width=5, state_len=4), - ), - "state_too_short": ("conv_state_len >= width - 1", dict(state_len=2)), - "indices_length": ("conv_state_indices must have shape", None), - "indices_int64": ("conv_state_indices.scalar_type", None), - "indices_non_contiguous": ("is_contiguous", None), - "cache_seqlens_int64": ("cache_seqlens.scalar_type", None), - "no_indices_oversize_pool": ("conv_state must have shape", dict(slots=None)), - "weight_shape": ("weight must have shape", None), - "conv_state_dtype": ("conv_state.scalar_type", None), - "weight_dtype": ("weight type must equal to input type", None), - "bias_dtype": ("bias.scalar_type", None), - "float64": ("input_type == at::ScalarType::Float", None), -} - - -@pytest.mark.parametrize("case", sorted(_REJECTIONS)) -def test_op_rejects_every_metadata_violation_loudly(case: str) -> None: - """Why the entry has no `is_valid`: every metadata violation probed raises.""" - fragment, overrides = _REJECTIONS[case] - args = _small(**(overrides or {})) - if case == "indices_length": - args[6] = args[6][:3] - elif case == "indices_int64": - args[6] = args[6].long() - elif case == "indices_non_contiguous": - args[6] = torch.arange(8, dtype=torch.int32, device="cuda")[::2] - elif case == "cache_seqlens_int64": - args[5] = torch.zeros(4, dtype=torch.int64, device="cuda") - elif case == "weight_shape": - args[2] = torch.randn(256, 4, dtype=torch.bfloat16, device="cuda") - elif case == "conv_state_dtype": - args[1] = args[1].float() - elif case == "weight_dtype": - args[2] = args[2].float() - elif case == "bias_dtype": - args[3] = torch.randn(512, dtype=torch.float32, device="cuda") - elif case == "float64": - args[0], args[1], args[2] = args[0].double(), args[1].double(), args[2].double() - with pytest.raises(RuntimeError, match=fragment.replace("(", r"\(").replace(")", r"\)")): + 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) 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 index 92fff412c4f4..6a8cf2fa2286 100644 --- 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 @@ -4,8 +4,9 @@ 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. "First call in the process" is a property of the process, so the cold-call tests -re-run this file as a child: `python `. +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 @@ -13,12 +14,12 @@ import sys import zlib from pathlib import Path -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.chunk_gated_delta_rule import ( # noqa: E501 + ENGINE_SLOTS, HEAD_K_DIM, HEAD_V_DIM, NUM_K_HEADS, @@ -49,7 +50,7 @@ def _compare(out: tuple, ref: tuple) -> None: 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: seven wrong references (no gating, beta = 1, q/k swapped, + 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. """ @@ -65,7 +66,6 @@ def _compare(out: tuple, ref: tuple) -> None: DEV = "cuda" -SCALE = HEAD_K_DIM**-0.5 def _l2norm(x: torch.Tensor) -> torch.Tensor: @@ -73,48 +73,35 @@ def _l2norm(x: torch.Tensor) -> torch.Tensor: return xf / torch.sqrt(xf.pow(2).sum(-1, keepdim=True) + 1e-6) -def _operands( - seqlens: Optional[tuple], - seed: int, - batch: int = 1, - num_tokens: Optional[int] = None, - num_k_heads: int = NUM_K_HEADS, - num_v_heads: int = NUM_V_HEADS, - normalize_qk: bool = True, -): +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) - if seqlens is None: - cu_seqlens, total = None, num_tokens - else: - total = sum(seqlens) - cu_seqlens = torch.tensor( - [0] + torch.tensor(seqlens).cumsum(0).tolist(), dtype=torch.long, device=DEV - ) - shape_qk = (batch, total, num_k_heads, HEAD_K_DIM) + 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) - if normalize_qk: - q, k = _l2norm(q), _l2norm(k) - q, k = q.to(torch.bfloat16), k.to(torch.bfloat16) + q, k = _l2norm(q).to(torch.bfloat16), _l2norm(k).to(torch.bfloat16) v = torch.randn( - batch, total, num_v_heads, HEAD_V_DIM, generator=gen, device=DEV, dtype=torch.bfloat16 + 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(batch, total, num_v_heads, generator=gen, device=DEV) + 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(batch, total, num_v_heads, generator=gen, device=DEV).sigmoid().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, num_v_heads: int = NUM_V_HEADS) -> torch.Tensor: - """An fp32 recurrent-state pool, `[slots, H, V, K]` (K innermost).""" +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 + 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: @@ -128,56 +115,23 @@ def _gate_pool(pool: torch.Tensor, ref: tuple) -> None: def _build(spec, seed): - """`__call__`'s 13 arguments for one spec, plus the wide output buffer.""" - q, k, v, g, beta, cu_seqlens = _operands( - spec.get("seqlens"), - seed, - batch=spec.get("batch", 1), - num_tokens=spec.get("num_tokens"), - normalize_qk=spec.get("normalize_qk", True), - ) - num_seqs = q.shape[0] if cu_seqlens is None else len(cu_seqlens) - 1 - initial, indices, inplace = None, None, False - if spec["initial"] == "indexed": - initial = _pool(spec["slots"], seed=seed + 1) - indices = torch.tensor( - spec.get("indices", tuple(range(num_seqs))), dtype=torch.int32, device=DEV - ) - inplace = True - elif spec["initial"] == "positional": - initial = _pool(num_seqs, seed=seed + 1) - wide, output = None, None - shape_o = (q.shape[0], q.shape[1], NUM_V_HEADS, HEAD_V_DIM) + """`__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) - elif spec.get("output") == "window": - wide = torch.zeros(1, 200, NUM_V_HEADS, HEAD_V_DIM, device=DEV, dtype=torch.bfloat16) - output = wide[:, 50 : 50 + q.shape[1]] - args = [ - q, - k, - v, - g, - beta, - spec.get("scale"), - initial, - indices, - inplace, - spec.get("output_final_state", False), - cu_seqlens, - spec.get("l2norm", False), - output, - ] - return args, wide + 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 - spec = cell.spec certifying(op) - args, wide = _build(spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) pool = args[6] - pool_before = None if pool is None else pool.clone() + pool_before = pool.clone() ref = op.reference(*args) o, final_state = op(*args) @@ -185,48 +139,10 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res _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" - if wide is not None: - assert torch.count_nonzero(wide[:, :50]) == 0 and torch.count_nonzero(wide[:, 146:]) == 0 - if spec["initial"] == "positional": - assert torch.equal(pool, pool_before), "a positional initial state is read-only" - if spec["initial"] == "indexed": - 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_op_rejects_these_domains_loudly() -> None: - """Why `is_valid` guards only `output`: every other refusal is the op's.""" - q, k, v, g, beta, cu = _operands((80,), seed=67) - idx1 = torch.zeros(1, dtype=torch.int32, device=DEV) - cases = [ - (AssertionError, "does not support float32", dict(q=q.float(), k=k.float(), v=v.float())), - (AssertionError, "", dict(k=k.half())), - (AssertionError, "beta must be of shape", dict(beta=beta[0])), - (ValueError, "batch size is expected to be 1", dict(q=q.expand(2, -1, -1, -1))), - ( - ValueError, - "number of initial-state indices", - dict( - initial_state=_pool(4, seed=671), - initial_state_indices=torch.tensor([0, 1], dtype=torch.int32, device=DEV), - inplace_indexed_state_update=True, - ), - ), - (ValueError, "number of initial states", dict(initial_state=_pool(3, seed=672))), - ( - ValueError, - "require inplace_indexed_state_update=True", - dict(initial_state=_pool(1, seed=673), initial_state_indices=idx1), - ), - ] - for exc, match, override in cases: - kwargs = dict(q=q, k=k, v=v, g=g, beta=beta, cu_seqlens=cu) - kwargs.update(override) - with pytest.raises(exc, match=match): - with validating(op): - op(**kwargs) + 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: @@ -246,46 +162,9 @@ def test_guard_refuses_a_noncontiguous_output() -> None: op(q, k, v, g, beta, cu_seqlens=cu, output=head_slice) -def test_guard_refuses_an_output_of_another_dtype() -> None: - """Unguarded, an fp32 buffer is filled but the returned bf16 tensor no - longer aliases it -- the allocation the buffer existed to avoid.""" - q, k, v, g, beta, cu = _operands((96,), seed=62) - buf = torch.zeros(1, 96, NUM_V_HEADS, HEAD_V_DIM, device=DEV, dtype=torch.float32) - o, _ = op(q, k, v, g, beta, cu_seqlens=cu, output=buf) - assert o.dtype == torch.bfloat16 and o.data_ptr() != buf.data_ptr(), "hazard gone" - with pytest.raises(AssertionError, match="output must have q's dtype"): - with validating(op): - op(q, k, v, g, beta, cu_seqlens=cu, output=buf) - - -def test_gate_discriminates_wrong_variants() -> None: - """Seven references wrong in one way each must all fail the `o` gate the - correct one passes. Positional path: no warm-up needed.""" - q, k, v, g, beta, cu = _operands((128, 64), seed=71) - initial = _pool(2, seed=711) - base = dict(q=q, k=k, v=v, g=g, beta=beta, scale=SCALE, initial_state=initial, cu_seqlens=cu) - o, final_state = op(**base, output_final_state=True) - ref = op.reference(**base, output_final_state=True) - _compare((o, final_state, initial), ref) - - variants = { - "no gating": dict(g=torch.zeros_like(g)), - "beta forced to 1": dict(beta=torch.ones_like(beta)), - "q and k swapped": dict(q=k, k=q), - "initial state dropped": dict(initial_state=None), - "scale left at 1": dict(scale=1.0), - "key heads reversed": dict(q=q.flip(2).contiguous(), k=k.flip(2).contiguous()), - "initial state [K, V]": dict(initial_state=initial.transpose(-1, -2).contiguous()), - } - for name, override in variants.items(): - wrong = op.reference(**{**base, **override}) - with pytest.raises(AssertionError): - _gate_output(wrong[0], ref) - - # --------------------------------------------------------------------------- -# The cold call. The cells above have already benchmarked these autotune keys by -# the time pytest reaches these, so each runs in a fresh interpreter. +# 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. @@ -306,9 +185,9 @@ def _state_kernel_autotuner(): raise AssertionError("no triton autotuner found around the chunked state kernel") -def _checkpoint_case(seed: int = 101): +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) + 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 @@ -316,10 +195,10 @@ def _checkpoint_case(seed: int = 101): return (q, k, v, g, beta, None, pool, indices, True, False, cu, False, buf) -def _child_cold_call() -> None: +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() + 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" @@ -327,24 +206,8 @@ def _child_cold_call() -> None: _gate_pool(args[6], ref) -def _child_cold_call_odd_heads() -> None: - """A head geometry (5 key / 20 value heads) nothing else uses.""" - autotuner = _state_kernel_autotuner() - assert not autotuner.cache, "the child must be the first indexed call in its process" - q, k, v, g, beta, cu = _operands((96,), seed=83, num_k_heads=5, num_v_heads=20) - pool = _pool(4, seed=831, num_v_heads=20) - indices = torch.tensor([1], dtype=torch.int32, device=DEV) - args = (q, k, v, g, beta, None, pool, indices, True, False, cu, False, None) - 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(pool, ref) - - _CHILD_MODES = { - "cold-call": _child_cold_call, - "cold-call-odd-heads": _child_cold_call_odd_heads, + "cold-call-bf16-pool": lambda: _child_cold_call(torch.bfloat16), } @@ -368,15 +231,11 @@ def _run_child(mode: str) -> None: ) -def test_cold_indexed_call_at_this_checkpoints_geometry() -> None: - """The first indexed call of a process at 16 / 48 heads, K = V = 128, in the - mixer's shape (256 + 96 packed, by slot, into `output=`), matches the reference.""" - _run_child("cold-call") - - -def test_cold_indexed_call_at_an_unused_head_geometry() -> None: - """The same at 5 / 20 heads, an autotune key nothing else here touches.""" - _run_child("cold-call-odd-heads") +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__": 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py index 5e011da0ac63..64926845e39c 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/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_cu_seqlens_from_seq_lens.py @@ -2,9 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 """GPU test for the cu_seqlens_from_seq_lens entry. -Integer arithmetic, gated bitwise. Beyond the cells: a near-miss vector is -monotonic and plausible, so the tests that matter are the ones proving the -vector actually splits a packed buffer and differs from every near miss. +Integer arithmetic, gated bitwise. A near-miss vector is monotonic and +plausible, so each cell also proves the vector actually splits a packed +buffer. """ import zlib @@ -12,9 +12,6 @@ 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 - BATCHES, -) 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 cu_seqlens_from_seq_lens as op, ) @@ -49,12 +46,6 @@ def _build(spec): buf = torch.full((len(lengths) + 9,), 99, dtype=dtype, device=DEV) buf[: len(lengths)] = _seq_lens(lengths, dtype) return buf[: len(lengths)] - if layout == "strided": - buf = torch.full((2 * len(lengths),), 99, dtype=dtype, device=DEV) - buf[::2] = _seq_lens(lengths, dtype) - view = buf[::2] - assert not view.is_contiguous() - return view return _seq_lens(lengths, dtype) @@ -81,61 +72,12 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res assert torch.equal(packed[bounds[s] : bounds[s + 1]], piece), f"sequence {s} mis-sliced" -def test_wrong_variants_all_differ() -> None: - """Every near-miss spelling is monotonic and plausible; none is this one.""" - lengths = BATCHES["mixed_p2"] - seq_lens = _seq_lens(lengths) - out = op(seq_lens, torch.int32) - ref = op.reference(seq_lens, torch.int32) - _compare(out, ref) - - inclusive = torch.cumsum(seq_lens, dim=0, dtype=torch.int32) - no_total = torch.cat([torch.zeros(1, dtype=torch.int32, device=DEV), inclusive[:-1]]) - right_padded = torch.cat([inclusive, inclusive[-1:]]) - # The one-row-per-generation form `Mamba2Metadata.prepare` builds: a real - # cumsum over the context prefix, then a constant plus arange. - num_ctx, num_gen = 2, len(lengths) - 2 - engine = torch.zeros(len(lengths) + 1, dtype=torch.int32, device=DEV) - torch.cumsum(seq_lens[:num_ctx], dim=0, dtype=torch.int32, out=engine[1 : num_ctx + 1]) - engine[num_ctx + 1 :] = engine[num_ctx] + torch.arange( - 1, num_gen + 1, dtype=torch.int32, device=DEV - ) - for name, wrong in { - "inclusive scan": inclusive, - "total dropped": no_total, - "right-padded": right_padded, - "one row per generation": engine, - "shifted": torch.roll(out, 1), - }.items(): - with pytest.raises(AssertionError): - _compare(wrong, ref) - assert wrong.shape != ref.shape or not torch.equal(wrong, ref), name - - # The metadata form agrees exactly when every generation request emits one - # row, which is why the mistake is invisible until speculation is on. - unspec = _seq_lens(list(lengths[:num_ctx]) + [1] * num_gen) - assert torch.equal(op(unspec, torch.int32), engine) - - -def test_result_is_fresh_and_deterministic() -> None: - seq_lens = _seq_lens(BATCHES["decode_p2"]) - results = [op(seq_lens, torch.int32) for _ in range(5)] - assert len({r.data_ptr() for r in results} | {seq_lens.data_ptr()}) == 6 - for other in results[1:]: - assert torch.equal(results[0], other) - - -def test_accumulation_happens_in_the_requested_dtype() -> None: - """The hard-limit row: int64 widens an int32 input; int32 wraps silently.""" - seq_lens = _seq_lens([2_000_000_000, 2_000_000_000], torch.int32) - assert op(seq_lens, torch.int64).tolist() == [0, 2_000_000_000, 4_000_000_000] - assert op(seq_lens, torch.int32).tolist() == [0, 2_000_000_000, -294_967_296] - - 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.""" - static_lens = _seq_lens([2] * 8) + 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): @@ -147,39 +89,17 @@ def test_cuda_graph_capture_and_replay() -> None: graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): static_out = op(static_lens, torch.int32) - for lengths in ([1] * 8, [2] * 8, [3, 1, 2, 1, 1, 1, 1, 1]): + 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_a_seq_lens_that_is_not_1d() -> None: - """Unguarded, a [1, 3] row vector returns a length-2 result, not length 4.""" - row = _seq_lens([7, 1, 13]).unsqueeze(0) - with pytest.warns(UserWarning): - unguarded = op(row, torch.int32) - assert unguarded.shape == (2,), "the hazard this guard covers is gone" - for bad in (row, torch.zeros(3, 2, dtype=torch.int32, device=DEV), _seq_lens([3])[0]): - with pytest.raises(AssertionError, match="1-D per-request row count"): - with validating(op): - op(bad, 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" - for dtype in (torch.float32, torch.bfloat16, torch.bool): - with pytest.raises(AssertionError, match="seq_lens must have an integer dtype"): - with validating(op): - op(torch.ones(3, dtype=dtype, device=DEV), torch.int32) - - -def test_guard_refuses_a_non_integer_result_dtype() -> None: - """Unguarded, a float `dtype` returns a float vector the consumers misread.""" - assert op(_seq_lens([1, 2, 3]), torch.float32).dtype is torch.float32 - for dtype in (torch.float32, torch.float64, torch.bool): - with pytest.raises(AssertionError, match="dtype must be an integer dtype"): - with validating(op): - op(_seq_lens([1, 2, 3]), dtype) + 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/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 index a7d49b0d7aa5..8908d558b5ad 100644 --- 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 @@ -1,17 +1,11 @@ # 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. - -The one engine op used here (`test_dilation_one_agrees_with_the_engine`) is an -independent oracle, not a reference: at dilation 1 this operator is exactly -`torch.ops.trtllm.causal_conv1d_fwd`, written from the other end. -""" +"""GPU test for the dilated_short_conv_paged entry.""" 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.ssm.dilated_short_conv_paged import ( # noqa: E501 ENGINE_SLOTS, @@ -51,41 +45,32 @@ def _compare( def _build(spec, seed): - """`__call__`'s positional arguments for one spec, freshly allocated. + """`__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"] - kernel_size, dilation = spec.get("kernel_size", KERNEL_SIZE), spec.get("dilation", DILATION) - dtype = spec.get("dtype", torch.bfloat16) - num_seqs, state_len = len(lengths), (kernel_size - 1) * dilation + num_seqs = len(lengths) def randn(*shape, scale=1.0): - return (scale * torch.randn(*shape, generator=g, device="cuda")).to(dtype) + 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=spec.get("qsl_dtype", torch.int32), device="cuda") - if spec.get("x_form", "channel_major") == "token_major": - x = randn(starts[-1], channels).t() - else: - x = randn(channels, starts[-1]) - weight = randn(channels, 1, kernel_size, scale=0.3) - bias = randn(channels, scale=0.1) if spec.get("bias", False) else None - pool = randn(spec.get("pool_slots", 8), channels, state_len) - if spec.get("no_indices", False): - cache_indices = None - else: - slots = spec.get("slots") or tuple(range(num_seqs)) - index_dtype = spec.get("index_dtype", torch.int32) - cache_indices = torch.tensor(slots, dtype=index_dtype, device="cuda") + 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 = None if initial == "none" else torch.tensor(initial, dtype=torch.bool, device="cuda") - silu, pad = spec.get("silu", True), spec.get("pad_slot_id", PAD_SLOT_ID) - return [x, weight, bias, pool, qsl, cache_indices, his, dilation, silu, pad] + 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): @@ -114,228 +99,11 @@ def _small(lengths, slots, seed, **spec): return _build(dict(channels=256, lengths=lengths, slots=slots, **spec), seed=seed) -def test_published_state_is_the_raw_input_tail() -> None: - """Independent of the reference, on both sides of the state length.""" - args = _small((STATE_LEN + 5, 2), (3, 1), seed=7) - _, pool = _run(args) - x, before = args[0], args[3] - end0, end1 = STATE_LEN + 5, STATE_LEN + 7 - assert torch.equal(pool[3], x[:, end0 - STATE_LEN : end0]) - assert torch.equal(pool[1], torch.cat([before[1][:, 2:], x[:, end0:end1]], dim=1)) - - -def test_impulse_response_pins_the_tap_map() -> None: - """Tap `k` lights output `t` exactly when the impulse sits at window column - `t + k * dilation` -- crossing the pool boundary, where a paged - implementation goes wrong.""" - channels, tokens = 8, 6 - qsl = torch.tensor([0, tokens], dtype=torch.int32, device="cuda") - cidx = torch.tensor([1], dtype=torch.int32, device="cuda") - his = torch.tensor([True], dtype=torch.bool, device="cuda") - for k in range(KERNEL_SIZE): - weight = torch.zeros(channels, 1, KERNEL_SIZE, device="cuda") - weight[:, 0, k] = 1.0 - for column in range(STATE_LEN + tokens): - pool = torch.zeros(2, channels, STATE_LEN, device="cuda") - x = torch.zeros(channels, tokens, device="cuda") - if column < STATE_LEN: - pool[1, :, column] = 1.0 - else: - x[:, column - STATE_LEN] = 1.0 - out = op(x, weight, None, pool, qsl, cidx, his, DILATION, False, PAD_SLOT_ID) - expected = torch.zeros(channels, tokens, device="cuda") - t = column - k * DILATION - if 0 <= t < tokens: - expected[:, t] = 1.0 - assert torch.equal(out, expected), (k, column) - - -def test_the_last_tap_is_the_current_token() -> None: - weight = torch.zeros(16, 1, KERNEL_SIZE, device="cuda") - weight[:, 0, KERNEL_SIZE - 1] = 1.0 - g = torch.Generator(device="cuda").manual_seed(10) - x = torch.randn(16, 8, generator=g, device="cuda") - pool = torch.randn(4, 16, STATE_LEN, generator=g, device="cuda") - qsl = torch.tensor([0, 5, 8], dtype=torch.int32, device="cuda") - cidx = torch.tensor([3, 0], dtype=torch.int32, device="cuda") - his = torch.tensor([True, True], dtype=torch.bool, device="cuda") - out = op(x, weight, None, pool, qsl, cidx, his, DILATION, False, PAD_SLOT_ID) - torch.testing.assert_close(out, x) - - -def test_the_carry_is_actually_read() -> None: - """Dropping the carry moves the first `state_len` outputs of each sequence - and nothing else -- invisible to any tolerance on a prefill batch.""" - args = _small((40, 40), (1, 2), seed=11) - with_carry, _ = _run(args) - without, _ = _run(args[:6] + [torch.zeros_like(args[6])] + args[7:]) - for begin, end in ((0, 40), (40, 80)): - head = with_carry[:, begin : begin + STATE_LEN] - without[:, begin : begin + STATE_LEN] - tail = with_carry[:, begin + STATE_LEN : end] - without[:, begin + STATE_LEN : end] - assert head.abs().max().item() > 0.05 - assert torch.equal(tail, torch.zeros_like(tail)) - - -def test_discrimination_against_the_wrong_variants() -> None: - """Five plausible mis-implementations, each `reference` on altered inputs. - - Measured max |deviation| from the entry (fp32, activation off, these - seeds): correct under 1e-5; taps reversed 4.89; dilation 1 over the same - taps 4.50; carry reversed 4.38; carry dropped 2.80; the batch as one - sequence 4.36 -- the gate at 1.0 sits five orders above the correct - result. The last two touch only nine columns per sequence. - """ - args = _small((11, 13, 7), (2, 0, 1), seed=12, dtype=torch.float32) - args[8] = False - x, weight, _, pool, qsl, cidx, his = args[:7] - out, _ = _run(args) - correct, _ = op.reference(*args) - assert (out - correct).abs().max().item() < 1e-5 - - def variant(**changes): - names = ( - "x", - "weight", - "bias", - "conv_states", - "query_start_loc", - "cache_indices", - "has_initial_state", - "dilation", - "silu_activation", - "pad_slot_id", - ) - kwargs = dict(zip(names, args)) - kwargs.update(changes) - return op.reference(**kwargs)[0] - - wrong = { - "flipped_taps": variant(weight=weight.flip(-1)), - "dilation_1": variant( - conv_states=pool[:, :, -(KERNEL_SIZE - 1) :].contiguous(), dilation=1 - ), - "reversed_state": variant(conv_states=pool.flip(-1)), - "no_carry": variant(has_initial_state=torch.zeros_like(his)), - "one_sequence": variant( - query_start_loc=torch.tensor([0, x.shape[1]], dtype=torch.int32, device="cuda"), - cache_indices=cidx[:1], - has_initial_state=his[:1], - ), - } - for name, value in wrong.items(): - deviation = (out - value).abs().max().item() - print(f"discrimination {name}: {deviation:.6g}") - assert deviation > 1.0, name - - -def test_chunked_prefill_matches_one_shot() -> None: - """The pool is the whole contract between steps: a chunk plus a remainder - is bitwise the one-shot call, including a remainder shorter than the state.""" - tokens = 64 - args = _small((tokens,), (2,), seed=13, bias=True) - x, weight, bias, pool, _, cidx, his = args[:7] - whole = pool.clone() - one_shot = op(x, weight, bias, whole, args[4], cidx, his, DILATION, True, PAD_SLOT_ID) - for split in (1, 2, STATE_LEN, tokens - 1): - chunked = pool.clone() - first = op( - x[:, :split], - weight, - bias, - chunked, - torch.tensor([0, split], dtype=torch.int32, device="cuda"), - cidx, - his, - DILATION, - True, - PAD_SLOT_ID, - ) - second = op( - x[:, split:].contiguous(), - weight, - bias, - chunked, - torch.tensor([0, tokens - split], dtype=torch.int32, device="cuda"), - cidx, - torch.ones_like(his), - DILATION, - True, - PAD_SLOT_ID, - ) - assert torch.equal(torch.cat([first, second], dim=1), one_shot), split - assert torch.equal(chunked, whole), split - - -def test_token_major_view_is_bitwise_the_contiguous_call() -> None: - spec = dict( - channels=256, lengths=(9, 1, 20), slots=(1, 3, 0), pool_slots=4, initial=(True, False, True) - ) - view_out, view_pool = _run(_build(dict(spec, x_form="token_major"), seed=14)) - args = _build(dict(spec, x_form="token_major"), seed=14) - args[0] = args[0].contiguous() - packed_out, packed_pool = _run(args) - assert torch.equal(view_out, packed_out) and torch.equal(view_pool, packed_pool) - - -def test_dilation_one_agrees_with_the_engine() -> None: - """An independent oracle on the ragged bookkeeping: fp32 activations within - the default band (2.4e-7), the published pool bitwise.""" - import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* - - args = _small((23, 1, 64, 1), (3, 1, 0, 2), seed=15, dilation=1, dtype=torch.float32) - x, weight, _, pool, qsl, cidx, his = args[:7] - engine_x, engine_pool = x.clone(), pool.clone() - torch.ops.trtllm.causal_conv1d_fwd( - engine_x, - weight.squeeze(1).contiguous(), - None, - engine_pool, - qsl, - cidx, - his, - True, - PAD_SLOT_ID, - None, - ) - mine, mine_pool = _run(args[:7] + [1, True, PAD_SLOT_ID]) - torch.testing.assert_close(mine, engine_x) - assert torch.equal(mine_pool, engine_pool) - - -def test_silu_off_is_the_bare_convolution() -> None: - args = _small((12, 3), (1, 0), seed=16, bias=True) - raw, _ = _run(args[:8] + [False, PAD_SLOT_ID]) - with_act, _ = _run(args) - assert torch.equal(with_act, F.silu(raw.float()).to(raw.dtype)) - - -def test_rounding_is_two_roundings() -> None: - """The conv rounds on its store and the SiLU rounds again; a single fused - rounding disagrees on a measurable fraction (0.2496 at these values).""" - args = _small((64, 64), (0, 1), seed=23) - out, _ = _run(args) - two, _ = op.reference(*args) - fp32 = [a.float() if i in (0, 1, 3) else a for i, a in enumerate(args)] - fp32[8] = False - acc, _ = op.reference(*fp32) - one = F.silu(acc).to(out.dtype) - assert torch.equal(out, two) - mismatches = (one != two).float().mean().item() - print(f"single-rounding disagrees on {mismatches:.4f} of the elements") - assert mismatches > 0.0 - - -def test_repeat_calls_are_bitwise_stable() -> None: - args = _small((33, 1), (2, 3), seed=24) - first = _run(args) - for _ in range(3): - again = _run(args) - assert torch.equal(first[0], again[0]) and torch.equal(first[1], again[1]) - - def test_cuda_graph_capture_and_replay() -> None: - """Captured once, replayed after the batch tensors' values change in place.""" - args = _small((4, 4, 4, 4), (0, 1, 2, 3), seed=25, pool_slots=6) + """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() @@ -351,9 +119,10 @@ def test_cuda_graph_capture_and_replay() -> None: captured = op(x, weight, None, live, qsl, cidx, his, DILATION, True, PAD_SLOT_ID) g = torch.Generator(device="cuda").manual_seed(26) - qsl.copy_(torch.tensor([0, 1, 2, 3, 16], dtype=torch.int32, device="cuda")) - cidx.copy_(torch.tensor([4, 3, 2, 1], dtype=torch.int32, device="cuda")) - his.copy_(torch.tensor([True, False, True, False], dtype=torch.bool, device="cuda")) + # 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() @@ -366,28 +135,13 @@ def test_cuda_graph_capture_and_replay() -> None: assert torch.equal(live, ref_pool) -@pytest.mark.parametrize( - "case,fragment", - [ - ("pool_dtype", "index_copy_(): self and source expected to have the same dtype"), - ("weight_2d", "weight should have at least three dimensions"), - ("weight_dtype", "expected scalar type BFloat16 but found Float"), - ("bias_dtype", "Input type (c10::BFloat16) and bias type (float) should be the same"), - ], -) -def test_mismatched_operands_raise(case: str, fragment: str) -> None: +def test_mismatched_operands_raise() -> None: """Refused loudly by the torch calls the entry is composed from -- why it - has no `is_valid`.""" + has no `is_valid`. Probed at a pool of another dtype (the index copy).""" args = _small((4, 4), None, seed=27) - if case == "pool_dtype": - args[3] = args[3].float() - elif case == "weight_2d": - args[1] = args[1].squeeze(1).contiguous() - elif case == "weight_dtype": - args[1] = args[1].float() - elif case == "bias_dtype": - args[2] = torch.zeros(args[0].shape[0], device="cuda") + 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 index 233386dcf207..d0e4fff2e5ad 100644 --- 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 @@ -30,12 +30,7 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: def _src(num_rows: int, num_cols: int, dtype: torch.dtype, seed: int) -> torch.Tensor: g = torch.Generator(device="cuda").manual_seed(seed) - base = torch.randn(num_rows, num_cols, generator=g, device="cuda") - if dtype in (torch.int32, torch.int8): - return (base * 10).to(dtype) - if dtype is torch.bool: - return base > 0 - return base.to(dtype) + return torch.randn(num_rows, num_cols, generator=g, device="cuda").to(dtype) def _build(spec, seed): @@ -60,38 +55,11 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res _compare(out, op.reference(src, n, start_col, width)) -def test_the_callable_refuses_a_non_unit_column_stride() -> None: - """Loud on its own, which is why `is_valid` does not restate it.""" - src = torch.randn(256, 256, dtype=torch.bfloat16, device="cuda") - with pytest.raises(AssertionError): - op(src.t(), 128, 0, 128) - with pytest.raises(AssertionError): - op(src[:, ::2], 128, 0, 64) - - -@pytest.mark.parametrize( - "n,start_col,width,shape", - [ - (128, 0, 256, (64, 256)), # 128 rows out of 64 - (64, 200, 128, (64, 256)), # columns [200, 328) out of 256 - (64, -1, 64, (64, 256)), - ], - ids=["rows_past_src", "cols_past_src", "negative_start"], -) -def test_guard_refuses_a_range_outside_src(n, start_col, width, shape) -> None: +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.""" - src = torch.randn(*shape, dtype=torch.bfloat16, device="cuda") + `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, n, start_col, width) - - -def test_guard_refuses_a_3d_src_the_callable_would_accept() -> None: - """`[rows, cols, 1]` has `stride(1) == 1`, passes the callable's assert, and - is read as a 2-D tensor through its row stride.""" - src = torch.randn(64, 256, 1, dtype=torch.bfloat16, device="cuda") - assert src.stride(1) == 1 - with pytest.raises(AssertionError, match="must be 2D"): - with validating(op): - op(src, 64, 0, 1) + 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 new file mode 100644 index 000000000000..96f88e6dbdaa --- /dev/null +++ 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 @@ -0,0 +1,108 @@ +# 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 + ) + 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 + ) + return dict( + 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" 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 index 38a6cf3f54d1..4022900963ef 100644 --- 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 @@ -29,7 +29,8 @@ assert torch.cuda.is_available(), "fused_gdn_gating requires a CUDA device" -#: The smallest softplus `beta` a cell drives; the tail error scales as 1 / beta. +#: 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 @@ -38,9 +39,9 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: 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 the certified extremes - (exp(A_log) <= 139.42, beta >= 0.5): atol = 4 * 2^-22 * 139.42 / 0.5 = - 2.66e-04. + `_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, @@ -55,9 +56,6 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: DEV = "cuda" SOFTPLUS_THRESHOLD = 20.0 -#: log(1 + exp(x)) in fp32 is exactly 0 once exp(x) <= 2**-24. Measured: the -#: largest x with g == 0 is -16.635532, the smallest with g != 0 -16.635530. -ZERO_BELOW = -24.0 * math.log(2.0) def _params(g, num_heads=NUM_V_HEADS, dtype=torch.float32): @@ -71,32 +69,21 @@ def _params(g, num_heads=NUM_V_HEADS, dtype=torch.float32): def _build(spec, seed): """`__call__`'s positional arguments for one spec.""" g = torch.Generator(device=DEV).manual_seed(seed) - rows, heads = spec["rows"], spec.get("num_heads", NUM_V_HEADS) - a_dtype = spec.get("a_dtype", torch.bfloat16) - layout = spec.get("layout", "contiguous") + rows, heads = spec["rows"], NUM_V_HEADS def randn(*shape): - return (torch.randn(*shape, generator=g, device=DEV) * spec.get("a_scale", 1.0)).to(a_dtype) + return torch.randn(*shape, generator=g, device=DEV).to(torch.bfloat16) if "fill" in spec: - a = torch.full((rows, heads), spec["fill"], dtype=a_dtype, device=DEV) - elif layout == "target_rows": + 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 - elif layout == "ba_column_half": - a = randn(rows, BA_WIDTH)[:, NUM_V_HEADS:] - assert a.stride() == (BA_WIDTH, 1) - elif layout == "ba_row_and_column": - a = randn(rows + 48, BA_WIDTH)[48:, NUM_V_HEADS:] - assert a.stride() == (BA_WIDTH, 1) else: a = randn(rows, heads) - A_log, dt_bias = _params(g, heads, spec.get("param_dtype", torch.float32)) - if "overlong" in spec: - pad = torch.full((spec["overlong"],), 9.0, dtype=A_log.dtype, device=DEV) - A_log, dt_bias = torch.cat([A_log, pad]), torch.cat([dt_bias, pad]) - return A_log, a, dt_bias, spec.get("beta", 1.0), spec.get("threshold", SOFTPLUS_THRESHOLD) + 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]) @@ -106,149 +93,23 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res g = op(*args) _compare(g, op.reference(*args)) assert g.is_contiguous() and torch.isfinite(g).all() - A_log, a, dt_bias = args[:3] - if not a.is_contiguous(): - # The row stride is consumed in place, not approximated. - assert torch.equal(g, op(A_log, a.contiguous(), dt_bias, *args[3:])) - if "overlong" in cell.spec: - heads = a.shape[1] - trimmed = op(A_log[:heads].clone(), a, dt_bias[:heads].clone(), *args[3:]) - assert torch.equal(g, trimmed), "elements past a.shape[1] must not be read" - - -def test_threshold_branch_is_live_for_this_checkpoint() -> None: - """dt_bias reaching 19.25 puts part of an ordinary batch past the default - threshold, so both softplus branches run in one launch.""" - args = _build(dict(rows=4096), seed=106) - A_log, a, dt_bias = args[:3] - over = ((a.double() + dt_bias.double()) > SOFTPLUS_THRESHOLD).double().mean().item() - assert 0.0 < over < 1.0, f"expected a mixed batch, got {over} of elements over the threshold" - _compare(op(*args), op.reference(*args)) - - -def test_unstabilised_softplus_tail() -> None: - """The kernel's `log(1 + exp(x))` is exactly 0 below -24 ln 2; the absolute - error that costs stays inside `_compare`, the relative error reaches 1.""" - heads = 8 - A_log = torch.zeros(heads, device=DEV) # exp(A_log) == 1, so g == -softplus - dt_bias = torch.zeros(heads, device=DEV) - x = torch.linspace(-40.0, 5.0, heads * 8192, device=DEV).view(-1, heads) - g = op(A_log, x, dt_bias) - ref = op.reference(A_log, x, dt_bias) - _compare(g, ref) - - assert (g[x <= ZERO_BELOW - 1e-3] == 0).all() - assert (g[x >= ZERO_BELOW + 1e-3] != 0).all() - exact = -torch.log1p(torch.exp(x.double())) - rel = (g.double() - exact).abs() / exact.abs().clamp_min(1e-300) - assert float(rel[x <= ZERO_BELOW].max()) == 1.0 - assert float(rel[x >= -9.0].max()) < 1e-3 - - -def test_determinism() -> None: - args = _build(dict(rows=256), seed=105) - assert torch.equal(op(*args), op(*args)) - - -def test_gate_discriminates_wrong_variants() -> None: - """Five plausible mis-implementations, each rejected by `_compare` -- and - by a margin: each lands at least 1e3 times the allowance away.""" - args = _build(dict(rows=256), seed=108) - A_log, a, dt_bias = args[:3] - g = op(*args) - ref = op.reference(*args) - _compare(g, ref) - - e = torch.exp(A_log.double()) - x = a.double() + dt_bias.double() - sp = torch.where(x <= SOFTPLUS_THRESHOLD, torch.log1p(torch.exp(x)), x) - sp_no_bias = torch.where( - a.double() <= SOFTPLUS_THRESHOLD, torch.log1p(torch.exp(a.double())), a.double() - ) - wrong = { - "dt_bias dropped": -e * sp_no_bias, - "leading minus dropped": e * sp, - "A_log without exp": -A_log.double() * sp, - "exp(-A_log)": -torch.exp(-A_log.double()) * sp, - "sigmoid for softplus": -e * torch.sigmoid(x), - } - allowance = 4.0 * APPROX_ULP * (math.exp(A_LOG_RANGE[1]) / 0.5 + ref.double().abs()) - correct = float(((g.double() - ref.double()).abs() / allowance).max()) - assert correct < 1.0, correct - for name, variant in wrong.items(): - with pytest.raises(AssertionError): - _compare(g, variant.float()) - ratio = float(((variant - ref.double()).abs() / allowance).max()) - print(f"{name}: {ratio:.3e}x the allowance (correct {correct:.3e})") - assert ratio > 1e3, f"{name} lands only {ratio}x outside the gate" - - -@pytest.mark.parametrize( - "shape,message", - [ - ((2, 4, NUM_V_HEADS), "too many values to unpack"), - ((NUM_V_HEADS,), "not enough values to unpack"), - ], -) -def test_op_rejects_a_non_2d_a(shape, message) -> None: - """Loud, from the op's own `batch, num_heads = a.shape` -- which is why the - guard reads `shape[-1]` / `stride(-1)` and does not preempt it.""" - A_log, dt_bias = _params(torch.Generator(device=DEV).manual_seed(109)) - a = torch.randn(*shape, device=DEV).to(torch.bfloat16) - with pytest.raises(ValueError, match=message): - with validating(op): - op(A_log, a, dt_bias) -def test_op_rejects_cpu_tensors() -> None: - A_log, dt_bias = _params(torch.Generator(device=DEV).manual_seed(110)) - a = torch.randn(4, NUM_V_HEADS, device=DEV).to(torch.bfloat16) - with pytest.raises(ValueError, match="Pointer argument cannot be accessed from Triton"): - op(A_log.cpu(), a.cpu(), dt_bias.cpu()) - - -def _guard_case(case): - """`(args, hazard)`: the violating call and what it silently computes instead.""" +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) - a = torch.randn(32, NUM_V_HEADS, generator=g, device=DEV).to(torch.bfloat16) - if case == "a_head_stride_2": - wide = torch.randn(32, BA_WIDTH, generator=g, device=DEV).to(torch.bfloat16) - bad = (A_log, wide[:, ::2], dt_bias) - # Read as `row * stride(0) + head`: the first 48 columns of each row. - return bad, (A_log, wide[:, :NUM_V_HEADS], dt_bias) - wide = torch.randn(2 * NUM_V_HEADS, generator=g, device=DEV) - if case == "A_log_stride_2": - return (wide[::2], a, dt_bias), (wide[:NUM_V_HEADS], a, dt_bias) - if case == "dt_bias_stride_2": - return (A_log, a, wide[::2]), (A_log, a, wide[:NUM_V_HEADS]) - # Short by 8: a view into a longer buffer, so the over-read lands in known memory. - if case == "A_log_short": - return (wide[:40], a, dt_bias), (wide[:NUM_V_HEADS], a, dt_bias) - return (A_log, a, wide[:40]), (A_log, a, wide[:NUM_V_HEADS]) - - -@pytest.mark.parametrize( - "case,match", - [ - ("a_head_stride_2", "head axis must be unit-stride"), - ("A_log_stride_2", "must be contiguous 1-D"), - ("dt_bias_stride_2", "must be contiguous 1-D"), - ("A_log_short", "at least a.shape"), - ("dt_bias_short", "at least a.shape"), - ], -) -def test_guard_refuses_what_the_kernel_reads_at_the_wrong_addresses(case, match) -> None: - """Unguarded, each call returns a finite `g` -- exactly the one for the - memory the kernel actually addressed, not the tensor it was handed; the - guard refuses it.""" - bad, addressed = _guard_case(case) + 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" - if "short" not in case: - with pytest.raises(AssertionError): - _compare(unguarded, op.reference(*bad)) - with pytest.raises(AssertionError, match=match): + 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 index 5cdbfe632874..01b35ad0d378 100644 --- 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 @@ -75,147 +75,69 @@ def _compare(out: tuple, ref: tuple) -> None: DTYPE = torch.bfloat16 HEADS = (NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM) L2_NORM_EPS = 1e-6 -#: The stock mixer's `z` half of `in_proj_qkvz`, and its `[b|a]` projection width. +#: The `z` columns of the input projection. Z_WIDTH = NUM_V_HEADS * HEAD_V_DIM -BA_WIDTH = 2 * NUM_V_HEADS +#: 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, then its keyword `beta_dtype`.""" + """`__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 - layout = spec.get("layout", "contiguous") def randn(*shape): return torch.randn(*shape, generator=g, device=DEV).to(DTYPE) - if layout == "production": - mixed = randn(n, QKV_DIM + Z_WIDTH)[:, :QKV_DIM] - prefill, decode = mixed[:n_p].t().contiguous(), mixed[n_p:] - ba = randn(n, BA_WIDTH) - b, a = ba[:, :NUM_V_HEADS], ba[:, NUM_V_HEADS:] - assert not decode.is_contiguous() and not a.is_contiguous() - elif layout == "target_mixed": - # `_gdn_block`: mixed_qkv = x @ w_qkv is [T, 10240]; the decode half is - # its row split, the prefill half the transposing extraction of the rest. - mixed = randn(n, QKV_DIM) - prefill, decode = mixed[:n_p].t().contiguous(), mixed[n_p:] - assert decode.stride() == (QKV_DIM, 1) and decode.storage_offset() == n_p * QKV_DIM - a, b = randn(n, NUM_V_HEADS), randn(n, NUM_V_HEADS) - elif layout == "prefill_view": - prefill = randn(n_p, QKV_DIM + Z_WIDTH)[:, :QKV_DIM].t() - assert prefill.stride() == (1, QKV_DIM + Z_WIDTH) - decode = None - a, b = randn(n, NUM_V_HEADS), randn(n, NUM_V_HEADS) - else: - prefill = randn(QKV_DIM, n_p) - decode = randn(n_d, QKV_DIM) if n_d else None - a, b = randn(n, NUM_V_HEADS), randn(n, NUM_V_HEADS) + # `_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) - gate_dtype = spec.get("gate_dtype", torch.float32) lo, hi = A_LOG_RANGE - A_log = (torch.rand(NUM_V_HEADS, generator=g, device=DEV) * (hi - lo) + lo).to(gate_dtype) + 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).to(gate_dtype) + 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).to(gate_dtype) - args = [prefill, decode, a, b, A_log, dt_bias, *HEADS, L2_NORM_EPS, SOFTPLUS_THRESHOLD] - return args, spec.get("beta_dtype", torch.float32) + 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, beta_dtype = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + 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, beta_dtype=beta_dtype) + 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, beta_dtype=beta_dtype)) - - -def test_g_gate_discriminates_wrong_variants() -> None: - """`g`'s floor is derived, so prove it still separates. Each variant is a - plausible way to get the gate wrong, and must land >1e4x outside it.""" - args, _ = _build(dict(num_prefill=96, num_decode=0), seed=106) - prefill, decode, a, b, A_log, dt_bias = args[:6] - out = op(*args) - ref = op.reference(*args) - _compare(out, ref) - g_ref = ref[3].double() - atol = 4 * 2.0**-22 * SOFTPLUS_THRESHOLD * math.exp(A_LOG_RANGE[1]) - observed = (out[3].double() - g_ref).abs().max().item() - assert observed < atol, f"correct result {observed} outside its own gate {atol}" - - x = a.double() + dt_bias.double() - softplus = torch.where(x <= SOFTPLUS_THRESHOLD, torch.log1p(torch.exp(x)), x) - - def g_of(*overrides): - return op.reference(*[*overrides, *args[len(overrides) :]])[3].double() - - variants = { - "sign flipped": -g_ref, - "dt_bias omitted": g_of(prefill, decode, a, b, A_log, torch.zeros_like(dt_bias)), - "A_log not exponentiated": (-A_log.double() * softplus).unsqueeze(0), - "softplus replaced by identity": (-torch.exp(A_log.double()) * x).unsqueeze(0), - "a and b swapped": g_of(prefill, decode, b, a), - } - for name, variant in variants.items(): - distance = (variant - g_ref).abs().max().item() - assert distance > 1e4 * atol, f"{name} lands {distance / atol:.1f}x the gate" - with pytest.raises(AssertionError): - _compare((*out[:3], variant.float(), out[4]), ref) - - -def test_qk_normalization_is_not_a_pass_through() -> None: - """The default bf16 band on q/k separates a missing norm and a q/k swap.""" - args, _ = _build(dict(num_prefill=96, num_decode=0), seed=107) - out = op(*args) - ref = op.reference(*args) - _compare(out, ref) - q_raw = args[0].t()[:, : NUM_K_HEADS * HEAD_K_DIM].reshape(1, -1, NUM_K_HEADS, HEAD_K_DIM) - for wrong in ((q_raw.contiguous(), *out[1:]), (out[1], *out[1:])): - with pytest.raises(AssertionError): - _compare(wrong, ref) + _compare(out, op.reference(*args)) def test_guard_refuses_strided_gate_parameters() -> None: - """Unguarded, a stride-2 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) - for position in (4, 5): - strided = torch.stack([args[position], torch.zeros_like(args[position])], dim=1)[:, 0] - assert strided.stride() == (2,) - bad = list(args) - bad[position] = 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) - - -def test_guard_refuses_short_gate_parameters_and_a_b() -> None: - """Each of these is an unmasked out-of-bounds read, so the hazard is not - replayed here; the contract's measurements are in the entry's note.""" - args, _ = _build(dict(num_prefill=32, num_decode=8), seed=109) - for position, match in ((4, "at least num_v_heads"), (5, "at least num_v_heads")): - bad = list(args) - bad[position] = args[position][: NUM_V_HEADS - 8].clone() - with pytest.raises(AssertionError, match=match): - with validating(op): - op(*bad) - for position in (2, 3): - bad = list(args) - # 32 rows cover the prefill block but not prefill + decode = 40. - bad[position] = args[position][:32].clone() - with pytest.raises(AssertionError, match="at least num_tokens=40"): - with validating(op): - op(*bad) + """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_cached_replay_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_cached_replay_update.py new file mode 100644 index 000000000000..99a60e6e074d --- /dev/null +++ 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_cached_replay_update.py @@ -0,0 +1,547 @@ +# 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_cached_replay_update entry. + +The op is stateful across verify steps: what one call leaves in the checkpoint +pool and the history double buffer is what the next call reads, after the +engine's own bookkeeping ran in between. So every cell drives several +consecutive verify steps the way the engine does -- work items before the +forward (partitioned batches), the op once per layer, then the all-layer +commit and `_advance_replay_state` with a random partial acceptance -- and +checks two things at every step: + +* each call against the entry's `reference` on that call's own inputs (the + output and all four state buffers), plus bitwise locks on everything the + call must not touch; +* every output against an *exact twin*: the entry's `reference` driven on the + same checkpoint and history held in fp32 and deep enough never to wrap, + advanced by `p += accepted` alone. The twin's state is the plain + 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. +""" + +from __future__ import annotations + +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.ssm.fused_recurrent_gated_delta_rule_cached_replay_update import ( # noqa: E501 + ALL_LAYER_COMMIT_MIN_BATCH, + HEAD_K_DIM, + HEAD_V_DIM, + HISTORY_SIZE, + NUM_K_HEADS, + 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 + fused_recurrent_gated_delta_rule_cached_replay_update as op, +) +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_triton +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + ReplayStateUpdateMetadata, + _advance_replay_state, +) + +__extra_import_path__ = ["../../../.."] +from _validating import ULP, certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), ( + "fused_recurrent_gated_delta_rule_cached_replay_update requires a CUDA device" +) + +DEV = "cuda" +H, HV, K, V = NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM +W = PACKED_QKV_WIDTH +HIST = HISTORY_SIZE +BF16_ULP = ULP[torch.bfloat16] + +#: This checkpoint's 36 GDN layers of `linear_attn.A_log` / `linear_attn.dt_bias` +#: (bf16, 48 each): mean, standard deviation, range. +A_LOG_STATS = (1.006, 1.487, (-7.125, 5.0625)) +DT_BIAS_STATS = (-4.411, 1.732, (-11.4375, 2.53125)) + +#: Element gate of the bf16 members, in bf16 ulps of the block they belong to. +_ELEMENT_ULPS = 4.0 +#: Relative-RMS gate of the output, in bf16 ulps. +_OUTPUT_RMS_ULPS = 1.5 +#: Element gate of an fp32 checkpoint state, relative to its (slot, head) scale. +_STATE_TOL = 2.0**-13 +#: Element gate of a bf16 checkpoint state, in bf16 ulps of its (slot, head) scale. +_BF16_STATE_ULPS = 2.0 +#: Gate of the stored cumulative log-decays: |dG| <= atol + rtol * |G|. +_G_ATOL, _G_RTOL = 2.0**-11, 2.0**-18 + + +def _gate(name, got, want, element, rms=None): + """`got` against `want` per block of the last two dims (a token's + `[heads, dim]`, or one head's `[V, K]` state): every element within + `element` of its block's absmax, NaN blocks of `want` excluded.""" + 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] + 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 + 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() + assert rel <= rms, f"{name}: relative RMS {rel:.3e} exceeds {rms:.3e}" + + +def _compare(out: tuple, ref: tuple) -> None: + """`(returned, ssm_states, old_u, old_k, old_G)` as the call left them, + against `reference`'s five; a member `None` on both sides is skipped (the + exact twin is compared on the output alone). + + The kernel is not fp32 throughout, and the bands say where. It rounds the + L2-normalized q and k to bf16 before every product (2^-9 relative each), + rounds the history-mixing coefficients to bf16 for its tensor-core dots, and + stores the output and the history's u and k in bf16. So the output, `old_u` + and `old_k` are gated per block -- one token's `[heads, dim]` -- at 4 bf16 + ulps of the block's absmax, which keeps a near-silent head from turning + rounding into relative noise, and the output's relative RMS at 1.5 ulps. An fp32 + checkpoint state is fp32 by construction: the commit splits the decayed + history update into a bf16 hi and lo part, ~16 significant bits per term, and + up to 16 terms each within 2^-17 give the 2^-13 gate, relative to one (slot, + head) state's absmax. A bf16 checkpoint is that fp32 commit rounded once, as the + reference's is, so the two stored values are equal or adjacent, and adjacent bf16 + values are at most 2 bf16 ulps of their head's absmax apart: the bf16 gate is one + representable step. `G` only ever enters as exp(G_t - G_j), so an absolute + error is a relative decay error: the kernel's softplus runs on approximate + transcendentals whose fp32 sum drops exp(x)'s low bits; atol 2^-11 is an + eighth of a bf16 ulp of the u it scales, and rtol 2^-18 covers the cumulative + sum's rounding where |G| reaches thousands on fast-forgetting heads. + """ + names = ("output", "ssm_states", "old_u", "old_k", "old_G") + assert len(out) == len(ref) == len(names) + for name, got, want in zip(names, out, ref): + assert (got is None) == (want is None), f"{name}: present on one side only" + if got is None: + continue + assert got.dtype == want.dtype and got.shape == want.shape, ( + name, + (got.dtype, tuple(got.shape)), + (want.dtype, tuple(want.shape)), + ) + if name == "old_G": + torch.testing.assert_close(got, want, rtol=_G_RTOL, atol=_G_ATOL, msg=name) + elif name == "ssm_states": + bf16 = got.dtype == torch.bfloat16 + _gate(name, got, want, _BF16_STATE_ULPS * BF16_ULP if bf16 else _STATE_TOL) + else: + rms = _OUTPUT_RMS_ULPS * BF16_ULP if name == "output" else None + _gate(name, got, want, _ELEMENT_ULPS * BF16_ULP, rms) + + +def _seed(why): + return zlib.crc32(why.encode()) + + +def _gating(gen, layers): + """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, 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] + return a_log.contiguous(), dt_bias.contiguous() + + +def _unit(x): + return (x / x.norm(dim=-1, keepdim=True)).to(torch.bfloat16) + + +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 + it: work items before the forward (`_build_replay_work_items_triton`, what + mamba metadata's prepare runs for a partitioned batch), the op per layer, + then `update_mamba_states`'s order -- the all-layer commit, then + `_advance_replay_state`. + + With `exact=True` it also carries the exact twin (see the module docstring). + """ + + def __init__(self, spec, seed, exact=True): + self.spec = spec + self.path, self.n, self.t = spec["path"], spec["n"], spec["t"] + 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 + + if 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: + chosen = torch.randperm(S, generator=self.cpu)[:n] + chosen = chosen.to(torch.int32) + live_rows = torch.ones(n, dtype=torch.bool) + dup = tuple(spec.get("dup", ())) + if dup: + chosen[list(dup)] = int(chosen[dup[0]]) + # The padding rows' shared slot is nobody's state: the bookkeeping never names it. + live_rows[list(dup)] = False + 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 + 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) + 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) + g = -torch.exp(self.a_log)[:, None, None, :, None] * F.softplus( + raw_a + self.dt_bias[:, None, None, :, None] + ) + self.old_dt.copy_(g.cumsum(-1)) + pnat = torch.randint(0, HIST + 1, (S,), generator=self.cpu, dtype=torch.int32) + live = self.live.long().cpu() + pnat[live[0]] = HIST # a wrap at step 0 + if live.numel() > 1: + pnat[live[1]] = 0 # an append at step 0 + self.pnat.copy_(pnat) + self.buf.copy_(torch.randint(0, 2, (S,), generator=self.cpu, dtype=torch.int32)) + self.meta = ReplayStateUpdateMetadata(self.pnat, self.buf, t, HIST) + self.work_items = self.n_writes = None + if self.path == "all_layer": + self.work_items = torch.zeros(n, 4, dtype=torch.int32, device=DEV) + self.n_writes = torch.zeros(1, dtype=torch.int32, device=DEV) + + 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) + 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] + + def _randn(self, *shape, dtype=torch.float32): + return torch.randn(*shape, generator=self.gen, device=DEV, dtype=dtype) + + 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) + return dict(q=q, k=k, v=v, a=a, b=b, packed=packed) + + def kwargs(self, layer, act): + return dict( + q=act["q"], + k=act["k"], + v=act["v"], + g=act["a"], + beta=act["b"], + ssm_states=self.pool[layer], + state_indices=self.state_indices, + old_u=self.old_x[layer], + old_k=self.old_B[layer], + old_G=self.old_dt[layer], + old_beta=self.old_dA[layer], + cache_buf_idx=self.buf, + prev_num_accepted_tokens=self.pnat, + history_size=HIST, + scale=None, + use_qk_l2norm_in_kernel=True, + A_log=self.a_log[layer], + dt_bias=self.dt_bias[layer], + launch_with_pdl=False, + replay_work_items=self.work_items, + n_writes=self.n_writes, + packed_qkv=act["packed"], + use_all_layer_commit=self.path == "all_layer", + output=None, + ) + + def exact_kwargs(self, layer, kw): + return dict( + kw, + ssm_states=self.x_pool[layer], + old_u=self.x_u[layer], + old_k=self.x_k[layer], + old_G=self.x_G[layer], + cache_buf_idx=self.x_buf, + prev_num_accepted_tokens=self.x_pnat, + history_size=self.depth, + replay_work_items=None, + n_writes=None, + use_all_layer_commit=False, + output=None, + ) + + def adopt_exact(self, layer, ref): + assert torch.equal(ref[1], self.x_pool[layer]), "the twin must never wrap" + self.x_u[layer].copy_(ref[2]) + self.x_k[layer].copy_(ref[3]) + self.x_G[layer].copy_(ref[4]) + + def state(self, layer): + return self.pool[layer], self.old_x[layer], self.old_B[layer], self.old_dt[layer] + + def prepare(self): + if self.path == "all_layer": + self.work_items.zero_() + _build_replay_work_items_triton( + self.live, self.pnat, self.buf, self.work_items, self.n_writes, self.t, HIST + ) + + def commit(self): + if self.path == "all_layer": + 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), + old_u=self.old_x, + old_k=self.old_B, + old_G=self.old_dt, + replay_work_items=self.work_items, + n_writes=self.n_writes, + history_size=HIST, + ) + + def draw_acceptance(self): + """Accepted tokens per live row, bonus token included, by `spec["accept"]`.""" + mode = self.spec.get("accept", "random") + rows = self.live.numel() + if mode == "full": + return torch.full((rows,), self.t, dtype=torch.int32, device=DEV) + assert mode == "random", mode + return torch.randint(1, self.t + 1, (rows,), generator=self.cpu, dtype=torch.int32).to(DEV) + + def accept(self, accepted): + _advance_replay_state(self.meta, self.live, accepted) + if self.exact: + self.x_pnat[self.live.long()] += accepted + + def writes(self): + """Where this step's calls append, by the documented rule, which slots wrap, and how + many append exactly up to the history's end: `[slots, 2, HIST]` mask, wrapping slot + ids, the `pnat + t == HIST` count.""" + s = self.live.long() + p, b = self.pnat[s].long(), self.buf[s].long() + wrap = p + self.t > HIST + half, start = torch.where(wrap, 1 - b, b), torch.where(wrap, 0, p) + mask = torch.zeros(self.slots, 2, HIST, dtype=torch.bool, device=DEV) + for t in range(self.t): + mask[s, half, start + t] = True + return mask, s[wrap], int((p + self.t == HIST).sum()) + + +def _check_duplicate_padding_slot(cell) -> None: + """The last rows of a batch share 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 those rows alone -- the padding rows change nothing about them. + The shared slot's state and the padding rows' outputs are unspecified.""" + bench = _Bench(cell.spec, seed=_seed(cell.why), exact=False) + dup = cell.spec["dup"] + unique = bench.n - len(dup) + assert dup == tuple(range(unique, bench.n)), "the padding rows trail the batch" + shared = int(bench.state_indices[unique]) + assert (bench.state_indices[unique:] == shared).all() + assert shared not in bench.live.tolist() + 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) + 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) + _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[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + if "dup" in cell.spec: + return _check_duplicate_padding_slot(cell) + bench = _Bench(cell.spec, seed=_seed(cell.why)) + wraps = appends = to_the_end = 0 + for _ in range(bench.steps): + mask, wrapped, at_end = bench.writes() + wraps += wrapped.numel() + to_the_end += at_end + appends += bench.live.numel() - wrapped.numel() + bench.prepare() + after_commit = [] + for layer in range(bench.layers): + act = bench.inputs() + kw = bench.kwargs(layer, act) + pool, u, k, G = bench.state(layer) + before = dict( + pool=pool.clone(), + u=u.clone(), + k=k.clone(), + G=G.clone(), + old_beta=bench.old_dA[layer].clone(), + ) + ref = op.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) + ) + got = op(**kw) + _compare((got, *bench.state(layer)), ref) + # The locks: in place, and nothing outside the documented writes. + pool, 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]), ( + "the pool moved where the call commits nothing (a partitioned call commits none)" + ) + 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. + _compare((got, *bench.state(layer)), committed) + _compare((got, None, None, None, None), (exact[0], None, None, None, None)) + bench.accept(bench.draw_acceptance()) + assert wraps > 0 and appends > 0, ( + f"the cell never wrapped or never appended ({wraps}, {appends})" + ) + if cell.spec.get("accept") == "full": + live = bench.live.numel() + assert wraps >= 3 * live and to_the_end >= 3 * live, ( + f"full acceptance should wrap every row at least 3 times and append up to the " + f"history's end as often ({wraps}, {to_the_end} over {live} rows)" + ) + + +# ── CUDA graphs ─────────────────────────────────────────────────────────── + + +def test_cuda_graph_replay_on_new_data_is_bitwise_eager() -> None: + """Two layers of 32 sequences x 8 tokens over an fp32 checkpoint captured with the + all-layer commit after them; replayed three times on new activations, with acceptance + in between, against the same steps run eagerly on a twin state.""" + n, t = 32, 8 + assert ALL_LAYER_COMMIT_MIN_BATCH == CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE + assert n >= ALL_LAYER_COMMIT_MIN_BATCH + spec = dict(path="all_layer", n=n, t=t, layers=2, slots=40, state_dtype=torch.float32) + seed = 7000 + 10 * n + t + eager, graph, warm = (_Bench(spec, seed, exact=False) for _ in range(3)) + acts = [warm.inputs() for _ in range(2)] + warm.prepare() + for layer in range(2): + op(**warm.kwargs(layer, acts[layer])) + warm.commit() + torch.cuda.synchronize() + + graph.prepare() + captured = torch.cuda.CUDAGraph() + with torch.cuda.graph(captured): + outs_graph = [op(**graph.kwargs(layer, acts[layer])) for layer in range(2)] + graph.commit() + assert torch.equal(graph.backing, eager.backing) and torch.equal(graph.old_x, eager.old_x) + + refill = torch.Generator(device=DEV).manual_seed(seed + 1) + 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=DEV, dtype=act[key].dtype) + ) + eager.prepare() + outs_eager = [op(**eager.kwargs(layer, acts[layer])) for layer in range(2)] + eager.commit() + graph.prepare() + 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) + + +# ── 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/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 index 980c4be80374..3b2041b63919 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_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 @@ -3,16 +3,14 @@ """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 sweep `P`, and the two state sinks -(write-back into the pool, per-step states into the rollback buffer) are -driven separately and together. +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. -Two tests answer "is this harness blind" rather than a correctness question: -the neighbouring `chunk_gated_delta_rule` is wrong on the first indexed +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, and `test_the_gate_rejects_a_doubly_applied_state` -shows `_compare` rejects exactly the damage that mechanism leaves. +process that has made none. """ from __future__ import annotations @@ -26,7 +24,6 @@ import pytest import torch -import tensorrt_llm._torch.modules.fla.fused_recurrent as _impl 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, @@ -74,67 +71,49 @@ def _compare(out: tuple, ref: tuple) -> None: DEV = "cuda" -REPEAT = NUM_V_HEADS // NUM_K_HEADS # 3 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 what the locks need. + """`__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 = spec["mode"], spec["n"] - lengths = spec.get("lengths") - batch, tokens = (n, spec["p"]) if lengths is None else (1, sum(lengths)) - total = batch * tokens + 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(batch, tokens, NUM_K_HEADS, HEAD_K_DIM) - k = mixed[:, kd : 2 * kd].view(batch, tokens, NUM_K_HEADS, HEAD_K_DIM) - v = mixed[:, 2 * kd :].view(batch, tokens, NUM_V_HEADS, HEAD_V_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(batch, tokens, NUM_V_HEADS) + g = g.view(n, tokens, NUM_V_HEADS) b = randn(total, NUM_V_HEADS, dtype=torch.bfloat16) - beta_kind = spec.get("beta", "fp32") - beta = (b.sigmoid() if beta_kind == "bf16" else b.float().sigmoid()).view( - batch, tokens, NUM_V_HEADS - ) - if beta_kind == "channel": - beta = beta[..., None].expand(batch, tokens, NUM_V_HEADS, HEAD_V_DIM).contiguous() + 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 = layered = None - if pool_slots is not None: - pool = randn(pool_slots, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM) * 0.05 - pool = pool.to(spec.get("pool_dtype", torch.float32)) - if spec.get("pool_view") == "layer_major": - layered = torch.zeros((3,) + pool.shape, dtype=pool.dtype, device=DEV) - layered[1].copy_(pool) - pool = layered[1] - assert pool.is_contiguous() + 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=spec.get("index_dtype", torch.int32), device=DEV) - cu_seqlens = None - if lengths is not None: - bounds = torch.tensor((0,) + tuple(lengths)).cumsum(0) - cu_seqlens = bounds.to(device=DEV, dtype=spec.get("cu_dtype", torch.int64)) + indices = torch.tensor(slots, dtype=torch.int32, device=DEV) buffer = cache_steps = None - if mode in ("verify", "both"): - rows = spec.get("cache_rows", SLOTS if pool_slots is None else pool_slots) + 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), @@ -143,17 +122,6 @@ def randn(*shape, dtype=torch.float32): device=DEV, ) - output = out_buf = None - if spec.get("output") == "prefix": - # A length-prefix slice of a longer row-major buffer, as the stock mixer - # passes it: contiguous, and the tail is not offered. - out_buf = torch.full( - (total + 5, NUM_V_HEADS, HEAD_V_DIM), -7.0, dtype=torch.bfloat16, device=DEV - ) - output = out_buf[:total].view(batch, tokens, NUM_V_HEADS, HEAD_V_DIM) - elif spec.get("output") == "fp32": - output = torch.zeros(batch, tokens, NUM_V_HEADS, HEAD_V_DIM, device=DEV) - kwargs = dict( q=q, k=k, @@ -164,16 +132,14 @@ def randn(*shape, dtype=torch.float32): initial_state_indices=indices, scale=None, use_qk_l2norm_in_kernel=True, - cu_seqlens=cu_seqlens, + cu_seqlens=None, disable_state_update=mode == "verify", disable_output_calculation=False, intermediate_states_buffer=buffer, cache_steps=cache_steps, - output=output, - ) - return kwargs, dict( - layered=layered, out_buf=out_buf, total=total, live=[s for s in slots if s >= 0] + output=None, ) + return kwargs, list(slots) def _seed(why): @@ -183,55 +149,23 @@ def _seed(why): @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, ctx = _build(cell.spec, seed=_seed(cell.why)) - pool, buffer, live = kw["initial_state_source"], kw["intermediate_states_buffer"], ctx["live"] - pool_before = None if pool is None else pool.clone() - buffer_before = None if buffer is None else buffer.clone() - layered_before = None if ctx["layered"] is None else ctx["layered"].clone() - tail_before = None if ctx["out_buf"] is None else ctx["out_buf"][ctx["total"] :].clone() + 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. - if kw["output"] is not None: - assert got.data_ptr() == kw["output"].data_ptr(), "output must be written in place" - if tail_before is not None: - assert torch.equal(ctx["out_buf"][ctx["total"] :], tail_before), "buffer tail touched" - if pool is not None: - if kw["disable_state_update"]: - assert torch.equal(pool, pool_before), "the pool must not move on a verify call" - else: - for s in live: - assert not torch.equal(pool[s], pool_before[s]), f"slot {s} was not written" - pool_before[live] = pool[live] - assert torch.equal(pool, pool_before), "a slot outside the indices was written" - if buffer is not None: - 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" - if layered_before is not None: - layered = ctx["layered"] - assert torch.equal(layered[0], layered_before[0]) and torch.equal( - layered[2], layered_before[2] - ), "the neighbouring layers must be untouched" - if cell.spec["mode"] == "both": - # Independent of the reference: the last cached step is the stored state. - assert torch.equal(buffer[live, cell.spec["p"] - 1], pool[live]) - - -# ── the mechanism the neighbouring entry fails on ───────────────────────── - - -def test_the_wrapped_kernel_carries_no_autotuner() -> None: - """Why the cold call is safe here: `triton.heuristics` + `triton.jit`, so - there is no trial launch to consume the in-place state.""" - kernel = _impl.fused_recurrent_gated_delta_rule_update_fwd_kernel - names = [cls.__name__ for cls in type(kernel).__mro__] - assert "Autotuner" not in names, f"the wrapped kernel gained an autotuner ({names})" - assert not hasattr(kernel, "configs"), "an autotuner config list appeared on the wrapped kernel" + 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)) @@ -298,171 +232,7 @@ def test_cold_call_in_a_fresh_process_is_correct() -> None: assert result["out_bitwise_equal"] and result["pool_bitwise_equal"], result -def test_the_gate_rejects_a_doubly_applied_state() -> None: - """The damage an autotuner trial launch leaves -- the state applied twice -- - is rejected by `_compare`, so a cold-call error of that shape cannot pass.""" - kw, _ = _build(_COLD_SPEC, seed=8) - pool = kw["initial_state_source"] - ref = op.reference(**kw) - _compare((op(**kw), pool, None), ref) # one application is right - - doubled = op(**kw) - out_err = (doubled.float() - ref[0].float()).abs().max().item() - pool_err = (pool - ref[1]).abs().max().item() - print(f"\ndoubly applied: output {out_err:.3e}, state {pool_err:.3e}") - assert out_err > 1e-2 and pool_err > 1e-2 - with pytest.raises(AssertionError): - _compare((doubled, pool, None), ref) - - -# ── properties beyond a single reference comparison ─────────────────────── - - -def test_varlen_form_matches_the_batched_form_bitwise() -> None: - """`cu_seqlens` at B = 1 is the same arithmetic as the uniform batch.""" - n, p = 3, 4 - kw, _ = _build(dict(mode="write_back", n=n, p=p, slots=(2, 6, 3)), seed=60) - pool_before = kw["initial_state_source"].clone() - batched = op(**kw) - total = n * p - flat = dict( - kw, - q=kw["q"].reshape(1, total, NUM_K_HEADS, HEAD_K_DIM), - k=kw["k"].reshape(1, total, NUM_K_HEADS, HEAD_K_DIM), - v=kw["v"].reshape(1, total, NUM_V_HEADS, HEAD_V_DIM), - g=kw["g"].reshape(1, total, NUM_V_HEADS), - beta=kw["beta"].reshape(1, total, NUM_V_HEADS), - initial_state_source=pool_before.clone(), - cu_seqlens=torch.arange(0, total + 1, p, device=DEV, dtype=torch.int64), - ) - varlen = op(**flat) - assert torch.equal(varlen.reshape(n, p, NUM_V_HEADS, HEAD_V_DIM), batched) - assert torch.equal(flat["initial_state_source"], kw["initial_state_source"]) - - -def test_scale_default_is_bitwise_head_dim_pow_minus_half() -> None: - kw, _ = _build(dict(mode="write_back", n=2, p=3, slots=(1, 2)), seed=66) - pool_a, pool_b = kw["initial_state_source"].clone(), kw["initial_state_source"].clone() - out_a = op(**dict(kw, initial_state_source=pool_a)) - out_b = op(**dict(kw, initial_state_source=pool_b, scale=HEAD_K_DIM**-0.5)) - assert torch.equal(out_a, out_b) and torch.equal(pool_a, pool_b) - - -def test_deterministic() -> None: - kw, _ = _build(dict(mode="both", n=3, p=4), seed=68) - pool0, buf0 = kw["initial_state_source"].clone(), kw["intermediate_states_buffer"].clone() - runs = [] - for _ in range(2): - pool, buf = pool0.clone(), buf0.clone() - out = op(**dict(kw, initial_state_source=pool, intermediate_states_buffer=buf)) - runs.append((out, pool, buf)) - assert all(torch.equal(x, y) for x, y in zip(*runs)) - - -def test_duplicate_slot_indices_race_on_the_write_back() -> None: - """Two sequences on one slot: the outputs stay right, the slot does not. - Nothing raises, so it is a caller obligation recorded as a measurement.""" - kw, _ = _build(dict(mode="write_back", n=2, p=3, slots=(3, 3)), seed=69) - pool = kw["initial_state_source"] - pool_before = pool.clone() - ref = op.reference(**kw) - got = op(**kw) - torch.testing.assert_close(got, ref[0]) - - def one(i, source): - rows = slice(i, i + 1) - return op.reference( - **dict( - kw, - q=kw["q"][rows], - k=kw["k"][rows], - v=kw["v"][rows], - g=kw["g"][rows], - beta=kw["beta"][rows], - initial_state_source=source, - initial_state_indices=kw["initial_state_indices"][rows], - ) - )[1] - - after_both = one(1, one(0, pool_before)) - drift = (pool[3] - after_both[3]).abs().max().item() - print(f"\nshared slot drift from a sequential application: {drift:.3e}") - assert drift > 1e-2 - - -# ── what the guard refuses, and what the op refuses itself ──────────────── - - -def _small(seed, **spec): - return _build(dict(dict(mode="write_back", n=2, p=3), **spec), seed=seed)[0] - - -def test_a_strided_output_clobbers_its_neighbour_and_is_refused() -> None: - """Unguarded, a last-dim slice of a doubly wide buffer is written with - packed arithmetic: the half never offered is overwritten.""" - kw = _small(70) - wide = torch.zeros(2, 3, NUM_V_HEADS, 2 * HEAD_V_DIM, dtype=torch.bfloat16, device=DEV) - strided = wide[..., :HEAD_V_DIM] - op(**dict(kw, initial_state_source=kw["initial_state_source"].clone(), output=strided)) - assert (wide[..., HEAD_V_DIM:] != 0).any(), "the hazard this guard covers is gone" - with pytest.raises(AssertionError, match="output must be contiguous"): - with validating(op): - op(**dict(kw, output=wide[..., :HEAD_V_DIM])) - - -def test_output_with_the_calculation_off_is_never_written_and_is_refused() -> None: - kw = _build(dict(mode="verify", n=2, p=3), seed=71)[0] - buf = torch.full((2, 3, NUM_V_HEADS, HEAD_V_DIM), -7.0, dtype=torch.bfloat16, device=DEV) - got = op(**dict(kw, disable_output_calculation=True, output=buf)) - assert tuple(got.shape) == (1, 1, 1, 1) and (buf == -7.0).all(), "the hazard is gone" - with pytest.raises(AssertionError, match="mutually exclusive"): - with validating(op): - op(**dict(kw, disable_output_calculation=True, output=buf)) - - -def test_a_strided_pool_loses_the_write_back_and_is_refused() -> None: - """Unguarded, a `[::2]` slot view is replaced by a contiguous copy: the - output is right, and the caller's pool comes back bitwise unchanged.""" - kw = _small(72) - backing = torch.zeros(2 * SLOTS, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM, device=DEV) - strided = backing[::2] - strided.copy_(kw["initial_state_source"]) - before = strided.clone() - ref = op.reference(**dict(kw, initial_state_source=before)) - got = op(**dict(kw, initial_state_source=strided)) - torch.testing.assert_close(got, ref[0]) - assert torch.equal(strided, before), "the hazard this guard covers is gone" - with pytest.raises(AssertionError, match="initial_state_source must be contiguous"): - with validating(op): - op(**dict(kw, initial_state_source=strided)) - - -def test_a_strided_step_buffer_loses_every_cached_state_and_is_refused() -> None: - """The same copy-in trap on the per-step buffer: unguarded, the caller's - buffer comes back bitwise unchanged and nothing raises.""" - kw = _build(dict(mode="verify", n=2, p=3), seed=73)[0] - backing = torch.full( - (2 * SLOTS, 3, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM), CACHE_SENTINEL, device=DEV - ) - strided = backing[::2] - op(**dict(kw, intermediate_states_buffer=strided)) - assert (backing == CACHE_SENTINEL).all(), "the hazard this guard covers is gone" - with pytest.raises(AssertionError, match="intermediate_states_buffer must be contiguous"): - with validating(op): - op(**dict(kw, intermediate_states_buffer=strided)) - - -def test_a_short_cache_steps_writes_into_the_next_row_and_is_refused() -> None: - """Unguarded, cache_steps 2 at P 3: row 0's last step lands in row 1, a - row the call never named.""" - kw = _build(dict(mode="verify", n=1, p=3), seed=74)[0] - short = torch.full((SLOTS, 2, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM), CACHE_SENTINEL, device=DEV) - op(**dict(kw, intermediate_states_buffer=short, cache_steps=2)) - assert not (short[1, 0] == CACHE_SENTINEL).all(), "the hazard this guard covers is gone" - for cache_steps in (2, None): - with pytest.raises(AssertionError, match="cache_steps must be at least"): - with validating(op): - op(**dict(kw, intermediate_states_buffer=short.clone(), cache_steps=cache_steps)) +# ── what the guard refuses ───────────────────────────────────────────── def test_cache_steps_off_the_buffer_step_axis_misplaces_rows_and_is_refused() -> None: @@ -478,103 +248,3 @@ def test_cache_steps_off_the_buffer_step_axis_misplaces_rows_and_is_refused() -> 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)) - - -def test_op_rejects_what_it_really_rejects() -> None: - """Three loud refusals; the guard admits each so the op's own error is what - surfaces.""" - from triton.compiler.errors import CompilationError - - kw = _small(76) - cu = torch.tensor([0, 3, 6], device=DEV, dtype=torch.int64) - with pytest.raises(ValueError, match="batch size is expected to be 1"): - with validating(op): - op(**dict(kw, cu_seqlens=cu)) - - flat = _build(dict(mode="write_back", n=1, p=6), seed=77)[0] - with pytest.raises(ValueError, match="number of initial states"): - with validating(op): - op(**dict(flat, cu_seqlens=cu)) - - with pytest.raises(CompilationError): - with validating(op): - op(**dict(kw, initial_state_source=None)) - - -def test_sequence_count_past_the_grid_cap_fails_loudly() -> None: - """gridDim.z is `N * HV`, capped at 65535: 1365 sequences at HV 48. Above - it the launch is refused, not truncated, and the context survives.""" - over = 1400 - assert over * NUM_V_HEADS > 65535 - kw = _build(dict(mode="write_back", n=over, p=1, slots=(0,)), seed=78)[0] - indices = torch.arange(over, device=DEV, dtype=torch.int32) % SLOTS - with pytest.raises(RuntimeError, match="invalid argument"): - op(**dict(kw, initial_state_indices=indices, disable_state_update=True)) - - kw, _ = _build(dict(mode="verify", n=2, p=2), seed=79) - ref = op.reference(**kw) - _compare((op(**kw), kw["initial_state_source"], kw["intermediate_states_buffer"]), ref) - - -# ── discrimination ──────────────────────────────────────────────────────── - - -def _wrong_variants(kw, pool_before): - """Plausible wrong recurrences, each expressed as `reference` driven on - transformed inputs -- so the rival is the same arithmetic, misapplied.""" - base = dict(kw, initial_state_source=pool_before) - n, p = kw["q"].shape[:2] - - def flip(t): - return t.flip(1) - - def per_token(t): - return t.reshape((n * p, 1) + tuple(t.shape[2:])) - - out = { - "no_gating": op.reference(**dict(base, g=torch.zeros_like(kw["g"])))[0], - "beta_one": op.reference(**dict(base, beta=torch.ones_like(kw["beta"])))[0], - "qk_swapped": op.reference(**dict(base, q=kw["k"], k=kw["q"]))[0], - "no_initial_state": op.reference( - **dict(base, initial_state_source=torch.zeros_like(pool_before)) - )[0], - "transposed_pool": op.reference( - **dict(base, initial_state_source=pool_before.transpose(-1, -2).contiguous()) - )[0], - # Tiled instead of interleaved: v-head hv reads q/k head hv % H. - "repeat_not_interleave": op.reference( - **dict(base, q=kw["q"].repeat(1, 1, REPEAT, 1), k=kw["k"].repeat(1, 1, REPEAT, 1)) - )[0], - "reversed_tokens": flip( - op.reference(**dict(base, **{x: flip(kw[x]) for x in ("q", "k", "v", "g", "beta")}))[0] - ), - # Every predicted token re-read from the same initial state. - "independent_tokens": op.reference( - **dict( - base, - **{x: per_token(kw[x]) for x in ("q", "k", "v", "g", "beta")}, - initial_state_indices=kw["initial_state_indices"].repeat_interleave(p), - disable_state_update=True, - ) - )[0].reshape(kw["v"].shape), - } - return out - - -def test_default_gate_separates_the_correct_reference_from_wrong_ones() -> None: - """`_compare` runs at torch's default band, so it is only as good as its - ability to reject a near miss. Two of these (`reversed_tokens`, - `independent_tokens`) are wrong only when P > 1 -- this entry's axis.""" - kw, _ = _build(dict(mode="write_back", n=3, p=4, slots=(2, 5, 0)), seed=80) - pool = kw["initial_state_source"] - pool_before = pool.clone() - ref = op.reference(**kw) - got = op(**kw) - _compare((got, pool, None), ref) - own = (got.float() - ref[0].float()).abs().max().item() - for name, wrong in _wrong_variants(kw, pool_before).items(): - gap = (wrong.float() - ref[0].float()).abs().max().item() - print(f"{name}: {gap:.3e} from the reference (the kernel: {own:.3e})") - assert gap > 1e-2, f"{name} is not a distinct reference (gap {gap:.3e})" - with pytest.raises(AssertionError): - _compare((got, pool, None), (wrong, ref[1], None)) 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 index 114df72f3257..b1039ed71fd6 100644 --- 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 @@ -2,9 +2,10 @@ # 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 fp32 recurrent pool -this target runs with keeps the FlashInfer branch out, which the first test -below proves. +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 @@ -13,7 +14,6 @@ import pytest import torch -import tensorrt_llm._torch.modules.fla.fused_sigmoid_gating_recurrent as _impl 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, ) @@ -35,7 +35,9 @@ def _compare( 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). Five wrong + 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 @@ -57,65 +59,59 @@ def _compare( 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 -_LAYERS = 3 -_TAIL = 5 + + +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 backing buffers some - layouts view. q/k/v are views over one packed `[T, 10240]` conv output, - exactly as the decode call reshapes its column splits.""" + """`__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) - layout = spec.get("layout", "varlen") - mixed = torch.randn(total, MIXED_QKV_WIDTH, generator=g, device=DEV, dtype=torch.bfloat16) + 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.get("normalize_qk", False): - for t in (q, k): - tf = t.float() - t.copy_(tf / (tf.pow(2).sum(-1, keepdim=True).sqrt() + 1e-6)) - if spec.get("ab", "packed") == "separate": + 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) - else: - ba = torch.randn(total, 2 * NUM_V_HEADS, generator=g, device=DEV, dtype=torch.bfloat16) - b, a = ba[:, :NUM_V_HEADS], ba[:, NUM_V_HEADS:] A_log = torch.randn(NUM_V_HEADS, generator=g, device=DEV) dt_bias = torch.randn(NUM_V_HEADS, generator=g, device=DEV) - extras = {} - slots = spec["slots"] - pool = None - if slots is not None: - lead = (_LAYERS, slots) if layout == "layer_view" else (slots,) - state_shape = (*lead, 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)) - if layout == "layer_view": - extras["cache"] = pool - pool = pool[1] + 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 layout == "batched": - n, t = len(seq_lens), seq_lens[0] - q, k = (x.reshape(n, t, NUM_K_HEADS, HEAD_K_DIM).contiguous() for x in (q, k)) - v = v.reshape(n, t, NUM_V_HEADS, HEAD_V_DIM).contiguous() - a, b = (x.reshape(n, t, NUM_V_HEADS).contiguous() for x in (a, b)) + if spec.get("layout", "varlen") == "rows": + q, k, v, a, b = _rows_form(mixed, a, b) cu_seqlens = None - output = None - if spec.get("output") == "prefix": - extras["buf"] = torch.full( - (1, total + _TAIL, NUM_V_HEADS, HEAD_V_DIM), -7.0, device=DEV, dtype=torch.bfloat16 - ) - output = extras["buf"][:, :total] args = [ A_log, a, @@ -129,133 +125,107 @@ def _build(spec, seed): pool, indices, None, - spec.get("l2norm", True), + True, cu_seqlens, - output, + None, ] - return args, extras + 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, extras = _build(spec, seed=zlib.crc32(cell.why.encode())) + args, _ = _build(spec, seed=zlib.crc32(cell.why.encode())) pool = args[9] - pool_before = None if pool is None else pool.clone() - cache_before = extras["cache"].clone() if "cache" in extras else None - tail_before = extras["buf"][:, -_TAIL:].clone() if "buf" in extras else None + pool_before = pool.clone() ref = op.reference(*args) out = op(*args) _compare((out, pool), ref) - if args[14] is not None: - assert out.data_ptr() == args[14].data_ptr(), "output= must be written in place" - assert torch.equal(extras["buf"][:, -_TAIL:], tail_before), "buffer tail written" - if cache_before is not None: - for other in (0, 2): - assert torch.equal(extras["cache"][other], cache_before[other]), other - if pool is not None: - # 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 - - -def test_certified_branch_is_the_vendored_triton_kernel() -> None: - """The fp32 pool keeps the FlashInfer arm out, armed both ways: its entry - point is replaced by one that raises, the fp32 call completes, and a bf16 - pool control in the same process does raise -- so the witness is not blind.""" - - def _boom(*args, **kwargs): - raise AssertionError("FlashInfer arm taken") - - original = _impl._flashinfer_gdn_decode - try: - _impl._flashinfer_gdn_decode = _boom - args, _ = _build(dict(seq_lens=(1,) * 4, slots=24), seed=100) - assert not _impl._can_use_flashinfer_gdn_decode(args[9], HEAD_K_DIM, HEAD_V_DIM, 4, 4) - assert op(*args).shape == (1, 4, NUM_V_HEADS, HEAD_V_DIM) - - args, _ = _build(dict(seq_lens=(1,) * 4, slots=24, pool_dtype=torch.bfloat16), seed=101) - assert _impl._can_use_flashinfer_gdn_decode(args[9], HEAD_K_DIM, HEAD_V_DIM, 4, 4) - with pytest.raises(AssertionError, match="FlashInfer arm taken"): - op(*args) - finally: - _impl._flashinfer_gdn_decode = original - - -def test_default_scale_and_repeated_calls_are_bitwise() -> None: - """`scale=None` equals an explicit `K ** -0.5` bitwise, and two identical - calls give identical output and state.""" - args, _ = _build(dict(seq_lens=(1,) * 4, slots=24), seed=95) - pool = args[9] - results = [] - for scale in (None, None, HEAD_K_DIM**-0.5): - call = list(args) - call[9], call[11] = pool.clone(), scale - results.append((op(*call), call[9])) - for o, p in results[1:]: - assert torch.equal(o, results[0][0]) and torch.equal(p, results[0][1]) - - -def test_guard_refuses_a_noncontiguous_output_buffer() -> None: - """Unguarded, the kernel writes `output` with packed arithmetic: a last-dim - slice of a doubly wide buffer is filled with misplaced values.""" - args, _ = _build(dict(seq_lens=(1,) * 4, slots=24), seed=97) + # 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) - wide = torch.zeros(1, 4, NUM_V_HEADS, 2 * HEAD_V_DIM, device=DEV, dtype=torch.bfloat16) - unguarded = list(args) - unguarded[14] = wide[..., :HEAD_V_DIM] - got = op(*unguarded) + 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, args[9]), ref) - assert torch.count_nonzero(wide[..., HEAD_V_DIM:]) > 0, "the hazard this guard covers is gone" - - strided = torch.zeros(1, 8, NUM_V_HEADS, HEAD_V_DIM, device=DEV, dtype=torch.bfloat16) - for bad in (wide[..., :HEAD_V_DIM], strided[:, ::2]): - guarded = list(args) - guarded[9], guarded[14] = pool_before.clone(), bad - with pytest.raises(AssertionError, match="output must be contiguous"): - with validating(op): - op(*guarded) - - -def test_wrong_variants_land_far_outside_the_gate() -> None: - """The default band separates right from wrong by > 100x the correct gap.""" - args, _ = _build(dict(seq_lens=(1,) * 4, slots=24), seed=98) - pool_before = args[9].clone() - ref = op.reference(*args) - got = op(*args) - _compare((got, args[9]), ref) - gap = (got.float() - ref[0].float()).abs().max().item() - - def ref_with(**overrides): - names = ("A_log", "a", "dt_bias", "q", "b", "pool", "l2norm") - positions = (0, 1, 2, 5, 8, 9, 12) - call = list(args) - call[9] = pool_before - for name, position in zip(names, positions): - if name in overrides: - call[position] = overrides[name] - return op.reference(*call)[0] - - q = args[5] - flipped = q[:, :, torch.arange(NUM_K_HEADS, device=DEV).flip(0), :] - wrong = { - "a and b swapped": ref_with(a=args[8], b=args[1]), - "l2norm omitted": ref_with(l2norm=False), - "pool (V, K) transposed": ref_with(pool=pool_before.transpose(2, 3).contiguous()), - "initial state ignored": ref_with(pool=torch.zeros_like(pool_before)), - "q head order reversed": ref_with(q=flipped), - } - for name, variant in wrong.items(): - distance = (got.float() - variant.float()).abs().max().item() - assert distance > 100 * gap, f"{name}: {distance} vs correct {gap}" - with pytest.raises(AssertionError): - _compare((got, args[9]), (variant, ref[1])) + _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 new file mode 100644 index 000000000000..b551a1ec8097 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_gdn_block.py @@ -0,0 +1,397 @@ +# 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 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` for decode, +`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, +) + +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) -> types.SimpleNamespace: + """The attributes `_gdn_block` reads off the core, with layer 0's pools.""" + 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 + 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), + ) + + +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) + si_all = buf[:m] + qsl_long = torch.arange(m + 1, dtype=torch.int64, device=DEV) + return (si_all, False, m, 1, None, None), qsl_long + + +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, qsl_long): + """The decode block with the Triton single-token recurrence.""" + 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(mixed, [KD, KD, VD], 1) + out = triton_decode( + A_log=L.a_log, + a=a, + dt_bias=L.dt_bias, + softplus_beta=M._SOFTPLUS_BETA, + softplus_threshold=M._SOFTPLUS_THRESHOLD, + q=torch.reshape(qv, [1, m, H, K]), + k=torch.reshape(kv, [1, m, H, K]), + v=torch.reshape(vv, [1, m, HV, V]), + b=b, + initial_state_source=core._bound.ssm[0], + initial_state_indices=si_all, + scale=None, + use_qk_l2norm_in_kernel=True, + cu_seqlens=qsl_long, + 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) + M.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() + + +def test_decode_block() -> None: + """A con-128 decode step; the graph-replay test below runs 32 rows.""" + m = 128 + core, L, x = _core(100 + m), _layer(200 + m), _x(m, 300 + m) + core_ref = _clone_core(core) + st, qsl_long = _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], qsl_long) + _check(y, y_ref, core, core_ref, f"decode m={m}") + named = set(_slots(m)) + unnamed = [s for s in range(SLOTS) if s not in named] + original = _core(100 + m) + assert torch.equal(core._bound.ssm[0][unnamed], original._bound.ssm[0][unnamed]), ( + "unnamed slots written" + ) + + +def test_prefill_block(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), _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") + + +def test_decode_block_replays_under_cuda_graph() -> 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), _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_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 new file mode 100644 index 000000000000..24ebe8c685a8 --- /dev/null +++ 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 @@ -0,0 +1,122 @@ +# 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 new file mode 100644 index 000000000000..adf92c5b6ba6 --- /dev/null +++ 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 @@ -0,0 +1,155 @@ +# 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 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 pool-dtype +guard. 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.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"] + 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=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" + + +def test_guards() -> None: + """A bf16 pool (the target sends this entry an fp32 one only) is refused.""" + args = _build(dict(rows=4), seed=5) + with validating(op), pytest.raises(AssertionError): + op(args[0], args[1].to(torch.bfloat16), *args[2:]) + + +# ---- 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 new file mode 100644 index 000000000000..e6a5f55bd730 --- /dev/null +++ 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 @@ -0,0 +1,103 @@ +# 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_prefill_conv_post.py b/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 new file mode 100644 index 000000000000..8482175c4a3a --- /dev/null +++ b/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 @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""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. +""" + +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.gdn_prefill_conv_post import ( # noqa: E501 + A_LOG_RANGE, + CONV_WIDTH, + DT_BIAS_RANGE, + HEAD_K_DIM, + HEAD_V_DIM, + 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 + gdn_prefill_conv_post as op, +) + +__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 +#: 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 _build(spec, seed): + """`__call__`'s positional arguments for one cell.""" + g = torch.Generator(device=DEV).manual_seed(seed) + lengths = spec["lengths"] + n_tok = sum(lengths) + + 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) + 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] + slots = spec.get("slots") or _engine_slots(len(lengths)) + bounds = [0] + for n in lengths: + bounds.append(bounds[-1] + n) + qsl = torch.tensor(bounds, dtype=torch.int32, device=DEV) + 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 + lo, hi_ = DT_BIAS_RANGE + dt_bias = torch.rand(NUM_V_HEADS, generator=g, device=DEV) * (hi_ - lo) + lo + return [ + x, + weight, + conv_states, + qsl, + si, + hi, + a, + b, + A_log, + dt_bias, + *HEADS, + PAD, + L2_NORM_EPS, + SOFTPLUS_THRESHOLD, + ] + + +def _compare(out, pool, want, what=""): + q, k, v, g, beta = out + q_r, k_r, v_r, g_r, beta_r, pool_r = want + for got, ref in zip(out, want[:5]): + assert got.dtype == ref.dtype and got.shape == ref.shape and got.is_contiguous() + for name, got, ref in (("q", q, q_r), ("k", k, k_r), ("v", v, v_r)): + d = (got.float() - ref.float()).abs() + flips = (d > 0).float().mean().item() if d.numel() else 0.0 + print( + f"{what} {name}: max abs {d.max().item() if d.numel() else 0:.3e}, " + f"differing fraction {flips:.2e}" + ) + torch.testing.assert_close(got, ref) + torch.testing.assert_close(g, g_r, rtol=8 * _APPROX_ULP, atol=G_ATOL) + torch.testing.assert_close(beta, beta_r) + assert torch.equal(pool, pool_r), "conv pool differs" + + +@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())) + want = op.reference(*args) + inputs = [t.clone() if isinstance(t, torch.Tensor) else t for t in args] + 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) 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 index 54fe1ac9e2d0..4adc8c137f6f 100644 --- 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 @@ -2,15 +2,13 @@ # 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; the discrimination test -is what shows that gate is not vacuous on these operands. +The op is a pure copy, so every comparison is bitwise. """ import zlib import pytest import torch -from triton.compiler.errors import CompilationError 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, @@ -45,133 +43,60 @@ def _compare(out: tuple, ref: tuple) -> None: DEV = "cuda" Q_WIDTH = NUM_K_HEADS * HEAD_K_DIM # 2048 V_WIDTH = NUM_V_HEADS * HEAD_V_DIM # 6144 -#: The `z` half of `in_proj_qkvz`, of which the stock mixer's `mixed_qkv` is a column slice. -Z_WIDTH = NUM_V_HEADS * HEAD_V_DIM def _build(spec, seed): - """`(mixed_qkv, conv_out or None)`: the input, and for a verify batch the + """`(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) - dtype = spec.get("dtype", torch.bfloat16) - heads = spec.get("heads", CKPT_HEADS) - qkv_dim = 2 * heads[0] * heads[1] + heads[2] * heads[3] - layout = spec.get("layout", "contiguous") def randn(*shape): - return torch.randn(*shape, generator=g, device=DEV).to(dtype) + return torch.randn(*shape, generator=g, device=DEV).to(torch.bfloat16) if "draft" in spec: b, p = spec["num_decodes"], spec["draft"] - if layout == "conv_dim_major": - conv_out = randn(b, qkv_dim, p) # contiguous; the flatten copies unless b == 1 - else: - # 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) - if layout != "conv_dim_major": - assert mixed.is_contiguous() and mixed.data_ptr() == conv_out.data_ptr() + # 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 - n = spec["num_tokens"] - if layout == "row_slice_of_projection": - mixed = randn(n + 9, qkv_dim + Z_WIDTH)[9:, :qkv_dim] - assert mixed.stride() == (qkv_dim + Z_WIDTH, 1) - elif layout == "column_stride_2": - mixed = randn(n, 2 * qkv_dim)[:, ::2] - assert mixed.stride() == (2 * qkv_dim, 2) - elif layout == "transposed_view": - mixed = randn(qkv_dim, n).t() - assert mixed.stride() == (1, n) - elif layout == "wide": - mixed = randn(n, qkv_dim + Z_WIDTH) - else: - mixed = randn(n, qkv_dim) - return mixed, None + 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 - heads = spec.get("heads", CKPT_HEADS) mixed, conv_out = _build(spec, seed=zlib.crc32(cell.why.encode())) snapshot = mixed.clone() - out = op(mixed, *heads) - _compare(out, op.reference(mixed, *heads)) + 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" - if mixed.numel(): # an empty allocation has no storage to tell apart - storages = {t.untyped_storage().data_ptr() for t in out} - assert len(storages) == 3 and mixed.untyped_storage().data_ptr() not in storages + 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"] - num_k_heads, head_k_dim, num_v_heads, head_v_dim = heads 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), + 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), + 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_draft_rows_are_independent() -> None: - """Every row carries a distinct constant, so a token-axis mix-up -- invisible - on rows that happen to agree -- cannot pass.""" - num_tokens = 6 * 4 - marks = torch.arange(1, num_tokens + 1, dtype=torch.float32, device=DEV).to(torch.bfloat16) - mixed = marks.reshape(num_tokens, 1).expand(num_tokens, QKV_DIM).contiguous() - q, k, v = op(mixed, *CKPT_HEADS) - for out, width in ((q, Q_WIDTH), (k, Q_WIDTH), (v, V_WIDTH)): - flat = out.reshape(num_tokens, -1) - assert flat.shape[1] == width - assert torch.equal(flat, marks.reshape(num_tokens, 1).expand(num_tokens, width)) - - -def test_bitwise_gate_discriminates_wrong_variants() -> None: - """Plausible ways to get the packing wrong; `_compare` rejects each.""" - g = torch.Generator(device=DEV).manual_seed(600) - n = 24 - mixed = torch.randn(n, QKV_DIM, generator=g, device=DEV).to(torch.bfloat16) - out = op(mixed, *CKPT_HEADS) - q_ref, k_ref, v_ref = op.reference(mixed, *CKPT_HEADS) - _compare(out, (q_ref, k_ref, v_ref)) - wrong = { - "q_k_swapped": (k_ref, q_ref, v_ref), - "v_from_q_offset": ( - q_ref, - k_ref, - mixed.narrow(1, 0, V_WIDTH).reshape(n, NUM_V_HEADS, HEAD_V_DIM).unsqueeze(0), - ), - "head_axis_transposed": (q_ref.transpose(2, 3).reshape(q_ref.shape), k_ref, v_ref), - "token_order_reversed": (q_ref.flip(1), k_ref, v_ref), - } - for name, bad in wrong.items(): - with pytest.raises(AssertionError): - _compare(out, bad) - print(f"{name}: rejected") - - -@pytest.mark.parametrize("heads", [(16, 128, 48, 64), (16, 64, 48, 128)]) -def test_op_rejects_head_dims_that_round_to_different_blocks(heads) -> None: - """Loud at compile time, before any launch -- so not an `is_valid` guard.""" - qkv_dim = 2 * heads[0] * heads[1] + heads[2] * heads[3] - mixed = torch.randn(13, qkv_dim, device=DEV).to(torch.bfloat16) - with pytest.raises(CompilationError, match="Mismatched type for dim_mask"): - op(mixed, *heads) - - 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. @@ -190,20 +115,3 @@ def test_a_short_row_reads_the_next_row_and_is_refused() -> None: with pytest.raises(AssertionError, match="must be at least qkv_dim=10240 wide"): with validating(op): op(mixed, *CKPT_HEADS) - - -def test_guard_refuses_over_declared_heads() -> None: - mixed = torch.randn(8, QKV_DIM, device=DEV).to(torch.bfloat16) - with pytest.raises(AssertionError, match="must be at least qkv_dim="): - with validating(op): - op(mixed, NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS + 16, HEAD_V_DIM) - - -def test_guard_refuses_a_three_dimensional_input() -> None: - """`[batch, draft, qkv_dim]` -- a caller who skipped the flatten. Unguarded - it faults with an illegal memory access (measured; not replayed here, it - poisons the process); the guard refuses it before the launch.""" - mixed = torch.randn(4, 2, QKV_DIM, device=DEV).to(torch.bfloat16) - with pytest.raises(AssertionError, match="must be at least qkv_dim="): - 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_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 index 84c95ef21b64..13a81044a508 100644 --- 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 @@ -47,45 +47,26 @@ def _compare( def _build(spec, seed): - """`__call__`'s arguments, plus the backing allocations a layout views.""" + """`__call__`'s arguments, plus the backing allocations a layer view views.""" g = torch.Generator(device=DEV).manual_seed(seed) slots = spec["slots"] - rec_shape, conv_shape = spec.get("rec_shape", REC_SHAPE), spec.get("conv_shape", CONV_SHAPE) rec_dtype = spec.get("rec_dtype", torch.float32) - layout = spec.get("layout", "plain") + layer_view = spec.get("layout", "plain") == "layer_view" def randn(*shape, dtype): return torch.randn(*shape, generator=g, device=DEV).to(dtype) backing = {} - if layout == "row_padding": - row = 1 - for d in rec_shape: - row *= d - backing["rec"] = randn(slots, 2, row, dtype=rec_dtype) - rec = backing["rec"][:, 0] - assert rec.stride() == (2 * row, 1) - elif layout == "layer_view": - backing["rec"] = randn(_LAYERS, slots, *rec_shape, dtype=rec_dtype) - rec = backing["rec"][1] + 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 = None - if spec.get("conv", True): - if layout == "layer_view": - backing["conv"] = randn(_LAYERS, slots, *conv_shape, dtype=torch.bfloat16) - conv = backing["conv"][1] - else: - conv = randn(slots, *conv_shape, dtype=torch.bfloat16) - - if spec["indices"] == "perm": - indices = torch.randperm(slots, generator=g, device=DEV).to(torch.int32) - flags = torch.rand(slots, generator=g, device=DEV) < 0.5 - else: - indices = torch.tensor( - spec["indices"], dtype=spec.get("index_dtype", torch.int32), device=DEV - ) - flags = torch.tensor(spec["flags"], dtype=spec.get("flag_dtype", torch.bool), device=DEV) + 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 @@ -111,75 +92,10 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res # 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") == "row_padding": - assert torch.equal(backing["rec"][:, 1], before["rec"][:, 1]), "padding lane written" 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) - if spec["indices"] == "perm": - # The cell discriminates "reset slot s" from "reset row i". - by_position = {i for i, f in enumerate(flags.tolist()) if not f} - assert reset != by_position, "the permutation failed to separate the two readings" - - -def test_the_following_kernel_sees_the_zeros() -> None: - """The launch is ordered on the current stream: no explicit sync is needed.""" - torch.manual_seed(5) - rec = torch.randn(4, *REC_SHAPE, dtype=torch.float32, device=DEV) - state_indices = torch.arange(4, dtype=torch.int32, device=DEV) - has_initial_states = torch.zeros(4, dtype=torch.bool, device=DEV) - for _ in range(20): - rec.normal_() - op(rec, state_indices, has_initial_states) - assert rec.sum().item() == 0.0 - - -@pytest.mark.parametrize( - "case,message", - [ - ("empty_pool", "recurrent state pool must have at least one cache line"), - ("slot_count_mismatch", "recurrent and convolution pools must have the same slot count"), - ("flag_shape_mismatch", "has_initial_states and state_indices must have matching shapes"), - ("recurrent_rows_noncontiguous", "recurrent state rows must be contiguous"), - ("recurrent_rows_overlapping", "recurrent state rows must not overlap"), - ("conv_rows_noncontiguous", "convolution state rows must be contiguous"), - ("conv_rows_overlapping", "convolution state rows must not overlap"), - ("cpu_pool", "Pointer argument cannot be accessed from Triton"), - ("cpu_indices", "Pointer argument cannot be accessed from Triton"), - ], -) -def test_op_rejects_layout_violations_loudly(case: str, message: str) -> None: - """Why `is_valid` guards only the index tensors: the rest raises itself.""" - torch.manual_seed(6) - rec = torch.randn(2, 4, dtype=torch.float32, device=DEV) - conv = torch.randn(2, 4, dtype=torch.bfloat16, device=DEV) - state_indices = torch.tensor([0], dtype=torch.int32, device=DEV) - has_initial_states = torch.tensor([False], device=DEV) - if case == "empty_pool": - rec, conv = torch.empty(0, 4, dtype=torch.float32, device=DEV), None - elif case == "slot_count_mismatch": - conv = torch.randn(3, 4, dtype=torch.bfloat16, device=DEV) - elif case == "flag_shape_mismatch": - has_initial_states = torch.tensor([False, True], device=DEV) - elif case == "recurrent_rows_noncontiguous": - rec = torch.randn(2, 8, dtype=torch.float32, device=DEV)[:, ::2] - elif case == "recurrent_rows_overlapping": - rec = torch.as_strided(torch.randn(6, dtype=torch.float32, device=DEV), (2, 4), (2, 1)) - conv = None - elif case == "conv_rows_noncontiguous": - conv = torch.randn(2, 8, dtype=torch.bfloat16, device=DEV)[:, ::2] - elif case == "conv_rows_overlapping": - conv = torch.as_strided(torch.randn(6, dtype=torch.bfloat16, device=DEV), (2, 4), (2, 1)) - elif case == "cpu_pool": - rec, conv = rec.cpu(), None - state_indices, has_initial_states = state_indices.cpu(), has_initial_states.cpu() - elif case == "cpu_indices": - conv = None - state_indices, has_initial_states = state_indices.cpu(), has_initial_states.cpu() - with pytest.raises(ValueError, match=message): - with validating(op): - op(rec, state_indices, has_initial_states, conv) def _pools(): @@ -189,50 +105,13 @@ def _pools(): return rec, conv -def test_guard_hazard_strided_index_tensors() -> None: - """Unguarded, a stride-2 view is read through its base pointer: indices - [0, 1, 2] reset [0, 1], and flags [T, T, T] reset row 1. No error.""" - good_indices = torch.tensor([0, 1, 2], dtype=torch.int32, device=DEV) - good_flags = torch.zeros(3, dtype=torch.bool, device=DEV) - strided_indices = torch.tensor([0, 9, 1, 9, 2, 9], dtype=torch.int32, device=DEV)[::2] - strided_flags = torch.tensor([True, False] * 3, device=DEV)[::2] - for indices, flags in ((strided_indices, good_flags), (good_indices, strided_flags)): - rec, conv = _pools() - # The reference reads the tensors by value, which is what a caller meant. - want = op.reference(rec, indices, flags, conv) - got = op(rec, indices, flags, conv) - with pytest.raises(AssertionError, match="pool differs"): - _compare(got, want) - - -def test_guard_hazard_two_d_index_tensors() -> None: - """Unguarded, a matching [2, 3] pair passes the op's shape check and only - its first shape[0] = 2 of 6 requests are considered.""" +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.arange(6, dtype=torch.int32, device=DEV).reshape(2, 3) - flags = torch.zeros(2, 3, dtype=torch.bool, device=DEV) - want = op.reference(rec, indices.reshape(-1), flags.reshape(-1), conv) - got = op(rec, indices, flags, conv) - with pytest.raises(AssertionError, match="pool differs"): - _compare(got, want) - - -@pytest.mark.parametrize("which", ["state_indices", "has_initial_states"]) -@pytest.mark.parametrize("how", ["strided", "two_d"]) -def test_guard_refuses_index_tensors_the_op_misreads(which: str, how: str) -> None: - rec, conv = _pools() - indices = torch.zeros(3, dtype=torch.int32, device=DEV) + indices = torch.tensor([0, 9, 1, 9, 2, 9], dtype=torch.int32, device=DEV)[::2] flags = torch.zeros(3, dtype=torch.bool, device=DEV) - if how == "strided": - bad_indices = torch.tensor([0, 9, 1, 9, 2, 9], dtype=torch.int32, device=DEV)[::2] - bad_flags = torch.zeros(6, dtype=torch.bool, device=DEV)[::2] - else: - bad_indices = torch.arange(6, dtype=torch.int32, device=DEV).reshape(2, 3) - bad_flags = torch.zeros(2, 3, dtype=torch.bool, device=DEV) - if which == "state_indices": - indices = bad_indices - else: - flags = bad_flags - with pytest.raises(AssertionError, match=f"{which} must be a contiguous 1-D tensor"): + 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 new file mode 100644 index 000000000000..6e5b42957ea4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_gdn_replay_verify.py @@ -0,0 +1,355 @@ +# 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.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 +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + ReplayStateUpdateMetadata, + _advance_replay_state, +) + +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, + ) + 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 new file mode 100644 index 000000000000..548008dacdf4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py @@ -0,0 +1,168 @@ +# 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 new file mode 100644 index 000000000000..2603aaf63e00 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_kv_range_probe.py @@ -0,0 +1,82 @@ +# 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 new file mode 100644 index 000000000000..c4aea680b2e9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_logits_widening.py @@ -0,0 +1,50 @@ +# 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)) 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__tp1/test_modeling_v2_spec_rollback.py new file mode 100644 index 000000000000..c56c098ff397 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_spec_rollback.py @@ -0,0 +1,423 @@ +# 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. +""" + +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 + 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 + 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, +) + +assert torch.cuda.is_available(), "the rollback handler test requires a CUDA device" + +DEV = "cuda" +EOS = 248044 +VOCAB = 248320 +NGRAM = 3 +KERNEL = 4 +DILATION = 3 +STATE_LEN = (KERNEL - 1) * DILATION +CHANNELS = 64 +TPB = 32 +INDEX_DIM = 128 +LAYERS = 3 +COMPRESS = 4 +SLOTS = 24 +PAGES = 48 +BLOCKS_PER_SEQ = 4 + + +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.""" + + def __init__(self, seed, ctx_lens, gen_starts, p, slots, dup_of=None, shared_page=False): + g = torch.Generator().manual_seed(seed) + self.ctx_lens = list(ctx_lens) + self.num_contexts = len(ctx_lens) + self.n = len(gen_starts) + self.p = p + self.slots = torch.tensor(slots, dtype=torch.int32, device=DEV) + assert self.slots.shape[0] == self.num_contexts + self.n + lengths = self.ctx_lens + [p] * self.n + starts = [0] * self.num_contexts + list(gen_starts) + # Distinct pages per request; a duplicate takes its original's row. + perm = torch.randperm(PAGES, generator=g) + table = torch.reshape(perm[: len(lengths) * BLOCKS_PER_SEQ], [len(lengths), -1]) + for dup, original in (dup_of or {}).items(): + table[self.num_contexts + dup] = table[self.num_contexts + original] + self.block_table = table.to(DEV, torch.int32) + req, logical = [], [] + for r, (length, start) in enumerate(zip(lengths, starts)): + req += [r] * length + logical += range(start, start + length) + self.rows = len(req) + self.prefix = sum(self.ctx_lens) + self.req_idx = torch.tensor(req, dtype=torch.int32, device=DEV) + self.logical = torch.tensor(logical, dtype=torch.int64, device=DEV) + self.carry = torch.tensor( + [False] * self.num_contexts + [True] * self.n, dtype=torch.bool, device=DEV + ) + self.lengths = lengths + self.qsl = _qsl(lengths) + + # Pre-step state: every pool and plane random, so a stray write shows. + self.ngram0 = torch.randint(0, VOCAB, (SLOTS, NGRAM - 1), generator=g).to(DEV) + self.conv0 = torch.randn(SLOTS, CHANNELS, STATE_LEN, generator=g).to(DEV, torch.bfloat16) + self.index0 = [ + torch.randn(PAGES, TPB, 1, INDEX_DIM, generator=g).to(DEV, torch.bfloat16) + for _ in range(LAYERS) + ] + 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) + self.x = torch.randn(self.rows, CHANNELS, generator=g).to(DEV, torch.bfloat16) + self.keys = [ + torch.randn(self.rows, 1, INDEX_DIM, generator=g).to(DEV, torch.bfloat16) + for _ in range(LAYERS) + ] + self.anchor_keys = [ + torch.randn(self.rows, 1, INDEX_DIM, generator=g).to(DEV, torch.bfloat16) + for _ in range(LAYERS) + ] + self.coords = torch.randint(0, 1 << 20, (self.rows, 3), generator=g).to(DEV, torch.int32) + + # Live state the step mutates. + self.ngram = self.ngram0.clone() + self.conv = self.conv0.clone() + if shared_page: + self.index, self.position = _shared_planes(self.index0, self.position0) + else: + self.index = [t.clone() for t in self.index0] + self.position = self.position0.clone() + self.accepted = torch.ones(self.num_contexts + self.n, dtype=torch.int32, device=DEV) + + def reset(self): + self.ngram.copy_(self.ngram0) + self.conv.copy_(self.conv0) + for live, pre in zip(self.index, self.index0): + live.copy_(pre) + self.position.copy_(self.position0) + + def set_accepted(self, gen_accepted): + assert len(gen_accepted) == self.n + assert all(1 <= a <= self.p for a in gen_accepted) + self.gen_accepted = list(gen_accepted) + values = [1] * self.num_contexts + list(gen_accepted) + self.accepted.copy_(torch.tensor(values, dtype=torch.int32)) + + +def _shared_planes(index0, position0): + """Live copies of the planes as views of one byte buffer per page: every layer's index-K + page, then the position page (what the QSA cache manager packs when an NVFP4 pool's data + and scale pages would otherwise absorb them).""" + index_bytes = TPB * INDEX_DIM * 2 + page = torch.zeros(PAGES, LAYERS * index_bytes + TPB * 3 * 4, dtype=torch.uint8, device=DEV) + index = [ + page[:, layer * index_bytes : (layer + 1) * index_bytes] + .view(torch.bfloat16) + .view(PAGES, TPB, 1, INDEX_DIM) + for layer in range(LAYERS) + ] + position = page[:, LAYERS * index_bytes :].view(torch.int32).view(PAGES, TPB, 3) + assert index[0].stride(0) == page.stride(0) // 2 and not index[0].is_contiguous() + for live, pre in zip(index, index0): + live.copy_(pre) + position.copy_(position0) + return index, position + + +def _qsl(lengths): + starts = [0] + for length in lengths: + starts.append(starts[-1] + length) + return torch.tensor(starts, dtype=torch.int64, device=DEV) + + +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.""" + 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 + anchor = ((logical + 1) % COMPRESS == 0).reshape(-1, 1, 1) + layer_index[pages, within] = torch.where(anchor, anchor_keys, layer_index[pages, within]) + + +def _verify_step(b, rollback): + """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.""" + rollback.discard() + qsl = b.qsl + rollback.begin( + num_contexts=b.num_contexts, + prefix_tokens=b.prefix, + num_gens=b.n, + rows_per_gen=b.p, + slots=b.slots[b.num_contexts :], + request_indices=b.req_idx[b.prefix :], + logical_positions=b.logical[b.prefix :], + block_table=b.block_table, + tokens_per_block=TPB, + ) + rollback.stage_ple_pools(b.ngram, b.conv) + window = ngram_token_history(b.ids, b.ngram, b.slots, b.carry, qsl, EOS, NGRAM) + rollback.stage_ple_rows(window, b.x) + dilated_short_conv_paged( + b.x.t(), b.weight, None, b.conv, qsl, b.slots, b.carry, DILATION, True, -1 + ) + for layer in range(LAYERS): + rollback.stage_index(b.index[layer], b.position) + _write_planes( + b, + b.index[layer], + b.position, + b.req_idx, + b.logical, + b.keys[layer], + b.anchor_keys[layer], + b.coords, + ) + rollback.commit_speculative_states(b.accepted, b.slots, b.num_contexts) + + +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.""" + keep = list(range(b.prefix)) + lengths = list(b.ctx_lens) + for s, a in enumerate(b.gen_accepted): + first = b.prefix + s * b.p + 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 + ) + for layer in range(LAYERS): + _write_planes( + b, + index[layer], + position, + b.req_idx[keep], + b.logical[keep], + b.keys[layer][keep], + b.anchor_keys[layer][keep], + b.coords[keep], + ) + return ngram, conv, index, position + + +def _bits(t): + return t.view(torch.int16) if t.dtype == torch.bfloat16 else t + + +def _assert_state(b, expected, what): + ngram, conv, index, position = expected + assert torch.equal(b.ngram, ngram), f"{what}: n-gram pool" + assert torch.equal(_bits(b.conv), _bits(conv)), f"{what}: conv pool" + for layer in range(LAYERS): + assert torch.equal(_bits(b.index[layer]), _bits(index[layer])), f"{what}: index {layer}" + 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. `decode` and `mixed` are the accepted- +# prefix check; `decode_p8` drives the shared-page and capture tests. +CASES = { + "decode": ((), (5, 30, 61, 100), 4, (1, 2, 4, 3)), + "mixed": ((7, 3), (12, 29, 95), 4, (2, 1, 3)), + # 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)), +} + + +def _slots(num): + return list(range(3, 3 + 2 * num, 2)) + + +@pytest.mark.parametrize("case", ["decode", "mixed"]) +def test_commit_matches_the_accepted_prefix(case): + 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, SpeculativeRollback()) + _assert_state(b, _reference(b), case) + + +def test_commit_over_the_strided_planes_of_a_shared_page(): + """The handler stages and restores the index and position planes by advanced indexing, so + planes that are strided views of one buffer (an NVFP4 KV cache's side planes) roll back + like contiguous ones, and the bytes of the other layers' pages in the same buffer stay as + the step wrote them.""" + case = "decode_p8" + ctx_lens, gen_starts, p, gen_accepted = CASES[case] + b = _Batch( + 13, ctx_lens, gen_starts, p, _slots(len(ctx_lens) + len(gen_starts)), shared_page=True + ) + b.set_accepted(gen_accepted) + _verify_step(b, SpeculativeRollback()) + _assert_state(b, _reference(b), case) + + +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) + _assert_state(b, (b.ngram0, b.conv0, b.index0, b.position0), "unstaged") + + +def test_discard_drops_a_stale_staging(): + """A staged step whose commit never ran is dropped by the next forward.""" + 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.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, + ) + rollback.discard() + rollback.commit_speculative_states(b.accepted, b.slots, 0) + _assert_state(b, (b.ngram0, b.conv0, b.index0, b.position0), "discarded") + + +def test_commit_is_sync_free(): + 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 = SpeculativeRollback() + torch.cuda.synchronize() + torch.cuda.set_sync_debug_mode("error") + try: + _verify_step(b, rollback) + finally: + torch.cuda.set_sync_debug_mode(0) + _assert_state(b, _reference(b), "sync-free") + + +def test_padding_duplicates_leave_the_unique_request_exact(): + # 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, SpeculativeRollback()) + # The unique request's slot is exactly its accepted prefix. + pages = b.block_table[0].long() + expected = _reference(_single(b, 0)) + assert torch.equal(b.ngram[3], expected[0][3]) + assert torch.equal(_bits(b.conv[3]), _bits(expected[1][3])) + for layer in range(LAYERS): + assert torch.equal(_bits(b.index[layer][pages]), _bits(expected[2][layer][pages])), ( + f"index {layer}" + ) + assert torch.equal(b.position[pages], expected[3][pages]) + + +def _single(b, s): + """`b` restricted to its generation request `s`, same pre-step state.""" + one = _Batch.__new__(_Batch) + one.__dict__.update(b.__dict__) + rows = torch.arange(b.prefix + s * b.p, b.prefix + (s + 1) * b.p, device=DEV) + one.ctx_lens, one.num_contexts, one.n, one.prefix = [], 0, 1, 0 + one.lengths = [b.p] + one.slots = b.slots[b.num_contexts + s : b.num_contexts + s + 1] + one.carry = b.carry[b.num_contexts + s : b.num_contexts + s + 1] + one.qsl = _qsl(one.lengths) + one.ids, one.x = b.ids[rows], b.x[rows] + one.keys = [k[rows] for k in b.keys] + one.anchor_keys = [k[rows] for k in b.anchor_keys] + one.coords = b.coords[rows] + one.logical = b.logical[rows] + one.req_idx = torch.zeros_like(b.req_idx[rows]) + one.block_table = b.block_table[b.num_contexts + s : b.num_contexts + s + 1] + one.gen_accepted = [b.gen_accepted[s]] + return one + + +#: 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, + (8, 7, 6, 5, 4, 3, 2, 1), + (2, 8, 1, 7, 3, 6, 4, 5), + (8,) * 8, + ], +} + + +@pytest.mark.parametrize("case", sorted(_REPLAYS)) +def test_capture_and_replay_follow_the_accepted_counts(case): + """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 = SpeculativeRollback() + b.set_accepted((p,) * len(gen_starts)) + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + _verify_step(b, rollback) + torch.cuda.current_stream().wait_stream(side) + b.reset() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _verify_step(b, rollback) + + for accepted in _REPLAYS[case]: + b.reset() + b.set_accepted(accepted) + graph.replay() + torch.cuda.synchronize() + _assert_state(b, _reference(b), f"replay {accepted}") 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 index 8724f9b5997d..21bdba2b0d0c 100644 --- 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 @@ -5,7 +5,7 @@ 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. +restating it. Beyond the cells: the one-patch-per-row guard and graph capture. """ import zlib @@ -13,7 +13,6 @@ 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.vision.patch_embed_conv3d import ( # noqa: E501 patch_embed_conv3d as op, ) @@ -73,170 +72,30 @@ def _operands( return pixel_values, weight, bias_t -def _build(spec, seed): - pixel_values, weight, bias = _operands( - spec["num_patches"], - dtype=spec["dtype"], - seed=seed, - in_channels=spec.get("in_channels", IN_CHANNELS), - kernel=spec.get("kernel", (KT, KH, KW)), - embed_dim=spec.get("embed_dim", EMBED_DIM), - bias=spec.get("bias", True), - ) - layout = spec.get("layout", "flat") - n, row = pixel_values.shape - if layout == "five_d": - pixel_values = pixel_values.view(-1, *weight.shape[1:]) - elif layout == "column_slice": - wide = torch.empty(n, 2 * row, device="cuda", dtype=pixel_values.dtype) - wide[:, :row] = pixel_values - pixel_values = wide[:, :row] - assert not pixel_values.is_contiguous() and pixel_values.stride(-1) == 1 - elif layout == "transposed": - pixel_values = pixel_values.t().contiguous().t() - assert pixel_values.stride(-1) != 1 - return pixel_values, weight, bias - - @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) - pixel_values, weight, bias = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + spec = cell.spec + pixel_values, weight, bias = _operands( + spec["num_patches"], dtype=spec["dtype"], seed=zlib.crc32(cell.why.encode()) + ) before = pixel_values.clone() - previous = torch.backends.cudnn.allow_tf32 - torch.backends.cudnn.allow_tf32 = False - try: - out = op(pixel_values, weight, bias) - finally: - torch.backends.cudnn.allow_tf32 = previous + 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" - if not pixel_values.is_contiguous() and pixel_values.stride(-1) == 1: - assert torch.equal(out, op(pixel_values.contiguous(), weight, bias)) _compare(out, op.reference(pixel_values, weight, bias)) -def test_five_dimensional_input_is_the_same_call() -> None: - pixel_values, weight, bias = _operands(64, seed=1) - flat = op(pixel_values, weight, bias) - explicit = op(pixel_values.view(-1, IN_CHANNELS, KT, KH, KW), weight, bias) - assert torch.equal(flat, explicit) - - -def test_zero_bias_is_the_identity() -> None: - pixel_values, weight, _ = _operands(128, seed=2, bias=False) - zero = torch.zeros(EMBED_DIM, dtype=weight.dtype, device="cuda") - assert torch.equal(op(pixel_values, weight, None), op(pixel_values, weight, zero)) - - -def test_rows_are_independent() -> None: - """Same shape, so cuDNN selects the same algorithm: asserted bitwise.""" - pixel_values, weight, bias = _operands(128, seed=5) - order = torch.randperm(128, generator=torch.Generator().manual_seed(0)).to("cuda") - out = op(pixel_values, weight, bias) - assert torch.equal(op(pixel_values[order].contiguous(), weight, bias), out[order]) - - -def test_wrong_variants_land_outside_the_gate() -> None: - """Three structural mis-derivations land hundreds of times outside - `_compare`'s atol; dropping the bias is bounded by `max|bias|` (~0.02 at the - tower's initializer scale) and is asserted to *be* the bias.""" - pixel_values, weight, bias = _operands(64, seed=6) - out = op(pixel_values, weight, bias).float() - ref = op.reference(pixel_values, weight, bias) - rms = ref.float().pow(2).mean().sqrt().item() - atol = max(1e-5, 4.0 * rms * 2.0**-8) - correct_dev = (out - ref.float()).abs().max().item() - rows = pixel_values.float() - - def _matmul(w): - with true_fp32_matmul(): - return rows @ w.reshape(w.shape[0], -1).float().t() + bias.float() - - structural = { - "row read as (T, C, Ph, Pw)": _matmul(weight.permute(0, 2, 1, 3, 4)), - "kernel flipped (true convolution)": _matmul(weight.flip(2, 3, 4)), - "patch rows and columns transposed": _matmul(weight.transpose(-1, -2)), - } - print(f"\n gate atol {atol:.5f}; correct deviation {correct_dev:.5f}") - for name, wrong in structural.items(): - dev = (out - wrong).abs().max().item() - print(f" {name:38s} dev {dev:9.5f} ({dev / atol:8.1f}x the gate)") - assert dev > 100.0 * atol, (name, dev, atol) - with pytest.raises(AssertionError): - _compare(out.to(ref.dtype), wrong.to(ref.dtype)) - - no_bias = (out - (ref.float() - bias.float())).abs().max().item() - print(f" {'bias dropped':38s} dev {no_bias:9.5f} ({no_bias / atol:8.1f}x the gate)") - assert no_bias > 5.0 * atol, (no_bias, atol) - assert abs(no_bias - bias.abs().max().item()) <= correct_dev - - -def test_guard_refuses_an_image_volume() -> None: - """The hazard, unguarded: a [1, C, D, H, W] volume with the element count of - 16 patches returns a [16, 1152] result of the right shape and dtype computed - from the wrong pixels.""" - g = torch.Generator(device="cuda").manual_seed(7) - volume = torch.empty(1, IN_CHANNELS, KT, 4 * KH, 4 * KW, device="cuda") - volume = volume.uniform_(-1.0, 1.0, generator=g).to(torch.bfloat16) - _, weight, bias = _operands(1, seed=7) - - reinterpreted = op(volume, weight, bias) - meant = ( - torch.nn.functional.conv3d(volume, weight, bias, stride=(KT, KH, KW)) - .flatten(2) - .squeeze(0) - .t() - ) - assert reinterpreted.shape == meant.shape == (16, EMBED_DIM) - dev = (reinterpreted.float() - meant.float()).abs().max().item() - print(f"\n unguarded volume reinterpretation deviates by {dev:.5f}") - assert dev > 0.1, dev - - with pytest.raises(AssertionError, match="each row must be exactly one patch"): - with validating(op): - op(volume, 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) - for bad in (doubled, pixel_values[:, :-1].contiguous(), pixel_values.reshape(-1)): - with pytest.raises(AssertionError, match="each row must be exactly one patch"): - with validating(op): - op(bad, weight, bias) - - -def test_dtype_and_device_mismatches_raise() -> None: - """Loud from torch itself, so the entry does not guard them.""" - pixel_values, weight, bias = _operands(8, seed=10) - with pytest.raises(RuntimeError, match="should be the same"): - op(pixel_values, weight.float(), bias.float()) - with pytest.raises(RuntimeError, match="should be the same"): - op(pixel_values, weight, bias.float()) - with pytest.raises(RuntimeError, match="should be the same"): - op(pixel_values, weight.cpu(), bias.cpu()) - - -def test_fp32_accumulates_in_tf32_unless_cudnn_is_told_otherwise() -> None: - pixel_values, weight, bias = _operands(256, dtype=torch.float32, seed=11) - ref = op.reference(pixel_values, weight, bias) - previous = torch.backends.cudnn.allow_tf32 - try: - torch.backends.cudnn.allow_tf32 = True - tf32_dev = (op(pixel_values, weight, bias) - ref).abs().max().item() - torch.backends.cudnn.allow_tf32 = False - exact = op(pixel_values, weight, bias) - finally: - torch.backends.cudnn.allow_tf32 = previous - exact_dev = (exact - ref).abs().max().item() - print(f"\n fp32 conv, cudnn TF32 on {tf32_dev:.3e}; off {exact_dev:.3e}") - _compare(exact, ref) - assert tf32_dev > 20.0 * exact_dev, (tf32_dev, exact_dev) + 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: 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 index 78611e4d4c96..3517e2f3d18e 100644 --- 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 @@ -7,7 +7,8 @@ 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), which is why fp32 cells gate on `op.reference`. +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 @@ -45,15 +46,14 @@ def _merge_order(x, height, width, merge): ) -def _ref_resize(table, grid_thw, merge, mode="bilinear", align_corners=True): +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: - kwargs = {} if mode == "nearest" else {"align_corners": align_corners} - resized = F.interpolate(src, size=(height, width), mode=mode, **kwargs) + 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] @@ -71,117 +71,8 @@ def test_certified_cells(cell, certifying): # noqa: F811 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"])) - if spec["dtype"] is not torch.float32: - # 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_exact_grid_is_the_table_itself_in_merge_order() -> None: - """At h = w = N every sample lands on a cell; asserted bitwise. This is what - pins the corner-aligned convention -- a cell-centre sampler lacks it.""" - table = _table(seed=1) - out = op(table, [(1, NUM_GRID, NUM_GRID)], MERGE) - direct = _merge_order(table.view(NUM_GRID, NUM_GRID, HIDDEN), NUM_GRID, NUM_GRID, MERGE) - assert torch.equal(out, direct) - - -def test_the_two_grid_axes_are_not_interchangeable() -> None: - """The bug a square test image hides: (1, 4, 8) and (1, 8, 4) differ by ~4.7.""" - table = _table(seed=2) - tall = op(table, [(1, 8, 4)], MERGE) - wide = op(table, [(1, 4, 8)], MERGE) - assert (tall.float() - wide.float()).abs().max().item() > 1.0 - - -def test_discrimination_against_the_wrong_operators() -> None: - """Six plausible mis-implementations at [2304, 1152] bf16, grid (1, 62, 90). - - Measured max deviation from the entry: correct 0.0147; axes transposed - 6.31; raster order 6.06; intra-block transpose 4.63; align_corners=False - 1.38; column-major table read 6.31; nearest neighbour 5.64. - """ - table = _table(seed=3) - grid = [(1, 62, 90)] - height, width = 62, 90 - out = op(table, grid, MERGE).double() - correct_dev = (out - _ref_resize(table, grid, MERGE)).abs().max().item() - - src = table.double().view(NUM_GRID, NUM_GRID, HIDDEN).permute(2, 0, 1).unsqueeze(0).contiguous() - raster = F.interpolate(src, size=(height, width), mode="bilinear", align_corners=True) - raster = raster[0].permute(1, 2, 0) - transposed = F.interpolate(src, size=(width, height), mode="bilinear", align_corners=True) - transposed = transposed[0].permute(1, 2, 0).transpose(0, 1).contiguous() - col_major = table.double().view(NUM_GRID, NUM_GRID, HIDDEN).permute(2, 1, 0).unsqueeze(0) - col_major = F.interpolate( - col_major.contiguous(), size=(height, width), mode="bilinear", align_corners=True - )[0].permute(1, 2, 0) - intra = raster.reshape(height // MERGE, MERGE, width // MERGE, MERGE, HIDDEN) - intra = intra.permute(0, 2, 3, 1, 4).reshape(height * width, HIDDEN) - - wrong = { - "transposed_axes": _merge_order(transposed, height, width, MERGE), - "raster_order": raster.reshape(height * width, HIDDEN), - "intra_block_transpose": intra, - "align_corners_false": _ref_resize(table, grid, MERGE, align_corners=False), - "column_major_table": _merge_order(col_major, height, width, MERGE), - "nearest_neighbour": _ref_resize(table, grid, MERGE, mode="nearest"), - } - print(f"\n correct {correct_dev:.6f}") - for name, variant in wrong.items(): - dev = (out - variant).abs().max().item() - print(f" {name:24s} {dev:.6f}") - assert dev > 20 * correct_dev, (name, dev, correct_dev) - with pytest.raises(AssertionError): - torch.testing.assert_close(out.to(table.dtype), variant.to(table.dtype)) - - -def test_correctly_rounded_and_not_a_bf16_accumulate_chain() -> None: - """fp32 accumulation, one rounding: within one bf16 ulp of the fp64 value - and nearer it than the bf16 multiply-then-add chain a literal HF - transcription computes.""" - table = _table(seed=4) - for grid in ([(1, 62, 90)], [(1, 64, 64)]): - _, height, width = grid[0] - out = op(table, grid, MERGE) - exact = op.reference(table.double(), grid, MERGE) - - h_pos = torch.linspace(0, NUM_GRID - 1, height, device="cuda") - w_pos = torch.linspace(0, NUM_GRID - 1, width, device="cuda") - 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, d_w = (h_pos - h_lo).unsqueeze(1), (w_pos - w_lo).unsqueeze(0) - grid_bf = table.view(NUM_GRID, NUM_GRID, HIDDEN) - chain = ( - grid_bf[h_lo][:, w_lo] * ((1 - d_h) * (1 - d_w)).bfloat16()[..., None] - + grid_bf[h_lo][:, w_hi] * ((1 - d_h) * d_w).bfloat16()[..., None] - + grid_bf[h_hi][:, w_lo] * (d_h * (1 - d_w)).bfloat16()[..., None] - + grid_bf[h_hi][:, w_hi] * (d_h * d_w).bfloat16()[..., None] - ) - chain = _merge_order(chain, height, width, MERGE) - - entry_dev = (out.double() - exact).abs().max().item() - chain_dev = (chain.double() - exact).abs().max().item() - print(f"\n {grid}: entry {entry_dev:.6f}, bf16 chain {chain_dev:.6f}") - assert entry_dev < chain_dev, (grid, entry_dev, chain_dev) - assert entry_dev <= 2.0**-8 * exact.abs().max().item(), (grid, entry_dev) - - -def test_a_batch_is_the_concatenation_of_its_images() -> None: - table = _table(seed=7) - grid = [(1, 4, 8), (2, 6, 4), (1, 48, 48), (1, 62, 90)] - batched = op(table, grid, MERGE) - one_by_one = torch.cat([op(table, [g], MERGE) for g in grid]) - assert torch.equal(batched, one_by_one) - - -def test_video_frames_repeat_the_spatial_embedding() -> None: - table = _table(seed=8) - frames, height, width = 3, 6, 10 - per_frame = op(table, [(frames, height, width)], MERGE).view(frames, height * width, HIDDEN) - for f in range(1, frames): - assert torch.equal(per_frame[f], per_frame[0]), f - assert torch.equal(per_frame[0], op(table, [(1, height, width)], 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: @@ -197,23 +88,6 @@ def test_guard_refuses_a_non_square_table() -> None: op(table, [(1, 4, 4)], MERGE) -def test_grid_not_divisible_by_the_merge_size_raises() -> None: - """Loud from torch, so not guarded.""" - table = _table(hidden=64, seed=11) - with pytest.raises(RuntimeError, match="is invalid for input of size"): - op(table, [(1, 5, 8)], MERGE) - with pytest.raises(RuntimeError, match="is invalid for input of size"): - op(table, [(1, 8, 5)], MERGE) - with pytest.raises(RuntimeError, match="is invalid for input of size"): - op(table, [(1, 2, 2)], 4) - - -def test_non_2d_table_raises() -> None: - table = torch.randn(NUM_POS, 4, 8, device="cuda", dtype=torch.bfloat16) - with pytest.raises(ValueError, match="weight has to be a 2D Tensor"): - 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) 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 new file mode 100644 index 000000000000..d8a1f9396844 --- /dev/null +++ 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 @@ -0,0 +1,165 @@ +# 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/test_modeling_v2_qwen4_exp_kv_quant.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_kv_quant.py new file mode 100644 index 000000000000..a699b909f1d7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_kv_quant.py @@ -0,0 +1,229 @@ +# 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/test_modeling_v2_qwen4_exp_moe_scalars.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py new file mode 100644 index 000000000000..27a677f5931a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_moe_scalars.py @@ -0,0 +1,74 @@ +# 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/test_modeling_v2_qwen4_exp_multimodal_embeds.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py new file mode 100644 index 000000000000..5a371f24bf36 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py @@ -0,0 +1,267 @@ +# 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, `mm_item_order`, 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, +) + +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) + return SimpleNamespace(multimodal_data=data, mm_item_order=order, 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": modality, "index": 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 mm_item_order"): + _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/test_modeling_v2_qwen4_exp_pool_contract.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py index d89a528d1417..20cdb53a3b25 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py @@ -9,15 +9,20 @@ 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 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, ) +pytestmark = pytest.mark.cpu_only + G = M._PoolGeometry( tokens_per_block=32, heads_kv=2, @@ -26,6 +31,7 @@ index_head_dim=128, conv_dim=8192, conv_kernel=4, + num_k_heads=16, num_v_heads=48, head_v_dim=128, head_k_dim=128, @@ -33,8 +39,15 @@ ple_state_len=9, ngram_context_len=3, dtype=torch.bfloat16, + kv_dtype=torch.bfloat16, + kv_width=256, ) +# The same pools under an fp8 KV cache: K/V are e4m3, every other plane keeps its type. +G_FP8 = G._replace(kv_dtype=torch.float8_e4m3fn) +# Under NVFP4 the K/V pool is uint8 e2m1 pairs, half a head wide. +G_NVFP4 = G._replace(kv_dtype=torch.uint8, kv_width=128) SLOTS = 6 +HISTORY = 16 def _t(*shape, dtype=torch.bfloat16): @@ -45,8 +58,17 @@ def _conv(**kw): return _t(SLOTS, G.conv_dim, G.conv_kernel - 1, **kw) -def _ssm(**kw): - return _t(SLOTS, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=torch.float32) +def _ssm(dtype=torch.float32): + return _t(SLOTS, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=dtype) + + +def _history(slots=SLOTS, history=HISTORY): + """One GDN layer's cached-replay history buffers (old_x, old_B, old_dt).""" + return ( + _t(slots, 2, history, G.num_v_heads, G.head_v_dim), + _t(slots, 2, history, G.num_k_heads, G.head_k_dim), + _t(slots, 2, G.num_v_heads, history, dtype=torch.float32), + ) def test_valid_pools_pass() -> None: @@ -54,7 +76,16 @@ def test_valid_pools_pass() -> None: M.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_recurrent_pools(G, _conv(), _ssm()) + M.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) + for dtype in (torch.float32, torch.bfloat16): + M.check_recurrent_pools(G, _conv(), _ssm(dtype)) M.check_ple_pools( G, _t(SLOTS, G.stream, G.ple_state_len), @@ -62,27 +93,30 @@ def test_valid_pools_pass() -> None: _conv(), ) M.check_step_buffers( - G, - _t(SLOTS, 2, G.conv_dim, G.conv_kernel - 1), - _t(SLOTS, 2, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=torch.float32), - _conv(), - _ssm(), - _t(SLOTS, dtype=torch.int32), + 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( + 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( - "plane, match", - [ - (None, "did not allocate"), - (_t(9, 32, 3), "torch.bfloat16"), - (_t(9, 16, 3, dtype=torch.int32), "expected"), - (_t(9, 32, 2, dtype=torch.int32), "expected"), - ], -) -def test_position_plane_rejects(plane, match) -> None: - with pytest.raises(AssertionError, match=match): - M.check_position_plane(G, plane) +def test_position_plane_rejects_the_wrong_dtype() -> None: + with pytest.raises(AssertionError, match="torch.bfloat16"): + M.check_position_plane(G, _t(9, 32, 3)) def test_attention_pools_reject_wrong_index_plane_and_kv_layout() -> None: @@ -101,6 +135,62 @@ def test_attention_pools_reject_wrong_index_plane_and_kv_layout() -> None: M.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: + """The attention op reads the pool through one `quant_mode`, so a pool of the other + storage would be read through the wrong format. The draft layer's bind runs the same check.""" + index = _t(9, 32, 1, 128) + 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) + with pytest.raises(AssertionError, match="paged K/V"): + M.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) + # 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) + # 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( + 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( + G, 3, index, _t(9, 2, G.heads_kv, G.tokens_per_block, 128, dtype=torch.uint8) + ) + + +def _strided(shape, slot_stride, dtype=torch.uint8): + """A view of `shape` whose slot stride is `slot_stride`, as a manager's role planes of a + lifecycle slot holding several layers' pages are.""" + pages = shape[0] * slot_stride + flat = torch.empty(pages, dtype=dtype, device="meta") + return flat.as_strided(shape, (slot_stride, *torch.empty(shape[1:], device="meta").stride())) + + +def test_kv_scale_planes_must_hold_the_data_pools_pages_per_slot() -> None: + """The attention op addresses the scale pages with the data pages' block table, so a + scale pool with another number of pages per slot would be read at other slots' scales. + The engine's trunk holds 24 data role pages per slot (12 layers x K, V): 196608 bytes + 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) + # 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) + 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)) + + +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)) + + def test_recurrent_pools_reject_wrong_key_dim_order_and_dtype() -> None: with pytest.raises(AssertionError, match="recurrent state pool"): M.check_recurrent_pools( @@ -108,6 +198,9 @@ def test_recurrent_pools_reject_wrong_key_dim_order_and_dtype() -> None: ) with pytest.raises(AssertionError, match="conv state pool"): M.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)) def test_ple_pools_reject_slot_mismatch_and_int32_history() -> None: @@ -123,19 +216,117 @@ def test_ple_pools_reject_slot_mismatch_and_int32_history() -> None: M.check_ple_pools(G, None, ngram, _conv()) -def test_step_buffers_reject_missing_dtype_drift_and_index() -> None: +def test_step_buffers_reject_missing_conv_rank_dtype_and_index() -> None: int_conv = _t(SLOTS, 2, G.conv_dim, G.conv_kernel - 1) - int_ssm = _t(SLOTS, 2, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=torch.float32) idx = _t(SLOTS, dtype=torch.int32) - with pytest.raises(AssertionError, match="no per-step"): - M.check_step_buffers(G, None, None, _conv(), _ssm(), idx) + with pytest.raises(AssertionError, match="no per-step conv"): + M.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(), int_ssm, _conv(), _ssm(), idx) - # The promotion kernel copies without converting, so the dtypes must match the pools. - with pytest.raises(AssertionError, match="per-step recurrent buffer"): - M.check_step_buffers(G, int_conv, int_ssm.to(torch.bfloat16), _conv(), _ssm(), idx) + M.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) with pytest.raises(AssertionError, match="intermediate_state_indices"): - M.check_step_buffers(G, int_conv, int_ssm, _conv(), _ssm(), idx.to(torch.int64)) + M.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, int_ssm, _conv(), _ssm(), None) + 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()) + + +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)) + 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): + with pytest.raises(AssertionError, match="replay_step_width"): + M.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()) + with pytest.raises(AssertionError, match="cache_buf_idx"): + M.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): + with pytest.raises(AssertionError, match="predicted_tokens_per_seq"): + M.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) + for rows in (1, 2, 4): + with pytest.raises(AssertionError, match="step width 3"): + M.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) + with pytest.raises(AssertionError, match="slots for a"): + M.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)) + with pytest.raises(AssertionError, match="old_B"): + M.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)) + with pytest.raises(AssertionError, match="double-buffered"): + M.check_replay_history( + G, + 3, + HISTORY, + SLOTS, + old_x, + old_B, + _t(SLOTS, 3, G.num_v_heads, HISTORY, dtype=torch.float32), + ) + + +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) + with pytest.raises(AssertionError, match="no row"): + M.check_shared_selection(shared, {4: 0}, 3, 512) + with pytest.raises(AssertionError, match="no row"): + M.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 diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_thop_attention_calls.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_thop_attention_calls.py new file mode 100644 index 000000000000..87584e317806 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_thop_attention_calls.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Every `thop_attention` call of the qwen4_exp target names the KV storage it runs over. + +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. + +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`. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest + +import tensorrt_llm._torch._experimental.modeling_v2 as _modeling_v2 + +pytestmark = pytest.mark.cpu_only + +_MODELING = ( + Path(_modeling_v2.__file__).resolve().parent + / "models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py" +) +_KV_STORAGE_ARGS = ("quant_mode", "kv_scale_orig_quant", "kv_scale_quant_orig") +_TREE = ast.parse(_MODELING.read_text()) + + +def _thop_attention_calls() -> list[ast.Call]: + return [ + node + 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: + 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}") + + +def _keywords(call: ast.Call) -> dict[str, ast.expr]: + return {kw.arg: kw.value for kw in call.keywords if kw.arg is not None} + + +def _unpacked(call: ast.Call) -> set[str]: + """Names of the `**name` arguments of `call`.""" + return { + kw.value.id for kw in call.keywords if kw.arg is None and isinstance(kw.value, ast.Name) + } + + +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)] + + +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 + for call in _KV_CACHE_CALLS: + keywords = _keywords(call) + for name in _KV_STORAGE_ARGS: + assert name in keywords, f"thop_attention call at line {call.lineno} omits {name}" + value = keywords[name] + # Read off the core, whose `FP8` / `None` resolution all three follow. + assert isinstance(value, ast.Attribute) and value.attr == name, ( + f"line {call.lineno}: {name}={ast.unparse(value)}, not the core's `{name}`" + ) + + +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} + assert constants + assert not constants & set(_KV_STORAGE_ARGS) + for call in _KV_CACHE_CALLS: + assert "_CALL_CONSTANTS" in _unpacked(call) + steps = [ + node + for node in ast.walk(_TREE) + if isinstance(node, ast.FunctionDef) and node.name == "step_args" + ] + assert len(steps) == 2, "one step_args per target (prefill, decode)" + for step in steps: + returned = [n for n in ast.walk(step) if isinstance(n, ast.Return)] + assert returned + for node in returned: + assert isinstance(node.value, ast.Call), "step_args returns dict(...)" + assert not {kw.arg for kw in node.value.keywords} & set(_KV_STORAGE_ARGS) + + +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} + 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 + for call in vision: + assert not set(_keywords(call)) & set(_KV_STORAGE_ARGS) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py new file mode 100644 index 000000000000..447ca2875a90 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py @@ -0,0 +1,203 @@ +# 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, +) + +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) + fields = {"py_multimodal_data": data, "py_mm_item_order": order, "prompt_len": prompt_len} + fields.update(overrides) + return SimpleNamespace(**fields) + + +def _admit(request): + _shell().validate_request(request) + + +def _item(modality, index): + return {"modality": modality, "index": 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 mm_item_order"): + _admit(_request(image=IMAGES, video=VIDEOS)) + with pytest.raises(ValueError, match="needs mm_item_order"): + _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=[{"modality": "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/test_modeling_v2_qwen4_exp_weight_table.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py index bdb99113b301..155db4b45143 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py @@ -20,10 +20,15 @@ __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( @@ -65,8 +70,9 @@ ) -def _core(mtp: bool): - """The configuration surface the table reads.""" +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), @@ -75,15 +81,17 @@ def _core(mtp: bool): ), vision_config=SimpleNamespace(**_VISION), mtp_enabled=mtp, + trunk_moe_backend=backend, + moe_prefill_copy=backend == "TRTLLM" if prefill_copy is None else prefill_copy, ) -@pytest.mark.parametrize("mtp", [False, True], ids=["no_mtp", "mtp"]) -def test_every_declared_weight_has_a_source(mtp): +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.""" - core = _core(mtp) + 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" @@ -97,9 +105,16 @@ def test_mtp_rows_appear_only_when_mtp_is_live(): on = W.MODEL_WEIGHTS.manifest(_core(True)) added = set(on) - set(off) assert set(off) <= set(on) - assert added and all(k.startswith(f"l{_TEXT['num_hidden_layers']}_mtp_") for k in added) - # The draft layer's own QSA indexer is never loaded; see weights.py. - assert not any("indexer" in str(src) for k in added for src, _, _ in on[k]) + 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: @@ -139,32 +154,34 @@ def hyper_connection(prefix): P(f"{prefix}_up", stream, rank) P(f"{prefix}_inject", t.hc_count, stream) - def gated_attention(prefix): - 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) + 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): - P(f"{prefix}sh_gu", 2 * shared, hidden) + # 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) - P(f"{prefix}sh_gate", 1, hidden) 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}_") - P(f"l{i}_idx_w", index_width, hidden) + 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: - P(f"l{i}_in_qkv_w", conv_dim, hidden) - P(f"l{i}_in_z_w", value_dim, hidden) - P(f"l{i}_in_a_w", t.linear_num_value_heads, hidden) - P(f"l{i}_in_b_w", t.linear_num_value_heads, hidden) + # [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) @@ -175,6 +192,11 @@ def shared_expert(prefix): 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}_") @@ -223,6 +245,9 @@ def shared_expert(prefix): 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) @@ -235,16 +260,142 @@ def shared_expert(prefix): return w -@pytest.mark.parametrize("mtp", [False, True], ids=["no_mtp", "mtp"]) -def test_the_table_declares_what_the_init_block_did(mtp): - core = _core(mtp) +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(mtp) + 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/test_modeling_v2_routing.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py index a526d0752a0a..b47b5836a1f8 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -31,7 +31,13 @@ from pydantic import ValidationError from transformers import PretrainedConfig -from tensorrt_llm._torch._experimental.modeling_v2._router_index import modeling_v2_resolve +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + NoTarget, + modeling_v2_resolve, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp import ( + routing as qwen4_exp_routing, +) from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_auto import AutoModelForCausalLM @@ -549,13 +555,23 @@ def _qwen4_exp_sm(_on_sm10x): @pytest.mark.usefixtures("_qwen4_exp_sm") @pytest.mark.parametrize( "spec_config", - [None, MTPDecodingConfig(max_draft_len=1), MTPDecodingConfig(max_draft_len=3)], - ids=["no-spec", "mtp1", "mtp3"], + [ + None, + MTPDecodingConfig(max_draft_len=1), + MTPDecodingConfig(max_draft_len=3), + MTPDecodingConfig(max_draft_len=7), + ], + ids=["no-spec", "mtp1", "mtp3", "mtp7"], ) -def test_qwen4_exp_tp1_matches(spec_config): +def test_qwen4_exp_tp1_matches(monkeypatch, spec_config): + # mtp7 needs the GDN cached replay, which TRTLLM_USE_GDN_REPLAY=0 turns off. + monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY", raising=False) config = _model_config(_qwen4_exp_config(), spec_config=spec_config) - target = modeling_v2_resolve(config, _llm_args("require")) - assert _qualified(target).startswith(_QWEN4_EXP_TARGET) + # The recurrent-state cache at either dtype the target serves; `auto` is bf16 here. + for dtype in ("float32", "bfloat16", "auto"): + kv_cache_config = KvCacheConfig(mamba_ssm_cache_dtype=dtype) + target = modeling_v2_resolve(config, _llm_args("require", kv_cache_config=kv_cache_config)) + assert _qualified(target).startswith(_QWEN4_EXP_TARGET), dtype def _mtp(max_draft_len=3, **attributes): @@ -572,14 +588,14 @@ def _mtp(max_draft_len=3, **attributes): "spec_config", [ _NGRAM, - MTPDecodingConfig(max_draft_len=4), + MTPDecodingConfig(max_draft_len=8), MTPDecodingConfig(max_draft_len=3, use_mtp_vanilla=True), _mtp(max_total_draft_tokens=5), _mtp(num_nextn_predict_layers=2), ], ids=[ "ngram", - "draft4", + "draft8", "vanilla-mtp", "max-total-draft", "two-predictors", @@ -606,9 +622,177 @@ def test_qwen4_exp_on_two_ranks_does_not_match(): @pytest.mark.usefixtures("_qwen4_exp_sm") -def test_qwen4_exp_with_an_fp8_kv_cache_does_not_match(): +def test_qwen4_exp_with_an_fp8_ds_mla_kv_cache_does_not_match(): + """fp8_ds_mla is FP8 to the loader but selects the packed sparse-MLA cache layout.""" config = _model_config(_qwen4_exp_config()) - args = dict(kv_cache_config=KvCacheConfig(dtype="fp8")) + args = dict(kv_cache_config=KvCacheConfig(dtype="fp8_ds_mla")) assert modeling_v2_resolve(config, _llm_args("auto", **args)) is None with pytest.raises(ValueError, match="KV cache"): modeling_v2_resolve(config, _llm_args("require", **args)) + + +# --- qwen4_exp: engine choices the constructor would refuse ----------------------- + + +def _declined(config, llm_args_kwargs=None, match=None): + """``auto`` falls back to the built-in model and ``require`` says why.""" + kwargs = llm_args_kwargs or {} + assert modeling_v2_resolve(config, _llm_args("auto", **kwargs)) is None + with pytest.raises(ValueError, match=match): + modeling_v2_resolve(config, _llm_args("require", **kwargs)) + + +def _claimed(config, llm_args_kwargs=None): + target = modeling_v2_resolve(config, _llm_args("require", **(llm_args_kwargs or {}))) + assert _qualified(target).startswith(_QWEN4_EXP_TARGET) + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +@pytest.mark.parametrize("algo", ["fp8", "nvfp4"]) +def test_qwen4_exp_quantized_kv_cache_dtypes_match(algo): + _claimed(_model_config(_qwen4_exp_config()), dict(kv_cache_config=KvCacheConfig(dtype=algo))) + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +@pytest.mark.parametrize("algo", [QuantAlgo.INT8, QuantAlgo.FP8_BLOCK_SCALES]) +def test_qwen4_exp_other_kv_cache_algorithms_are_declined(algo): + choices = qwen4_exp_routing.engine_choices( + _model_config(_qwen4_exp_config()), _llm_args("auto") + ) + with pytest.raises(NoTarget, match="KV cache"): + qwen4_exp_routing._check_engine(None, algo, choices) + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +def test_qwen4_exp_piecewise_prefill_cuda_graph_falls_back(): + config = _model_config(_qwen4_exp_config()) + _declined(config, dict(prefill_cuda_graph_backend="piecewise"), match="piecewise") + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +@pytest.mark.parametrize("backend", ["disabled", "breakable"]) +def test_qwen4_exp_matches_without_or_with_the_breakable_prefill_cuda_graph(backend): + _claimed( + _model_config(_qwen4_exp_config()), + dict(prefill_cuda_graph_backend=backend), + ) + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +@pytest.mark.parametrize("backend", ["CUTLASS", "TRITON", "DEEPGEMM", "VANILLA"]) +def test_qwen4_exp_unimplemented_moe_backend_falls_back(backend): + config = replace(_model_config(_qwen4_exp_config()), moe_backend_requested=backend) + _declined(config, match="moe_config.backend") + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +@pytest.mark.parametrize("backend", ["AUTO", "TRTLLM", "CUTEDSL"]) +def test_qwen4_exp_implemented_moe_backend_matches(backend): + _claimed(replace(_model_config(_qwen4_exp_config()), moe_backend_requested=backend)) + + +@pytest.mark.parametrize( + "spec_kwargs", + [ + dict(max_concurrency=8), + dict(draft_len_schedule={4: 3, 8: 1}), + dict(acceptance_rate_threshold=0.5), + ], + ids=["max-concurrency", "draft-len-schedule", "acceptance-rate"], +) +@pytest.mark.usefixtures("_qwen4_exp_sm") +def test_qwen4_exp_runtime_draft_length_falls_back(spec_kwargs): + """The verify step runs at one width; a draft length that changes at runtime + would trip the width assert in the first reduced step.""" + spec = MTPDecodingConfig(max_draft_len=3, **spec_kwargs) + _declined(_model_config(_qwen4_exp_config(), spec_config=spec), match="draft length") + + +@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.""" + monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "0") + config = _model_config( + _qwen4_exp_config(), spec_config=MTPDecodingConfig(max_draft_len=max_draft_len) + ) + kv_cache_config = KvCacheConfig(mamba_ssm_cache_dtype="float32") + _declined(config, dict(kv_cache_config=kv_cache_config), match="TRTLLM_USE_GDN_REPLAY=0") + monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY") + _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): + """The CuteDSL expert entries capture up to 128 rows (32 requests of 4 verify rows).""" + monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY", raising=False) + kv_cache_config = KvCacheConfig(mamba_ssm_cache_dtype="float32") + spec = MTPDecodingConfig(max_draft_len=max_draft_len) + cutedsl = replace( + _model_config(_qwen4_exp_config(), spec_config=spec), moe_backend_requested="CUTEDSL" + ) + _declined(cutedsl, dict(kv_cache_config=kv_cache_config), match="CUTEDSL") + _claimed( + _model_config(_qwen4_exp_config(), spec_config=spec), dict(kv_cache_config=kv_cache_config) + ) + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +def test_qwen4_exp_cutedsl_matches_up_to_three_drafted_tokens(monkeypatch): + monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY", raising=False) + spec = MTPDecodingConfig(max_draft_len=3) + cutedsl = replace( + _model_config(_qwen4_exp_config(), spec_config=spec), moe_backend_requested="CUTEDSL" + ) + _claimed(cutedsl, dict(kv_cache_config=KvCacheConfig(mamba_ssm_cache_dtype="float32"))) + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +@pytest.mark.parametrize("dtype", ["float32", "bfloat16", "auto"]) +def test_qwen4_exp_mtp_matches_a_float32_or_bfloat16_state_cache(monkeypatch, dtype): + """GDN cached replay runs on either checkpoint dtype; `auto` resolves to bf16 here.""" + monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY", raising=False) + config = _model_config(_qwen4_exp_config(), spec_config=MTPDecodingConfig(max_draft_len=3)) + _claimed(config, dict(kv_cache_config=KvCacheConfig(mamba_ssm_cache_dtype=dtype))) + + +@pytest.mark.parametrize( + "spec_config", [None, MTPDecodingConfig(max_draft_len=3)], ids=["no-spec", "mtp3"] +) +@pytest.mark.usefixtures("_qwen4_exp_sm") +def test_qwen4_exp_declines_a_float16_state_cache(monkeypatch, spec_config): + """The engine's replay gate accepts fp16; the target's pool contract does not.""" + monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY", raising=False) + config = _model_config(_qwen4_exp_config(), spec_config=spec_config) + kv_cache_config = KvCacheConfig(mamba_ssm_cache_dtype="float16") + _declined(config, dict(kv_cache_config=kv_cache_config), match="recurrent-state cache") + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +def test_qwen4_exp_mtp_on_a_float16_state_cache_declines_with_replay_off(monkeypatch): + monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "0") + config = _model_config(_qwen4_exp_config(), spec_config=MTPDecodingConfig(max_draft_len=3)) + kv_cache_config = KvCacheConfig(mamba_ssm_cache_dtype="float16") + _declined(config, dict(kv_cache_config=kv_cache_config), match="recurrent-state cache") + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +def test_qwen4_exp_mtp_on_a_bf16_state_cache_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) + + +@pytest.mark.usefixtures("_qwen4_exp_sm") +def test_qwen4_exp_no_speculation_matches_the_default_state_cache(monkeypatch): + monkeypatch.delenv("TRTLLM_USE_GDN_REPLAY", raising=False) + _claimed(_model_config(_qwen4_exp_config())) diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 43a80f733f87..ea9a23ef8f68 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -2014,7 +2014,8 @@ def test_vbws_rejects_decreasing_beam_width_array(beam_width_array: list[int], assert request.py_beam_width == max(beam_width_array) # Exercise the real beam admission checks with unrelated cache, transport, - # token-range, sampler, and budget validation disabled in the executor stub. + # token-range, sampler, model, and budget validation disabled in the + # executor stub. executor = types.SimpleNamespace( _is_kv_manager_v2=True, max_beam_width=request.py_beam_width, @@ -2023,6 +2024,7 @@ def test_vbws_rejects_decreasing_beam_width_array(beam_width_array: list[int], _configure_csa2_remote_tail=lambda _request: None, _validate_token_id_range=lambda _request: None, sampler=types.SimpleNamespace(validate_request=lambda _request: None), + model_engine=types.SimpleNamespace(model=None), _validate_request_budget=lambda _request: None, ) validate = functools.partial(