diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 833ca696496b..77bc45f49e75 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -391,8 +391,7 @@ # experiment with named owners, not a standing domain. /tensorrt_llm/_torch/_experimental/modeling_v2 @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju /tests/unittest/_torch/modeling_v2 @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju -/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v3.py @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju -/tests/integration/defs/accuracy/test_modeling_v2_gpt_oss.py @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju +/tests/integration/defs/accuracy/modeling_v2 @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju # The two catalog categories whose contracts state kernel behaviour the # attention and MoE owners are the authority on; co-owned rather than reassigned. /tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention @tianyuxbear @Wanli-Jiang @WeiHaocheng @litaotju @xxi-nv @yuxianq diff --git a/jenkins/scripts/cbts/rules/modeling_v2_rule.py b/jenkins/scripts/cbts/rules/modeling_v2_rule.py index ea0a6597fab9..33b4e4927dde 100644 --- a/jenkins/scripts/cbts/rules/modeling_v2_rule.py +++ b/jenkins/scripts/cbts/rules/modeling_v2_rule.py @@ -25,9 +25,10 @@ `_MV2_ENTRY_PATTERNS`. Those markers are exact by construction rather than by luck: every unit test file in the subtree is named `test_modeling_v2_*` precisely so it cannot collide with the upstream -test of the same op, and the accuracy files follow the same prefix. So -there is no substring that could claim an unrelated entry, and no -`mtp_nextn=0`-style carve-out is needed. +test of the same op, and the accuracy gates live in their own +`accuracy/modeling_v2/` directory. So there is no substring that could +claim an unrelated entry, and no `mtp_nextn=0`-style carve-out is +needed. Outward fallback: not needed, and that is a property of the design rather than an accident. Nothing beyond the resolver module is imported @@ -71,14 +72,16 @@ # to TestsDefRule; the two scopes combine via _TESTSONLY_FAMILY. _MV2_SRC_PREFIXES: tuple[str, ...] = ("tensorrt_llm/_torch/_experimental/modeling_v2/",) -# Substrings that mark a test entry as modeling_v2. Both are unambiguous: +# Substrings that mark a test entry as modeling_v2. All are unambiguous: # - "unittest/_torch/modeling_v2/" → the op-level catalog matrix, taken # as whole-directory entries (one on l0_gb300, one on # l0_gb300_multi_gpus for the 4-rank collectives) -# - "test_modeling_v2_" → the accuracy gates, and any future unit file -# named by the subtree's own convention +# - "accuracy/modeling_v2/" → the per-model accuracy gates +# - "test_modeling_v2_" → any future unit file named by the subtree's +# own convention _MV2_ENTRY_PATTERNS: tuple[str, ...] = ( "unittest/_torch/modeling_v2/", + "accuracy/modeling_v2/", "test_modeling_v2_", ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md index c8c91dfcb665..481c86552981 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md @@ -218,7 +218,7 @@ widening a tolerance, and each is written up in its own entry: by the op, not certified here). Both targets are gated by the accuracy suite's modeling_v2 files, -`tests/integration/defs/accuracy/test_modeling_v2_{gpt_oss,deepseek_v3}.py`, +`tests/integration/defs/accuracy/modeling_v2/test_{gpt_oss,deepseek_v3}.py`, with `modeling_v2="require"` in force, which is what rules out the built-in implementation having been measured instead. gpt-oss-120b / sm_103 / tp1 runs gsm8k; deepseek-r1-0528-nvfp4 / sm_103 / dep4 runs gsm8k for the identity diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index cdb288f7220f..b0b6e78ef8d8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -51,6 +51,9 @@ MODELING_V2_ROUTERS = { "GptOssForCausalLM": "models.gpt_oss.routing", "DeepseekV3ForCausalLM": "models.deepseek_v3.routing", + # Qwen3.8-Flash-Next publishes ``Qwen4ExpForConditionalGeneration``, and the + # config loader's composite branch writes the same name back. + "Qwen4ExpForConditionalGeneration": "models.qwen4_exp.routing", } diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/ops/attention/fused_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py similarity index 56% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/ops/attention/fused_qk_norm_rope.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py index 4242a08264ce..c2f7af53d702 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/ops/attention/fused_qk_norm_rope.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/fused_qk_norm_rope.py @@ -17,15 +17,22 @@ class FusedQkNormRope(OpWrapper): The v heads are not touched. `qkv` is one `[num_tokens, (Hq + Hk + Hv) * head_dim]` buffer, which is the layout the fused projection already - produces, so nothing is split or copied. + produces, so nothing is split or copied. The norm is `rmsnorm(x) * w`, or + the Gemma form `rmsnorm(x) * (1 + w)` with `use_gemma`; it covers the whole + head, the rotation only its first `rotary_dim` dims. With `use_mrope`, + `position_ids` is `[3, T]` and rotary pair `j` takes its position from row 1 + when `j % 3 == 1 and j < 3 * mrope_section1`, from row 2 when + `j % 3 == 2 and j < 3 * mrope_section2`, from row 0 otherwise -- the + interleaved mRoPE layout. The caller owns the QKV projection before the + call and the KV-cache append and attention after it. """ ARCHS = frozenset({Arch.SM_103}) - # gpt-oss-120b is the only caller, and it asks for the RoPE half only: - # `is_qk_norm=False`, neox pairing, and YaRN-blended frequencies. 64 query - # heads, 8 key and 8 value, head_dim 64, all of which the cells carry - # because the op reads them as scalars rather than deriving them. + # gpt-oss-120b asks for the RoPE half only: `is_qk_norm=False`, neox + # pairing, and YaRN-blended frequencies. 64 query heads, 8 key and 8 value, + # head_dim 64, all of which the cells carry because the op reads them as + # scalars rather than deriving them. _GPT_OSS = dict( num_heads_q=64, num_heads_k=8, @@ -43,6 +50,30 @@ class FusedQkNormRope(OpWrapper): attention_factor=1.0, ) + # qwen4_exp (Qwen3.8-Flash-Next) runs Gemma QK norm and an unscaled rope + # on the first 64 dims (32 pairs), base 1e7, at two geometries. Its + # mrope_section is [11, 11, 10], passed as sections 11 / 11. Positions are + # drawn uniformly below `max_position` (a test axis, not an op argument): + # the kernel's fast sincos loses phase as `position * freq` grows, so the + # position range is part of what a cell certifies. + _QWEN_ROPE = dict( + rotary_dim=64, + base=1e7, + is_neox=True, + is_qk_norm=True, + use_gemma=True, + max_position=32768, + ) + _QWEN_MROPE = dict(use_mrope=True, mrope_section1=11, mrope_section2=11, **_QWEN_ROPE) + # Main attention (trunk and MTP): 24 query heads, 2 kv, head_dim 256. + _QWEN_MAIN = dict(num_heads_q=24, num_heads_k=2, num_heads_v=2, head_dim=256, **_QWEN_MROPE) + # The QSA index plane's query rope: 4 index heads plus the one k and one v + # slot of its scratch layout, over index_head_dim 128 (== 2 * rotary_dim, + # so the partial factor is not reapplied). + _QWEN_INDEX_HEADS = dict(num_heads_q=4, num_heads_k=1, num_heads_v=1, head_dim=128) + _QWEN_INDEX_MROPE = dict(**_QWEN_INDEX_HEADS, **_QWEN_MROPE) + _QWEN_INDEX_PLAIN = dict(**_QWEN_INDEX_HEADS, **_QWEN_ROPE) + CELLS: tuple[Cell, ...] = ( Cell(why="decode: one token", spec=dict(tokens=1, **_GPT_OSS)), Cell(why="prefill-shaped batch", spec=dict(tokens=2048, **_GPT_OSS)), @@ -50,19 +81,94 @@ class FusedQkNormRope(OpWrapper): why="a token count that is not a multiple of any tile", spec=dict(tokens=37, **_GPT_OSS), ), + Cell( + why="qwen4_exp main rope, decode: one token, Gemma norm and interleaved mRoPE " + "with three distinct position rows", + spec=dict(tokens=1, **_QWEN_MAIN), + ), + Cell( + why="qwen4_exp main rope, 19 tokens: the size its wiring and saturation tests run at", + spec=dict(tokens=19, **_QWEN_MAIN), + ), + Cell( + why="qwen4_exp main rope, 8192-token prefill: the Gemma one-ulp tail, under " + "GEMMA_PREFILL_ATOL", + spec=dict(tokens=8192, gemma_prefill_band=True, **_QWEN_MAIN), + ), + 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", + spec=dict(tokens=1, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 2 tokens: the second token's heads " + "straddle a block boundary", + spec=dict(tokens=2, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 5 tokens: the last block runs one warp", + spec=dict(tokens=5, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 17 tokens: one past a power of two", + spec=dict(tokens=17, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 19 tokens: the size its wiring and " + "saturation tests run at", + spec=dict(tokens=19, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 33 tokens: one past a power of two", + spec=dict(tokens=33, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 65 tokens: one past a power of two", + spec=dict(tokens=65, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (mRoPE), 8192-token prefill: the index site sees the " + "whole packed batch; under GEMMA_PREFILL_ATOL", + spec=dict(tokens=8192, gemma_prefill_band=True, **_QWEN_INDEX_MROPE), + ), + Cell( + why="qwen4_exp index rope (plain), decode: the form the index plane runs " + "today, row 0 alone as a flat [T] position_ids", + spec=dict(tokens=1, **_QWEN_INDEX_PLAIN), + ), + Cell( + why="qwen4_exp index rope (plain), 19 tokens: an odd count across blocks", + spec=dict(tokens=19, **_QWEN_INDEX_PLAIN), + ), + Cell( + why="qwen4_exp index rope (plain), 8192-token prefill: under GEMMA_PREFILL_ATOL", + spec=dict(tokens=8192, gemma_prefill_band=True, **_QWEN_INDEX_PLAIN), + ), ) note = """ - Three of this op's paths are uncertified because no shipped target reaches - them, and the first is the one its own name advertises: + Certified by caller: gpt-oss runs `is_qk_norm=False` with YaRN; qwen4_exp + runs `is_qk_norm=True, use_gemma=True`, unscaled, with and without + interleaved mRoPE. Uncertified because no shipped target reaches them: + `is_qk_norm=True` with `use_gemma=False` (the plain-weight norm the op's + name advertises), `is_neox=False`, YaRN together with the norm or with + mRoPE. A target that wants one adds the cell rather than assuming the path + works because the argument exists. - * `is_qk_norm=True` -- gpt-oss passes False, so the norm half runs in - `reference` but never against the kernel; - * `is_neox=False`, the interleaved pairing; - * `use_mrope`, `use_gemma`. + mRoPE preconditions. `position_ids` is a contiguous int32 `[3, T]`; the op + rejects a strided view loudly, so an engine's `[3, 1, max_num_tokens]` + position buffer has to be copied, not only reshaped. A HuggingFace + `mrope_section` `[s0, s1, s2]` summing to `rotary_dim / 2` is passed as + `mrope_section1 = s1`, `mrope_section2 = s2`; `s0` is implied. Only the + interleaved layout exists here -- the contiguous-section variant is a + different operator, measured ~500x out when wired in its place. Sections + saturate: at 32 pairs `mrope_section2` 10, 11 and 12 are bit-identical (the + largest `j % 3 == 2` below 32 is 29), while `mrope_section1` 10 moves pair + 31 back to row 0. Three equal rows reduce bitwise to the plain + `use_mrope=False` call on that row, so a text batch cannot see a wrong + split; an image span is where one shows. - Adding a target that wants any of them means adding the cell, not assuming - the path works because the argument exists. + Metadata consumed: none. """ def __call__( 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 29803a4aa2c8..7ec0b60bfe18 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,14 +26,72 @@ #: quant_mode bit for an fp8 KV cache. QM_FP8_KV_CACHE = 128 +#: 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 +#: prefill). Pool sizes are chosen, not dumped: 128 pages of 32 tokens and 4 +#: block-offset rows, larger than any cell's batch. +_QWEN4_DENSE = dict( + input_type=STANDARD, + num_heads=24, + num_kv_heads=2, + head_size=256, + v_head_dim=256, + sinks=False, + window=None, + quant_mode=0, + mask_type=1, + tokens_per_block=32, + num_blocks=128, + max_num_requests=4, + use_paged_context_fmha=True, +) + +#: qwen4_exp's vision tower: the no-KV-cache mode, one context sequence per +#: frame, MHA 16/16 at head_size 72 (hidden 1152 over 16 heads), padding mask. +#: `max_num_requests` is the segment count, as the target passes it. +_QWEN4_VISION = dict( + input_type=STANDARD, + kv_cache="none", + num_heads=16, + num_kv_heads=16, + head_size=72, + v_head_dim=72, + sinks=False, + window=None, + quant_mode=0, + mask_type=0, +) + +#: 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( + input_type=STANDARD, + num_heads=24, + num_kv_heads=4, + head_size=256, + v_head_dim=256, + sinks=False, + window=None, + quant_mode=QM_FP8_KV_CACHE, + kv_scale=1.0, + mask_type=1, + tokens_per_block=32, + num_blocks=128, + max_num_requests=4, +) + class ThopAttention(OpWrapper): """Append this step's K/V to the paged cache, then attend over it. The most load-bearing entry in the catalog and the only one wrapping a - pybind binding rather than a registered torch op. Both shipped targets go - through it, by different paths: gpt-oss on the standard path with attention - sinks and an alternating sliding window, deepseek on MLA's two phases. + pybind binding rather than a registered torch op. Three targets go through + 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. 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 @@ -42,10 +100,10 @@ class ThopAttention(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # The four shapes the shipped targets reach, which is four different code - # paths rather than four sizes. gpt-oss alternates its two window classes - # inside one engine, so both of its cells are live every step; deepseek runs - # context and generation against the same pool. + # gpt-oss alternates its two window classes inside one engine, so both of + # its cells are live every step; deepseek runs context and generation + # against the same pool. qwen4_exp's cells give `cached` (tokens already in + # the pool) and `new` (this call's q rows) per sequence, contexts first. CELLS: tuple[Cell, ...] = ( Cell( why="gpt-oss sliding layer: GQA 64q/8kv d64, sinks, window 128", @@ -99,6 +157,253 @@ class ThopAttention(OpWrapper): quant_mode=QM_FP8_KV_CACHE, ), ), + Cell( + why="qwen4_exp prefill: GQA 24q/2kv d256, bf16 pool, paged context; " + "a fresh 70-row sequence beside 33 rows over a 48-token cached prefix", + spec=dict( + _QWEN4_DENSE, + cached=(0, 48), + new=(70, 33), + num_contexts=2, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="qwen4_exp decode P=1: one row per generation sequence over " + "page-crossing histories (70, 35)", + spec=dict( + _QWEN4_DENSE, + cached=(70, 35), + new=(1, 1), + num_contexts=0, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="qwen4_exp MTP verify P=2 (max_draft_len 1): draft row t attends " + "[0, L-P+t], bottom-right causal within the block", + spec=dict( + _QWEN4_DENSE, + cached=(70, 35), + new=(2, 2), + num_contexts=0, + predicted_tokens_per_seq=2, + ), + ), + Cell( + why="qwen4_exp MTP verify P=3 (max_draft_len 2): a separate decode " + "compile, P reaches the kernel as maxSeqLenQ", + spec=dict( + _QWEN4_DENSE, + cached=(70, 35), + new=(3, 3), + num_contexts=0, + predicted_tokens_per_seq=3, + ), + ), + Cell( + why="qwen4_exp MTP verify P=4 (max_draft_len 3): the target's ceiling", + spec=dict( + _QWEN4_DENSE, + cached=(70, 35), + new=(4, 4), + num_contexts=0, + predicted_tokens_per_seq=4, + ), + ), + 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", + spec=dict( + _QWEN4_DENSE, + cached=(0, 70, 35), + new=(20, 2, 2), + num_contexts=1, + predicted_tokens_per_seq=2, + ), + ), + Cell( + why="qwen4_exp hybrid pool: row 47 of the 48-row two-pool mapping " + "(pool 0, layer-in-pool 11), P=2 verify", + spec=dict( + _QWEN4_DENSE, + pool_layout="hybrid48", + local_layer_idx=47, + cached=(64, 25), + new=(2, 2), + num_contexts=0, + predicted_tokens_per_seq=2, + ), + ), + Cell( + why="48-layer single pool, identity mapping: row 47 = layer-in-pool 47, " + "24 rows over a 40-token cached prefix", + spec=dict( + _QWEN4_DENSE, + pool_layout="identity48", + local_layer_idx=47, + cached=(40,), + new=(24,), + num_contexts=1, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="qwen4_exp vision tower: no KV cache, MHA 16/16 d72, padding mask, " + "frames of 27/9/41 rows each attending only to itself", + spec=dict(_QWEN4_VISION, seg_lens=(27, 9, 41)), + ), + Cell( + why="encoder-only, d72 causal: measured, not what the tower passes", + spec=dict(_QWEN4_VISION, mask_type=1, seg_lens=(27, 9, 41)), + ), + Cell( + why="encoder-only, d128 padding: the tower's other admitted width", + spec=dict(_QWEN4_VISION, head_size=128, v_head_dim=128, seg_lens=(27, 9, 41)), + ), + Cell( + why="encoder-only, d128 causal", + spec=dict( + _QWEN4_VISION, head_size=128, v_head_dim=128, mask_type=1, seg_lens=(27, 9, 41) + ), + ), + Cell( + why="encoder-only, a single one-token frame", + spec=dict(_QWEN4_VISION, seg_lens=(1,)), + ), + Cell( + why="mmmu serving shape: one 5580-row frame, the 62x90 patch grid", + spec=dict(_QWEN4_VISION, seg_lens=(5580,)), + ), + Cell( + why="mmmu serving shape: one 8160-row frame, the largest single-segment call", + spec=dict(_QWEN4_VISION, seg_lens=(8160,)), + ), + Cell( + why="mmmu serving shape: 2204 rows in four frames", + spec=dict(_QWEN4_VISION, seg_lens=(728, 728, 460, 288)), + ), + Cell( + why="mmmu serving shape: 1892 rows in five frames, the most segments per call", + spec=dict(_QWEN4_VISION, seg_lens=(480, 480, 400, 352, 180)), + ), + Cell( + why="mmmu serving shape: 9008 rows in two frames", + spec=dict(_QWEN4_VISION, seg_lens=(8000, 1008)), + ), + Cell( + why="mmmu serving shape: 22144 rows in two frames, the largest call; its " + "12544-row frame is the longest segment", + spec=dict(_QWEN4_VISION, seg_lens=(12544, 9600)), + ), + Cell( + why="fp8 pool packed prefill, GQA 24q/4kv d256: the FMHA reads the " + "bf16 q rows, not the pool", + spec=dict( + _FP8_24_4_256, + use_paged_context_fmha=False, + cached=(0, 0), + new=(48, 17), + num_contexts=2, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="fp8 pool padding-mask prefill, 24/4/256, context-only batch", + spec=dict( + _FP8_24_4_256, + mask_type=0, + use_paged_context_fmha=False, + cached=(0, 0), + new=(31, 9), + num_contexts=2, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="fp8 pool decode, 24/4/256: q, K, V and the softmax probabilities " + "are e4m3 operands; histories 200 and 48 span pages and KV tiles", + spec=dict( + _FP8_24_4_256, + use_paged_context_fmha=False, + cached=(200, 48), + new=(1, 1), + num_contexts=0, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="fp8 pool paged-context prefill, 24/4/256: the FMHA reads the " + "pool, so its operands are e4m3", + spec=dict( + _FP8_24_4_256, + use_paged_context_fmha=True, + cached=(0, 48), + new=(48, 33), + num_contexts=2, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="fp8 pool paged-context prefill at GQA 32q/8kv d128", + spec=dict( + _FP8_24_4_256, + num_heads=32, + num_kv_heads=8, + head_size=128, + v_head_dim=128, + use_paged_context_fmha=True, + cached=(0, 48), + new=(48, 33), + num_contexts=2, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="fp8 pool MTP verify P=2, 24/4/256", + spec=dict( + _FP8_24_4_256, + use_paged_context_fmha=True, + cached=(70, 35), + new=(2, 2), + num_contexts=0, + predicted_tokens_per_seq=2, + ), + ), + Cell( + why="fp8 pool MTP verify P=3, 24/4/256", + spec=dict( + _FP8_24_4_256, + use_paged_context_fmha=True, + cached=(70, 35), + new=(3, 3), + num_contexts=0, + predicted_tokens_per_seq=3, + ), + ), + Cell( + why="fp8 pool MTP verify P=4, 24/4/256", + spec=dict( + _FP8_24_4_256, + use_paged_context_fmha=True, + cached=(70, 35), + new=(4, 4), + num_contexts=0, + predicted_tokens_per_seq=4, + ), + ), + Cell( + why="fp8 pool mixed step, 24/4/256: a 20-row context ahead of two P=3 " + "blocks over histories 31 and 64", + spec=dict( + _FP8_24_4_256, + use_paged_context_fmha=True, + cached=(0, 31, 64), + new=(20, 3, 3), + num_contexts=1, + predicted_tokens_per_seq=3, + ), + ), ) note = """ @@ -106,7 +411,13 @@ class ThopAttention(OpWrapper): sliding-window and full attention, and KVCacheManagerV2 gives each window class its own layer group, so its mapping carries two pool ids; the op reads the pool column itself and this entry passes the mapping through untouched. - A third pool is outside what was measured. + A third pool is outside what was measured. qwen4_exp's hybrid tower is the + other two-pool shape: a 48-row mapping whose full-attention rows 3, 7, ..., + 47 sit on pool 0 at layer-in-pool 0..11 and whose 36 recurrent rows name + pool 1. `local_layer_idx` is the mapping *row* -- the layer's + `layer_offsets` entry -- and the pool-base shift comes from the row's + layer column; an index computed by counting attention layers addresses a + sibling's slabs and returns plausible output. `reference` covers the masked-softmax half. The append half -- that this step's K/V land in the right pool slots -- is what the paged-cache entries @@ -116,6 +427,45 @@ class its own layer group, so its mapping carries two pool ids; the op reads The softmax statistics the MLA chunked-prefill path writes are a second output; `softmax_stats` states them separately so `reference` stays one thing. + + predicted_tokens_per_seq (P) is the q-row count of every generation + sequence, token-major within it; the append writes all P and + `sequence_length` must count them. P alone makes the within-block mask + bottom-right causal: every `spec_decoding_*` tensor stays None and + `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. + + 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. + 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. + 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. + + The sparse / skip-softmax arguments are bitwise inert in the standard + configuration (nine off-default settings at 24/2/256): passing them does + not buy sparse attention. + + The kv head count moves the compiled decode kernel at fixed Hq and D + (24/2/256 runs a VarSeqQ16 variant, 24/4/256 VarSeqQ8), so a geometry that + differs from a cell only in `num_kv_heads` is not covered by it. + + Generation kernels JIT through NVRTC and find their headers with `pip show + tensorrt_llm`: a `pip` on PATH that cannot see the install fails every + generation call, and the first failure poisons the process ("Jitify fatal + error: Uninitialized"). """ def __call__( @@ -378,6 +728,7 @@ def reference( window: Optional[int] = None, e4m3_inputs: bool = False, out_scale: float = 1.0, + e4m3_probabilities: bool = False, ) -> torch.Tensor: """Masked softmax attention over explicit per-sequence K/V, in fp32. @@ -389,7 +740,9 @@ def reference( Causal is bottom-right aligned -- query `i` of a sequence sits at absolute position `kv_len - q_len + i`, which is what makes a decode - step's single query attend to everything before it. + step's single query attend to everything before it, and a P-token draft + block's row `t` attend `[0, kv_len - P + t]`. With no KV cache, `k` and + `v` are the call's own packed rows and `q_lens == kv_lens`. `attention_sinks` adds one per-head logit to the softmax *denominator* and drops it from the numerator, so the attention weights sum to less @@ -404,6 +757,17 @@ def reference( `e4m3_inputs` rounds q, k and v through e4m3, which is what the fp8 path quantizes them to; `out_scale` is the second of that path's two dequantization factors, the first riding in `softmax_scale`. + + `e4m3_probabilities` rounds the softmax probabilities too, the value + product's other operand on the standard configuration's fp8 kernels: + `exp(score - row max)` scaled by e4m3's 448, rounded, unscaled, over a + row sum kept in fp32. Measured as what those kernels compute: 0.22 bf16 + ulp RMS where the KV range fits one tile, 2.4 on a 200-token decode + (the kernel rounds against each tile's running max, this against the + row's), against 5-6 with the probabilities left in fp32 and 8 with them + rounded unscaled. At kv scale 2.0 it is bitwise exact on a one-tile + decode; at 1.5 nothing tried fits (4.2 ulp RMS at best). Not modelled + with sinks. """ rep = num_heads // num_kv_heads total_q = sum(q_lens) @@ -433,7 +797,13 @@ def reference( keep = self._mask(q_len, kv_len, causal, window, q.device) scores = scores.masked_fill(~keep.unsqueeze(0), float("-inf")) - if attention_sinks is None: + if e4m3_probabilities: + assert attention_sinks is None, "e4m3 probabilities are not modelled with sinks" + weights = torch.exp(scores - scores.max(dim=-1, keepdim=True).values) + e4m3_max = torch.finfo(torch.float8_e4m3fn).max + rounded = (weights * e4m3_max).to(torch.float8_e4m3fn).float() / e4m3_max + probs = rounded / weights.sum(dim=-1, keepdim=True) + elif attention_sinks is None: probs = torch.softmax(scores, dim=-1) else: sink = attention_sinks.float().view(num_heads, 1) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/__init__.py new file mode 100644 index 000000000000..5e898b4e1168 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Mixture-of-experts entries.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/moe/fp4_block_scale_moe_runner.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fp4_block_scale_moe_runner.py similarity index 83% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/moe/fp4_block_scale_moe_runner.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fp4_block_scale_moe_runner.py index d07905a94666..bfa0b0644755 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/moe/fp4_block_scale_moe_runner.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/moe/fp4_block_scale_moe_runner.py @@ -99,8 +99,9 @@ class Fp4BlockScaleMoeRunner(OpWrapper): The caller owns top-k: `topk_ids` carries global expert ids and `topk_weights` the combine weights, neither renormalized here. This rank answers for `[local_expert_offset, local_expert_offset + local_num_experts)` - and contributes zero for every other id, so an expert-parallel split is four - calls whose outputs sum -- which is what the dep4 target does. + and contributes zero for every other id, so an expert-parallel split is N + calls whose outputs sum -- which is what the deepseek dep4 target does; + qwen4_exp makes one call over its whole expert stack. With `do_finalize=True` the result is a one-element list holding the combined `[T, H]` bf16 rows, or an empty tensor when `output` was given. @@ -116,6 +117,22 @@ class Fp4BlockScaleMoeRunner(OpWrapper): # single-rank one. _R1 = dict(hidden=7168, intermediate=2048, num_experts=256, top_k=8) + # qwen4_exp's (Qwen3.8-Flash-Next) routed geometry: many small experts, a + # top_k past 8, an intermediate width of 640, all 512 experts local. Its + # router renormalizes the top-10 softmax weights to sum to 1 before the + # call (`routing="renormalized"`, a test axis rather than an op argument; + # the R1 cells draw uniform [0, 1) weights). Nothing is padded: + # H % 256, 2I % 128, H/16 % 4 and I/16 % 4 are all zero. + _QWEN4_EXP = dict( + hidden=2560, + intermediate=640, + num_experts=512, + top_k=10, + local_expert_offset=0, + local_num_experts=512, + routing="renormalized", + ) + CELLS: tuple[Cell, ...] = ( Cell( why="dep4's own call: the 64-wide expert window at offset 0, R1 geometry", @@ -129,6 +146,25 @@ class Fp4BlockScaleMoeRunner(OpWrapper): why="the single-rank shape: all 256 experts local, and one token, the decode row count", spec=dict(tokens=1, local_expert_offset=0, local_num_experts=256, **_R1), ), + Cell( + why="qwen4_exp decode: one token over 512 experts, top-10 renormalized, H 2560 / I 640", + spec=dict(tokens=1, **_QWEN4_EXP), + ), + Cell( + why="qwen4_exp at 4 rows: the last count the pre-routed path's static-block " + "permutation kernel takes", + spec=dict(tokens=4, **_QWEN4_EXP), + ), + Cell( + why="qwen4_exp at 5 rows: the first count on the single-cluster permutation " + "kernel (the dynamic-block one stops at 256 experts)", + spec=dict(tokens=5, **_QWEN4_EXP), + ), + Cell( + why="qwen4_exp prefill: 8192 rows, the single-cluster kernel's cap and the top " + "of the certified column", + spec=dict(tokens=8192, **_QWEN4_EXP), + ), ) note = """ @@ -137,12 +173,29 @@ class Fp4BlockScaleMoeRunner(OpWrapper): `topk_group` and `routed_scaling_factor` are inert -- they are passed because the signature takes them, not because they do anything here. + That holds at 512 experts / top-10 too, where it would matter if it were + read: which `(num_experts, top_k)` pairs the routing kernels are compiled + for is a per-policy tier list, and a pair with no tier is a logged error, + not a raise. The pre-routed path substitutes its own preprocess/postprocess + pair before dispatch, and `routing_method_type` 0, 2, 4, 5, 6 were measured + bitwise identical to 1 there. It picks its permutation kernel by row count + instead: a static-block kernel up to 4 rows, a dynamic-block one up to 16 + rows only at <= 256 experts, a single-cluster one up to 8192 rows. + `act_type` 0 (SwiGlu) only. `gemm1_bias`, `gemm1_alpha`, `gemm1_beta`, - `gemm1_clamp_limit` and `gemm2_bias` are `None`. + `gemm1_clamp_limit` and `gemm2_bias` are `None`. `hidden_states_scale` is + the linear (unswizzled) e4m3 block-scale layout, one scale per 16 elements + of each row. + + At 512 experts only the whole stack is certified: no expert-parallel window + of it, `do_finalize=False` or a caller-supplied `output` was measured there. `tune_max_num_tokens` is an autotuner bucket hint and was measured inert; it is not a token count. `T` is `hidden_states.shape[0]`, and 8192 is the - top of the certified column. + top of the certified column. A warm autotuner cache was measured bitwise + inert against a cold one at the R1 stack, its windows and the 512-expert + stack; `local_num_experts` separates the tuner's cache entries, so each + layout warms its own. The caller owns the weight layout entirely: FC1 and FC2 weights are pre-shuffled and their scales additionally 128x4 swizzled. Nothing here diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/__init__.py new file mode 100644 index 000000000000..7d0c28eefd77 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Quantization entries.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/quantization/fp4_quantize.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py similarity index 88% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/quantization/fp4_quantize.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py index 69e43cdc819b..b0c0e864d886 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/quantization/fp4_quantize.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/quantization/fp4_quantize.py @@ -47,10 +47,14 @@ class Fp4Quantize(OpWrapper): ARCHS = frozenset({Arch.SM_103}) - # deepseek-r1 is the only caller and reaches this op from both layouts: the - # swizzled one feeding nvfp4_gemm in the dense and shared-expert MLPs, the - # linear one feeding the routed-expert path. Widths are its hidden size and - # its intermediate. + # deepseek-r1 reaches this op from both layouts: the swizzled one feeding + # nvfp4_gemm in the dense and shared-expert MLPs, the linear one feeding the + # routed-expert path, at its hidden size and its intermediate. qwen3.8 + # (qwen4_exp) calls it only for its routed experts: linear scales at hidden + # 2560. A linear buffer has no row padding, so at a row count that is not a + # multiple of 128 it is shorter than the swizzled one; at a multiple of 128 + # the two are the same length and only the cells' byte order tells them + # apart. CELLS: tuple[Cell, ...] = ( Cell( why="decode, hidden 7168, swizzled scales -- the dense MLP's input", @@ -72,6 +76,18 @@ class Fp4Quantize(OpWrapper): why="a row count that is not a multiple of the swizzle's 128-row tile", spec=dict(shape=(37, 7168), dtype=torch.bfloat16, vec=16, swizzled=True), ), + Cell( + why="qwen3.8 decode: 1 token, hidden 2560, linear -- its routed-expert input", + spec=dict(shape=(1, 2560), dtype=torch.bfloat16, vec=16, swizzled=False), + ), + Cell( + why="qwen3.8, 7 rows at hidden 2560, linear: several rows, no row padding", + spec=dict(shape=(7, 2560), dtype=torch.bfloat16, vec=16, swizzled=False), + ), + Cell( + why="qwen3.8 prefill at max_num_tokens 8192, hidden 2560, linear", + spec=dict(shape=(8192, 2560), dtype=torch.bfloat16, vec=16, swizzled=False), + ), ) note = """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/modeling.py index f84c25f519f1..8ae250b7afe0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/modeling.py @@ -27,12 +27,18 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.reducescatter import reducescatter from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.nvfp4_gemm import nvfp4_gemm +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.norm.flashinfer_fused_add_rmsnorm import ( # noqa: E501 flashinfer_fused_add_rmsnorm, ) from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( flashinfer_rmsnorm, ) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) 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 @@ -46,10 +52,8 @@ ) from .ops.attention.mla_rope_generation import mla_rope_generation from .ops.gemm.bmm_out import bmm_out -from .ops.moe.fp4_block_scale_moe_runner import fp4_block_scale_moe_runner # noqa: E501 from .ops.moe.fused_moe import fused_moe from .ops.moe.noaux_tc_op import noaux_tc_op -from .ops.quantization.fp4_quantize import fp4_quantize _CACHED_CTX_FIELDS = ( "enable_context_mla_with_cached_kv", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/modeling.py index 1e6e28db6868..01124bdf0a4c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/modeling.py @@ -17,6 +17,9 @@ from transformers import PretrainedConfig from tensorrt_llm._torch._experimental.modeling_v2._router_index import step_contract_enabled +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import ( thop_attention, ) @@ -33,7 +36,6 @@ from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM from . import weights as _weights -from .ops.attention.fused_qk_norm_rope import fused_qk_norm_rope from .ops.moe.mxe4m3_mxe2m1_block_scale_moe_runner import ( # noqa: E501 ACT_TYPE_SWIGLU, ROUTING_METHOD_RENORMALIZE, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/__init__.py new file mode 100644 index 000000000000..c0df4d8f5d62 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Qwen4ExpForConditionalGeneration targets.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/__init__.py new file mode 100644 index 000000000000..8fd2a7f0c45b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Qwen3.8-Flash-Next-NVFP4 / sm_103 / tp1.""" 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 new file mode 100644 index 000000000000..f71b25463bdd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py @@ -0,0 +1,4025 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ModelingV2 target: Qwen3.8-Flash-Next-NVFP4 / sm_103 / tp1 — self-contained modeling code. + +The core holds the static structure (weights, derived operands, per-manager pool +views); `PrefillTarget` (any batch with context rows) and `DecodeTarget` +(generation-only) run a step over it, passing every argument to every op at the +call site. + +The checkpoint is the composite Qwen4-Exp VLM export. The MTP module is built only +when the engine has a ``speculative_config`` (else ``mtp.*`` stays in +``weights.py``'s non-load set). The vision tower is always built; a text-only +request never reaches it. Every speculation-dependent change is gated on a host +int that is 1 when speculation is off. + +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. +* **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. +* **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. + +Integration facts this file depends on: + +1. The engine rebuilds the composite config and normalizes its text tower + (``config_utils.load_pretrained_config``). The shell adopts that text config as + ``model_config.pretrained_config`` before building the core, because the + recurrent-state and PLE pool geometry are read from it. +2. ``position_ids`` arrives in the mRoPE layout (``[3, 1, T]``) even for text-only + batches, because the normalized text config carries ``rope_scaling['type'] == + 'mrope'``. The main and MTP rope read all three rows; the QSA index plane takes + coordinate 0 only (see ``_QSA_MROPE_SECTION``). +3. The engine owns the multimodal input processor (resolved off the checkpoint's + ``architectures[0]``; modeling_v2 has no hook), so it is the built-in + ``Qwen3VLInputProcessorBase``. This target asserts every ``multimodal_data`` key + where it is read. + +**Host offload.** By default the PLE table lives in host memory the device +addresses (`weights._host_resident_empty`), and the gather reads it unchanged; +``TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD=0`` restores the resident declaration. Output +is identical, so only the device-memory delta sees it. The gather stays +device-driven because row indices are produced on device inside a captured decode +forward, which cannot express a host readback. +""" + +import math +import os +from typing import NamedTuple + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tensorrt_llm._torch._experimental.modeling_v2._router_index import step_contract_enabled +from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.flashinfer_silu_and_mul import ( # noqa: E501 + flashinfer_silu_and_mul, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import ( + 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.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.speculative import get_spec_worker, should_use_separate_draft_kv_cache +from tensorrt_llm._torch.utils import is_gdn_replay_enabled + +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.attention.qsa_expand_block_indices import qsa_expand_block_indices +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_sparse_gqa import qsa_paged_sparse_gqa +from .ops.attention.qsa_prefill_compress import qsa_prefill_compress +from .ops.attention.qsa_topk_select import qsa_topk_select +from .ops.attention.vision_rope_2d import vision_rope_2d +from .ops.embedding.ngram_hash import ngram_hash +from .ops.hc.qwen4_exp_gated_residual import qwen4_exp_gated_residual +from .ops.hc.qwen4_exp_post_inject import qwen4_exp_post_inject +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.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.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 +from .ops.ssm.causal_conv1d_triton_update import causal_conv1d_triton_update +from .ops.ssm.causal_conv1d_update import causal_conv1d_update +from .ops.ssm.chunk_gated_delta_rule import chunk_gated_delta_rule +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.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_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.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 +from .ops.vision.pos_embed_interpolate import pos_embed_interpolate + +# The GPU architecture this target is certified for; a direct instantiation could +# bypass routing. +_SM = (10, 3) + +# QSA backend metadata buffers rebuilt by prepare(). Their absence means the engine +# did not select the QSA metadata and the index plane would never be written. +_QSA_FIELDS = ( + "qsa_block_table", + "qsa_req_idx_per_token", + "qsa_logical_positions", + "qsa_sequence_lengths", + "qsa_visible_blocks", + "qsa_topk_indices", + "qsa_topk_row_starts", +) + +# Fields read off `attn_metadata.mamba_metadata`. Absence is silent: a smoke run +# passes with no recurrent state at all. +_MAMBA_FIELDS = ( + "state_indices", + "has_initial_states", + "query_start_loc", + "query_start_loc_long", +) + + +# 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. +_CALL_CONSTANTS = dict( + output_sf=None, + k=None, + v=None, + is_fused_qkv=True, + update_kv_cache=True, + attention_input_type=0, + is_mla_enable=False, + mask_type=1, + q_scaling=1.0, + attention_sinks=None, + # 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, + rotary_cos_sin=None, + rope_dim=0, + rope_base=10000.0, + rope_scale_type=0, + rope_scale=1.0, + rope_short_m_scale=1.0, + rope_long_m_scale=1.0, + rope_max_positions=1024, + rope_original_max_positions=1024, + out_scale=None, + latent_cache=None, + q_pe=None, + q_lora_rank=None, + kv_lora_rank=None, + qk_nope_head_dim=None, + qk_rope_head_dim=None, + v_head_dim=None, + rope_append=None, + chunked_prefill_buffer_batch_size=1, + attention_chunk_size=None, + softmax_stats_tensor=None, + sparse_kv_indices=None, + sparse_kv_offsets=None, + sparse_attn_indices=None, + sparse_attn_offsets=None, + sparse_attn_indices_block_size=0, + mrope_rotary_cos_sin=None, + mrope_position_deltas=None, + helix_position_offsets=None, + helix_is_inactive_rank=None, + cross_kv=None, + relative_attention_bias=None, + relative_attention_max_distance=0, + # MLA-only surface, unused here. + kv_norm_weight=None, + kv_norm_eps=1e-6, + skip_correction_threshold=0.0, +) + +# ------------------------------------------------------------ vision ------ + +#: Tower rope base (`vision_config` has no rope field); unrelated to the text 1e7. +_VISION_ROPE_THETA = 10000.0 + +#: LayerNorm epsilon of `norm1`, `norm2`, `merger.norm`: a literal in the reference. +_VISION_NORM_EPS = 1e-6 + +#: `mask_type` 0 = padding mask: bidirectional within a batch segment. Causal (1) +#: would make a patch depend on raster order. +_VISION_MASK_PADDING = 0 + +# Per-call constants of the tower's attention: the encoder-only mode of +# `attention/thop_attention`. Omitting `kv_cache_block_offsets`, +# `host_kv_cache_pool_pointers` and `host_kv_cache_pool_mapping` selects that mode. +# `update_kv_cache=False` is rejected ("KV cache update cannot be disabled now"), +# so it stays True. +_VISION_CALL_CONSTANTS = dict( + k=None, + v=None, + output_sf=None, + kv_cache_block_offsets=None, + host_kv_cache_pool_pointers=None, + host_kv_cache_pool_mapping=None, + tokens_per_block=None, + use_paged_context_fmha=False, + is_fused_qkv=True, + update_kv_cache=True, + predicted_tokens_per_seq=1, + # No pool, so no layer to address. + local_layer_idx=0, + max_context_q_len_override=None, + cache_indirection=None, + kv_scale_orig_quant=None, + kv_scale_quant_orig=None, + out_scale=None, + quant_mode=0, + beam_width=1, + mask_type=_VISION_MASK_PADDING, + q_scaling=1.0, + # Rotation is applied outside (`ops/attention/vision_rope_2d`); rope_* inert. + position_embedding_type=0, + rotary_inv_freq=None, + rotary_cos_sin=None, + rope_dim=0, + rope_base=10000.0, + rope_scale_type=0, + rope_scale=1.0, + rope_short_m_scale=1.0, + rope_long_m_scale=1.0, + rope_max_positions=1024, + rope_original_max_positions=1024, + attention_input_type=0, + is_mla_enable=False, + chunked_prefill_buffer_batch_size=1, + attention_sinks=None, + block_ids_per_seq=None, + latent_cache=None, + q_pe=None, + q_lora_rank=None, + kv_lora_rank=None, + qk_nope_head_dim=None, + qk_rope_head_dim=None, + v_head_dim=None, + rope_append=None, + attention_chunk_size=None, + softmax_stats_tensor=None, + is_spec_decoding_enabled=False, + use_spec_decoding=False, + is_spec_dec_tree=False, + spec_decoding_generation_lengths=None, + spec_decoding_position_offsets_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, + mrope_rotary_cos_sin=None, + mrope_position_deltas=None, + helix_position_offsets=None, + helix_is_inactive_rank=None, + is_cross=False, + cross_kv=None, + relative_attention_bias=None, + relative_attention_max_distance=0, + kv_norm_weight=None, + kv_norm_eps=1e-6, + skip_correction_threshold=0.0, +) + +#: Head geometries the encoder-only mode is certified at: MHA, `head_size` 72 or +#: 128, context rows only. `head_size` 72 is context-phase only ("unsupported +#: head_size 72" in generation). +_VISION_CERTIFIED_HEADS = {(16, 16, 72), (16, 16, 128)} + +#: Head widths `ops/attention/vision_rope_2d` is certified at. +_VISION_ROPE_HEAD_DIMS = (40, 64, 72, 128) + +# 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 +# kernel launch. +_ROUTING_METHOD_RENORM = 1 +_ACT_TYPE_SWIGLU = 0 + +# Skip sentinel the conv ops and paged short conv compare `cache_indices` against +# (the runtime's mamba `PAD_SLOT_ID`). Never a substitution: the scheduler's slot +# vector gives every row a real slot, and `ops/ple/ngram_token_history` reads its +# carry via `index_select`, so a -1 there is a device-side assert that poisons the +# process. +_PAD_SLOT_ID = -1 + +# Gated-DeltaNet gating form: the only pair `fused_sigmoid_gating_delta_rule_update` +# is certified at; matches `fused_gdn_post_conv`'s default so prefill and decode +# compute the same `g`. +_SOFTPLUS_BETA = 1.0 +_SOFTPLUS_THRESHOLD = 20.0 + +# Gated-DeltaNet Q/K L2-norm epsilon: a literal, not `rms_norm_eps`. +_L2_NORM_EPS = 1e-6 + +# PLE gate magnitude floor before its signed square root (a model constant). +_PLE_CLAMP_MIN = 1e-6 + +# --------------------------------------------------------------- host offload +# +# Where the PLE n-gram table lives (module docstring, *Host offload*). No engine +# arg owns this, hence an env var, read once at construction (the table is +# allocated once, during `load_weights`). Default on: 47.68 GiB otherwise comes +# out of the KV pool. +_PLE_HOST_OFFLOAD_ENV = "TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD" + +#: Spellings of "off"; anything else, including empty, leaves offload on. +_PLE_HOST_OFFLOAD_OFF = frozenset({"0", "false", "no", "off"}) + + +def _ple_host_offload_enabled() -> bool: + """Whether the n-gram table is allocated in host memory.""" + return os.environ.get(_PLE_HOST_OFFLOAD_ENV, "1").strip().lower() not in _PLE_HOST_OFFLOAD_OFF + + +# QSA's fixed constants (backend `constants.py`): one index KV head, three position +# axes, index head exactly twice the rotary width (128 == 2 * 64). +_QSA_INDEX_KV_HEADS = 1 +_QSA_POSITION_AXES = 3 +_QSA_INDEX_HEAD_TO_ROTARY_RATIO = 2 + +# The QSA index plane takes coordinate 0 alone, a deliberate departure from the +# checkpoint. `qsa_prefill_compress` (key side, `mrope_section`) and +# `fused_qk_norm_rope` on `index_scratch` (query side, `use_mrope`) must agree: +# both read the three coordinates or neither, else scores mix two position +# encodings. Both-on is certified (`mrope_section` (11, 11, 10)); None is a +# modeling choice. No cost on text; on an image batch the index plane sees a frame +# at one position, changing block selection only above `dense_seq_len_threshold`. +# To flip, set the config's `(s0, s1, s2)` and `_QSA_INDEX_USE_MROPE = True`. +_QSA_MROPE_SECTION = None +_QSA_INDEX_USE_MROPE = False + +# ---------------------------------------------------------------- MTP ------ + +#: The checkpoint's MTP module count (`mtp_num_hidden_layers`), replayed +#: `max_draft_len` times; NOT the draft length. +_MTP_NUM_LAYERS = 1 + +#: How the `[T, hidden]` embedding branch enters the `[T, hc_count * hidden]` MTP +#: stream; nothing in the checkpoint or config says. True: added into all four +#: streams; False: stream 0 only. A wrong choice yields bit-correct text with every +#: draft rejected; only the acceptance gate can tell. +_MTP_EMBED_BLOCK_ALL_STREAMS = True + +#: FP8 block-scale block edge, on both weight-scale axes and the activation row. +_FP8_BLOCK = _weights.FP8_BLOCK + +#: 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 + + +# ----------------------------------------------- KV cache manager binding --- +# +# The ops are plain functions over raw tensors, and the pool views taken from the +# KV cache manager are `TensorWrapper`s over raw device addresses that do not keep +# memory alive, so a view is only valid for the manager that produced it. `_Bound` +# is everything derived from one manager; the core replaces it as a unit +# (`_bind_cache_manager`) and the hot path pays one identity comparison. +# +# The `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. + + +class _PoolGeometry(NamedTuple): + """The pool shapes this target expects, taken from its config and the metadata.""" + + tokens_per_block: int + heads_kv: int + head_dim: int + index_kv_heads: int + index_head_dim: int + conv_dim: int + conv_kernel: int + num_v_heads: int + head_v_dim: int + head_k_dim: int + stream: int + ple_state_len: int + ngram_context_len: int + dtype: torch.dtype # the activation dtype + + +class _Bound(NamedTuple): + """Everything derived from one KV cache manager.""" + + cm: object + pool_idx: dict # attention layer -> pool offset (what thop_attention addresses) + index_cache: dict # attention layer -> QSA index plane + position_cache: torch.Tensor # shared QSA position plane + kv: dict # attention layer -> (K, V) views + 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. + int_idx: torch.Tensor | None = None + int_conv: dict | None = None + int_ssm: dict | None = None + + +def _require(cm, names, what: str) -> None: + for name in names: + assert hasattr(cm, name), ( + f"the KV cache manager {type(cm).__name__} does not expose {name}; this target {what}" + ) + + +def _check_pool(t, what, lead, tail, dtype=None) -> None: + """`t` has `lead` leading dims of any size, then exactly `tail`, at `dtype`.""" + assert t is not None, f"{what}: the runtime did not allocate it" + expected = ["*"] * lead + [str(n) for n in tail] + assert t.dim() == lead + len(tail) and tuple(t.shape[lead:]) == tuple(tail), ( + f"{what} is {tuple(t.shape)}, expected [{', '.join(expected)}]" + ) + assert dtype is None or t.dtype == dtype, f"{what} is {t.dtype}, expected {dtype}" + + +def check_position_plane(g: _PoolGeometry, position) -> None: + shape = (g.tokens_per_block, _QSA_POSITION_AXES) + _check_pool(position, "QSA position plane", 1, shape, torch.int32) + + +def check_attention_pools(g: _PoolGeometry, layer: int, index_cache, pool) -> None: + """QSA index plane and paged K/V for one attention layer.""" + index_shape = (g.tokens_per_block, g.index_kv_heads, g.index_head_dim) + _check_pool(index_cache, f"layer {layer}: QSA index plane", 1, index_shape, g.dtype) + # The sparse path needs adjacent HND K/V role planes: [pages, 2, H, block, D]. + kv_shape = (g.heads_kv, g.tokens_per_block, g.head_dim) + _check_pool(pool, f"layer {layer}: paged K/V", 2, kv_shape) + assert pool.shape[1] == 2, f"layer {layer}: paged K/V has {pool.shape[1]} role planes, not 2" + + +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)) + + +def check_ple_pools(g: _PoolGeometry, ple_conv, ple_ngram, conv) -> None: + _check_pool(ple_conv, "PLE conv state", 1, (g.stream, g.ple_state_len), g.dtype) + # The id-history entry writes int64 into it; an int32 pool raises. + _check_pool(ple_ngram, "PLE n-gram context pool", 1, (g.ngram_context_len,), torch.int64) + # The forward addresses all three with one scheduler slot vector. + assert ple_conv.shape[0] == conv.shape[0] == ple_ngram.shape[0], ( + "the PLE and Gated-DeltaNet pools must share one slot space" + ) + + +def check_step_buffers(g: _PoolGeometry, int_conv, 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" + ) + _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" + ) + + +class Qwen4ExpTp1ModelingV2(DecoderModel): + def __init__(self, model_config: ModelConfig, composite_config: PretrainedConfig): + """`composite_config` is passed in because the text sub-config adopted as + `model_config.pretrained_config` carries neither the vision half nor the + media token ids.""" + super().__init__(model_config) + cfg = model_config.pretrained_config + assert cfg is not None + + # This target IS this geometry — assert, never adapt. + assert torch.cuda.get_device_capability() == _SM, ( + f"target certified on sm_{_SM[0]}{_SM[1]}, running on " + f"sm_{''.join(map(str, torch.cuda.get_device_capability()))}" + ) + assert model_config.mapping.tp_size == 1, "tp1 target" + assert model_config.mapping.pp_size == 1, "tp1 target" + assert not model_config.mapping.enable_attention_dp, "tp1 target" + + # `cfg` is the engine-normalized Qwen4ExpTextConfig (pool geometry is read + # from it). + assert getattr(cfg, "model_type", None) == "qwen4_exp_text", ( + f"the core must be built on the normalized text config; got " + f"model_type {getattr(cfg, 'model_type', None)!r}" + ) + + dt = model_config.torch_dtype + assert dt == torch.bfloat16, f"bf16 target, engine resolved {dt}" + assert cfg.torch_dtype == torch.bfloat16, cfg.torch_dtype + assert not cfg.tie_word_embeddings, "untied lm_head" + assert not cfg.attention_bias, "no q/k/v/o bias anywhere" + assert cfg.hidden_act == "silu", "SwiGLU over a silu gate" + + self.num_layers = cfg.num_hidden_layers + self.hidden = cfg.hidden_size + self.vocab = cfg.vocab_size + self.eps = cfg.rms_norm_eps + + # `layer_types` is authoritative; `full_attention_interval` is the runtime + # helper's fallback, so both must agree. + layer_types = list(cfg.layer_types) + assert len(layer_types) == self.num_layers, (len(layer_types), self.num_layers) + assert set(layer_types) <= {"full_attention", "linear_attention"}, set(layer_types) + self.is_attn = [t == "full_attention" for t in layer_types] + interval = cfg.full_attention_interval + assert self.is_attn == [i % interval == interval - 1 for i in range(self.num_layers)], ( + "layer_types disagrees with full_attention_interval" + ) + self.attn_layers = [i for i, a in enumerate(self.is_attn) if a] + self.gdn_layers = [i for i, a in enumerate(self.is_attn) if not a] + assert self.attn_layers and self.gdn_layers, "this target is the hybrid assembly" + + # The checkpoint picks the speculative mode (`MTP_EAGLE_ONE_MODEL`, one module + # replayed `max_draft_len` times); the module is built only when the engine + # asked for speculation. + spec_config = getattr(model_config, "spec_config", None) + self.mtp_enabled = spec_config is not None + declared = getattr(cfg, "mtp_num_hidden_layers", None) + assert declared in (None, _MTP_NUM_LAYERS), ( + f"this target implements the one-layer MTP module this checkpoint " + f"ships; the text config declares mtp_num_hidden_layers {declared}" + ) + # The module shares the trunk's embedding and the shell's lm_head. + assert not getattr(cfg, "mtp_use_dedicated_embeddings", False), ( + "this checkpoint's MTP module shares the decoder's embedding " + "(`mtp_use_dedicated_embeddings: false`) and ships no " + "`mtp.embed_tokens` tensor; a dedicated-embedding module is a " + "different assembly" + ) + # The nested `mtp` block, when kept: one `full_attention` layer at the trunk's + # rope base. + mtp_block = getattr(cfg, "mtp", None) + self.mtp_theta = None + if mtp_block is not None: + assert isinstance(mtp_block, dict), ( + f"the text config's `mtp` block is {type(mtp_block).__name__}; " + "this target reads it as the published mapping" + ) + assert mtp_block.get("num_hidden_layers", _MTP_NUM_LAYERS) == _MTP_NUM_LAYERS, ( + f"mtp.num_hidden_layers is {mtp_block.get('num_hidden_layers')}" + ) + assert list(mtp_block.get("layer_types", ["full_attention"])) == ["full_attention"], ( + f"the MTP module's layer_types is {mtp_block.get('layer_types')}; this " + "target implements the full-attention form, which is the only " + "one whose weights the checkpoint carries" + ) + # `null` means the trunk's own output. + assert mtp_block.get("mtp_use_hidden_state_from_layer") is None, ( + f"the MTP module asks for the hidden state of layer " + f"{mtp_block.get('mtp_use_hidden_state_from_layer')}; this " + "target hands it the end of the stack, which is what the " + "runtime passes and all this file can reach" + ) + if "rope_theta" in mtp_block: + self.mtp_theta = float(mtp_block["rope_theta"]) + + # ---- hyper-connections ------------------------------------------ + self.hc_count = cfg.hc_count + self.hc_lowrank = cfg.hc_lowrank + assert self.hc_count >= 2, ( + f"the hyper-connection residual carries hc_count streams; this " + f"config declares {self.hc_count}, which is not a multi-stream " + f"residual and is a different assembly" + ) + self.stream = self.hc_count * self.hidden + + # ---- full attention --------------------------------------------- + self.heads_q = cfg.num_attention_heads + self.heads_kv = cfg.num_key_value_heads + self.head_dim = cfg.head_dim + assert self.head_dim in (64, 128, 256), "fused_qk_norm_rope head_dim set" + assert self.heads_q % self.heads_kv == 0, "thop_attention GQA rule" + self.q_width = self.heads_q * self.head_dim + self.kv_width = self.heads_kv * self.head_dim + self.qkv_width = self.q_width + 2 * self.kv_width + # The per-head interleaved output gate makes q_proj twice as wide as q. No + # config flag exists, so it is derived from the weight (asserted in weights.py); + # the per-head split below would otherwise address the wrong columns. + assert not hasattr(cfg, "attn_output_gate"), ( + "this config gained an attn_output_gate flag; the gate is derived " + "from q_proj's width here, and the two must be reconciled" + ) + + # ---- RoPE -------------------------------------------------------- + # Partial, neox half-split pairs, no scaling; `rope_theta` is inside + # `rope_parameters` (mirrored as `rope_scaling`). + rope = getattr(cfg, "rope_parameters", None) or getattr(cfg, "rope_scaling", None) + assert isinstance(rope, dict), f"rope parameters must be a dict, got {rope!r}" + self.theta = float(rope["rope_theta"]) + assert self.theta > 0.0, "rope theta" + partial = float(rope.get("partial_rotary_factor", cfg.partial_rotary_factor)) + self.rotary_dim = int(self.head_dim * partial) + assert self.rotary_dim % 2 == 0 and 0 < self.rotary_dim <= self.head_dim, ( + f"rotary_dim {self.rotary_dim} must be even and at most head_dim " + f"{self.head_dim}; the op enforces only the evenness half" + ) + # The mRoPE declaration is what makes the engine hand this forward a + # [3, 1, T] position buffer. + self.mrope_declared = str(rope.get("type", rope.get("rope_type"))) == "mrope" + # `fused_qk_norm_rope` implements the interleaved mRoPE layout only (pair `j` + # reads coordinate 1 when `j % 3 == 1 and j < 3*s1`, 2 when `j % 3 == 2 and + # j < 3*s2`, else 0); the contiguous variant is a different operator. A text + # gate cannot see the split. + assert rope.get("mrope_interleaved") is True, ( + f"this target passes the interleaved mRoPE layout, which is the only " + f"one attention/fused_qk_norm_rope implements; the config declares " + f"mrope_interleaved={rope.get('mrope_interleaved')!r}" + ) + section = list(rope["mrope_section"]) + assert len(section) == _QSA_POSITION_AXES, ( + f"mrope_section must carry one entry per position coordinate; got {section}" + ) + rotary_pairs = self.rotary_dim // 2 + assert sum(section) == rotary_pairs, ( + f"mrope_section {section} sums to {sum(section)} but rotary_dim " + f"{self.rotary_dim} makes {rotary_pairs} rotation pairs; the split " + f"partitions pairs, not dimensions" + ) + self.mrope_section1 = int(section[1]) + self.mrope_section2 = int(section[2]) + # Row-1 section is load-bearing here; row-2 saturates (10, 11, 12 are + # bit-identical at 32 pairs). `use_gemma` with `use_mrope` is certified only + # at the geometry asserted below. + assert ( + self.heads_q, + self.heads_kv, + self.head_dim, + self.rotary_dim, + self.theta, + ) == (24, 2, 256, 64, 1e7), ( + f"the main attention rope runs use_gemma together with use_mrope, " + f"which attention/fused_qk_norm_rope certifies only at " + f"(24, 2, 256, 64, 1e7); this config gives " + f"{(self.heads_q, self.heads_kv, self.head_dim, self.rotary_dim, self.theta)}" + ) + # No YaRN or other scaling. + assert "factor" not in rope, ( + f"this target builds a plain (unscaled) rope; rope carries {rope!r}" + ) + self.max_positions = cfg.max_position_embeddings + + # ---- QSA sparse index -------------------------------------------- + self.index_heads = cfg.indexer_n_heads + self.index_kv_heads = cfg.indexer_kv_heads + self.index_head_dim = cfg.indexer_head_dim + self.token_topk = cfg.indexer_budget + self.compress_ratio = cfg.indexer_compress_ratio + assert self.index_kv_heads == _QSA_INDEX_KV_HEADS, ( + f"QSA requires exactly one index KV head; this config declares {self.index_kv_heads}" + ) + assert self.compress_ratio >= 2 and self.token_topk % self.compress_ratio == 0, ( + f"QSA selection geometry: token_topk {self.token_topk} must divide " + f"by compress_ratio {self.compress_ratio}, which must be at least 2" + ) + # Fused index kernels: both powers of two, index head twice the rotary width. + assert self.index_head_dim == self.rotary_dim * _QSA_INDEX_HEAD_TO_ROTARY_RATIO, ( + f"the fused QSA index kernels require index_head_dim " + f"({self.index_head_dim}) == 2 * rotary_dim ({self.rotary_dim})" + ) + assert self.index_head_dim & (self.index_head_dim - 1) == 0, "index_head_dim power of two" + assert self.compress_ratio & (self.compress_ratio - 1) == 0, "compress_ratio power of two" + self.block_topk = self.token_topk // self.compress_ratio + self.expanded_topk = self.token_topk + self.compress_ratio - 1 + # Dense-attention cutoff: QSA resolves max(token_topk, seq_len_threshold); + # checked against the engine's value below because the CUDA-graph runner + # splits its short and long decode families on it. + self.dense_threshold = self.token_topk + + # ---- the engine's sparse configuration, checked at construction ----- + # The index/position side caches, `qsa_*` metadata and sparse kernels all hang + # off `sparse_attention_config`; checking here fails at cold start rather than + # at the first forward. + sparse_cfg = getattr(model_config, "sparse_attention_config", None) + assert getattr(sparse_cfg, "algorithm", None) == "qsa", ( + f"this target is a QSA sparse-attention assembly and the engine " + f"resolved sparse_attention_config={sparse_cfg!r}. QSA is this " + f"architecture's own model default, so an engine that reaches here " + f"without it dropped the default on the way in rather than being " + f"asked for dense attention" + ) + # The engine's own lowering on this core's config: a mismatch with the side + # cache sizing would otherwise be a silent mis-index. + qsa_params = sparse_cfg.to_sparse_params(pretrained_config=cfg) + assert ( + qsa_params.index_n_heads, + qsa_params.index_kv_heads, + qsa_params.index_head_dim, + qsa_params.token_topk, + qsa_params.compress_ratio, + ) == ( + self.index_heads, + self.index_kv_heads, + self.index_head_dim, + self.token_topk, + self.compress_ratio, + ), ( + f"the engine lowered the QSA index geometry to {qsa_params!r} and " + f"this target derived heads {self.index_heads}, kv heads " + f"{self.index_kv_heads}, head_dim {self.index_head_dim}, token_topk " + f"{self.token_topk}, compress_ratio {self.compress_ratio}; the side " + f"caches are sized from the engine's copy and the kernels are " + f"called with this one" + ) + assert qsa_params.dense_seq_len_threshold == self.dense_threshold, ( + f"the engine's resolved QSA cutoff is " + f"{qsa_params.dense_seq_len_threshold} and this forward branches at " + f"{self.dense_threshold}. The runner captures its short decode " + f"graph family at the engine's number, so a disagreement freezes " + f"the dense body into the graph that long batches replay — which " + f"is finite, plausible output. A serving config that raises " + f"seq_len_threshold above token_topk is the way to get here, and " + f"it is a configuration this target does not implement" + ) + 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 + self.softmax_scale = 1.0 / math.sqrt(self.head_dim) + + # ---- gated deltanet ---------------------------------------------- + self.num_k_heads = cfg.linear_num_key_heads + self.num_v_heads = cfg.linear_num_value_heads + self.head_k_dim = cfg.linear_key_head_dim + self.head_v_dim = cfg.linear_value_head_dim + self.conv_kernel = cfg.linear_conv_kernel_dim + assert self.num_v_heads % self.num_k_heads == 0, "GDN key-head broadcast" + assert self.head_k_dim == 128 and self.head_v_dim == 128, ( + f"the delta-rule entries certify head dims 128/128; this config " + f"declares {self.head_k_dim}/{self.head_v_dim}" + ) + assert self.conv_kernel in (2, 3, 4), "causal_conv1d width domain" + self.key_dim = self.num_k_heads * self.head_k_dim + self.value_dim = self.num_v_heads * self.head_v_dim + self.conv_dim = 2 * self.key_dim + self.value_dim + # Gate activation of the gated RMS norm (`silu` or `sigmoid`). + gate_type = getattr(cfg, "output_gate_type", None) + assert gate_type == "sigmoid", ( + f"gate activation {gate_type!r} is not the sigmoid form this " + f"target passes to rms_norm_gated_token_major" + ) + + # ---- MoE ---------------------------------------------------------- + self.num_experts = cfg.num_experts + self.topk = cfg.num_experts_per_tok + self.moe_inter = cfg.moe_intermediate_size + self.shared_inter = cfg.shared_expert_intermediate_size + assert cfg.norm_topk_prob, "renorm_moe_routing_op always renormalizes" + assert self.num_experts <= 512 and self.topk <= 16, ( + f"renorm_moe_routing_op takes at most 512 experts and top-16; this " + f"config declares {self.num_experts} / {self.topk}" + ) + # Every layer is MoE; the dense `intermediate_size` is dead config. + assert getattr(cfg, "decoder_sparse_step", 1) == 1, "every layer must be MoE" + assert not getattr(cfg, "mlp_only_layers", []), "every layer must be MoE" + # trtllm-gen NVFP4 runner operand-width rules and fp4_quantize's block rule. + assert self.hidden % 256 == 0, "fp4_block_scale_moe_runner hidden alignment" + assert (2 * self.moe_inter) % 128 == 0, "fp4_block_scale_moe_runner 2I alignment" + assert (self.hidden // _SF_VEC) % 4 == 0 and (self.moe_inter // _SF_VEC) % 4 == 0, ( + "the 128x4 block-scale swizzle needs both scale widths to be a multiple of 4" + ) + # flashinfer_silu_and_mul's precondition is on the half width. + assert (self.shared_inter * 2) % 16 == 0, "flashinfer_silu_and_mul half alignment" + + # ---- PLE ----------------------------------------------------------- + ple_ids = list(getattr(cfg, "ple_layer_ids", None) or []) + assert len(ple_ids) == 1, ( + f"this target implements the single PLE decoder layer this " + f"checkpoint ships; the config declares ple_layer_ids={ple_ids}" + ) + # One-based in the config. + self.ple_layer = ple_ids[0] - 1 + assert 0 <= self.ple_layer < self.num_layers, self.ple_layer + assert not self.is_attn[self.ple_layer], ( + f"the PLE lifecycle state rides the recurrent pools, so its layer " + f"must be a linear-attention layer; layer {self.ple_layer} is not" + ) + self.ngram_size = cfg.ngram_size + self.heads_per_ngram = cfg.heads_per_ngram + self.ngram_heads = self.heads_per_ngram * (self.ngram_size - 1) + self.ple_embed_dim = cfg.ple_embed_dim + self.ple_head_dim = self.ple_embed_dim // self.ngram_heads + assert self.ple_head_dim * self.ngram_heads == self.ple_embed_dim, ( + f"{self.ngram_heads} n-gram heads must tile ple_embed_dim {self.ple_embed_dim}" + ) + self.ple_conv_kernel = cfg.ple_conv_kernel_size + # Length of the runtime's paged conv state. + self.ple_dilation = self.ngram_size + self.ple_state_len = (self.ple_conv_kernel - 1) * self.ple_dilation + self.ngram_context_len = self.ngram_size - 1 + # Shard count and padding rule are config; the table height is not, so the + # table is declared at load time (`weights.py`). + self.ple_shards = cfg.split_ngram_parts + self.ngram_vocab_divisor = cfg.make_ngram_vocab_size_divisible_by + assert self.ple_shards > 0 and self.ngram_vocab_divisor > 0, ( + self.ple_shards, + self.ngram_vocab_divisor, + ) + # Fresh-sequence seed for the n-gram id history and the hash's segment token + # (element 0 when a list, as in the HF reference). + eos = cfg.eos_token_id + self.eos_token_id = int(eos[0] if isinstance(eos, (list, tuple)) else eos) + + # ---- the vision tower ------------------------------------------------ + # A 27-block ViT (no KV cache, deepstack or quantization) feeding a 2x2 patch + # merger. Constants the config omits: `_VISION_ROPE_THETA`, + # `_VISION_NORM_EPS`, and the merger's erf GELU. + vcfg = getattr(composite_config, "vision_config", None) + assert vcfg is not None, ( + "this target implements the composite export's vision tower; the " + "composite config carries no vision_config" + ) + assert getattr(vcfg, "model_type", None) == "qwen4_exp_vision", ( + f"vision_config.model_type is {getattr(vcfg, 'model_type', None)!r}; " + "this target is the Qwen4-Exp vision tower" + ) + self.v_depth = vcfg.depth + self.v_hidden = vcfg.hidden_size + self.v_heads = vcfg.num_heads + self.v_inter = vcfg.intermediate_size + self.v_merge = vcfg.spatial_merge_size + self.v_patch = vcfg.patch_size + self.v_tpatch = vcfg.temporal_patch_size + self.v_channels = vcfg.in_channels + self.v_pos_rows = vcfg.num_position_embeddings + assert self.v_hidden % self.v_heads == 0, (self.v_hidden, self.v_heads) + self.v_head_dim = self.v_hidden // self.v_heads + # vision_rope_2d needs a head divisible by 4 (two axes, one rotation pair + # each). + assert self.v_head_dim in _VISION_ROPE_HEAD_DIMS, ( + f"ops/attention/vision_rope_2d is certified at head widths " + f"{_VISION_ROPE_HEAD_DIMS}; this tower's is {self.v_head_dim}" + ) + assert (self.v_heads, self.v_heads, self.v_head_dim) in _VISION_CERTIFIED_HEADS, ( + f"the tower's attention is the encoder-only (no KV cache) mode of " + f"attention/thop_attention, certified at {sorted(_VISION_CERTIFIED_HEADS)}; " + f"this tower is {(self.v_heads, self.v_heads, self.v_head_dim)}" + ) + # `hidden_act` names the block MLP's activation only; the merger runs the erf + # `nn.GELU()`. Applying it to both is wrong by under one bf16 ulp, so there + # are two entries (`gelu_tanh`, `gelu_erf`). + assert vcfg.hidden_act == "gelu_pytorch_tanh", ( + f"the vision blocks' MLP activation is {vcfg.hidden_act!r}; this " + "target calls ops/activation/gelu_tanh for it" + ) + # No deepstack (`routing.py` refuses it; this is the second lock). + assert not list(getattr(vcfg, "deepstack_visual_indexes", []) or []), ( + f"deepstack_visual_indexes={list(vcfg.deepstack_visual_indexes)} asks " + "for per-layer visual feature injection, which this target does not " + "implement" + ) + assert vcfg.out_hidden_size == self.hidden, ( + f"the merger projects into the language model's hidden width; " + f"out_hidden_size {vcfg.out_hidden_size} against hidden {self.hidden}" + ) + # pos_embed_interpolate asserts a square table; fail before the first image. + side = int(math.isqrt(self.v_pos_rows)) + assert side * side == self.v_pos_rows, ( + f"num_position_embeddings {self.v_pos_rows} is not a perfect square, " + "so it is not a square position grid" + ) + self.v_grid_side = side + # `spatial_merge_size ** 2` patches become one token at the merger. + self.v_merge_dim = self.v_hidden * self.v_merge * self.v_merge + # One flat pixel row per patch, `(in_channels, kt, kh, kw)` channel-slowest. + self.v_pixel_width = self.v_channels * self.v_tpatch * self.v_patch * self.v_patch + # Placeholder token ids live on the composite config's top level. + self.image_token_id = int(composite_config.image_token_id) + self.video_token_id = int(composite_config.video_token_id) + assert self.image_token_id < self.vocab and self.video_token_id < self.vocab, ( + f"the media placeholders ({self.image_token_id}, {self.video_token_id}) " + f"are in-vocabulary against vocab_size {self.vocab}, so the engine " + "cannot find them with its out-of-vocabulary fallback and this target " + "must publish them as `multimodal_token_ids`" + ) + + # ---- KV pool dtype -------------------------------------------------- + # An unsupported KV dtype makes the QSA backend silently fall back to dense + # (`QSA_SPARSE_KV_CACHE_DTYPES`). + 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 + # The weight table reads the tower's geometry off this. + self.vision_config = vcfg + + self.mtp_layer_idx = self.num_layers + if self.mtp_theta is None: + self.mtp_theta = self.theta + assert self.mtp_theta > 0.0, self.mtp_theta + if self.mtp_enabled: + mode = spec_config.spec_dec_mode + assert mode.is_mtp_eagle_one_model(), ( + f"this target implements the one-model MTP-Eagle draft loop, " + f"where one module is replayed; the engine resolved " + f"spec_dec_mode {mode!r}" + ) + resolved = getattr(spec_config, "num_nextn_predict_layers", None) + assert resolved == _MTP_NUM_LAYERS, ( + f"the engine resolved num_nextn_predict_layers={resolved} from " + f"the checkpoint; this target carries the weights of " + f"{_MTP_NUM_LAYERS} MTP module" + ) + # `max_draft_len` becomes `predicted_tokens_per_seq` on every attention + # call, certified at 1..4, so attention binds. + 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" + ) + # A linear tree keeps the Blackwell spec-dec mask surface inert. + assert not getattr(spec_config, "use_dynamic_tree", False), ( + "dynamic-tree drafting is not implemented by this target" + ) + assert getattr(spec_config, "eagle_choices", None) is None, ( + "static-tree drafting is not implemented by this target" + ) + assert spec_config.max_total_draft_tokens == max_draft, ( + f"a linear tree mirrors max_total_draft_tokens from " + f"max_draft_len; the engine resolved " + f"{spec_config.max_total_draft_tokens} against {max_draft}" + ) + + # The 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 FP8 block-scale runner's own shape rules. + assert self.hidden % _FP8_BLOCK == 0 and self.moe_inter % _FP8_BLOCK == 0, ( + f"fp8_block_scale_moe_runner needs hidden ({self.hidden}) and " + f"intermediate ({self.moe_inter}) to be multiples of {_FP8_BLOCK}" + ) + assert self.topk < self.num_experts and self.topk <= 32, ( + f"the no-groups renormalize routing kernel takes 0 < top_k <= 32 " + f"and top_k < num_experts; this config declares {self.topk} of " + f"{self.num_experts}" + ) + + # The dtype the table is declared at in weights.py. + self.ple_table_dtype = torch.float8_e4m3fn + # Resolved at construction; read by `weights._declare_ngram_table`. + self.ple_host_offload = _ple_host_offload_enabled() + # CPU allocation behind the offloaded table. The `ple_table` parameter is a + # device-address view that owns none of those pages; dropping this reference + # frees memory a kernel still points into. None when resident. + self._ple_table_host = None + + # Every weight but the n-gram table; meta-init intercepts torch.empty. + self.w = _weights.MODEL_WEIGHTS.declare(_weights.MODEL_WEIGHTS.dims(self)) + + # Derived once after load (`post_load_weights`). + self._moe_scalars: list | None = None + self._gdn_gating: dict | None = None + self._cos_sin: torch.Tensor | None = None + # MTP module operands; None with speculation off. + self._mtp: _MtpWeights | None = None + self._prefill: "PrefillTarget | None" = None + self._decode: "DecodeTarget | 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 + + # ---------------------------------------------------------------- load --- + + def _norm_slices(self, key: str) -> tuple: + """Split a grouped-norm weight into its per-group slices (taken once).""" + return tuple(torch.split(self.w[key], self.hidden, dim=0)) + + def 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(), + ) + ) + self._moe_scalars = moe_scalars + # Both delta-rule entries index A_log and dt_bias as `ptr + head` and certify + # fp32; the checkpoint is bf16. + self._gdn_gating = { + i: ( + w[f"l{i}_A_log"].float().contiguous(), + w[f"l{i}_dt_bias"].float().contiguous(), + ) + for i in self.gdn_layers + } + # Reconcile the table height against the checkpoint's id-space tensors: an + # index past the end is a device-side assert that poisons the process. + table_rows = w["ple_table"].shape[0] + used_rows = int(w["ple_offs"][-1]) + int(w["ple_vocab"][-1]) + assert used_rows <= table_rows, ( + f"the n-gram table has {table_rows} rows but its per-head offsets " + f"and vocabulary sizes address {used_rows}" + ) + assert table_rows % self.ngram_vocab_divisor == 0, ( + f"the n-gram table's {table_rows} rows are not a multiple of " + f"make_ngram_vocab_size_divisible_by={self.ngram_vocab_divisor}" + ) + assert table_rows - used_rows < self.ngram_vocab_divisor, ( + f"the n-gram table carries {table_rows - used_rows} padding rows, " + f"more than one divisor of {self.ngram_vocab_divisor}: the shards " + f"do not concatenate to the padded id space this checkpoint " + f"declares" + ) + # 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._report_ple_residency() + if self.mtp_enabled: + self._mtp = _mtp_views(self) + + self._prefill = PrefillTarget(self) + self._decode = DecodeTarget(self) + self.vision_tower = VisionTower(self) + + def _report_ple_residency(self) -> None: + """Check where the n-gram table ended up, and report it once. + + Offload is output-identical, so a stray `.to()` or a driver that declined + the mapping would pass every smoke test; the pointer's memory kind read + back from the driver is the signal that cannot be faked. The report is + printed here, before the KV pool is sized from what is left.""" + table = self.w["ple_table"] + kind = _weights.memory_kind(table.data_ptr()) + expected = ( + _weights.CUDA_MEMORY_TYPE_HOST + if self.ple_host_offload + else _weights.CUDA_MEMORY_TYPE_DEVICE + ) + assert kind == expected, ( + f"the n-gram table was declared with {_PLE_HOST_OFFLOAD_ENV}=" + f"{'1' if self.ple_host_offload else '0'}, so its pages must be " + f"{'host' if self.ple_host_offload else 'device'} memory, but the " + f"driver reports cudaMemoryType {kind} for the pointer the forward " + f"gathers from. The offload is output-identical, so nothing " + f"downstream of this would notice" + ) + mib = 1024.0 * 1024.0 + table_mib = table.numel() * table.element_size() / mib + free, total = torch.cuda.mem_get_info() + print( + f"ple n-gram table: {tuple(table.shape)} {table.dtype}, " + f"{table_mib:.1f} MiB, " + f"{'HOST-resident (device-addressable)' if self.ple_host_offload else 'DEVICE-resident'}" + f" [{_PLE_HOST_OFFLOAD_ENV}=" + f"{os.environ.get(_PLE_HOST_OFFLOAD_ENV, '1 (default)')}]; " + f"after load: allocated {torch.cuda.memory_allocated() / mib:.1f} MiB, " + f"reserved {torch.cuda.memory_reserved() / mib:.1f} MiB, " + f"driver-used {(total - free) / mib:.1f} MiB of {total / mib:.1f} MiB", + flush=True, + ) + + # ------------------------------------------------------------ fail-fast --- + + def _pool_geometry(self, tokens_per_block: int) -> _PoolGeometry: + return _PoolGeometry( + tokens_per_block=tokens_per_block, + heads_kv=self.heads_kv, + head_dim=self.head_dim, + index_kv_heads=self.index_kv_heads, + index_head_dim=self.index_head_dim, + conv_dim=self.conv_dim, + conv_kernel=self.conv_kernel, + num_v_heads=self.num_v_heads, + head_v_dim=self.head_v_dim, + head_k_dim=self.head_k_dim, + stream=self.stream, + ple_state_len=self.ple_state_len, + ngram_context_len=self.ngram_context_len, + dtype=self.w["embed"].dtype, + ) + + def _bind_cache_manager(self, md, position_ids) -> None: + """Bind this module to one KV cache manager: fail-fast on the contract, then + take every pool view from it. + + Called once per manager, not per model: with capacity estimation on, the + engine builds a throwaway manager, warms up, shuts it down and builds the + real one on the same module, and a view held across the swap dangles + silently. The result is built in locals and installed as one `_Bound`, so a + failed bind leaves the previous binding untouched. The contract checks + re-run on every bind.""" + # Checked first: some fields below are properties that assert on a + # manager-less metadata instead of raising AttributeError. + cm = md.kv_cache_manager + assert cm is not None, "a paged KV cache is required" + # A re-bind during capture would bake two managers' pointers into one graph; + # the engine releases graphs before its own swap, so this is an ordering + # violation. + assert not (self._binds and torch.cuda.is_current_stream_capturing()), ( + f"the KV cache manager changed to {type(cm).__name__} during CUDA graph capture" + ) + self._binds += 1 + + missing = [name for name in _QSA_FIELDS if not hasattr(md, name)] + assert not missing, ( + f"metadata fields missing: {missing}; the engine did not select QSA " + "metadata for this model" + ) + + # `position_ids` is the int32 [3, 1, T] buffer the engine builds on an mRoPE + # config even for text-only batches; row 0 is the text position. + assert position_ids.dtype == torch.int32, position_ids.dtype + assert self.mrope_declared, "this target takes the mRoPE position layout" + assert position_ids.dim() == 3 and position_ids.shape[0] == _QSA_POSITION_AXES, ( + f"expected the [{_QSA_POSITION_AXES}, 1, T] mRoPE position buffer, got " + f"{tuple(position_ids.shape)}" + ) + + # thop_attention addresses the pool by the layer's pool offset, not its + # ordinal among attention layers; counting them aliases layers with finite, + # plausible output. + pool_idx = {i: cm.layer_offsets[i] for i in self.attn_layers} + pools = {row[0] for row in md.host_kv_cache_pool_mapping.tolist()} + assert pools <= {0, 1}, ( + f"KV addressing over {len(pools)} pools is not certified " + f"(thop_attention's cells cover one and two); layer->pool ids {sorted(pools)}" + ) + + # No sliding window: both targets pass `max_seq_len` as the window. The vector + # is indexed by pool offset and recurrent layers carry a large negative + # sentinel, so only attention offsets are read. + windows = cm.max_attention_window_vec + attn_windows = {windows[off] for off in pool_idx.values()} + assert attn_windows <= {None, md.max_seq_len}, ( + f"the attention layers' pool windows are {sorted(map(repr, attn_windows))} " + f"but this target passes max_seq_len {md.max_seq_len} as the window; " + "a sliding window is not implemented here" + ) + + g = self._pool_geometry(md.tokens_per_block) + _require( + cm, + ("get_index_k_buffer", "get_qsa_position_buffer", "get_buffers"), + "reads the QSA index plane, the shared position plane and the paged K/V", + ) + # `qsa_topk_select` checks only that the buffer is at least `index_topk` wide, + # then addresses it as flat `index_topk`-wide rows ignoring the real stride, so + # a width mismatch silently mis-writes every row but the first. + assert md.qsa_topk_indices.shape[1] == self.block_topk, ( + f"the engine sized the QSA selection buffer {md.qsa_topk_indices.shape[1]} " + f"wide and this target selects {self.block_topk} blocks per row" + ) + position_cache = cm.get_qsa_position_buffer() + check_position_plane(g, position_cache) + index_cache, kv = {}, {} + for i in self.attn_layers: + index_cache[i] = cm.get_index_k_buffer(i) + pool = cm.get_buffers(i, kv_layout="HND") + check_attention_pools(g, i, index_cache[i], pool) + kv[i] = (pool[:, 0], pool[:, 1]) + + # A missed hybrid predicate would still build and pass a smoke run with 36 + # layers carrying no state. + mm = md.mamba_metadata + assert mm is not None and mm is not False, ( + "attn_metadata.mamba_metadata is absent: the runtime allocated no " + f"recurrent-state pool, so the {len(self.gdn_layers)} linear-attention " + "layers would run with no state" + ) + mm_missing = [name for name in _MAMBA_FIELDS if not hasattr(mm, name)] + assert not mm_missing, f"mamba metadata fields missing: {mm_missing}" + # `ple_layer_cache` exists only on the V2 hybrid manager; under V1 the PLE + # state silently does not exist. + _require( + cm, + ("get_conv_states", "get_ssm_states", "ple_layer_cache"), + "reads the recurrent, conv and PLE pools", + ) + conv = {i: cm.get_conv_states(i) for i in self.gdn_layers} + ssm = {i: cm.get_ssm_states(i) for i in self.gdn_layers} + conv0, ssm0 = conv[self.gdn_layers[0]], ssm[self.gdn_layers[0]] + check_recurrent_pools(g, conv0, ssm0) + + ple_cache = cm.ple_layer_cache(self.ple_layer) + assert ple_cache is not None, ( + f"no PLE lifecycle state for decoder layer {self.ple_layer}: the runtime " + "allocated neither the dilated conv state nor the n-gram id history" + ) + ple_conv, ple_ngram = ple_cache + check_ple_pools(g, ple_conv, ple_ngram, conv0) + + assert md.effective_beam_width == 1, "beam search is not implemented" + assert md.cache_indirection is None, "beam search is not implemented" + assert not md.is_cross, "cross attention is not implemented" + # With speculation on, the Blackwell trtllm-gen spec-dec mask machinery is off + # for a linear tree at sm_103: every `spec_decoding_*` tensor stays None and + # `predicted_tokens_per_seq` alone produces the within-block mask. + assert not md.is_spec_decoding_enabled, ( + "the attention metadata reports is_spec_decoding_enabled; on sm_103 a " + "linear-tree MTP forces that machinery off" + ) + bound = _Bound( + cm, pool_idx, index_cache, position_cache, kv, conv, ssm, (ple_conv, ple_ngram) + ) + if self.mtp_enabled: + bound = bound._replace(**self._bind_step_buffers(cm, g, conv0, ssm0)) + self._bound = bound + + def _bind_step_buffers(self, cm, g: _PoolGeometry, 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.""" + 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", + ) + # `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) + + # -------------------------------------------------------------- forward --- + + def _select_target(self, attn_metadata) -> "PrefillTarget | DecodeTarget": + """Which target runs this step: decode if `num_contexts == 0`, else prefill. + + A draft loop rewrites `num_contexts` in place between draft steps, so the + draft step routes through here per call.""" + return self._decode if attn_metadata.num_contexts == 0 else self._prefill + + def 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, + ): + """Check what holds for every step, bind the cache manager, then route.""" + assert self._prefill is not None and self._decode is not None, ( + "load_weights must run before forward" + ) + assert position_ids is not None + assert isinstance(attn_metadata, TrtllmAttentionMetadata) + # Unimplemented inputs must fail loudly, not be dropped. + assert lora_params is None, "LoRA is not implemented by this target" + # `input_ids` is required 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. + 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" + ) + # `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`). + if self._bound is None or attn_metadata.kv_cache_manager is not self._bound.cm: + self._bind_cache_manager(attn_metadata, position_ids) + return self._select_target(attn_metadata).forward( + attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + +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_up: torch.Tensor + ahc_inject: torch.Tensor + mhc_norm: tuple + mhc_down: torch.Tensor + mhc_up: torch.Tensor + mhc_inject: torch.Tensor + q: torch.Tensor | None + k: torch.Tensor | None + v: 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 + out: torch.Tensor | None + conv_w: torch.Tensor | None + a_log: torch.Tensor | None + dt_bias: torch.Tensor | None + gnorm: torch.Tensor | None + router: torch.Tensor + moe_gscale: torch.Tensor + fc1_w: torch.Tensor + fc1_s: torch.Tensor + fc2_w: torch.Tensor + fc2_s: torch.Tensor + out1_scale: torch.Tensor + out1_gate_scale: torch.Tensor + out2_scale: torch.Tensor + sh_gu: torch.Tensor + sh_dn: torch.Tensor + sh_gate: torch.Tensor + + +def _layer_views(core: Qwen4ExpTp1ModelingV2) -> list[_LayerWeights]: + """Every trunk layer's `_LayerWeights`, read off the realized weights.""" + w = core.w + layers = [] + for i in range(core.num_layers): + attn: dict = dict( + q=None, + k=None, + v=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, + out=None, + conv_w=None, + a_log=None, + dt_bias=None, + gnorm=None, + ) + 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(), + 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(), + out=w[f"l{i}_out_w"].t(), + # [conv_dim, 1, width] -> [conv_dim, width]: a view (unit width stride). + conv_w=w[f"l{i}_conv_w"].reshape(core.conv_dim, core.conv_kernel), + a_log=a_log, + dt_bias=dt_bias, + gnorm=w[f"l{i}_gnorm"], + ) + scalars = core._moe_scalars[i] + 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"], + **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(), + sh_dn=w[f"l{i}_sh_dn"].t(), + sh_gate=w[f"l{i}_sh_gate"].t(), + ) + ) + return layers + + +class _StackWeights(NamedTuple): + """The weights outside the layer loop: the stack-final mixer (this checkpoint's + final norm) and the PLE side path.""" + + fhc_norm: tuple + fhc_down: torch.Tensor + fhc_up: torch.Tensor + ple_key: torch.Tensor + ple_val: torch.Tensor + ple_nk: tuple + ple_nq: tuple + ple_nc: tuple + ple_conv_w: torch.Tensor + ple_mult: torch.Tensor + ple_vocab: torch.Tensor + ple_offs: torch.Tensor + ple_tscale: torch.Tensor + + +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"], + 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"), + ple_conv_w=w["ple_conv_w"], + ple_mult=w["ple_mult"], + ple_vocab=w["ple_vocab"], + ple_offs=w["ple_offs"], + # `dequantize_e4m3_per_tensor` needs the scale at the input's rank (3), all 1s. + ple_tscale=w["ple_tscale"].reshape(1, 1, 1), + ) + + +class _MtpWeights(NamedTuple): + """The MTP module's weights (keyed at its layer index, `num_hidden_layers`).""" + + pre_e_norm: torch.Tensor + pre_h_norm: tuple + fc_e: torch.Tensor + fc_h: torch.Tensor + ahc_norm: tuple + ahc_down: torch.Tensor + ahc_up: torch.Tensor + ahc_inject: torch.Tensor + mhc_norm: tuple + mhc_down: torch.Tensor + mhc_up: torch.Tensor + mhc_inject: torch.Tensor + q: torch.Tensor + k: torch.Tensor + v: torch.Tensor + o: torch.Tensor + q_norm: torch.Tensor + k_norm: torch.Tensor + router: torch.Tensor + fc1_w: torch.Tensor + fc1_s: torch.Tensor + fc2_w: torch.Tensor + fc2_s: torch.Tensor + sh_gu: torch.Tensor + sh_dn: torch.Tensor + sh_gate: torch.Tensor + fhc_norm: tuple + fhc_down: torch.Tensor + fhc_up: torch.Tensor + + +def _mtp_views(core: Qwen4ExpTp1ModelingV2) -> _MtpWeights: + w = core.w + p = f"l{core.mtp_layer_idx}_mtp_" + return _MtpWeights( + pre_e_norm=w[f"{p}pre_e_norm"], + pre_h_norm=core._norm_slices(f"{p}pre_h_norm"), + fc_e=w[f"{p}fc_e"].t(), + fc_h=w[f"{p}fc_h"].t(), + ahc_norm=core._norm_slices(f"{p}ahc_norm"), + ahc_down=w[f"{p}ahc_down"], + ahc_up=w[f"{p}ahc_up"], + ahc_inject=w[f"{p}ahc_inject"], + mhc_norm=core._norm_slices(f"{p}mhc_norm"), + mhc_down=w[f"{p}mhc_down"], + mhc_up=w[f"{p}mhc_up"], + mhc_inject=w[f"{p}mhc_inject"], + q=w[f"{p}q_w"].t(), + k=w[f"{p}k_w"].t(), + v=w[f"{p}v_w"].t(), + o=w[f"{p}o_w"].t(), + q_norm=w[f"{p}q_norm"], + k_norm=w[f"{p}k_norm"], + 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_dn=w[f"{p}sh_dn"].t(), + sh_gate=w[f"{p}sh_gate"].t(), + fhc_norm=core._norm_slices(f"{p}fhc_norm"), + fhc_down=w[f"{p}fhc_down"], + fhc_up=w[f"{p}fhc_up"], + ) + + +class PrefillTarget: + """Any batch with context rows (`num_contexts > 0`), possibly with generation rows + beside them; the recurrent layers' decode-only branches never run here.""" + + def __init__(self, core: Qwen4ExpTp1ModelingV2) -> None: + """Take the per-layer weight views once; MTP operands are read off `core._mtp`.""" + self.core = core + # Off unless TRTLLM_MODELING_V2_VALIDATE asks; each target owns its flag (the + # two `step_args` projections differ). + self._contract_pending = step_contract_enabled() + self._layers = _layer_views(core) + self._stack = _stack_views(core) + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): calling + `step_args` surfaces an upstream metadata rename here, not mid-forward.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md: TrtllmAttentionMetadata) -> dict: + """Project the prepared metadata onto thop_attention's batch state, once per + forward. Tensors are engine-owned persistent buffers refreshed in place; + Python ints are per-capture constants.""" + return dict( + sequence_length=md.kv_lens_cuda_runtime, + host_past_key_value_lengths=md.kv_lens_runtime, + host_total_kv_lens=md.host_total_kv_lens, + context_lengths=md.prompt_lens_cuda_runtime, + host_context_lengths=md.prompt_lens_cpu_runtime, + host_request_types=md.host_request_types_runtime, + kv_cache_block_offsets=md.kv_cache_block_offsets, + host_kv_cache_pool_pointers=md.host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping=md.host_kv_cache_pool_mapping, + workspace_=md.effective_workspace, + tokens_per_block=md.tokens_per_block, + max_num_requests=md.max_num_requests, + max_context_length=md.max_context_length, + max_seq_len=md.max_seq_len, + # No sliding window: the full admitted length on every attention layer. + attention_window_size=md.max_seq_len, + num_contexts=md.num_contexts, + num_ctx_tokens=md.num_ctx_tokens, + trtllm_gen_jit_warmup=md.trtllm_gen_jit_warmup, + use_paged_context_fmha=md.use_paged_context_fmha, + beam_width=md.effective_beam_width, + cache_indirection=md.cache_indirection, + block_ids_per_seq=md.block_ids_per_seq, + max_context_q_len_override=md.max_context_q_len_override, + is_cross=md.is_cross, + is_spec_decoding_enabled=md.is_spec_decoding_enabled, + use_spec_decoding=md.use_spec_decoding, + is_spec_dec_tree=md.is_spec_dec_tree, + spec_decoding_generation_lengths=md.spec_decoding_generation_lengths, + spec_decoding_position_offsets_for_cpp=md.spec_decoding_position_offsets_for_cpp, + spec_decoding_packed_mask=md.spec_decoding_packed_mask, + spec_decoding_bl_tree_mask_offset=md.spec_decoding_bl_tree_mask_offset, + spec_decoding_bl_tree_mask=md.spec_decoding_bl_tree_mask, + spec_decoding_target_max_draft_tokens=md.max_total_draft_tokens, + spec_bl_tree_first_sparse_mask_offset_kv=md.spec_bl_tree_first_sparse_mask_offset_kv, + num_sparse_topk=md.num_sparse_topk, + flash_mla_tile_scheduler_metadata=md.flash_mla_tile_scheduler_metadata, + flash_mla_num_splits=md.flash_mla_num_splits, + max_num_sequences=md.max_num_sequences, + force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask, + ) + + # --------------------------------------------------------------- blocks --- + + def _grouped_norm(self, x: torch.Tensor, weights: tuple) -> torch.Tensor: + """`Qwen4ExpTextRMSNorm(group_size=hidden)`: an independent RMS per `hidden`-wide + group, then `(1 + w)` over the full width; one gemma-RMS-norm per column + slice (row stride certified by that entry) plus a concat.""" + core = self.core + parts = torch.split(x, [core.hidden] * core.hc_count, 1) + return torch.cat( + [flashinfer_gemma_rmsnorm(part, w, core.eps) for w, part in zip(weights, parts)], 1 + ) + + def _moe_block(self, x: 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.""" + core = self.core + # The routing op reads logits from data_ptr() ignoring 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) + + def _attention_block(self, x, i, L, qkv, index_scratch, attn_out, qsa, step): + """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.""" + core = self.core + ( + req_idx, + logical, + seq_lens_q, + visible_blocks, + topk_indices, + 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, + ) + + # ---- 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], + 1, + )[1] + qsa_paged_index_store( + token_k, + coords, + index_cache, + core._bound.position_cache, + req_idx, + logical, + 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, + [num_tokens, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim], + ), + [core.index_heads, core.index_kv_heads, core.index_kv_heads], + 1, + )[0] + qsa_prefill_compress( + logical_positions=logical, + request_indices=req_idx, + block_table=block_table, + index_cache=index_cache, + position_cache=core._bound.position_cache, + k_norm_weight=L.idx_kn, + cos_sin=core._cos_sin, + eps=core.eps, + tokens_per_block=tokens_per_block, + compress_ratio=core.compress_ratio, + mrope_section=_QSA_MROPE_SECTION, + ) + + if sparse: + k_cache, v_cache = 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 + ) + 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( + 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, + # 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], + ) + else: + out = thop_attention( + q=qkv, + output=attn_out, + local_layer_idx=core._bound.pool_idx[i], + num_heads=core.heads_q, + num_kv_heads=core.heads_kv, + head_size=core.head_dim, + quant_mode=core.quant_mode, + **step, + **_CALL_CONSTANTS, + ) + fused_sigmoid_mul_inplace(out, gate) + return cublas_mm(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: + """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.""" + 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). + x3 = torch.transpose(torch.reshape(mq_d, [n, p, core.conv_dim]), 1, 2) + conv_out = causal_conv1d_triton_update( + x=x3, + conv_state=core._bound.conv[i], + weight=L.conv_w, + bias=None, + activation="silu", + conv_state_indices=si_d, + pad_slot_id=_PAD_SLOT_ID, + # Snapshots, not the accepted-token roll. + num_accepted_tokens=None, + intermediate_conv_window=core._bound.int_conv[i], + intermediate_state_indices=int_idx, + launch_dependent_kernels=False, + ) + # Back to token-major (a view). + mq_out = torch.reshape(torch.transpose(conv_out, 1, 2), [n * p, core.conv_dim]) + q, k, v = pack_gdn_decode_qkv( + mq_out, core.num_k_heads, core.head_k_dim, core.num_v_heads, core.head_v_dim + ) + g = torch.reshape( + fused_gdn_gating( + A_log=L.a_log, + a=a_d, + dt_bias=L.dt_bias, + beta=_SOFTPLUS_BETA, + threshold=_SOFTPLUS_THRESHOLD, + ), + [n, p, core.num_v_heads], + ) + beta = torch.reshape(torch.sigmoid(b_d), [n, p, core.num_v_heads]) + # Initial states are gathered from the live pool into a batch-scoped copy. The + # per-step buffer is indexed by batch row (`intermediate_state_indices[:n]`), + # not slot id: slot ids would write rows nothing promotes, or past the end. + src = torch.reshape( + nn.functional.embedding(si_d, torch.reshape(ssm_states, [ssm_states.shape[0], -1])), + [n, core.num_v_heads, core.head_v_dim, core.head_k_dim], + ) + out = fused_recurrent_gated_delta_rule_update( + torch.reshape(q, [n, p, core.num_k_heads, core.head_k_dim]), + torch.reshape(k, [n, p, core.num_k_heads, core.head_k_dim]), + torch.reshape(v, [n, p, core.num_v_heads, core.head_v_dim]), + g, + beta, + initial_state_source=src, + initial_state_indices=int_idx, + scale=None, + use_qk_l2norm_in_kernel=True, + cu_seqlens=None, + disable_state_update=True, + disable_output_calculation=False, + intermediate_states_buffer=core._bound.int_ssm[i], + cache_steps=p, + output=None, + ) + return torch.reshape(out, [1, n * p, core.num_v_heads, core.head_v_dim]) + + def _gdn_block(self, x, i, L, 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.""" + core = self.core + ( + 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, + ) = 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) + + conv_states = core._bound.conv[i] + ssm_states = core._bound.ssm[i] + + # 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) + # Prefill conv wants channel-major, decode token-major. + mq_p_t = extract_transpose_prefill_slice(mq_p, num_prefill_tokens, 0, core.conv_dim) + causal_conv1d_fwd( + x=mq_p_t, + weight=L.conv_w, + bias=None, + conv_states=conv_states, + query_start_loc=qsl_p, + cache_indices=si_p, + has_initial_state=hi_p, + silu_activation=True, + pad_slot_id=_PAD_SLOT_ID, + out=None, + ) + if is_verify: + # A context row advances its slot as on the base path; a generation + # row is a draft chain of unknown accepted length, so it goes to the + # rollback buffers. + a_p, a_d = torch.split(a, [num_prefill_tokens, num_decode_tokens], 0) + b_p, b_d = torch.split(b, [num_prefill_tokens, num_decode_tokens], 0) + q, k, v, g, beta = fused_gdn_post_conv( + prefill=mq_p_t, + decode=None, + a=a_p, + b=b_p, + 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, + l2_norm_eps=_L2_NORM_EPS, + softplus_threshold=_SOFTPLUS_THRESHOLD, + ) + # q and k are already L2-normalized by the post-conv call. + out_p, _ = 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, + use_qk_l2norm_in_kernel=False, + output=None, + ) + out_d = self._gdn_verify( + mq_d, a_d, b_d, ssm_states, i, L, si_d, int_idx, num_decodes, gen_p + ) + return self._gdn_epilogue(torch.cat([out_p, out_d], 1), z, L, num_tokens) + 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, + ) + q, k, v, g, beta = fused_gdn_post_conv( + prefill=mq_p_t, + decode=mq_d, + 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, + l2_norm_eps=_L2_NORM_EPS, + softplus_threshold=_SOFTPLUS_THRESHOLD, + ) + else: + mq_t = extract_transpose_prefill_slice(mixed_qkv, num_tokens, 0, core.conv_dim) + causal_conv1d_fwd( + x=mq_t, + weight=L.conv_w, + bias=None, + conv_states=conv_states, + query_start_loc=qsl_all, + cache_indices=si_all, + has_initial_state=hi_all, + silu_activation=True, + pad_slot_id=_PAD_SLOT_ID, + out=None, + ) + q, k, v, g, beta = fused_gdn_post_conv( + prefill=mq_t, + decode=None, + 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, + l2_norm_eps=_L2_NORM_EPS, + softplus_threshold=_SOFTPLUS_THRESHOLD, + ) + # q and k are already L2-normalized; the recurrence must not repeat it. + out, _ = 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, + use_qk_l2norm_in_kernel=False, + output=None, + ) + 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.""" + core = self.core + # Sigmoid gate and a plain (non-gemma) weight (`output_gate_type`). + o = rms_norm_gated_token_major( + torch.reshape(out, [-1, core.head_v_dim]), + torch.reshape(z, [num_tokens, core.num_v_heads, core.head_v_dim]), + weight=L.gnorm, + eps=core.eps, + fp8_scale=None, + gate_activation="sigmoid", + ) + 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. + + Per-request stateful: ids become one `ngram_size`-wide window per token + (predecessors from the paged id history); each window hashes to one row per + head; rows gather fp8 vectors from one table and dequantize by one per-tensor + scale; key/value projections and the gate scale one value into each stream; + a dilated depthwise short conv over a paged state is added back. A sequence + carrying nothing is seeded with `eos_token_id`, not 0: id 0 is a real token + the hash cannot distinguish.""" + 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 + ) + # 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 torch.add(gated, torch.transpose(convolved, 0, 1)) + + # -------------------------------------------------------------- forward --- + + def forward( + self, + 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 + md = attn_metadata + self._check_step_contract(md) + step = self.step_args(md) + spec_metadata = kwargs.get("spec_metadata") + 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, + ) + 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. + 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. + gen_p = num_decode_tokens // num_decodes if num_decodes else 1 + assert gen_p * num_decodes == num_decode_tokens, ( + f"{num_decode_tokens} generation rows do not divide over " + f"{num_decodes} generation sequences; a ragged per-sequence draft " + "length can be expressed neither to the attention call nor to the " + "multi-token recurrence" + ) + assert core.mtp_enabled or gen_p == 1, ( + f"{num_decode_tokens} generation rows over {num_decodes} generation " + "sequences with speculation off; the base path implements one query " + "token per generation sequence, and both the attention call's " + "within-block mask and the recurrent path's cumulative-length " + "vector assume it" + ) + # `gen_p > 1` is load-bearing: verify parks a generation row in the rollback + # buffers, which `update_mamba_states` promotes only if the worker drafts. + # With `runtime_draft_len == 0` (a `draft_len_schedule` can do this) it skips + # drafting before the promotion, leaving every recurrent layer one token + # stale; that row is the always-accepted golden token, so the ordinary decode + # path is right. + if spec_metadata is not None and num_decodes > 0: + declared_p = spec_metadata.runtime_draft_len + 1 + assert gen_p == declared_p, ( + f"the batch carries {gen_p} query rows per generation sequence " + f"but the worker's runtime_draft_len says {declared_p}; the " + "per-step conv snapshots and the recurrent rollback buffer are " + "sized from the second and written from the first" + ) + # Per-iteration counterpart of the core's `max_draft_len` bound. + assert 1 <= gen_p <= 4, ( + f"predicted_tokens_per_seq {gen_p} is outside the 1..4 range " + f"attention/thop_attention certifies" + ) + is_verify = spec_metadata is not None and num_decodes > 0 and gen_p > 1 + step["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) + hi_p = hi_all[:num_prefills] + # `query_start_loc` covers the context prefix; `_long` the whole batch. + 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 + int_idx = None + qsl_long_p = 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] + gdn_state = ( + num_prefill_tokens, + num_decode_tokens, + si_p, + si_all, + si_d, + hi_p, + hi_all, + qsl_p, + qsl_all, + qsl_long, + is_verify, + num_decodes, + gen_p, + int_idx, + qsl_long_p, + ) + + # PLE carry flags must cover the whole batch, but `has_initial_states` is + # filled for the context prefix only (generation entries are stale); every + # generation sequence carries, so pad with True (host-int sized, graph-safe). + ple_carry = nn.functional.pad(hi_p, [0, num_decodes], mode="constant", value=1.0) + # Under speculation `Mamba2Metadata.prepare`'s `query_start_loc_long` has one + # row per generation sequence, so with `gen_p > 1` it is short and both PLE + # stateful calls are silently mis-segmented; rebuild it from `seq_lens_cuda`, + # which still holds the true counts during the trunk forward. + if gen_p == 1: + ple_qsl = qsl_long + else: + ple_qsl = cu_seqlens_from_seq_lens(md.seq_lens_cuda[:batch], torch.int64) + ple_state = (core._bound.ple_state[0], core._bound.ple_state[1], si_all, ple_carry, ple_qsl) + + # Dense/sparse decision once per forward, from a host-side metadata read (as + # QSA's backend does): a per-capture constant, matching the engine's separate + # short/long decode graph families. + sparse = batch > 0 and int(md.kv_lens_runtime[:batch].max()) > core.dense_threshold + + qsa = ( + md.qsa_req_idx_per_token[:num_tokens], + md.qsa_logical_positions[:num_tokens], + md.qsa_sequence_lengths[:num_tokens], + md.qsa_visible_blocks[:num_tokens], + md.qsa_topk_indices[:num_tokens], + md.qsa_topk_row_starts[:num_tokens], + md.qsa_block_table, + md.tokens_per_block, + pos, + pos3, + coords, + sparse, + ) + + # Scratch shared by all attention layers (each fully rewrites them). + 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 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): + 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 core.is_attn[i]: + block_out = self._attention_block( + mixed, i, L, qkv, index_scratch, attn_out, qsa, step + ) + 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 + ) + 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 + ) + if spec_metadata is None: + return final + # `lm_head` reads the `[T, hidden]` mix; the MTP module reads the + # `[T, hc_count * hidden]` stream it was collapsed from. + return final, hc + + # ------------------------------------------------------------- the draft --- + + def _mtp_moe_block(self, x, dt) -> torch.Tensor: + """The draft layer's MoE: FP8 block-scale routed experts plus a bf16 shared expert. + + The scale buffer comes back transposed `[hidden/128, tokens]` and is passed + straight through; never reshape it. At `tokens == hidden/128` (20, an + ordinary decode batch) the shapes coincide, so nothing would catch a + re-materialization. Pre-routed; `routing_bias` must be `None` (the entry + refuses one).""" + core = self.core + M = core._mtp + logits = cublas_mm(x, M.router) + topk_ids, topk_w = renorm_moe_routing_op(logits, topk=core.topk, output_dtype=dt) + xq, xsf = fp8_quantize_1x128(x, use_ue8m0=_MTP_MOE_USE_UE8M0) + routed = fp8_block_scale_moe_runner( + routing_logits=None, + routing_bias=None, + hidden_states=xq, + hidden_states_scale=xsf, + gemm1_weights=M.fc1_w, + gemm1_weights_scale=M.fc1_s, + gemm2_weights=M.fc2_w, + gemm2_weights_scale=M.fc2_s, + num_experts=core.num_experts, + top_k=core.topk, + num_fused_shared_experts=None, + n_group=None, + topk_group=None, + intermediate_size=core.moe_inter, + local_expert_offset=0, + local_num_experts=core.num_experts, + routed_scaling_factor=None, + routing_method_type=_ROUTING_METHOD_RENORM, + topk_weights=topk_w, + topk_ids=topk_ids, + act_type=_ACT_TYPE_SWIGLU, + gemm1_clamp_limit=None, + ) + shared = cublas_mm(flashinfer_silu_and_mul(cublas_mm(x, M.sh_gu)), M.sh_dn) + gate_logits = cublas_mm(x, M.sh_gate) + return fused_sigmoid_gate_mul_add(routed, gate_logits, shared) + + def mtp_forward( + self, + embed_tokens: torch.Tensor, + all_rank_num_tokens, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + **kwargs, + ) -> torch.Tensor: + """One draft step whose metadata still carries context rows. + + `spec_metadata` and `all_rank_num_tokens` are accepted and unused.""" + core = self.core + assert isinstance(attn_metadata, TrtllmAttentionMetadata) + assert core._bound is not None, ( + "the trunk's cache-manager bind has not run; the shell calls the " + "core before the worker, so this cannot be reached first" + ) + md = attn_metadata + rows = hidden_states.shape[0] + assert input_ids.shape[0] == rows, ( + f"the draft step's {input_ids.shape[0]} token ids and {rows} " + "hidden-state rows must describe the same tokens" + ) + assert hidden_states.shape[1] == core.stream, ( + f"the MTP module reads the trunk's {core.stream}-wide four-stream " + f"residual; the runtime handed it {hidden_states.shape[1]} columns" + ) + dt = hidden_states.dtype + dev = hidden_states.device + + # Projected on every call: the worker rewrites `attn_metadata` in place + # between draft step 0 and 1+ (`_seq_lens` to 1, `num_contexts` and + # `num_ctx_tokens` to 0), so a first-call value is wrong on the rest. + step = self.step_args(md) + M = core._mtp + num_ctx = md.num_ctx_tokens + gen_rows = rows - num_ctx + gen_seqs = md.num_seqs - md.num_contexts + p = gen_rows // gen_seqs if gen_seqs else 1 + assert p * gen_seqs == gen_rows, ( + f"{gen_rows} generation rows do not divide over {gen_seqs} " + "generation sequences; a ragged per-sequence draft length cannot be " + "expressed to the attention call" + ) + + # The KV cache written here is not the trunk's: with a separate draft KV + # manager the worker swaps `attn_metadata.kv_cache_manager` and + # `kv_cache_block_offsets` for the draft loop. The layer index must be taken + # from the current manager every call; caching the trunk's addresses another + # layer's pages with finite, plausible drafts. + cm = md.kv_cache_manager + assert cm is not None, "a paged KV cache is required" + offsets = cm.layer_offsets + assert core.mtp_layer_idx in offsets, ( + f"the MTP module's attention is global layer index " + f"{core.mtp_layer_idx}, which the cache manager " + f"{type(cm).__name__} does not map (it maps {len(offsets)} layers, " + f"is_draft={getattr(cm, 'is_draft', None)}); the engine allocated " + "no KV layer for the draft module" + ) + layer_idx = offsets[core.mtp_layer_idx] + window = cm.max_attention_window_vec[layer_idx] + assert window is None or window == md.max_seq_len, ( + f"the MTP layer's pool window is {window!r} but this call passes " + f"max_seq_len {md.max_seq_len} as the attention window; a sliding " + "window is not implemented here" + ) + + # Copy into a fresh contiguous `[3, N]`: a sliced view keeps the source row + # stride and the op rejects it under `use_mrope`. + assert position_ids.shape[0] == _QSA_POSITION_AXES, ( + f"expected the mRoPE position layout with {_QSA_POSITION_AXES} " + f"coordinate rows, got {tuple(position_ids.shape)}" + ) + pos3 = torch.empty( + [_QSA_POSITION_AXES, position_ids.numel() // _QSA_POSITION_AXES], + dtype=torch.int32, + device=position_ids.device, + ) + pos3.copy_(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1])) + + # ---- the front end -------------------------------------------------- + e = nn.functional.embedding(input_ids, embed_tokens) + x_e = cublas_mm(flashinfer_gemma_rmsnorm(e, M.pre_e_norm, core.eps), M.fc_e) + # `fc_hidden` is one `[hidden, hidden]` matrix applied to each of the four + # streams (the checkpoint's norm statistics say they are distinct streams). + h_n = self._grouped_norm(hidden_states, M.pre_h_norm) + x_h = torch.reshape( + cublas_mm(torch.reshape(h_n, [rows * core.hc_count, core.hidden]), M.fc_h), + [rows, core.stream], + ) + if _MTP_EMBED_BLOCK_ALL_STREAMS: + x = torch.add(x_h, torch.cat([x_e] * core.hc_count, -1)) + else: + x = torch.add( + x_h, + nn.functional.pad(x_e, [0, core.stream - core.hidden], mode="constant", value=0.0), + ) + + # ---- the decoder layer ---------------------------------------------- + mixed, inject = qwen4_exp_gated_residual( + self._grouped_norm(x, M.ahc_norm), M.ahc_down, M.ahc_up, M.ahc_inject, core.hc_count + ) + qkv = torch.empty([rows, core.qkv_width], dtype=dt, device=dev) + attn_out = torch.empty([rows, core.q_width], dtype=dt, device=dev) + qg = cublas_mm(mixed, M.q) + k = cublas_mm(mixed, M.k) + v = cublas_mm(mixed, M.v) + qkv_h = torch.reshape(qkv, [rows, core.heads_q + 2 * core.heads_kv, core.head_dim]) + q_dst, k_dst, v_dst = torch.split(qkv_h, [core.heads_q, core.heads_kv, core.heads_kv], 1) + q_src, gate = torch.split( + torch.reshape(qg, [rows, core.heads_q, 2 * core.head_dim]), + [core.head_dim, core.head_dim], + 2, + ) + q_dst.copy_(q_src) + k_dst.copy_(torch.reshape(k, [rows, core.heads_kv, core.head_dim])) + v_dst.copy_(torch.reshape(v, [rows, core.heads_kv, core.head_dim])) + # The trunk rope's certified cell at the module's own base (`mtp.rope_theta`). + fused_qk_norm_rope( + qkv, + num_heads_q=core.heads_q, + num_heads_k=core.heads_kv, + num_heads_v=core.heads_kv, + head_dim=core.head_dim, + rotary_dim=core.rotary_dim, + eps=core.eps, + q_weight=M.q_norm, + k_weight=M.k_norm, + base=core.mtp_theta, + is_neox=True, + position_ids=pos3, + is_qk_norm=True, + use_gemma=True, + use_mrope=True, + mrope_section1=core.mrope_section1, + mrope_section2=core.mrope_section2, + ) + # 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, + ) + fused_sigmoid_mul_inplace(attn_out, gate) + x = qwen4_exp_post_inject(x, cublas_mm(attn_out, M.o), inject) + + mixed, inject = qwen4_exp_gated_residual( + self._grouped_norm(x, M.mhc_norm), M.mhc_down, M.mhc_up, M.mhc_inject, core.hc_count + ) + return qwen4_exp_post_inject(x, self._mtp_moe_block(mixed, dt), inject) + + +class DecodeTarget: + """Generation-only batches (no context rows): drops the context-row recurrent + ops and adds the fused single-token recurrence.""" + + def __init__(self, core: Qwen4ExpTp1ModelingV2) -> None: + """Take the per-layer weight views once; MTP operands are read off `core._mtp`.""" + self.core = core + # See PrefillTarget.__init__. + self._contract_pending = step_contract_enabled() + self._layers = _layer_views(core) + self._stack = _stack_views(core) + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast; see PrefillTarget._check_step_contract.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md: TrtllmAttentionMetadata) -> dict: + """See PrefillTarget.step_args. Context counts are stated, not read: routing + guarantees them, and a decode-only CUDA graph would freeze whatever it read.""" + return dict( + sequence_length=md.kv_lens_cuda_runtime, + host_past_key_value_lengths=md.kv_lens_runtime, + host_total_kv_lens=md.host_total_kv_lens, + context_lengths=md.prompt_lens_cuda_runtime, + host_context_lengths=md.prompt_lens_cpu_runtime, + host_request_types=md.host_request_types_runtime, + kv_cache_block_offsets=md.kv_cache_block_offsets, + host_kv_cache_pool_pointers=md.host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping=md.host_kv_cache_pool_mapping, + workspace_=md.effective_workspace, + tokens_per_block=md.tokens_per_block, + max_num_requests=md.max_num_requests, + max_context_length=md.max_context_length, + max_seq_len=md.max_seq_len, + # No sliding window: the full admitted length on every attention layer. + attention_window_size=md.max_seq_len, + num_contexts=0, + num_ctx_tokens=0, + trtllm_gen_jit_warmup=md.trtllm_gen_jit_warmup, + use_paged_context_fmha=md.use_paged_context_fmha, + beam_width=md.effective_beam_width, + cache_indirection=md.cache_indirection, + block_ids_per_seq=md.block_ids_per_seq, + max_context_q_len_override=md.max_context_q_len_override, + is_cross=md.is_cross, + is_spec_decoding_enabled=md.is_spec_decoding_enabled, + use_spec_decoding=md.use_spec_decoding, + is_spec_dec_tree=md.is_spec_dec_tree, + spec_decoding_generation_lengths=md.spec_decoding_generation_lengths, + spec_decoding_position_offsets_for_cpp=md.spec_decoding_position_offsets_for_cpp, + spec_decoding_packed_mask=md.spec_decoding_packed_mask, + spec_decoding_bl_tree_mask_offset=md.spec_decoding_bl_tree_mask_offset, + spec_decoding_bl_tree_mask=md.spec_decoding_bl_tree_mask, + spec_decoding_target_max_draft_tokens=md.max_total_draft_tokens, + spec_bl_tree_first_sparse_mask_offset_kv=md.spec_bl_tree_first_sparse_mask_offset_kv, + num_sparse_topk=md.num_sparse_topk, + flash_mla_tile_scheduler_metadata=md.flash_mla_tile_scheduler_metadata, + flash_mla_num_splits=md.flash_mla_num_splits, + max_num_sequences=md.max_num_sequences, + force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask, + ) + + # --------------------------------------------------------------- blocks --- + + def _grouped_norm(self, x: torch.Tensor, weights: tuple) -> torch.Tensor: + """See `PrefillTarget._grouped_norm`.""" + core = self.core + parts = torch.split(x, [core.hidden] * core.hc_count, 1) + return torch.cat( + [flashinfer_gemma_rmsnorm(part, w, core.eps) for w, part in zip(weights, parts)], 1 + ) + + def _moe_block(self, x: torch.Tensor, L: _LayerWeights, dt: torch.dtype) -> torch.Tensor: + """See `PrefillTarget._moe_block`.""" + 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`.""" + core = self.core + ( + req_idx, + logical, + seq_lens_q, + visible_blocks, + topk_indices, + 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 (see PrefillTarget). + use_mrope=True, + mrope_section1=core.mrope_section1, + mrope_section2=core.mrope_section2, + ) + + # ---- 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, + req_idx, + logical, + block_table, + tokens_per_block, + ) + # 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_index = torch.split( + torch.reshape( + scratch_h, + [num_tokens, core.index_heads + 2 * core.index_kv_heads, core.index_head_dim], + ), + [core.index_heads, core.index_kv_heads, core.index_kv_heads], + 1, + )[0] + qsa_prefill_compress( + logical_positions=logical, + request_indices=req_idx, + block_table=block_table, + index_cache=index_cache, + position_cache=core._bound.position_cache, + k_norm_weight=L.idx_kn, + cos_sin=core._cos_sin, + eps=core.eps, + tokens_per_block=tokens_per_block, + compress_ratio=core.compress_ratio, + mrope_section=_QSA_MROPE_SECTION, + ) + + if sparse: + 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 + ) + 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, + ) + # `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( + 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: + out = thop_attention( + q=qkv, + output=attn_out, + local_layer_idx=core._bound.pool_idx[i], + num_heads=core.heads_q, + num_kv_heads=core.heads_kv, + head_size=core.head_dim, + quant_mode=core.quant_mode, + **step, + **_CALL_CONSTANTS, + ) + fused_sigmoid_mul_inplace(out, gate) + return cublas_mm(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: + """Recurrent half of a speculative verify step; see + `PrefillTarget._gdn_verify`.""" + core = self.core + # [n*p, C] -> [n, C, p]: the multi-token conv reads channel-major. + x3 = torch.transpose(torch.reshape(mq_d, [n, p, core.conv_dim]), 1, 2) + conv_out = causal_conv1d_triton_update( + x=x3, + conv_state=core._bound.conv[i], + weight=L.conv_w, + bias=None, + activation="silu", + conv_state_indices=si_d, + pad_slot_id=_PAD_SLOT_ID, + # Snapshots, not the accepted-token roll. + num_accepted_tokens=None, + intermediate_conv_window=core._bound.int_conv[i], + intermediate_state_indices=int_idx, + launch_dependent_kernels=False, + ) + # Back to token-major (a view). + mq_out = torch.reshape(torch.transpose(conv_out, 1, 2), [n * p, core.conv_dim]) + q, k, v = pack_gdn_decode_qkv( + mq_out, core.num_k_heads, core.head_k_dim, core.num_v_heads, core.head_v_dim + ) + g = torch.reshape( + fused_gdn_gating( + A_log=L.a_log, + a=a_d, + dt_bias=L.dt_bias, + beta=_SOFTPLUS_BETA, + threshold=_SOFTPLUS_THRESHOLD, + ), + [n, p, core.num_v_heads], + ) + beta = torch.reshape(torch.sigmoid(b_d), [n, p, core.num_v_heads]) + # Initial states are gathered from the live pool; the per-step buffer is + # indexed by batch row (`intermediate_state_indices[:n]`), not slot id. + src = torch.reshape( + nn.functional.embedding(si_d, torch.reshape(ssm_states, [ssm_states.shape[0], -1])), + [n, core.num_v_heads, core.head_v_dim, core.head_k_dim], + ) + out = fused_recurrent_gated_delta_rule_update( + torch.reshape(q, [n, p, core.num_k_heads, core.head_k_dim]), + torch.reshape(k, [n, p, core.num_k_heads, core.head_k_dim]), + torch.reshape(v, [n, p, core.num_v_heads, core.head_v_dim]), + g, + beta, + initial_state_source=src, + initial_state_indices=int_idx, + scale=None, + use_qk_l2norm_in_kernel=True, + cu_seqlens=None, + disable_state_update=True, + disable_output_calculation=False, + intermediate_states_buffer=core._bound.int_ssm[i], + cache_steps=p, + output=None, + ) + return torch.reshape(out, [1, n * p, core.num_v_heads, core.head_v_dim]) + + def _gdn_block(self, x, i, L, gdn_state) -> 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.""" + core = self.core + si_all, qsl_long, is_verify, num_decodes, gen_p, int_idx = 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] + if is_verify: + out = self._gdn_verify( + mixed_qkv, a, b, ssm_states, i, L, si_all, int_idx, num_decodes, gen_p + ) + else: + causal_conv1d_update( + x=mixed_qkv, + conv_state=core._bound.conv[i], + weight=L.conv_w, + bias=None, + silu_activation=True, + cache_seqlens=None, + conv_state_indices=si_all, + pad_slot_id=_PAD_SLOT_ID, + ) + # 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) + out = fused_sigmoid_gating_delta_rule_update( + A_log=L.a_log, + a=a, + 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, + initial_state_source=ssm_states, + initial_state_indices=si_all, + scale=None, + use_qk_l2norm_in_kernel=True, + cu_seqlens=qsl_long, + output=None, + ) + 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.""" + 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]), + weight=L.gnorm, + eps=core.eps, + 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`.""" + 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 + ) + # 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 torch.add(gated, torch.transpose(convolved, 0, 1)) + + # -------------------------------------------------------------- forward --- + + def forward( + self, + 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 + md = attn_metadata + self._check_step_contract(md) + step = self.step_args(md) + spec_metadata = kwargs.get("spec_metadata") + mm = md.mamba_metadata + + # 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]) + 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, + ) + 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; every sequence is a generation sequence. + batch = md.num_seqs + # Host-int check (graph-safe); see PrefillTarget.forward. + assert num_tokens == md.num_tokens, ( + f"the engine handed {num_tokens} activation rows for a batch its " + f"metadata describes as {md.num_tokens} tokens " + f"(padded_num_tokens={getattr(md, 'padded_num_tokens', None)}). A " + "prefill/piecewise CUDA-graph backend pads the activation buffer " + "past the batch and expects the model to slice it back; this " + "target does not implement that configuration. Decode CUDA graphs " + "pad with real dummy requests instead, so they are in the metadata " + "and this assert is silent for them." + ) + # Rows per generation sequence: 1 without speculation, else + # `runtime_draft_len + 1`. + gen_p = num_tokens // batch if batch else 1 + assert gen_p * batch == num_tokens, ( + f"{num_tokens} generation rows do not divide over {batch} generation " + "sequences; a ragged per-sequence draft length can be expressed " + "neither to the attention call nor to the multi-token recurrence" + ) + assert core.mtp_enabled or gen_p == 1, ( + f"{num_tokens} generation rows over {batch} generation sequences with " + "speculation off; the base path implements one query token per " + "generation sequence" + ) + # `gen_p > 1` is load-bearing for verify (see PrefillTarget.forward). + if spec_metadata is not None and batch > 0: + declared_p = spec_metadata.runtime_draft_len + 1 + assert gen_p == declared_p, ( + f"the batch carries {gen_p} query rows per generation sequence " + f"but the worker's runtime_draft_len says {declared_p}; the " + "per-step conv snapshots and the recurrent rollback buffer are " + "sized from the second and written from the first" + ) + assert 1 <= gen_p <= 4, ( + f"predicted_tokens_per_seq {gen_p} is outside the 1..4 range " + f"attention/thop_attention certifies" + ) + 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 + 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) + + # Every generation sequence carries its n-gram history; the engine's + # `has_initial_states` covers the context prefix only (empty here). + ple_carry = torch.ones( + [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). + if gen_p == 1: + ple_qsl = qsl_long + else: + ple_qsl = cu_seqlens_from_seq_lens(md.seq_lens_cuda[:batch], torch.int64) + ple_state = (core._bound.ple_state[0], core._bound.ple_state[1], si_all, ple_carry, ple_qsl) + + # Dense/sparse decision once per forward, from a host-side metadata read: a + # per-capture constant matching the engine's short/long decode graph families. + sparse = batch > 0 and int(md.kv_lens_runtime[:batch].max()) > core.dense_threshold + + qsa = ( + md.qsa_req_idx_per_token[:num_tokens], + md.qsa_logical_positions[:num_tokens], + md.qsa_sequence_lengths[:num_tokens], + md.qsa_visible_blocks[:num_tokens], + md.qsa_topk_indices[:num_tokens], + md.qsa_topk_row_starts[:num_tokens], + md.qsa_block_table, + md.tokens_per_block, + pos, + pos3, + coords, + sparse, + ) + + # Scratch shared by all attention layers (each fully rewrites them). + 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 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): + 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 core.is_attn[i]: + block_out = self._attention_block( + mixed, i, L, qkv, index_scratch, attn_out, qsa, step + ) + 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 + ) + 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 + ) + if spec_metadata is None: + return final + # `lm_head` reads the `[T, hidden]` mix; the MTP module reads the + # `[T, hc_count * hidden]` stream it was collapsed from. + return final, hc + + # ------------------------------------------------------------- the draft --- + + def _mtp_moe_block(self, x, dt) -> torch.Tensor: + """FP8 block-scale draft MoE; see `PrefillTarget._mtp_moe_block`.""" + core = self.core + M = core._mtp + logits = cublas_mm(x, M.router) + topk_ids, topk_w = renorm_moe_routing_op(logits, topk=core.topk, output_dtype=dt) + xq, xsf = fp8_quantize_1x128(x, use_ue8m0=_MTP_MOE_USE_UE8M0) + routed = fp8_block_scale_moe_runner( + routing_logits=None, + routing_bias=None, + hidden_states=xq, + hidden_states_scale=xsf, + gemm1_weights=M.fc1_w, + gemm1_weights_scale=M.fc1_s, + gemm2_weights=M.fc2_w, + gemm2_weights_scale=M.fc2_s, + num_experts=core.num_experts, + top_k=core.topk, + num_fused_shared_experts=None, + n_group=None, + topk_group=None, + intermediate_size=core.moe_inter, + local_expert_offset=0, + local_num_experts=core.num_experts, + routed_scaling_factor=None, + routing_method_type=_ROUTING_METHOD_RENORM, + topk_weights=topk_w, + topk_ids=topk_ids, + act_type=_ACT_TYPE_SWIGLU, + gemm1_clamp_limit=None, + ) + shared = cublas_mm(flashinfer_silu_and_mul(cublas_mm(x, M.sh_gu)), M.sh_dn) + gate_logits = cublas_mm(x, M.sh_gate) + return fused_sigmoid_gate_mul_add(routed, gate_logits, shared) + + def mtp_forward( + self, + embed_tokens: torch.Tensor, + all_rank_num_tokens, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + **kwargs, + ) -> torch.Tensor: + """One draft step over generation rows only (every step past the first in an + engine forward, and the first when the batch had no context). + + `spec_metadata` and `all_rank_num_tokens` are accepted and unused.""" + core = self.core + assert isinstance(attn_metadata, TrtllmAttentionMetadata) + assert core._bound is not None, ( + "the trunk's cache-manager bind has not run; the shell calls the " + "core before the worker, so this cannot be reached first" + ) + md = attn_metadata + rows = hidden_states.shape[0] + assert input_ids.shape[0] == rows, ( + f"the draft step's {input_ids.shape[0]} token ids and {rows} " + "hidden-state rows must describe the same tokens" + ) + assert hidden_states.shape[1] == core.stream, ( + f"the MTP module reads the trunk's {core.stream}-wide four-stream " + f"residual; the runtime handed it {hidden_states.shape[1]} columns" + ) + dt = hidden_states.dtype + dev = hidden_states.device + + # Projected on every call: the worker rewrites `attn_metadata` in place + # between draft steps (see PrefillTarget.mtp_forward). + step = self.step_args(md) + M = core._mtp + gen_seqs = md.num_seqs + p = rows // gen_seqs if gen_seqs else 1 + assert p * gen_seqs == rows, ( + f"{rows} generation rows do not divide over {gen_seqs} " + "generation sequences; a ragged per-sequence draft length cannot be " + "expressed to the attention call" + ) + + # The KV cache written here is the draft manager's, swapped in by the worker + # per call (see PrefillTarget.mtp_forward); take the layer index from the + # current manager every time. + cm = md.kv_cache_manager + assert cm is not None, "a paged KV cache is required" + offsets = cm.layer_offsets + assert core.mtp_layer_idx in offsets, ( + f"the MTP module's attention is global layer index " + f"{core.mtp_layer_idx}, which the cache manager " + f"{type(cm).__name__} does not map (it maps {len(offsets)} layers, " + f"is_draft={getattr(cm, 'is_draft', None)}); the engine allocated " + "no KV layer for the draft module" + ) + layer_idx = offsets[core.mtp_layer_idx] + window = cm.max_attention_window_vec[layer_idx] + assert window is None or window == md.max_seq_len, ( + f"the MTP layer's pool window is {window!r} but this call passes " + f"max_seq_len {md.max_seq_len} as the attention window; a sliding " + "window is not implemented here" + ) + + # Copy into a fresh contiguous `[3, N]`: a sliced view keeps the source row + # stride and the op rejects it under `use_mrope`. + assert position_ids.shape[0] == _QSA_POSITION_AXES, ( + f"expected the mRoPE position layout with {_QSA_POSITION_AXES} " + f"coordinate rows, got {tuple(position_ids.shape)}" + ) + pos3 = torch.empty( + [_QSA_POSITION_AXES, position_ids.numel() // _QSA_POSITION_AXES], + dtype=torch.int32, + device=position_ids.device, + ) + pos3.copy_(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1])) + + # ---- the front end -------------------------------------------------- + e = nn.functional.embedding(input_ids, embed_tokens) + x_e = cublas_mm(flashinfer_gemma_rmsnorm(e, M.pre_e_norm, core.eps), M.fc_e) + # `fc_hidden` is one `[hidden, hidden]` matrix applied to each of the four + # streams (the checkpoint's norm statistics say they are distinct streams). + h_n = self._grouped_norm(hidden_states, M.pre_h_norm) + x_h = torch.reshape( + cublas_mm(torch.reshape(h_n, [rows * core.hc_count, core.hidden]), M.fc_h), + [rows, core.stream], + ) + if _MTP_EMBED_BLOCK_ALL_STREAMS: + x = torch.add(x_h, torch.cat([x_e] * core.hc_count, -1)) + else: + x = torch.add( + x_h, + nn.functional.pad(x_e, [0, core.stream - core.hidden], mode="constant", value=0.0), + ) + + # ---- the decoder layer ---------------------------------------------- + mixed, inject = qwen4_exp_gated_residual( + self._grouped_norm(x, M.ahc_norm), M.ahc_down, M.ahc_up, M.ahc_inject, core.hc_count + ) + qkv = torch.empty([rows, core.qkv_width], dtype=dt, device=dev) + attn_out = torch.empty([rows, core.q_width], dtype=dt, device=dev) + qg = cublas_mm(mixed, M.q) + k = cublas_mm(mixed, M.k) + v = cublas_mm(mixed, M.v) + qkv_h = torch.reshape(qkv, [rows, core.heads_q + 2 * core.heads_kv, core.head_dim]) + q_dst, k_dst, v_dst = torch.split(qkv_h, [core.heads_q, core.heads_kv, core.heads_kv], 1) + q_src, gate = torch.split( + torch.reshape(qg, [rows, core.heads_q, 2 * core.head_dim]), + [core.head_dim, core.head_dim], + 2, + ) + q_dst.copy_(q_src) + k_dst.copy_(torch.reshape(k, [rows, core.heads_kv, core.head_dim])) + v_dst.copy_(torch.reshape(v, [rows, core.heads_kv, core.head_dim])) + # The trunk rope's certified cell at the module's own base (`mtp.rope_theta`). + fused_qk_norm_rope( + qkv, + num_heads_q=core.heads_q, + num_heads_k=core.heads_kv, + num_heads_v=core.heads_kv, + head_dim=core.head_dim, + rotary_dim=core.rotary_dim, + eps=core.eps, + q_weight=M.q_norm, + k_weight=M.k_norm, + base=core.mtp_theta, + is_neox=True, + position_ids=pos3, + is_qk_norm=True, + use_gemma=True, + use_mrope=True, + mrope_section1=core.mrope_section1, + mrope_section2=core.mrope_section2, + ) + # 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, + ) + fused_sigmoid_mul_inplace(attn_out, gate) + x = qwen4_exp_post_inject(x, cublas_mm(attn_out, M.o), inject) + + mixed, inject = qwen4_exp_gated_residual( + self._grouped_norm(x, M.mhc_norm), M.mhc_down, M.mhc_up, M.mhc_inject, core.hc_count + ) + return qwen4_exp_post_inject(x, self._mtp_moe_block(mixed, dt), inject) + + +class _VisionBlockWeights(NamedTuple): + """One vision block's weights (`.t()` GEMM views). LayerNorm has a bias, unlike + every norm in the language tower.""" + + n1_w: torch.Tensor + n1_b: torch.Tensor + qkv: torch.Tensor + qkv_b: torch.Tensor + proj: torch.Tensor + proj_b: torch.Tensor + n2_w: torch.Tensor + n2_b: torch.Tensor + fc1: torch.Tensor + fc1_b: torch.Tensor + fc2: torch.Tensor + fc2_b: torch.Tensor + + +class _VisionWeights(NamedTuple): + """The tower's weights: patch embedding, position table, blocks and patch merger.""" + + patch_w: torch.Tensor + patch_b: torch.Tensor + pos: torch.Tensor + blocks: list + m_n_w: torch.Tensor + m_n_b: torch.Tensor + m_fc1: torch.Tensor + m_fc1_b: torch.Tensor + m_fc2: torch.Tensor + m_fc2_b: torch.Tensor + + +def _vision_views(core: Qwen4ExpTp1ModelingV2) -> _VisionWeights: + w = core.w + return _VisionWeights( + # Not transposed: the entry takes the checkpoint's `[out, in, kt, kh, kw]`. + patch_w=w["vis_patch_w"], + patch_b=w["vis_patch_b"], + pos=w["vis_pos"], + blocks=[ + _VisionBlockWeights( + n1_w=w[f"l{b}_vis_n1_w"], + n1_b=w[f"l{b}_vis_n1_b"], + # Row layout (q heads, k, v) is already thop_attention's packed form. + qkv=w[f"l{b}_vis_qkv_w"].t(), + qkv_b=w[f"l{b}_vis_qkv_b"], + proj=w[f"l{b}_vis_proj_w"].t(), + proj_b=w[f"l{b}_vis_proj_b"], + n2_w=w[f"l{b}_vis_n2_w"], + n2_b=w[f"l{b}_vis_n2_b"], + fc1=w[f"l{b}_vis_fc1_w"].t(), + fc1_b=w[f"l{b}_vis_fc1_b"], + fc2=w[f"l{b}_vis_fc2_w"].t(), + fc2_b=w[f"l{b}_vis_fc2_b"], + ) + for b in range(core.v_depth) + ], + m_n_w=w["vis_m_n_w"], + m_n_b=w["vis_m_n_b"], + m_fc1=w["vis_m_fc1_w"].t(), + m_fc1_b=w["vis_m_fc1_b"], + m_fc2=w["vis_m_fc2_w"].t(), + m_fc2_b=w["vis_m_fc2_b"], + ) + + +class VisionTower: + """One request's images 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.""" + + def __init__(self, core: Qwen4ExpTp1ModelingV2) -> None: + self.core = core + self._w = _vision_views(core) + # Persistent scratch: `thop_attention` grows `workspace_` in place; never + # the engine's own workspace. + self._ws: torch.Tensor | None = None + + def __call__(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: + """`pixel_values` is `[sum(t*h*w), in_channels*kt*kh*kw]`, one flat patch per + row, channel-slowest. `grid_thw` is one `(frames, height, width)` row per + image and stays on the host (no synchronization). Returns + `[sum(t*h*w) / merge**2, out_hidden_size]`, one row per placeholder.""" + core = self.core + dt, dev = pixel_values.dtype, pixel_values.device + assert dt == core.w["embed"].dtype, ( + f"the tower runs at the activation dtype {core.w['embed'].dtype}; the " + f"runtime handed pixel_values as {dt}. The input processor's `dtype` " + "property is what decides this, and no cast happens here" + ) + assert pixel_values.dim() == 2 and pixel_values.shape[1] == core.v_pixel_width, ( + f"pixel_values must be [patches, {core.v_pixel_width}] flat patch rows; " + f"got {tuple(pixel_values.shape)}" + ) + 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: + 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]) + 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}" + ) + 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. + host_lens = torch.tensor(seg_lens, dtype=torch.int32) + host_types = torch.zeros([num_segments], dtype=torch.int32) + host_total = torch.tensor([num_patches, 0], dtype=torch.int32) + dev_lens = host_lens.to(dev) + if self._ws is None: + self._ws = torch.empty([0], dtype=torch.int8, device=dev) + + # ---- patch embedding + absolute position ----------------------------- + V = self._w + h = torch.add( + patch_embed_conv3d(pixel_values, V.patch_w, V.patch_b), + pos_embed_interpolate(V.pos, grid, core.v_merge), + ) + attn_out = torch.empty([num_patches, core.v_hidden], dtype=dt, device=dev) + + for B in V.blocks: + qkv = cublas_mm(layer_norm(h, B.n1_w, B.n1_b, _VISION_NORM_EPS), B.qkv, B.qkv_b) + qkv_h = torch.reshape(qkv, [num_patches, 3 * core.v_heads, core.v_head_dim]) + q_view, k_view, _ = torch.split(qkv_h, [core.v_heads] * 3, 1) + q_rot, k_rot = vision_rope_2d( + q_view, k_view, grid_thw, core.v_merge, theta=_VISION_ROPE_THETA + ) + q_view.copy_(q_rot) + k_view.copy_(k_rot) + attn_out = thop_attention( + q=qkv, + output=attn_out, + workspace_=self._ws, + sequence_length=dev_lens, + host_past_key_value_lengths=host_lens, + host_total_kv_lens=host_total, + context_lengths=dev_lens, + host_context_lengths=host_lens, + host_request_types=host_types, + max_num_requests=num_segments, + max_context_length=max_seg, + max_seq_len=max_seg, + attention_window_size=max_seg, + num_contexts=num_segments, + num_ctx_tokens=num_patches, + num_heads=core.v_heads, + num_kv_heads=core.v_heads, + head_size=core.v_head_dim, + **_VISION_CALL_CONSTANTS, + ) + h = torch.add(h, cublas_mm(attn_out, B.proj, B.proj_b)) + mlp = cublas_mm(layer_norm(h, B.n2_w, B.n2_b, _VISION_NORM_EPS), B.fc1, B.fc1_b) + # Non-gated MLP with the tanh GELU (not the merger's erf). + h = torch.add(h, cublas_mm(gelu_tanh(mlp), B.fc2, B.fc2_b)) + + # ---- the patch merger ------------------------------------------------ + # The norm runs over `hidden` before the 4-to-1 shuffle; the reshape is legal + # only because both order-bearing steps emitted in spatial-merge order. + merged = torch.reshape( + layer_norm(h, V.m_n_w, V.m_n_b, _VISION_NORM_EPS), [-1, core.v_merge_dim] + ) + # The erf GELU (`nn.GELU()`); within one bf16 ulp of tanh, hence a separate entry. + return cublas_mm(gelu_erf(cublas_mm(merged, V.m_fc1, V.m_fc1_b)), V.m_fc2, V.m_fc2_b) + + +class MTPLayer: + """The checkpoint's MTP module, replayed once per draft step: one more decoder + layer (two hyper-connections, gated GQA, 512-expert MoE) behind a two-branch + front end (`fc_embedding` over the NEXT token's embedding, `fc_hidden` over the + hidden state per stream). It returns `x` un-collapsed; `shared_head` owns the + module's mixer. The step runs on the targets (`mtp_forward`) over `core._mtp`. + + The MoE is FP8 128x128 weight blocks (the trunk's is NVFP4); re-quantizing it + would change the checkpoint's meaning, visible only as lower acceptance. + + Attention is 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. + + Not an `nn.Module`: the weights live in the core's `ParameterDict`, and a second + registration would break the weight load's both-directions coverage assert.""" + + def __init__(self, core: Qwen4ExpTp1ModelingV2, logits_processor) -> None: + self.core = core + self.logits_processor = logits_processor + + def __call__(self, *args, **kwargs) -> torch.Tensor: + return self.forward(*args, **kwargs) + + def shared_head( + self, hidden_states, lm_head, attn_metadata, return_context_logits + ) -> torch.Tensor: + """The module's own head: its stack-final hyper-connection mixer (a full gated + mixer with its own parameters, not a plain RMSNorm), then the trunk's + `lm_head` through the shell's `logits_processor`.""" + core = self.core + M = core._mtp + parts = torch.split(hidden_states, [core.hidden] * core.hc_count, 1) + normed = torch.cat( + [flashinfer_gemma_rmsnorm(part, w, core.eps) for w, part in zip(M.fhc_norm, parts)], + 1, + ) + final, _ = qwen4_exp_gated_residual(normed, M.fhc_down, M.fhc_up, None, core.hc_count) + return self.logits_processor.forward(final, lm_head, attn_metadata, return_context_logits) + + def forward( + self, + embed_tokens: torch.Tensor, + all_rank_num_tokens, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + **kwargs, + ) -> torch.Tensor: + """Route the draft step by the phase this call's metadata shows, per call: the + worker rewrites `attn_metadata` between draft steps, so one engine step's + draft calls can cross from prefill to decode.""" + return self.core._select_target(attn_metadata).mtp_forward( + embed_tokens, + all_rank_num_tokens, + input_ids, + position_ids, + hidden_states, + attn_metadata, + **kwargs, + ) + + +class DraftModel: + """The container the spec worker reaches the drafter through. The runtime + inspects only `mtp_layers` (`[0]`: MTP-Eagle replays one layer), + `embed_tokens`, `lm_head`, and `model.d2t` (correctly absent: shared + vocabulary). + + `embed_tokens` resolves through the core on every read: this container is + built while parameters are meta tensors, and materialization replaces the + tensor objects, so a captured reference stays on meta.""" + + def __init__(self, core: Qwen4ExpTp1ModelingV2, lm_head, logits_processor) -> None: + self.core = core + self.mtp_layers = [MTPLayer(core, logits_processor)] + self.lm_head = lm_head + + @property + def embed_tokens(self) -> torch.Tensor: + return self.core.w["embed"] + + +class ModelingV2Qwen38FlashNextNvfp4Sm103Tp1( + DecoderModelForCausalLM[Qwen4ExpTp1ModelingV2, PretrainedConfig] +): + @classmethod + def get_preferred_kv_cache_manager_version(cls, pretrained_config=None) -> str: + """Prefer KV cache manager V2: PLE state (`ple_layer_cache`) and QSA's index and + position planes exist only there. Under V1 the model builds and passes a + smoke run with the PLE layer stateless. The engine resolves this from the + built-in class; the declaration serves a direct instantiation.""" + return "V2" + + def __init__(self, model_config: ModelConfig): + cfg = model_config.pretrained_config + assert cfg is not None + + # Adopt the engine-normalized text tower (`Qwen4ExpTextConfig`) as the model's + # config: everything that sizes the pools reads + # `model_engine.model.model_config.pretrained_config`, and the composite + # parent mirrors text geometry through a pass that silently skips keys the + # parent already defaults. The swap is in place because the engine keeps its + # own reference to this `ModelConfig`. + text = getattr(cfg, "text_config", None) + assert text is not None, ( + "this target is the language tower of the composite Qwen4-Exp " + "export; pretrained_config carries no text_config" + ) + assert getattr(text, "architectures", None) == ["Qwen4ExpForCausalLM"], ( + f"the text tower must be the normalized Qwen4-Exp text config; got " + f"architectures {getattr(text, 'architectures', None)!r}" + ) + assert text.torch_dtype is not None, "the text tower must declare a dtype" + # Captured before the swap: the core needs the vision half and media token ids. + composite = cfg + model_config.pretrained_config = text + + super().__init__( + Qwen4ExpTp1ModelingV2(model_config, composite), + config=model_config, + hidden_size=text.hidden_size, + vocab_size=text.vocab_size, + ) + # Media placeholder ids for the executor; built lazily (`__init__` runs under + # meta init). + self._mm_token_ids_value = (self.model.image_token_id, self.model.video_token_id) + self._mm_token_ids_cache: torch.Tensor | None = None + # `lm_head` is outside the quantization scheme: a dense bf16 `LMHead`; a + # quantized head would register a different parameter set. + assert self.lm_head.weight.dtype == torch.bfloat16, ( + f"expected a bf16 lm_head; the shell built {self.lm_head.weight.dtype}, " + "which means the engine resolved a quantized output head for a " + "checkpoint that stores it unquantized" + ) + + # The speculative branch, built only when `speculative_config` asked for it. + # Written out rather than inherited: the in-tree one-engine shell builds its + # drafter via a module-level function dispatching on `model_type` (still the + # upstream family name), which would construct trtllm's own MTP layer. + self.spec_config = getattr(model_config, "spec_config", None) + self.draft_model = None + self.spec_worker = None + # No separate draft checkpoint; the loader reads this, so None is stated. + self.draft_config = None + if self.spec_config is not None: + assert self.model.mtp_enabled, "core built without the MTP module" + assert hasattr(self, "logits_processor"), ( + "the inherited shell no longer exposes logits_processor; the " + "draft head and the shell's own gather both project through it" + ) + self.draft_model = DraftModel(self.model, self.lm_head, self.logits_processor) + # Must not be left at `get_spec_worker`'s default False: the resource + # manager decides separately (`_should_create_separate_draft_kv_cache`) and + # builds the target manager from target layer masks alone, so left False + # the MTP layer is handed a manager with no `layer_offsets` entry for it. + self.use_separate_draft_kv_cache = should_use_separate_draft_kv_cache(self.spec_config) + self.spec_worker = get_spec_worker( + self.spec_config, + model_config, + model_config.mapping, + use_separate_draft_kv_cache=self.use_separate_draft_kv_cache, + ) + assert self.spec_worker is not None, ( + f"the engine resolved no spec worker for spec_dec_mode " + f"{self.spec_config.spec_dec_mode!r}" + ) + # Caches the draft->target vocab offset map (correctly None here). + self.spec_worker.set_draft_model(self.draft_model) + + # ----------------------------------------------------------- multimodal --- + # Names the executor reads off a model class by `getattr`; each is silently + # optional and its absence costs correctness rather than raising. + + @property + def multimodal_data_device_paths(self) -> list[str]: + """Tensors under ``multimodal_data`` the executor moves to CUDA. Absent, it moves + everything, putting `image_grid_thw` on the device where + `pos_embed_interpolate` and `vision_rope_2d` would each pay a D2H sync.""" + return [ + "image.pixel_values", + "mrope_config.mrope_position_ids", + "mrope_config.mrope_position_deltas", + ] + + @property + def multimodal_token_ids(self) -> torch.Tensor: + """Placeholder ids whose rows the vision embeddings replace. Absent, the + executor uses `input_ids >= vocab_size`, but this checkpoint's placeholders + (248056, 248057) are in vocabulary, so nothing is selected and the engine + generates fluent text while blind. Built on first read (meta init); not a + parameter or buffer.""" + if self._mm_token_ids_cache is None: + self._mm_token_ids_cache = torch.tensor(self._mm_token_ids_value, dtype=torch.int32) + return self._mm_token_ids_cache + + def _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. + + 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.""" + 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: + 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" + ) + 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: + # Chunked prefill or reused prefix: splice only this chunk's slice. + start = int(runtime.num_cached_mm_tokens) + count = int(runtime.num_mm_tokens_in_chunk) + total = embeds.shape[0] + assert int(runtime.total_embeds_in_request) == total, ( + f"the executor counted {runtime.total_embeds_in_request} " + f"placeholder rows for this request and the tower produced " + f"{total}; the two are derived from the same grid 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: + return None + vision = rows[0] if len(rows) == 1 else torch.cat(rows, 0) + mm_token_indices = kwargs.get("mm_token_indices") + assert mm_token_indices is not None, ( + "the executor shipped multimodal params but no mm_token_indices; " + "that pairing is what `multimodal_token_ids` on this class buys, and " + "without it there is no sync-free way to find the placeholder rows" + ) + assert mm_token_indices.shape[0] == vision.shape[0], ( + f"the batch has {mm_token_indices.shape[0]} placeholder rows and the " + f"tower produced {vision.shape[0]} embedding rows" + ) + 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 + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: torch.IntTensor | None = None, + position_ids: torch.IntTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + return_context_logits: bool = False, + spec_metadata=None, + lora_params: dict | None = None, + **kwargs, + ): + """One engine step. + + Without a spec worker this delegates to the base forward with every argument + it was handed (`resource_manager` stays in `**kwargs` deliberately: naming + it would drop it from what the base receives). With one, the shell owns the + logits gather (one-model modes set `without_logits`) and the worker needs + ungathered hidden states; `position_ids` is passed in the engine's layout + (flattening gives silently wrong draft positions). The vision 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 + if self.spec_worker is None: + assert spec_metadata is None, ( + "spec_metadata arrived without a spec worker; this engine was " + "built without a speculative_config" + ) + return super().forward( + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + return_context_logits, + spec_metadata, + lora_params, + **kwargs, + ) + assert spec_metadata is not None, ( + "the spec worker is built but the engine passed no spec_metadata" + ) + # `final` is the `[T, hidden]` mix for `lm_head`, `hc` the stream the MTP reads. + final, hc = self.model( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + lora_params=lora_params, + spec_metadata=spec_metadata, + **kwargs, + ) + # `gather_ids`: one row per context request, `runtime_draft_len + 1` per + # generation request. + logits = self.logits_processor.forward( + nn.functional.embedding(spec_metadata.gather_ids, final), + self.lm_head, + attn_metadata, + True, + ) + return self.spec_worker( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hc, + logits=logits, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=self.draft_model, + resource_manager=kwargs.get("resource_manager"), + ) + + def load_weights(self, weights, *args, **kwargs): + _weights.MODEL_WEIGHTS.load(self, weights) + + def post_load_weights(self): + super().post_load_weights() + self.model.post_load_weights() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/__init__.py new file mode 100644 index 000000000000..1059445198db --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ops used only by this target. Shared ops stay in the catalog.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/quantization/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/ops/quantization/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/__init__.py 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 new file mode 100644 index 000000000000..4982ae3f00fe --- /dev/null +++ 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 @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""In-place sigmoid output gating: ``attention_output *= sigmoid(gate)``.""" + +from __future__ import annotations + +import torch + +# Registered by the module that defines it: after `import +# tensorrt_llm._torch.custom_ops` alone, torch.ops.trtllm has no +# `fused_sigmoid_mul_inplace` (measured). +import tensorrt_llm._torch.attention.kernels.fused_qk_norm_rope_gate # noqa: F401 +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The checkpoint's gated attention: 24 q heads of 256, q projection [T, 24, 512]. +_CKPT = dict(heads=24, head_dim=256, dtype=torch.bfloat16, gate="interleaved") + + +class FusedSigmoidMulInplace(OpWrapper): + """`attention_output[t, i] *= sigmoid(gate[t, i // D, i % D])` in place, `D = gate.shape[-1]`. + + Returns `attention_output`, the tensor it wrote. Both operands are widened to + fp32 and the product is rounded once to `attention_output.dtype`. `gate` is + read through its own token and head strides, so the gate half of an + interleaved `[T, heads, 2 * D]` q projection is passed as a strided view with + no repack (bitwise equal to its compacted copy, measured). A 2-D `gate` of + `attention_output`'s shape is one head spanning the row. + + Fusion boundary: sigmoid and multiply only. The caller owns producing both + operands and everything downstream (the o projection). `gate` is read-only. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode: 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, + ), + ), + Cell( + why="fp16 operands", + spec=dict(tokens=33, heads=4, head_dim=128, dtype=torch.float16, gate="packed"), + ), + 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), + ), + ) + + note = """ + Metadata consumed: none. One Triton kernel, JIT-compiled on the first call, + no autotuner: grid `(T, cdiv(hidden, block))`, `block = min(next_pow2(hidden), + 1024 if T < 1024 else 2048)`, 4 warps. + + Preconditions, every one enforced loudly by the op's own bare `assert` + (an AssertionError with no message): `attention_output` 2-D with + `stride(-1) == 1` (row stride free); `gate` 2-D of `attention_output`'s + shape, or 3-D with `shape[0] == T` and `shape[1] * shape[2] == hidden`, + `stride(-1) == 1`. A CPU tensor raises ValueError from Triton. 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). + + Numerics: bf16/fp16 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. + + Hard limits + T * row stride token offset computed in int64 boundary 2^63 not probed + (largest run 2048 rows) + + Not certified: CUDA-graph capture, torch.compile. + """ + + def __call__(self, attention_output: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + torch.ops.trtllm.fused_sigmoid_mul_inplace(attention_output, gate) + return attention_output + + def reference(self, attention_output: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + """What `attention_output` holds after the call; neither input is mutated. + + `gate.reshape(T, hidden)` is the per-head `(t, i // D, i % D)` read, for + the 2-D and 3-D forms alike. + """ + num_tokens, hidden = attention_output.shape + g = gate.reshape(num_tokens, hidden).float() + return (attention_output.float() * torch.sigmoid(g)).to(attention_output.dtype) + + +fused_sigmoid_mul_inplace = FusedSigmoidMulInplace() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/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 new file mode 100644 index 000000000000..7bfbe1ee84a6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_erf.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Non-gated exact (erf) GELU -- what `torch.nn.GELU()` is -- the vision patch +merger's activation, which no config key names. + +Tier-3 composed: one `torch.nn.functional.gelu(x)` call, `approximate="none"`. +`torch.ops.trtllm.moe_gelu` computes this function bitwise but is not a +substitute: its `tile_idx_to_mn_limit` / `num_non_exiting_tiles` operands must +be int32 *device* tensors (on the host it raises nothing and returns garbage, +nan, or an illegal memory access), so wrapping it means launching two fills +before the call; it also rejects fp32 and non-2-D inputs. +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: `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) + + +class GeluErf(OpWrapper): + """`0.5 * x * (1 + erf(x / sqrt(2)))`, elementwise, any shape, new tensor + out in `x`'s dtype. + + Fusion boundary: the activation only. `merger.linear_fc1` before it and + `merger.linear_fc2` after it are the caller's; there is no gate half and no + multiply. This is the erf GELU and not the tanh one (`GeluTanh`, the block + MLPs' `hidden_act`); both appear in this one tower. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # 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". + 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), + ), + ) + ) + + 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). + + 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. + + 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. + + Not certified: float8 operands. No index-typed argument; numel past 2^31 + not probed. + """ + + def __call__(self, x: torch.Tensor) -> torch.Tensor: + return F.gelu(x) + + def reference(self, x: torch.Tensor) -> torch.Tensor: + """The closed form in fp64 from primitive ops, rounded once to `x.dtype`.""" + x64 = x.double() + return (0.5 * x64 * (1.0 + torch.erf(x64 / _SQRT_2))).to(x.dtype) + + +gelu_erf = GeluErf() 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 new file mode 100644 index 000000000000..2f6c5955f1dc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/gelu_tanh.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Non-gated tanh GELU -- HuggingFace's `gelu_pytorch_tanh` -- the vision +tower's block-MLP activation. + +Tier-3 composed: one `torch.nn.functional.gelu(x, approximate="tanh")` call. +No registered op computes a plain non-gated GELU over an unquantized tensor: +`flashinfer_gelu_tanh_and_mul` is the gated form (and silently returns a +half-width `gelu(a) * b` when handed this tower's non-gated activand), +`moe_gelu` is the erf GELU behind MoE tile metadata, and the +`cute_dsl_nvfp4_dense_gemm_gelu*` ops are NVFP4 GEMM epilogues. +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The block MLP's `linear_fc1` is `[4304, 1152]`: one width, nothing split. +INTERMEDIATE_SIZE = 4304 + +_SQRT_2_OVER_PI = math.sqrt(2.0 / math.pi) + + +class GeluTanh(OpWrapper): + """`0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x**3)))`, elementwise, + any shape, new tensor out in `x`'s dtype. + + Fusion boundary: the activation only. `linear_fc1` before it and + `linear_fc2` after it are the caller's; there is no gate half and no + multiply. This is the tanh GELU and not the erf one (`GeluErf`, the patch + merger's); which one a module wants is a property of the checkpoint. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # 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". + 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), + ), + ) + ) + + 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). + + 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. + + 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. + + Not certified: float8 operands. No index-typed argument; numel past 2^31 + not probed. + """ + + def __call__(self, x: torch.Tensor) -> torch.Tensor: + return F.gelu(x, approximate="tanh") + + def reference(self, x: torch.Tensor) -> torch.Tensor: + """The closed form in fp64 from primitive ops, rounded once to `x.dtype`.""" + x64 = x.double() + inner = _SQRT_2_OVER_PI * (x64 + 0.044715 * x64 * x64 * x64) + return (0.5 * x64 * (1.0 + torch.tanh(inner))).to(x.dtype) + + +gelu_tanh = GeluTanh() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/__init__.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/ops/attention/__init__.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/__init__.py 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 new file mode 100644 index 000000000000..3bf02f3b3eac --- /dev/null +++ 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 @@ -0,0 +1,225 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA Top-K group ids expanded into raw KV token indices, plus the causal tail.""" + +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 geometry: `indexer_compress_ratio` and +#: `indexer_budget`, so `block_topk` 512 and an output 2051 wide. +_CKPT = dict(compress_ratio=4, token_topk=2048) + + +def _spec(batch, **overrides) -> dict: + """A cell spec at this checkpoint's geometry and the engine's metadata dtypes. + + `batch` is one of: + ("sweep", n) n rows at 4096 + 13 i, length = position + 1 + ("serving", requests, context_lens) one packed step: context chunks, then + two MTP rows per remaining request + ("rows", positions) explicit positions, length = position + 1 + """ + 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. + + Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels. + triton_expand_qsa_block_indices`: one Triton launch, + `_expand_qsa_block_indices_kernel`. Per row (one per query token), with + `final_topk = token_topk + compress_ratio - 1`: + + 1. column `c < token_topk` holds member `c % cr` of group + `block_indices[row, c // cr]`; `-1` if that group is `-1` or the + token is `>= sequence_lengths[row]`; + 2. the last `(query_positions[row] + 1) % cr` visible tokens -- the + group the indexer never scored -- go in the columns right after the + selected groups, i.e. from column `(# non-negative groups) * cr`; + 3. everything else is `-1`. + + Selection order is kept (nothing sorted, compacted or de-duplicated). + `query_positions` feeds the tail only; the expansion is masked by + `sequence_lengths` alone, so the causal bound on the *selection* is the + producer's. Returns a fresh contiguous int32 `[rows, final_topk]`; inputs + are read-only. + + Fusion boundary: index arithmetic and the tail only. Scoring, Top-K, the + per-row causal group bound and the geometry are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + # ---- decode rows: the column-split launch below 128 rows ------------- + 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))), + ), + Cell( + why="warm-up capture: 8191 rows of one request from 0", + spec=_spec(("serving", 1, (8191,))), + ), + 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 = """ + Metadata consumed: none -- every input is an argument. The target passes + `qsa_topk_indices[:rows]` (int32, written by `qsa_topk_select`), + `qsa_logical_positions` (int64) and `qsa_sequence_lengths` (int32: the + 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): + - `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 + 12..15 lost. `qsa_topk_select` emits packed rows. + - selected groups are complete and visible (`g < (pos + 1) // cr`). An + incomplete group comes back with `-1` holes in place and its real + tokens repeated in the tail; a group past the query's horizon but + 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. + + 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. + + Hard limits (read from the kernel, not probed): the row offset is + `program_id * stride(0)` in int32, so `rows * stride(0) >= 2^31` (4 Mi + rows at width 512) would wrap. Certified extent: 8192 rows x 2051 columns. + + Not certified: multiple devices, `torch.compile`, rows above 8192. + """ + + def __call__( + self, + block_indices: torch.Tensor, + query_positions: torch.Tensor, + sequence_lengths: torch.Tensor, + *, + compress_ratio: int, + token_topk: int, + ) -> torch.Tensor: + return _qsa_kernels.triton_expand_qsa_block_indices( + block_indices, + query_positions, + sequence_lengths, + compress_ratio=compress_ratio, + token_topk=token_topk, + ) + + def reference( + self, + block_indices: torch.Tensor, + query_positions: torch.Tensor, + sequence_lengths: torch.Tensor, + *, + compress_ratio: int, + token_topk: int, + ) -> torch.Tensor: + """The expansion as the contract states it, in int64 torch.""" + rows, block_topk = block_indices.shape + final_topk = token_topk + compress_ratio - 1 + device = block_indices.device + blocks = block_indices.long() + lengths = sequence_lengths.long() + offsets = torch.arange(compress_ratio, device=device) + + expanded = (blocks[:, :, None] * compress_ratio + offsets).reshape(rows, token_topk) + chosen = (blocks >= 0).repeat_interleave(compress_ratio, dim=1) + keep = chosen & (expanded < lengths[:, None]) + out = torch.full((rows, final_topk), -1, dtype=torch.int64, device=device) + out[:, :token_topk] = torch.where(keep, expanded, -1) + + first_free = (blocks >= 0).sum(dim=1) * compress_ratio + visible = query_positions.long() + 1 + tail_start = visible // compress_ratio * compress_ratio + row = torch.arange(rows, device=device) + for offset in range(compress_ratio - 1): + column = first_free + offset + token = tail_start + offset + live = (offset < visible - tail_start) & (token < lengths) & (column < final_topk) + out[row[live], column[live]] = token[live] + return out.to(torch.int32) + + +qsa_expand_block_indices = QsaExpandBlockIndices() 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 new file mode 100644 index 000000000000..d0dbbfa1c66a --- /dev/null +++ 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 @@ -0,0 +1,474 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA compressed-block index scoring, read straight out of the paged index-K side cache.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# Approved exception: QSA's selection path is Triton called as Python, with no +# torch.ops boundary. The engine's own sparse hook (`select_qsa_paged_tokens` +# in `_torch/attention/backends/sparse/qsa/module.py`) calls this same callable. +from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import ( + triton_qsa_paged_index_scores as _triton_qsa_paged_index_scores, +) + +_CHECKPOINT = dict( + heads=4, + head_dim=128, + tokens_per_block=32, + compress_ratio=4, + dtype=torch.bfloat16, + # `qsa_logical_positions` is int64, `qsa_req_idx_per_token` int32. + positions_dtype=torch.int64, + request_dtype=torch.int32, + # The target's q is a head slice of its packed [T, 4 + 1 + 1, 128] index + # scratch buffer, so its row stride is 768 rather than 512. + q_layout="scratch_slice", +) + + +def _spec(**overrides) -> dict: + """A cell spec: this checkpoint's geometry unless the cell says otherwise.""" + return {**_CHECKPOINT, **overrides} + + +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) +) + + +class QsaPagedIndexScores(OpWrapper): + """Score every causally visible compressed block of each query row. + + Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels. + triton_qsa_paged_index_scores` (one Triton launch after allocating the + output). For row `r` and compressed block `b` of request + `request_indices[r]`:: + + anchor = b * compress_ratio + compress_ratio - 1 + key_b = index_cache[block_table[req, anchor // tpb], anchor % tpb, 0] + out[r, b] = sum_h relu(q[r, h] . key_b) * index_head_dim ** -0.5 + + The block's key is whatever the compression step left at the group's LAST + position; nothing is pooled here. Row `r` sees `(query_positions[r] + 1) // + compress_ratio` blocks (complete groups only). Columns at or past that are + `-inf` when `only_visible_blocks` is False and **not written** -- the fresh + allocation's garbage -- when it is True. The column count is the block + table's width, `block_table.shape[1] * tpb // compress_ratio`. + + `context_rows` is a layout declaration ("consecutive rows mostly share a + request"), not a phase one: with more than 256 rows and an index-head count + that divides 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. + + Fusion boundary: the table walk, the gather, the dots, relu, head sum, + scale and causal bound are inside. Producing `q` (norm, RoPE), writing and + compressing the index cache, and everything after the scores are not. + Returns a fresh fp32 `[rows, columns]`; nothing is written in place. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + # ---- 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", + spec=_spec( + layout="decode", + gen_positions=(7000, 5000, 4200, 2100), + rows_per_request=2, + slots=32, + pages=_WIDTH, + only_visible_blocks=True, + context_rows=False, + branch="row_wise", + ), + ), + 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", + spec=_spec( + layout="decode", + gen_positions=tuple(100 + (499 * g) % 8000 for g in range(32)), + rows_per_request=2, + slots=32, + pages=_WIDTH, + only_visible_blocks=True, + context_rows=False, + 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( + layout="capture", + lengths=(1024, 512, 2048, 256, 1500, 700, 1151, 1000), + slots=8, + pages=_WIDTH, + only_visible_blocks=True, + context_rows=True, + branch="query_tile", + ), + ), + 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 = """ + Metadata consumed: none -- every input is an argument. What the target + passes: `block_table` is `QSAAttentionMetadata.qsa_block_table` (raw + pool-page ids, int32, -1 past each sequence); `query_positions` / + `request_indices` are `qsa_logical_positions` / `qsa_req_idx_per_token`; + `index_cache` is the layer's bf16 index-K plane (`get_index_k_buffer`). + + Preconditions (not checked): every block-table entry the causal bound + reaches is a valid page, and `request_indices[r]` is a row of the table -- + the loads are masked by the bound, not by the pool size, so a violation is + an out-of-bounds read (not probed). The compression step must already have + run: an uncompressed anchor scores as a raw index key. `rows == 0` returns + an empty `[0, columns]` tensor without launching. + + The op validates its own metadata loudly (10 `ValueError`s, pinned by the + test), 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 + 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. + """ + + def __call__( + self, + q: torch.Tensor, + index_cache: torch.Tensor, + block_table: torch.Tensor, + query_positions: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + compress_ratio: int, + only_visible_blocks: bool, + context_rows: bool, + ) -> torch.Tensor: + return _triton_qsa_paged_index_scores( + q=q, + index_cache=index_cache, + block_table=block_table, + query_positions=query_positions, + request_indices=request_indices, + tokens_per_block=tokens_per_block, + compress_ratio=compress_ratio, + only_visible_blocks=only_visible_blocks, + context_rows=context_rows, + ) + + def reference( + self, + q: torch.Tensor, + index_cache: torch.Tensor, + block_table: torch.Tensor, + query_positions: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + compress_ratio: int, + only_visible_blocks: bool, + context_rows: bool, + ) -> torch.Tensor: + """fp32 scores, `-inf` past each row's bound -- or NaN there, meaning + "unspecified", when `only_visible_blocks` is True. `context_rows` is + ignored: it picks a kernel, not a result. + + Walked in row chunks only to bound the gathered keys, which reach + 8.7 GiB in fp32 at 8191 rows x 2080 columns. + """ + rows, _, head_dim = q.shape + num_columns = block_table.shape[1] * tokens_per_block // compress_ratio + blocks = torch.arange(num_columns, device=q.device) + anchors = blocks * compress_ratio + compress_ratio - 1 + logical_pages, in_page = anchors // tokens_per_block, anchors % tokens_per_block + table = block_table.long() + requests = request_indices.long() + scores = torch.empty(rows, num_columns, dtype=torch.float32, device=q.device) + chunk = max(1, 2**26 // max(num_columns * head_dim, 1)) + for lo in range(0, rows, chunk): + hi = min(lo + chunk, rows) + # Pages past a sequence are -1; their columns are masked below. + physical = table[requests[lo:hi]][:, logical_pages].clamp_min(0) + keys = index_cache[physical, in_page.expand(hi - lo, num_columns)][:, :, 0, :] + per_head = torch.einsum("rhd,rbd->rbh", q[lo:hi].float(), keys.float()) + scores[lo:hi] = torch.relu(per_head).sum(-1) * (head_dim**-0.5) + visible = blocks[None, :] < ((query_positions.long() + 1) // compress_ratio)[:, None] + tail = float("nan") if only_visible_blocks else float("-inf") + return scores.masked_fill(~visible, tail) + + +qsa_paged_index_scores = QsaPagedIndexScores() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_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 new file mode 100644 index 000000000000..341287f03b95 --- /dev/null +++ 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 @@ -0,0 +1,268 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA index side-cache store: raw index keys and position coordinates into their pages.""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: 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 +#: buffer; `slots` x `width` is `qsa_block_table`. `pages` is the pool. +_CKPT = dict( + hkv=1, + head_dim=128, + index_heads=4, + tokens_per_block=32, + axes=3, + slots=32, + width=260, + key_dtype=torch.bfloat16, + request_dtype=torch.int32, + position_dtype=torch.int64, + key_layout="packed", + coords_layout="transposed", + meta_layout="contiguous", + cache_layout="plain", + table_layout="dense", +) + +_DECODE = ((0, 1), (1, 1), (31, 1), (32, 1), (33, 1), (63, 1), (64, 1), (95, 1)) + + +def _spec(segments, pages, **overrides) -> dict: + """`segments` is `(first position, rows)` per request, in row order.""" + 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. + + For row t, with `page = block_table[request_indices[t], logical_positions[t] + // tokens_per_block]` and `slot = logical_positions[t] % tokens_per_block`:: + + 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. + + 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. + + Fusion boundary: page lookup, slot arithmetic, both scatters. The index + projection, the query's norm and rotation, group compression, scoring, + selection, sparse attention and the main K/V pool are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + 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), + ), + Cell( + why="int32 requests and int32 positions", + spec=_spec(_DECODE, pages=32, position_dtype=torch.int32), + ), + Cell( + why="int64 requests and int64 positions", + spec=_spec(_DECODE, pages=32, request_dtype=torch.int64), + ), + Cell(why="zero rows: both planes unchanged", spec=_spec((), pages=8)), + ) + + note = """ + Metadata consumed: none -- every input is an argument. The target passes + `qsa_req_idx_per_token` (int32) and `qsa_logical_positions` (int64) + sliced to the step's rows, `qsa_block_table` unsliced, the layer's + `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. + """ + + def __call__( + 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, + ) -> 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 + ) + return index_cache, position_cache + + def reference( + 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, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Both planes after the store, as fresh tensors, built row by row. + + Pages and slots are computed in Python and each write is a basic + (integer) index -- none of the advanced indexing `__call__` is made of. + """ + expected_index, expected_position = index_cache.clone(), position_cache.clone() + table = block_table.tolist() + axes = position_coordinates.shape[-1] + 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" + ) + 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" + ) + + +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 new file mode 100644 index 000000000000..06796397a5f4 --- /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.py @@ -0,0 +1,269 @@ +# 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: this step's rows into the HND main pool.""" + +from __future__ import annotations + +from typing import Tuple + +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 main K/V geometry and the engine's QSA metadata. `heads_q` +#: places k and v as head slices of the packed `[T, 24 + 2 + 2, 256]` QKV +#: buffer; `slots` x `width` is `qsa_block_table` (max_num_sequences x +#: max_blocks_per_seq). `pages` is the pool, independent of the batch. +_CKPT = dict( + hkv=2, + head_dim=256, + heads_q=24, + tokens_per_block=32, + slots=32, + width=260, + row_dtype=torch.bfloat16, + cache_dtype=torch.bfloat16, + request_dtype=torch.int32, + position_dtype=torch.int64, + k_layout="qkv_slice", + table_layout="dense", + hole=None, +) + +_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} + + +def _bits(tensor: torch.Tensor) -> torch.Tensor: + """The same storage as integers: fp8 has no comparison ops, and -0.0 != 0.0.""" + as_int = {1: torch.uint8, 2: torch.int16, 4: torch.int32}[tensor.element_size()] + return tensor.view(as_int) + + +class QsaPagedKvStore(OpWrapper): + """Scatter every row's K and V into its request's page of the HND pool. + + Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels. + triton_qsa_paged_kv_store`: one Triton launch over `(rows, Hkv)`. For row t:: + + page = block_table[request_indices[t], logical_positions[t] // tokens_per_block] + slot = logical_positions[t] % tokens_per_block + k_cache[page, :, slot, :] = k[t] v_cache[page, :, slot, :] = v[t] + + `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. + + 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 + index plane is never touched. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + 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"), + ), + Cell( + why="skip: negative logical position", + spec=_spec(_DECODE, pages=64, hole="negative_position"), + ), + Cell( + why="skip: page column past the table", + spec=_spec(_DECODE, pages=64, hole="column_past_table"), + ), + Cell(why="zero rows: returns before launching", spec=_spec((), pages=4)), + ) + + note = """ + Metadata consumed: none -- every input is an argument. The target passes + `qsa_req_idx_per_token` (int32, the row's batch index) and + `qsa_logical_positions` (int64) sliced to the step's rows, + `qsa_block_table` unsliced (raw pool page ids, `-1` unallocated), and the + layer pool's K and V role views of `get_buffers(i, kv_layout="HND")`. + + Preconditions: + - destination slots are unique across the batch: two rows naming one + (page, slot) both store, unordered. The engine's mapping cannot + collide. + - `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. + + 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. + + 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. + + Not certified: duplicate destinations, multiple devices, a scaled fp8 + pool (there is no scale argument). + """ + + def __call__( + self, + k: torch.Tensor, + v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + request_indices: torch.Tensor, + logical_positions: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + _qsa_kernels.triton_qsa_paged_kv_store( + k=k, + v=v, + k_cache=k_cache, + v_cache=v_cache, + request_indices=request_indices, + logical_positions=logical_positions, + block_table=block_table, + tokens_per_block=tokens_per_block, + ) + return k_cache, v_cache + + def reference( + self, + k: torch.Tensor, + v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + request_indices: torch.Tensor, + logical_positions: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Both caches after the scatter, as fresh tensors; the inputs are not touched.""" + expected_k, expected_v = k_cache.clone(), v_cache.clone() + requests = request_indices.long() + positions = logical_positions.long() + columns = positions.div(tokens_per_block, rounding_mode="floor") + table_rows, table_width = block_table.shape + addressable = ( + (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() + addressable &= (pages >= 0) & (pages < k_cache.shape[0]) + page, slot = pages[addressable], (positions % tokens_per_block)[addressable] + expected_k[page, :, slot, :] = k[addressable].to(expected_k.dtype) + expected_v[page, :, slot, :] = v[addressable].to(expected_v.dtype) + return expected_k, expected_v + + def is_valid( + self, + k: torch.Tensor, + v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + request_indices: torch.Tensor, + logical_positions: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + ) -> None: + # The kernel reads these two as flat unit-stride buffers. Measured: a + # column of a 2-wide buffer is read through its base pointer and the + # rows land at other rows' slots, with nothing raised. + assert request_indices.is_contiguous(), "request_indices must be a contiguous 1-D tensor" + assert logical_positions.is_contiguous(), ( + "logical_positions must be a contiguous 1-D tensor" + ) + + +qsa_paged_kv_store = QsaPagedKvStore() 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 new file mode 100644 index 000000000000..c1af92df810c --- /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.py @@ -0,0 +1,323 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA sparse GQA: attention over an explicitly selected set of paged K/V tokens.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# Approved exception: QSA's sparse path never reaches the thop attention op's +# sparse argument group, so there is no torch.ops surface for this computation. +from tensorrt_llm._torch.attention.backends.sparse.qsa.kernels import ( + triton_qsa_paged_sparse_gqa as _triton_qsa_paged_sparse_gqa, +) + +_CHECKPOINT = dict( + q_heads=24, + kv_heads=2, + head_dim=256, + tokens_per_block=32, + compress_ratio=4, + token_topk=2048, + cache_dtype=torch.bfloat16, + q_dtype=torch.bfloat16, + # 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: 260 pages per slot, -1 past each sequence. + pages=260, +) + +_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. + """ + 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}) + + +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. + + Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels. + triton_qsa_paged_sparse_gqa`. For row `r`, query head `h`, KV head + `h // (Hq // Hkv)` (contiguous groups, not interleaved), over the columns + `j` of `selected_tokens[r]` that name a resident token:: + + t = selected_tokens[r, j] # logical position + page = block_table[request_indices[r], t // tokens_per_block] + s_j = (q[r, h] . k_cache[page, kv, t % tpb]) * softmax_scale + out[r, h] = sum_j softmax(s)_j * v_cache[page, kv, t % tpb] + + A negative index, a page past the table's width, a negative or + out-of-pool table entry, and a request index outside the table's rows all + contribute nothing; a row left with nothing returns zeros. Indices are not + deduplicated: a token named twice is attended twice. Returns a fresh, + contiguous `[rows, Hq, D]` tensor of `q.dtype`, even for a strided `q`. + + Two kernels, chosen inside the call: with `query_positions` and + `compress_ratio` given and `rows * Hkv > 512`, a causal-prefix prefill + kernel; otherwise split-K, plus a merge kernel when it used more than one + split. Each cell names the launches it was measured on. + + **The prefill pair is a bound, not a mask.** The prefill kernel reads only + `min(visible // cr, block_topk) * cr + visible % cr` leading columns of a + row (`visible = query_positions[r] + 1`, rounded up to a 64-wide tile) and + nothing after them. That is exactly what QSA's own top-k + expansion emit, + packed at the front; a selection with valid entries past the bound is + silently truncated (measured: a planted index moving the right answer by + 0.77 moved the op's by nothing). The split-K kernel ignores the pair. + + 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. + """ + + ARCHS = frozenset({Arch.SM_103}) + + 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), + ), + # ---- 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), + ), + 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), + ), + ) + + note = """ + Metadata consumed: none -- every input is an argument. What the target + passes: `block_table` is `qsa_block_table`, physical page ids into dim 0 + of the caches (the *unscaled* table; a reference assuming logical == + physical lands 306x outside tolerance); `selected_tokens` is + `qsa_expand_block_indices`' output, logical positions, `-1`-padded; + `k_cache` / `v_cache` are the layer pool's role planes. + + Preconditions: + - Whenever the prefill kernel can run, each selection row is packed at + the front with exactly the bound's count of valid entries (see the + class docstring). QSA's producer -- `qsa_topk_select` with zero + `row_starts` and `row_ends = (pos + 1) // cr`, then the expansion -- + was measured to satisfy it at 300, 8191 and 8192 rows, including with + NaN/inf planted in the scores. Checking it needs the contents, so it + is not an `is_valid` guard. + - The prefill pair is a property of the launch: passed because some + rows are context, it puts the generation rows on the prefill kernel + too, and the bound applies to them (asserted with a planted index). + - There is no KV scale argument: an fp8 pool is read at face value. + - `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. + + 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. + + 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. + """ + + def __call__( + self, + *, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + selected_tokens: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + softmax_scale: float, + query_positions: Optional[torch.Tensor] = None, + compress_ratio: Optional[int] = None, + ) -> torch.Tensor: + return _triton_qsa_paged_sparse_gqa( + q=q, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + selected_tokens=selected_tokens, + request_indices=request_indices, + tokens_per_block=tokens_per_block, + softmax_scale=softmax_scale, + query_positions=query_positions, + compress_ratio=compress_ratio, + ) + + def reference( + self, + *, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_table: torch.Tensor, + selected_tokens: torch.Tensor, + request_indices: torch.Tensor, + tokens_per_block: int, + softmax_scale: float, + query_positions: Optional[torch.Tensor] = None, + compress_ratio: Optional[int] = None, + ) -> torch.Tensor: + """fp32 sparse GQA over every resident selected token, rounded once. + + `query_positions` and `compress_ratio` are ignored: on a prefix-packed + selection the prefill kernel's bound reads exactly the valid entries. + Walked 32 rows at a time to bound the gathered K/V. + """ + rows, num_q_heads, head_dim = q.shape + num_pages, num_kv_heads = k_cache.shape[0], k_cache.shape[1] + groups = num_q_heads // num_kv_heads + num_requests, table_width = block_table.shape + requests = request_indices.long() + out = torch.empty(rows, num_q_heads, head_dim, dtype=torch.float32, device=q.device) + for start in range(0, rows, 32): + stop = min(start + 32, rows) + chunk = selected_tokens[start:stop].long() + logical_page = chunk.clamp_min(0) // tokens_per_block + in_page = chunk.clamp_min(0) % tokens_per_block + request = requests[start:stop] + owned = (request >= 0) & (request < num_requests) + physical = block_table[ + request.clamp(0, num_requests - 1)[:, None].expand_as(logical_page), + logical_page.clamp(max=table_width - 1), + ].long() + resident = ( + (chunk >= 0) + & owned[:, None] + & (logical_page < table_width) + & (physical >= 0) + & (physical < num_pages) + ) + physical = physical.clamp(0, num_pages - 1) + keys = k_cache[physical, :, in_page, :].float() + values = v_cache[physical, :, in_page, :].float() + grouped = q[start:stop].reshape(-1, num_kv_heads, groups, head_dim).float() + scores = torch.einsum("bhgd,bkhd->bhgk", grouped, keys) * softmax_scale + scores = scores.masked_fill(~resident[:, None, None, :], float("-inf")) + # A row with nothing resident softmaxes over all -inf: NaN, i.e. zeros. + weights = torch.nan_to_num(torch.softmax(scores, dim=-1), nan=0.0) + pooled = torch.einsum("bhgk,bkhd->bhgd", weights, values) + out[start:stop] = pooled.reshape(stop - start, num_q_heads, head_dim) + return out.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_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 new file mode 100644 index 000000000000..ced33a120477 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_prefill_compress.py @@ -0,0 +1,308 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA block compression of the paged index-key side cache: completed groups to anchors.""" + +from __future__ import annotations + +from typing import Optional, Tuple + +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 + +#: The mRoPE sections the checkpoint declares; the target passes `None` today +#: (`_QSA_MROPE_SECTION`) and these are certified for when it flips. +MROPE_SECTION = (11, 11, 10) + +#: This checkpoint's index geometry and the engine's QSA metadata: index head +#: 128, rotary width 64 (32 cos/sin pairs), compress ratio 4, rms eps 1e-6, +#: theta 1e7; `slots` x `width` is `qsa_block_table`. `pages` is the pool. +_CKPT = dict( + head_dim=128, + tokens_per_block=32, + slots=32, + width=260, + compress_ratio=4, + eps=1e-6, + 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)) + + +def _spec(segments, pages, **overrides) -> dict: + """`segments` is `(first position, rows)` per request, in row order.""" + 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).""" + return _spec(segments, pages, mrope_section=MROPE_SECTION, **overrides) + + +class QsaPrefillCompress(OpWrapper): + """Replace each completed group's anchor slot with its pooled, normed, rotated key. + + Wraps `tensorrt_llm._torch.attention.backends.sparse.qsa.kernels. + triton_qsa_prefill_compress`: one Triton launch, grid `(rows,)`. A row + with `(L + 1) % cr == 0` (L = its logical position) is the anchor of the + group `L - cr + 1 .. L`; every other row does nothing. Each member is + resolved through `block_table` on its own, so groups may straddle pages:: + + pooled = bf16(sum_f32(raw keys of the group) / cr) + y = f32(pooled) * rsqrt(mean(f32(pooled)^2) + eps) * (1 + k_norm_weight) + y = cos_sin.dtype(y) + (t, h, w) = position_cache at the group's FIRST member + pair q takes h if q % 3 == 1 and q < 3 * s1, w if q % 3 == 2 and + q < 3 * s2, else t (t for every pair when mrope_section is None) + NeoX-rotate y[:D/2] by cos_sin[pos[q], q], cos_sin[pos[q], R + q]; + y[D/2:] passes through (R = cos_sin.shape[1] // 2 = D / 4) + index_cache[anchor] = bf16(y) + + Returns `index_cache`, written in place at anchor slots only; everything + else, including the non-anchor members' raw keys, is untouched. + + Not idempotent: the anchor slot is the group's last input and its only + output, so a second call over the same rows re-pools a compressed anchor + (measured: all 48 of 48 anchors moved, up to 2.94). Two rows naming one + anchor in a launch are harmless. + + Fusion boundary: pool, Gemma norm, position gather, rotation, store. The + raw-key and position stores (`qsa_paged_index_store`, which must run + first), the table and metadata, and the main K/V pool are the caller's. + """ + + 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), + ), + 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)), + Cell( + why="int32 positions, int32 requests", + spec=_old(((0, 64),) * 2, pages=4, logical_dtype=torch.int32), + ), + 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 = """ + Metadata consumed: none -- every input is an argument. The target passes + `qsa_logical_positions` (int64) and `qsa_req_idx_per_token` (int32) + sliced to the step's rows, `qsa_block_table` unsliced, the layer's + `get_index_k_buffer(i)` and the shared `get_qsa_position_buffer()`, the + indexer's `k_layernorm` weight (bf16 [128]) and eps, a float32 + `[max_positions, 64]` cos|sin table built at load, and `mrope_section` + None. It runs on every step, decode and prefill rows alike. + + Preconditions: + - every member's raw key and the first member's coordinates are in the + caches before the call; this call stores neither. + - every logical position is non-negative. `L = -1` satisfies the anchor + test and addresses positions -4..-1: measured, an illegal memory + access that poisons the process. A value, so not guarded; pad with 0 + (measured an exact no-op). + - `logical_positions`, `request_indices`, `k_norm_weight` and the last + dim of `cos_sin` are read with an assumed unit stride while the + callable checks only shape: each was measured accepted and wrong + (max abs 4.72, 6.20, 0.80, 7.16). `is_valid` refuses them. + The callable rejects loudly: 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` + or `compress_ratio`, a norm weight of the wrong length, CPU operands, and + `mrope_section` that is not 3 entries summing to `R`. `rows == 0` returns + before launching. + + `cos_sin`'s dtype is the rotation's precision: bf16 is accepted and was + measured 1.4% of elements off by up to 1.56e-2 -- not certified. Text + rows reduce mRoPE to plain RoPE bit for bit. Plain `@triton.jit`, no + autotune; CUDA-graph capturable at a fixed row count. + + 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. + + 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. + """ + + def __call__( + self, + *, + logical_positions: torch.Tensor, + request_indices: torch.Tensor, + block_table: torch.Tensor, + index_cache: torch.Tensor, + position_cache: torch.Tensor, + k_norm_weight: torch.Tensor, + cos_sin: torch.Tensor, + eps: float, + tokens_per_block: int, + compress_ratio: int, + mrope_section: Optional[Tuple[int, int, int]], + ) -> torch.Tensor: + _qsa_kernels.triton_qsa_prefill_compress( + logical_positions=logical_positions, + request_indices=request_indices, + block_table=block_table, + index_cache=index_cache, + position_cache=position_cache, + k_norm_weight=k_norm_weight, + cos_sin=cos_sin, + eps=eps, + tokens_per_block=tokens_per_block, + compress_ratio=compress_ratio, + mrope_section=mrope_section, + ) + return index_cache + + def reference( + self, + *, + logical_positions: torch.Tensor, + request_indices: torch.Tensor, + block_table: torch.Tensor, + index_cache: torch.Tensor, + position_cache: torch.Tensor, + k_norm_weight: torch.Tensor, + cos_sin: torch.Tensor, + eps: float, + tokens_per_block: int, + compress_ratio: int, + mrope_section: Optional[Tuple[int, int, int]], + ) -> torch.Tensor: + """The index cache after the call, as a fresh tensor, from the pre-call state.""" + out = index_cache.clone() + logical = logical_positions.long() + anchored = (logical + 1) % compress_ratio == 0 + anchor = logical[anchored] + if anchor.numel() == 0: + return out + device = index_cache.device + request = request_indices.long()[anchored] + members = ( + anchor[:, None] - (compress_ratio - 1) + torch.arange(compress_ratio, device=device) + ) + pages = block_table.long()[request[:, None], members // tokens_per_block] + slots = members % tokens_per_block + + head_dim = index_cache.shape[3] + raw = index_cache[pages, slots, 0].float() # [anchors, cr, D] + pooled = (raw.sum(dim=1) / compress_ratio).to(torch.bfloat16).float() + scale = torch.rsqrt(pooled.mul(pooled).sum(dim=-1, keepdim=True) / head_dim + eps) + normed = (pooled * scale * (k_norm_weight.float() + 1.0)).to(cos_sin.dtype) + + t, h, w = position_cache[pages[:, 0], slots[:, 0], :3].long().unbind(-1) + pairs_n = cos_sin.shape[1] // 2 + pairs = torch.arange(pairs_n, device=device) + if mrope_section is None: + pos = t[:, None].expand(-1, pairs_n) + else: + pos = torch.where( + (pairs % 3 == 1) & (pairs < 3 * mrope_section[1]), + h[:, None], + torch.where( + (pairs % 3 == 2) & (pairs < 3 * mrope_section[2]), w[:, None], t[:, None] + ), + ) + cos, sin = cos_sin[pos, pairs], cos_sin[pos, pairs_n + pairs] + first, second = normed[:, :pairs_n], normed[:, pairs_n : 2 * pairs_n] + rotated = torch.cat( + [first * cos - second * sin, second * cos + first * sin, normed[:, 2 * pairs_n :]], + dim=-1, + ) + out[pages[:, -1], slots[:, -1], 0] = rotated.to(torch.bfloat16) + return out + + def is_valid( + self, + *, + logical_positions: torch.Tensor, + request_indices: torch.Tensor, + block_table: torch.Tensor, + index_cache: torch.Tensor, + position_cache: torch.Tensor, + k_norm_weight: torch.Tensor, + cos_sin: torch.Tensor, + eps: float, + tokens_per_block: int, + compress_ratio: int, + mrope_section: Optional[Tuple[int, int, int]], + ) -> None: + # Read by the kernel with an assumed unit stride; the callable checks + # only their shapes. Each measured accepted and silently wrong. + assert logical_positions.is_contiguous(), "logical_positions must be contiguous" + assert request_indices.is_contiguous(), "request_indices must be contiguous" + assert k_norm_weight.is_contiguous(), "k_norm_weight must be contiguous" + # `ndim != 2` is the callable's own loud error; leave it to the callable. + assert cos_sin.ndim != 2 or cos_sin.stride(1) == 1, ( + "cos_sin must have unit stride along its last dimension" + ) + + +qsa_prefill_compress = QsaPrefillCompress() 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 new file mode 100644 index 000000000000..1903a9cc6f32 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_topk_select.py @@ -0,0 +1,285 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""QSA block selection: an exact per-row top-k over a caller-supplied column window.""" + +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 + +#: This checkpoint's block budget: `indexer_budget` 2048 tokens over +#: `indexer_compress_ratio` 4. The op's own default, 2048, is the *token* budget. +_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 + + +class TopkExpectation(NamedTuple): + """What a correct selection must satisfy. + + The op's output order is unspecified, so `reference` cannot return the + indices themselves; it returns each row's expected top values plus the + frame the test needs to read the op's offsets back as values. + """ + + #: `[rows, index_topk]` fp32: the window's top values, descending, `-inf`-padded. + values: torch.Tensor + #: The logits as fp32 with every out-of-window column at `-inf`. + logits: torch.Tensor + #: int64 `row_starts`, to lift window-local offsets to columns. + row_starts: torch.Tensor + #: int64 window length per row, `max(row_ends - row_starts, 0)`. + window: torch.Tensor + + +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 + `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 + is written, so `indices` needs no pre-fill. Returns `indices`. + + Exact, not approximate: the kernel resolves the threshold over the full 32 + bits of the fp32 pattern. The **order of the offsets is unspecified** and + not stable between identical launches (candidates are appended through a + shared-memory atomic; measured at `[8, 2048]`, `k=512`: four launches, four + permutations of one selected set). Ties are broken arbitrarily within the + tied value. + + Fusion boundary: selection only. Scoring, the causal bound the caller + encodes in `row_ends`, and the block-to-token expansion are outside. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The target calls this with `index_topk = 512` over the `[rows, 2080]` + # score matrix `qsa_paged_index_scores` returns with + # `only_visible_blocks=True` -- so every column at or past a row's + # `row_end` is allocator garbage. Those cells poison it with NaN: a NaN + # inside the window would be selected (see `note`), so they prove the + # 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", + spec=dict( + rows=1, + cols=_SCORE_COLUMNS, + index_topk=_BLOCK_TOPK, + windows=("positions", (7777,)), + logits="relu", + tail="nan", + ), + ), + Cell( + why="decode MTP: one request's two rows (7777, 7778), NaN tail", + 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, + cols=_SCORE_COLUMNS, + index_topk=_BLOCK_TOPK, + windows=("positions", (7777, 7778, 4000, 4001, 1500, 1501, 300, 301)), + logits="relu", + tail="nan", + ), + ), + Cell( + why="context: 8191 rows at positions 0..8190 (warm-up capture)", + spec=dict( + rows=8191, + cols=_SCORE_COLUMNS, + index_topk=_BLOCK_TOPK, + windows=("chunk", 0), + logits="relu", + 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 = """ + Metadata consumed: none. Stateless; every input is an argument. + + Preconditions (the op checks none of these): + - `0 <= row_starts[r] <= row_ends[r] <= logits.shape[1]`. An empty window + is legal and yields an all-`-1` row; a `row_end` past the last column + reads out of bounds (not probed). + - Logits inside the window are finite. `-inf` is never selected, but a + NaN ranks above every finite value and IS selected: NaN in columns 0-3 + of a 32-wide window with `k=8` returned all four. 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 + 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). + + Hard limits: not probed. `indices` is addressed as `row * index_topk + + slot`; the largest certified extent is 12416 x 64. + """ + + def __call__( + self, + logits: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + indices: torch.Tensor, + index_topk: int = 2048, + ) -> torch.Tensor: + torch.ops.trtllm.indexer_topk_prefill(logits, row_starts, row_ends, indices, index_topk) + return indices + + def reference( + self, + logits: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + indices: torch.Tensor, + index_topk: int = 2048, + ) -> TopkExpectation: + """Each window's top values, sorted. `indices` (the output buffer) is ignored.""" + num_rows, num_cols = logits.shape + 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) + masked = logits.float().masked_fill(~in_window, float("-inf")) + values = torch.sort(masked, dim=1, descending=True).values + if num_cols < index_topk: + values = torch.nn.functional.pad( + values, (0, index_topk - num_cols), value=float("-inf") + ) + window = (row_ends.long() - row_starts.long()).clamp_min(0) + return TopkExpectation(values[:, :index_topk], masked, row_starts.long(), window) + + def is_valid( + self, + logits: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + indices: torch.Tensor, + index_topk: int = 2048, + ) -> None: + # The op checks only `indices.size(1) >= index_topk`, then writes a flat + # run of `index_topk`-wide rows and ignores the real row stride. Measured + # with a [4, 12] buffer at k=8: row 1's answer landed across row 0's + # columns 8-11 and its own 0-3, and the last row was never written. + assert indices.dim() == 2 and indices.shape[1] == index_topk, ( + f"indices must be [num_rows, index_topk]; got {tuple(indices.shape)} " + f"for index_topk={index_topk}" + ) + assert indices.is_contiguous(), "indices must be contiguous" + + +qsa_topk_select = QsaTopkSelect() 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 new file mode 100644 index 000000000000..559ba69efc9a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/vision_rope_2d.py @@ -0,0 +1,305 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Vision-tower 2-D axial rotary embedding (patch row / patch column), no KV +cache. + +Not the text tower's mRoPE: that rotates 64 of a 256-wide head against three +interleaved position rows (`mrope_section`, `rope_theta` 1e7); this rotates the +whole 72-wide vision head against two spatial rows from the image grid, 18 row +angles then 18 column angles, `theta` 1e4. Neither reads the other's config. + +Tier-3 composed from native torch (angle table, index gather, cos/sin, two +multiply-adds). Every registered rope op takes 1-D positions; +`flashinfer_apply_rope_with_cos_sin_cache_inplace` with an explicit per-token +table would sidestep that, but at `head_size=72` it raises "misaligned +address" and poisons the CUDA context. +""" + +from __future__ import annotations + +from typing import List, Sequence, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: This checkpoint: 1152 / 16 heads = a 72-wide head, spatial_merge_size 2, +#: theta 1e4 (`model.visual.blocks.N.attn.qkv.weight` is [3456, 1152]). +NUM_HEADS = 16 +HEAD_DIM = 72 +MERGE = 2 +THETA = 10000.0 + +#: (frames, height, width) in patch units. 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) + + +def _rotate_half(x: torch.Tensor, half: int) -> torch.Tensor: + """NeoX-style half rotation: `[x0, x1] -> [-x1, x0]` over the last dim.""" + return torch.cat((-x[..., half:], x[..., :half]), dim=-1) + + +def _spatial_merge_positions(grid: Sequence[Sequence[int]], merge: int) -> List[Tuple[int, int]]: + """Per-token `(row, col)` by four nested loops, in spatial-merge order.""" + positions: List[Tuple[int, int]] = [] + for frames, height, width in grid: + frame = [ + (block_row * merge + r, block_col * merge + c) + for block_row in range(height // merge) + for block_col in range(width // merge) + for r in range(merge) + for c in range(merge) + ] + positions.extend(frame * frames) + return positions + + +class VisionRope2d(OpWrapper): + """Rotate every head of `q` and `k` by its patch's 2-D axial angles; return + fresh `(q_out, k_out)` in the inputs' shapes and dtypes. + + rope_dim = head_dim // 2; axis_dim = rope_dim // 2 + inv_freq[j] = theta ** (-2j / rope_dim) j < axis_dim + angle[t] = [row(t) * inv_freq, col(t) * inv_freq] height half first + out[.., i] = x[i] cos a_i - x[i + rope_dim] sin a_i + out[.., i + rope_dim] = x[i + rope_dim] cos a_i + x[i] sin a_i + + NeoX half-split pairing `(i, i + rope_dim)`, not GPT-J `(2j, 2j+1)`. Tokens + are in spatial-merge order (whole `merge x merge` blocks row-major, + row-major inside a block) -- what the patch merger consumes; a raster walk + is silently wrong. A video's frames repeat the spatial grid: no temporal + axis. The angle table is built in fp64 and rounded once to fp32; the + rotation runs in fp32 and rounds once on the store. + + Fusion boundary: the angle table and the rotation. The qkv projection, the + unrotated `v`, the per-frame attention segmentation and the attention are + the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: grid, dtype; `heads` / `heads_k` / `head_dim` (default 16 / heads / + # 72); `theta` (default 1e4); `grid_dtype` / `grid_device` for `grid_thw` + # (default int32 on the host); `layout` "separate" (default) or + # "packed_qkv" -- q and k as the `split` views of one [T, 3*heads, D] + # projection, which is how the tower calls it. + CELLS: tuple[Cell, ...] = ( + 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=GRIDS["serving_frame_5580"], dtype=torch.bfloat16, layout="packed_qkv" + ), + ), + 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 = """ + Metadata consumed: none. `grid_thw` is read with `.tolist()`: a device + tensor forces a sync (int32/int64, host/device all bitwise equal). `q` and + `k` need last-dim stride 1 only; they are not mutated. + + `head_dim` must be divisible by 4; otherwise the angle table's width + mismatches the head and the broadcast raises (loud, so not guarded). + + 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`. + + 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. + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + grid_thw: torch.Tensor, + spatial_merge_size: int, + theta: float = 10000.0, + ) -> Tuple[torch.Tensor, torch.Tensor]: + num_tokens, _, head_dim = q.shape + # Angles per token: half the head (the other half is each pair's + # partner); each axis owns half of those, height first. + rope_dim = head_dim // 2 + + grid: List[List[int]] = grid_thw.tolist() + + device = q.device + # fp64 then one rounding to fp32: the largest angle is (side - 1) + # radians, and fp32 throughout rounds twice there. + inv_freq = 1.0 / ( + theta ** (torch.arange(0, rope_dim, 2, dtype=torch.float64, device=device) / rope_dim) + ) + max_side = max(max(height, width) for _, height, width in grid) + freq_table = torch.outer( + torch.arange(max_side, dtype=torch.float64, device=device), inv_freq + ).to(torch.float32) + + merge = spatial_merge_size + intra = torch.arange(merge, device=device) + pos_ids = torch.empty((num_tokens, 2), dtype=torch.long, device=device) + offset = 0 + for frames, height, width in grid: + blocks_h, blocks_w = height // merge, width // merge + rows = ( + ( + torch.arange(blocks_h, device=device)[:, None, None, None] * merge + + intra[None, None, :, None] + ) + .expand(blocks_h, blocks_w, merge, merge) + .reshape(-1) + ) + cols = ( + ( + torch.arange(blocks_w, device=device)[None, :, None, None] * merge + + intra[None, None, None, :] + ) + .expand(blocks_h, blocks_w, merge, merge) + .reshape(-1) + ) + coords = torch.stack((rows, cols), dim=-1) + if frames > 1: + coords = coords.repeat(frames, 1) + pos_ids[offset : offset + coords.shape[0]] = coords + offset += coords.shape[0] + + # [tokens, 2, axis_dim] -> [tokens, rope_dim]: height angles, then width. + angles = freq_table[pos_ids].flatten(1) + duplicated = torch.cat((angles, angles), dim=-1) + cos = duplicated.cos().unsqueeze(-2) + sin = duplicated.sin().unsqueeze(-2) + + q_fp32, k_fp32 = q.float(), k.float() + q_out = (q_fp32 * cos + _rotate_half(q_fp32, rope_dim) * sin).to(q.dtype) + k_out = (k_fp32 * cos + _rotate_half(k_fp32, rope_dim) * sin).to(k.dtype) + return q_out, k_out + + def reference( + self, + q: torch.Tensor, + k: torch.Tensor, + grid_thw: torch.Tensor, + spatial_merge_size: int, + theta: float = 10000.0, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Positions from nested loops, angles from Python-float `inv_freq` times + integer positions in fp64, rotation as a complex multiply in fp32.""" + rope_dim = q.shape[-1] // 2 + axis_dim = rope_dim // 2 + positions = _spatial_merge_positions(grid_thw.tolist(), spatial_merge_size) + inverse = torch.tensor( + [theta ** (-(2.0 * j) / rope_dim) for j in range(axis_dim)], dtype=torch.float64 + ) + pos = torch.tensor(positions, dtype=torch.float64).reshape(-1, 2) + angles = (pos[:, :, None] * inverse).flatten(1).to(device=q.device, dtype=torch.float32) + rotor = torch.polar(torch.ones_like(angles), angles).unsqueeze(-2) + + def rotate(x: torch.Tensor) -> torch.Tensor: + fp32 = x.float() + rotated = torch.complex(fp32[..., :rope_dim], fp32[..., rope_dim:]) * rotor + return torch.cat((rotated.real, rotated.imag), dim=-1).to(x.dtype) + + return rotate(q), rotate(k) + + def is_valid( + self, + q: torch.Tensor, + k: torch.Tensor, + grid_thw: torch.Tensor, + spatial_merge_size: int, + theta: float = 10000.0, + ) -> None: + # Each of these returns a full-size result without an error: a side + # that does not divide fills only the whole blocks' rows of the + # position table and leaves the rest uninitialized (four tokens of a + # 4x5 grid read back as unrotated); a grid describing fewer tokens than + # q carries leaves the tail uninitialized the same way; a one-token k + # broadcasts against every token's angles. + assert q.dim() == 3 and k.dim() == 3, ( + f"q and k must be [tokens, heads, head_dim]; got q {tuple(q.shape)}, k {tuple(k.shape)}" + ) + assert q.shape[0] == k.shape[0] and q.shape[-1] == k.shape[-1], ( + f"q and k must agree on tokens and head_dim; got q {tuple(q.shape)}, k {tuple(k.shape)}" + ) + total = 0 + for frames, height, width in grid_thw.tolist(): + assert height % spatial_merge_size == 0 and width % spatial_merge_size == 0, ( + f"grid sides must be divisible by spatial_merge_size={spatial_merge_size}; " + f"got height={height}, width={width}" + ) + total += frames * height * width + assert total == q.shape[0], f"grid_thw describes {total} tokens but q/k carry {q.shape[0]}" + + +vision_rope_2d = VisionRope2d() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/__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/embedding/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/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 new file mode 100644 index 000000000000..21190a07f730 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/embedding/ngram_hash.py @@ -0,0 +1,261 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""PLE n-gram hash: a token-id history becomes one global embedding row per head. + +Tier-3 composed from native torch integer ops (~30 small launches); exact, so +gated bitwise. Nothing registered under ``torch.ops.trtllm`` hashes n-grams. +``EngramHashProvider.compute_hashes`` in +``tensorrt_llm/_torch/modules/engram/engram.py`` is the same family and a +different function on three axes: it hashes *compressed* ids (a tokenizer +table this checkpoint does not ship), derives its multipliers from a seed +(six ``(seed, layer)`` combinations reproduced none of this checkpoint's +stored ``layer_multipliers``), and pads instead of resetting at every EOS. Its +moduli rule does agree. The built-in Qwen4-Exp implementation has its own, +not wrapped here. +""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: This checkpoint's `ple_embedding` buffers: 3-gram, 8 heads per order. +CKPT_MULTIPLIERS = (23703573157769, 20109073645365, 8052911324071) +CKPT_HEAD_VOCAB_SIZES = ( + 20000003, 20000023, 20000033, 20000047, 20000059, 20000063, 20000069, 20000077, + 20000081, 20000093, 20000107, 20000147, 20000153, 20000159, 20000161, 20000171, +) # fmt: skip +CKPT_HEAD_OFFSETS = ( + 0, 20000003, 40000026, 60000059, 80000106, 100000165, 120000228, 140000297, + 160000374, 180000455, 200000548, 220000655, 240000802, 260000955, 280001114, 300001275, +) # fmt: skip +CKPT_EOS_TOKEN_ID = 248044 +CKPT_VOCAB_SIZE = 248320 +#: `ngram_size - 1`: the cached ids ahead of the new tokens. +CKPT_CONTEXT_LEN = len(CKPT_MULTIPLIERS) - 1 + + +def _rows(batch: int, length: int, query: int, why: str, **spec) -> Cell: + return Cell(why=why, spec=dict(batch=batch, length=length, num_query_tokens=query, **spec)) + + +class NgramHash(OpWrapper): + """For a query token at history position `t` and head `h` of the `n`-gram + block (`m = layer_multipliers`, `^` int64 XOR, `%` non-negative):: + + mixed[t] = (c[t] * m[0]) ^ (c[t - 1] * m[1]) ^ ... ^ (c[t - n + 1] * m[n - 1]) + row[t, h] = mixed[t] % head_vocab_sizes[h] + head_offsets[h] + + `c[t-k]` is `token_history[t-k]` when `t - k` lies in `t`'s segment and + `eos_token_id` otherwise; a segment starts one past the most recent EOS + **strictly before** `t` (an EOS closes its segment, so the token after it + hashes as a sequence's first). Shifts before position 0 read EOS too. + Heads run `[2-gram block | ... | ngram_size-gram block]`, `heads_per_ngram + = head_vocab_sizes.numel() // (ngram_size - 1)` each, `ngram_size = + layer_multipliers.numel()`. + + Returns a fresh contiguous int64 `[batch, num_query_tokens, ngram_heads]` + for the **last** `num_query_tokens` positions; the leading ones only give + the first query token its context. Indices are **global** -- the offset is + applied, one concatenated table is addressed. No input is modified. + + Fusion boundary: segment scan, shifted streams, mix, modulo and offset. + The id history (here built by `ngram_token_history`) and the table gather + are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: batch, length, num_query_tokens; history is uniform random 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). + 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 = """ + Metadata consumed: none; rows are independent, and within a row the only + structure seen is `eos_token_id` -- packing two sequences into one row is + safe only when the first ends in EOS. Call site: `ngram_hash(window, + multipliers, head_vocab, head_offsets, eos_token_id, 1)` on the + `[T, 3]` int64 window, reshaped to `[T, 16]` by the caller. + + 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. + + 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 + in range and inside the table. The tail slice is taken before the mix + (bitwise equal to slicing after), and chunked prefill with the + `ngram_size - 1` carried ids reproduces one-shot bitwise. + + CUDA-graph capture certified, bitwise equal to eager. Cost, not gated: + 0.387 ms for 4096 x 16, 0.370 ms for an 8-row decode -- launch-bound. + + Hard limits + token id max_id * max(layer_multipliers) < 2^63 boundary 389114 for + this checkpoint's multipliers probed above it: accepted, no error, + indices still inside the table (int64 wraps silently). The largest + legal id 248319 gives 5.886e18. Not guarded: data, would sync. + """ + + def __call__( + self, + token_history: torch.Tensor, + layer_multipliers: torch.Tensor, + head_vocab_sizes: torch.Tensor, + head_offsets: torch.Tensor, + eos_token_id: int, + num_query_tokens: int, + ) -> torch.Tensor: + ids = token_history.long() + batch_size, history_len = ids.shape + ngram_size = layer_multipliers.numel() + total_heads = head_vocab_sizes.numel() + heads_per_ngram = total_heads // (ngram_size - 1) + + # Segment bookkeeping over the FULL history: the most recent EOS + # strictly before each position starts that position's segment. + positions = torch.arange(history_len, device=ids.device, dtype=torch.long) + minus_one = positions.new_full((), -1) + eos_positions = torch.where(ids == eos_token_id, positions, minus_one) + previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values + previous_eos = torch.cat( + [eos_positions.new_full((batch_size, 1), -1), previous_eos_inclusive[:, :-1]], + dim=1, + ) + position_in_segment = positions.unsqueeze(0) - (previous_eos + 1) + + # Elementwise along the token axis from here on, so slicing to the + # query tail now is bitwise equal to slicing after the modulo. + keep_from = history_len - num_query_tokens + eos_fill = ids.new_full((), eos_token_id) + shifted = [ids[:, keep_from:]] + for shift in range(1, ngram_size): + source_positions = positions - shift + gather_positions = source_positions.clamp_min(0).unsqueeze(0).expand(batch_size, -1) + gathered = ids.gather(1, gather_positions) + valid = (position_in_segment >= shift) & (source_positions.unsqueeze(0) >= 0) + shifted.append(torch.where(valid, gathered, eos_fill)[:, keep_from:]) + + blocks = [] + for ngram in range(2, ngram_size + 1): + low = (ngram - 2) * heads_per_ngram + high = low + heads_per_ngram + mixed = shifted[0] * layer_multipliers[0] + for position in range(1, ngram): + mixed = torch.bitwise_xor(mixed, shifted[position] * layer_multipliers[position]) + blocks.append( + torch.remainder(mixed.unsqueeze(-1), head_vocab_sizes[low:high].view(1, 1, -1)) + + head_offsets[low:high].view(1, 1, -1) + ) + return torch.cat(blocks, dim=-1) + + def reference( + self, + token_history: torch.Tensor, + layer_multipliers: torch.Tensor, + head_vocab_sizes: torch.Tensor, + head_offsets: torch.Tensor, + eos_token_id: int, + num_query_tokens: int, + ) -> torch.Tensor: + """Scalar Python-int arithmetic, one token and one head at a time, the + segment start tracked by walking the row; no torch op computes it.""" + multipliers = layer_multipliers.tolist() + sizes = head_vocab_sizes.tolist() + offsets = head_offsets.tolist() + ngram_size = len(multipliers) + heads_per_ngram = len(sizes) // (ngram_size - 1) + out = [] + for row in token_history.tolist(): + segment_start, last_eos = [], -1 + for position, token in enumerate(row): + segment_start.append(last_eos + 1) + if token == eos_token_id: + last_eos = position + row_out = [] + for position in range(len(row) - num_query_tokens, len(row)): + context = [ + row[position - shift] + if position - segment_start[position] >= shift and position - shift >= 0 + else eos_token_id + for shift in range(ngram_size) + ] + 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] + low = (ngram - 2) * heads_per_ngram + heads.extend( + mixed % sizes[low + h] + offsets[low + h] for h in range(heads_per_ngram) + ) + row_out.append(heads) + out.append(row_out) + return torch.tensor(out, dtype=torch.long, device=token_history.device).reshape( + token_history.shape[0], num_query_tokens, len(sizes) + ) + + def is_valid( + self, + token_history: torch.Tensor, + layer_multipliers: torch.Tensor, + head_vocab_sizes: torch.Tensor, + head_offsets: torch.Tensor, + eos_token_id: int, + num_query_tokens: int, + ) -> None: + # Each measured to be accepted and to return plausible in-range rows. + history_len = token_history.shape[-1] + ngram_size = layer_multipliers.numel() + total_heads = head_vocab_sizes.numel() + assert 0 <= num_query_tokens <= history_len, ( + f"num_query_tokens must be in [0, history_len]; got {num_query_tokens} " + f"against a history of {history_len} -- a value outside that range is " + f"silently truncated to a shorter output rather than rejected" + ) + assert layer_multipliers.dtype == torch.int64, ( + f"layer_multipliers must be int64; got {layer_multipliers.dtype} -- a " + f"narrower dtype is accepted and silently hashes with truncated " + f"multipliers" + ) + assert ngram_size < 2 or total_heads % (ngram_size - 1) == 0, ( + f"head_vocab_sizes.numel()={total_heads} must be a multiple of " + f"ngram_size-1={ngram_size - 1} -- the blocks are equal width, and a " + f"remainder is silently dropped rather than rejected" + ) + assert head_offsets.numel() == total_heads, ( + f"head_offsets and head_vocab_sizes must have the same length; got " + f"{head_offsets.numel()} and {total_heads} -- an over-long head_offsets " + f"is silently read for its first {total_heads} entries only" + ) + + +ngram_hash = NgramHash() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_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 new file mode 100644 index 000000000000..fc016c9add86 --- /dev/null +++ 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 @@ -0,0 +1,195 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Qwen4-Exp hyper-connection gating: low-rank stream mixer plus block-injection +weights. + +Tier-3 composed from native torch (three ``F.linear`` GEMMs and four elementwise +passes); it launches no TensorRT-LLM kernel. The ``mhc_*`` op family +(``mhc_gemm_sqrsum_fma``, ``mhc_big_fuse``, ``mhc_post_mapping``) computes a +structurally different hyper-connection -- a per-stream *scalar* gate from one +linear rather than a rank-320 SiLU bottleneck, a weighted *sum* rather than a +mean, the *raw* stream gated rather than the normalized one -- and its fused +members do not instantiate hidden 2560. The built-in Qwen4-Exp implementation +has fused ops for this (``trtllm::qwen4_exp_hc_*``), not wrapped here; the +semantics follow the HuggingFace ``Qwen4ExpTextGatedResidual`` minus its norm. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: This checkpoint's geometry: 4 streams of 2560, rank-320 mixer bottleneck. +CKPT_HC_COUNT = 4 +CKPT_HIDDEN = 2560 +CKPT_LOWRANK = 320 + + +def gate_ratio(out: torch.Tensor, ref: torch.Tensor) -> float: + """Worst element's distance from `ref` as a fraction of the test's band: + `rtol = eps`, `atol = eps * max|ref|`, both of `ref`'s dtype. Below 1 is + inside.""" + eps = torch.finfo(ref.dtype).eps + ref64 = ref.double() + atol = eps * ref64.abs().max().item() + return ((out.double() - ref64).abs() / (atol + eps * ref64.abs())).max().item() + + +class Qwen4ExpGatedResidual(OpWrapper): + """Collapse the `hc_count`-stream residual to one stream, and weight the + block's injection. With `C = hc_count`, `H = width // C` and `normed` the + `[..., C*H]` stream **after** its grouped norm, viewed stream-major as + `[..., C, H]`:: + + low = silu(normed @ down^T / C) # [..., lowrank] + gate = sigmoid(low @ up^T) # [..., C*H] + mixed = mean over streams of (gate * normed) # [..., H] + inject = 2 * sigmoid(normed @ block_inject^T / C) # [..., C] + + Returns fresh contiguous `(mixed, inject)` in the activation dtype; + `inject` is `None` exactly when `block_inject_weight` is `None` (the + stack-final mixer), and `mixed` is then bitwise what it is with one. + + Both divisions by `C` sit on the pre-activation logit; `mixed` gates the + **normed** stream; `inject` carries a factor of 2 (a zero logit is an + identity injection). GEMMs run in the activation dtype (fp32 accumulate), + the pointwise chain in fp32, each result rounded once. + + Fusion boundary: both projections, both activations, the gating, the + stream mean and the injection projection are inside. The grouped gemma + norm producing `normed`, and the accumulate `qwen4_exp_post_inject` (into + the **pre-norm** stream the caller still holds), are not. Weights are used + as stored (`nn.Linear` layout), nothing is cached. + """ + + 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. + 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), + ), + ) + ) + + note = """ + Metadata consumed: none; every row independent. Call sites: twice per + decoder layer (attention and MoE hyper-connections, 96 calls) plus the + stack-final mixer with `block_inject_weight=None` -- whose `mixed` *is* + the checkpoint's final norm -- and three more in the MTP layer (two + blocks and `shared_head`'s own mixer). + + Preconditions, all loud (each measured to raise `RuntimeError` before + producing anything): the width divisible by `hc_count` (unflatten), the + weights' input width equal to it and their dtype equal to the + activation's (`F.linear`). + + float32 activations are accepted and not certified: `F.linear` on fp32 + follows torch's TF32 default here, ~10 mantissa bits. A literal bf16 + transcription of the pointwise chain is ~2x this entry's error (0.0449 vs + 0.0227 on `mixed` with the real layer-0 weights). + + Operand range: with the real layer-0 weights the normed stream reaches + |30.1|, the logits 47.8 / 67.5, `inject` spans 0.430 .. 1.578. At up to + 4096x that stream every output stays finite, `|mixed| <= max|normed|`, + `inject` in `[0, 2]`; agreement with float64 degrades there because the + sigmoid saturates (a property of the function, outside the checkpoint's + range). + + 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. + + 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). + """ + + def __call__( + self, + hyper_input_normed: torch.Tensor, + input_mix_weight_down: torch.Tensor, + input_mix_weight_up: torch.Tensor, + block_inject_weight: Optional[torch.Tensor], + hc_count: int, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + activation_dtype = hyper_input_normed.dtype + hidden_size = hyper_input_normed.shape[-1] // hc_count + + # GEMMs in the activation dtype: a bf16 product accumulates in fp32 with + # no ambient state, an fp32 one would follow torch's TF32 default. + low_rank = F.linear(hyper_input_normed, input_mix_weight_down) + low_rank = F.silu(low_rank.float() / hc_count).to(activation_dtype) + + mix_weight = F.linear(low_rank, input_mix_weight_up) + mix_weight = torch.sigmoid(mix_weight.float()).unflatten(-1, (hc_count, hidden_size)) + streams = hyper_input_normed.float().unflatten(-1, (hc_count, hidden_size)) + mixed_input = (mix_weight * streams).mean(dim=-2).to(activation_dtype) + + injection_weights = None + if block_inject_weight is not None: + inject_logits = F.linear(hyper_input_normed, block_inject_weight) + injection_weights = 2.0 * torch.sigmoid(inject_logits.float() / hc_count) + injection_weights = injection_weights.to(activation_dtype) + + return mixed_input, injection_weights + + def reference( + self, + hyper_input_normed: torch.Tensor, + input_mix_weight_down: torch.Tensor, + input_mix_weight_up: torch.Tensor, + block_inject_weight: Optional[torch.Tensor], + hc_count: int, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """float64 throughout -- an fp32 GEMM here would be TF32 and the less + accurate side -- rounded once to the activation dtype.""" + dtype = hyper_input_normed.dtype + hidden = hyper_input_normed.shape[-1] // hc_count + x = hyper_input_normed.double() + low = F.silu((x @ input_mix_weight_down.double().t()) / hc_count) + gate = torch.sigmoid(low @ input_mix_weight_up.double().t()).unflatten( + -1, (hc_count, hidden) + ) + mixed = (gate * x.unflatten(-1, (hc_count, hidden))).mean(dim=-2).to(dtype) + inject = None + if block_inject_weight is not None: + inject = 2.0 * torch.sigmoid((x @ block_inject_weight.double().t()) / hc_count) + inject = inject.to(dtype) + return mixed, inject + + +qwen4_exp_gated_residual = Qwen4ExpGatedResidual() 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 new file mode 100644 index 000000000000..679d4ce50521 --- /dev/null +++ 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 @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hyper-connection post-block injection: a per-stream gated accumulate of one +block output into the multi-stream residual. + +Tier-3 composed from native torch: one kernel-launching call +(``torch.addcmul``) plus view glue. No ``torch.ops.trtllm.*`` op computes this +accumulate with an *identity* cross-stream mix: ``mhc_post_mapping`` applies a +Sinkhorn doubly-stochastic mix across streams, fixes the stream count at 4 and +instantiates its MMA paths for hidden sizes this checkpoint does not use. The +built-in Qwen4-Exp implementation has its own fused op for it, not wrapped here. +""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: This checkpoint's hyper-connection geometry: 4 streams of 2560, packed +#: stream-major into a 10240-wide residual. +CKPT_HC_COUNT = 4 +CKPT_HIDDEN = 2560 + + +class Qwen4ExpPostInject(OpWrapper): + """`out[..., s * H + i] = hyper_input[..., s * H + i] + block_out[..., i] * + inject[..., s]`, with `H = block_out.shape[-1]` and `hc_count = + inject.shape[-1]`: one block output accumulated into **every** stream, each + scaled by its own gate. Returns a fresh contiguous tensor of `hyper_input`'s + shape; no input is mutated. + + `hyper_input` is the **pre-norm** residual, packed **stream-major** (last + dim is `(hc_count, H)`, not `(H, hc_count)` -- both have the same shape, so + a stream-minor caller is silently wrong). The cross-stream mix is the + identity: no stream's value enters another. + + Fusion boundary: the accumulate only. The grouped norm, the input mixer + that produced the block's input and `inject` (`qwen4_exp_gated_residual`), + and the block itself are the caller's. + """ + + 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"), + ) + ) + + ( + 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) + ) + ) + + note = """ + Metadata consumed: none. Call sites: twice per decoder layer (after the + attention/GDN block and after the MoE) and twice in the MTP layer, with + `inject` the `[T, 4]` weights `qwen4_exp_gated_residual` returned for that + block -- in `[0, 2]`, because they are `2 * sigmoid(.)`. + + Preconditions, all loud except the guarded one: the packed width must equal + `inject.shape[-1] * block_out.shape[-1]` (the unflatten raises); all three + on one CUDA device (raises). The leading dims must be **equal**, not merely + broadcastable -- a 1-D `inject` or `block_out` was measured to be accepted + and to broadcast one row over every token (`is_valid`). + + Mixed dtypes are not rejected, they promote: an fp32 `inject` beside a + bf16 stream returns an fp32 result -- correct values, twice the bytes. + Recorded, not guarded. + + Rounding: `addcmul` forms product and sum in fp32 opmath and rounds once + on the store. Measured bit-exact against an fp32 reference rounded once at + 1..8192 tokens; the bf16-multiply-then-bf16-add form differs by 0.0625. + Wrong operators at `[64, 10240]` bf16, max deviation from the bit-exact + result: stream-minor packing 10.34, no gate 3.77, gate on the residual 5.89. + + CUDA-graph capture certified; the output is a fresh allocation, so under + capture it comes from the graph's private pool. Cost, not gated: 17 us at + 64 tokens, 53 us at 2048, 213 us at 8192 (~1.78 TB/s) -- the stride-0 + broadcast operands take a non-vectorized path. An in-place form was not + measured. + + Not certified: float8 operands. Hard limits: not probed (the largest + measured operand is 8192 x 10240 = 2^26.3 elements). + """ + + def __call__( + self, + hyper_input: torch.Tensor, + block_out: torch.Tensor, + inject: torch.Tensor, + ) -> torch.Tensor: + streams = hyper_input.unflatten(-1, (inject.shape[-1], block_out.shape[-1])) + return torch.addcmul(streams, block_out.unsqueeze(-2), inject.unsqueeze(-1)).flatten(-2) + + def reference( + self, + hyper_input: torch.Tensor, + block_out: torch.Tensor, + inject: torch.Tensor, + ) -> torch.Tensor: + """fp32 product and sum, rounded once to the promoted dtype.""" + hc_count, hidden = inject.shape[-1], block_out.shape[-1] + streams = hyper_input.float().unflatten(-1, (hc_count, hidden)) + gated = block_out.float().unsqueeze(-2) * inject.float().unsqueeze(-1) + out_dtype = torch.promote_types( + torch.promote_types(hyper_input.dtype, block_out.dtype), inject.dtype + ) + return (streams + gated).flatten(-2).to(out_dtype) + + def is_valid( + self, + hyper_input: torch.Tensor, + block_out: torch.Tensor, + inject: torch.Tensor, + ) -> None: + # Measured: a 1-D `inject` or `block_out` is accepted and returns a full + # `[8, 10240]` result computed from one row shared by every token. + assert hyper_input.shape[:-1] == block_out.shape[:-1] == inject.shape[:-1], ( + f"leading dims must match exactly; got hyper_input {tuple(hyper_input.shape)}, " + f"block_out {tuple(block_out.shape)}, inject {tuple(inject.shape)}" + ) + + +qwen4_exp_post_inject = Qwen4ExpPostInject() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__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 new file mode 100644 index 000000000000..fcb1f0eac64a --- /dev/null +++ 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 @@ -0,0 +1,412 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DeepSeek-style FP8 block-scale mixture-of-experts layer: given top-k + +grouped FC1 GEMM + SwiGLU + FP8 requantization + grouped FC2 GEMM + +routing-weighted combine, in one trtllm-gen call.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +#: Block edge on every scale axis: 1x128 activations, 128x128 weights. +BLOCK = 128 + +#: Largest finite e4m3 magnitude. +E4M3_MAX = 448.0 + +#: The activation kernel's floor on a block amax (`kDsActAmaxEpsilon` in +#: trtllmGenKernels/blockScaleMoe/DevKernel.cu). Only an all-zero block reaches it. +ACT_AMAX_FLOOR = 1.0e-10 + + +def dequantize_block_weight(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """One expert's `[M, K]` fp32 weight from e4m3 codes and `[M/128, K/128]` scales.""" + m, k = codes.shape + blocks = codes.float().reshape(m // BLOCK, BLOCK, k // BLOCK, BLOCK) + return (blocks * scales.float().reshape(m // BLOCK, 1, k // BLOCK, 1)).reshape(m, k) + + +def requantize_1x128(x: torch.Tensor) -> torch.Tensor: + """The e4m3 round trip the kernel applies to both intermediates. + + Per row and per 128 consecutive columns: `scale = max(amax, 1e-10) / 448`, + then round-to-nearest-even into e4m3 saturating at +-448. The plain + `amax / 448` scale, not the round-up power of two the MX members use. + """ + shape = x.shape + blocks = x.reshape(*shape[:-1], shape[-1] // BLOCK, BLOCK) + scale = blocks.abs().amax(dim=-1, keepdim=True).clamp(min=ACT_AMAX_FLOOR) / E4M3_MAX + codes = (blocks / scale).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn) + return (codes.float() * scale).reshape(shape) + + +def _geom(hidden, inter, num_experts, top_k, offset=0, local=None): + """A cell's MoE geometry; the window defaults to every expert local.""" + return dict( + hidden=hidden, + inter=inter, + num_experts=num_experts, + top_k=top_k, + local_expert_offset=offset, + local_num_experts=num_experts if local is None else local, + ) + + +#: The MTP draft layer's routed experts: H 2560, I 640, 512 experts, top-10, +#: every expert local (tp1). +_CKPT = _geom(2560, 640, 512, 10) + + +class Fp8BlockScaleMoeRunner(OpWrapper): + """One FP8 block-scale MoE layer over e4m3 activations, entered pre-routed. + + Per token `t` and slot `j`, with `A[t]` the fp32 activations + (`hidden_states[t, k] * hidden_states_scale[k // 128, t]`) and `W_*` the + dequantized 128x128-block weights of local expert `e = topk_ids[t, j] - + local_expert_offset` (slots outside `[0, local_num_experts)` contribute + nothing): + + up, gate = requant(A[t] @ [W_up; W_gate].T).split(I) + gate, up = min(gate, lim), clamp(up, -lim, lim) # gemm1_clamp_limit + out[t] += topk_weights[t, j] * (requant(silu(gate) * up) @ W_down.T) + + accumulated in fp32 and stored bf16, `requant` being `requantize_1x128`. + FC1 rows are `[up ; gate]` -- up first, the engine loader's `[w3; w1]`. + + Fusion boundary: permutation, both GEMMs with on-the-fly dequantization, + both intermediate requantizations, the clamped SwiGLU and the weighted + combine are inside. The router GEMM and top-k, quantizing the bf16 hidden + states (`fp8_quantize_1x128(x, use_ue8m0=False)`), the shared expert, any + all-reduce and the residual are the caller's. + + `hidden_states_scale` is `[H/128, T]` -- transposed, exactly what + `fp8_quantize_1x128` returns on this architecture. Pass it straight through: + at `T == H/128` (20 tokens at H 2560) its transpose has the same shape and + passes every check, and is 41-43 ulp RMS wrong. + + Returns a fresh `[T, H]` bf16 tensor, or an empty `[0]` bf16 tensor when + `output` is given (the result is then in `output`). + """ + + 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. + 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)), + ), + ) + + 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. + + Preconditions. `hidden_states` is `[T, H]` e4m3 and `hidden_states_scale` + `[H/128, T]` fp32 (see the class docstring); weights are the checkpoint's + per-expert tensors in plain row-major layout, no shuffle and no swizzle: + `gemm1_weights` `[E, 2I, H]` e4m3 rows `[up ; gate]`, `gemm1_weights_scale` + `[E, 2I/128, H/128]` fp32, `gemm2_weights` `[E, H, I]`, `gemm2_weights_scale` + `[E, H/128, I/128]`; `topk_ids` `[T, top_k]` int32 and `topk_weights` + `[T, top_k]` bf16, used as given (nothing is renormalized). H and I are + multiples of 128, `intermediate_size == gemm1_weights.shape[1] // 2`, + `0 < top_k < num_experts`, `top_k <= 32`. The op raises on each of those. + Supply distinct ids per row: a repeated id counted once at 16 tokens and + twice at 256 (the token count picks the routing kernel). + + The window: local slot i answers for global id `local_expert_offset + i`. + Nothing checks the window against `num_experts`; 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. + 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). + + Inert on the certified path, bitwise: `routing_method_type` 0/1/2/4 (the + pre-routed path forces the routing postprocess), `n_group`/`topk_group` + (1, 1), `routed_scaling_factor` None/1.0/2.5, `tune_max_num_tokens`, + `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 + 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, + which is the only path the C++ runner's small-tile tactic exclusion guards. + + The registered fake returns `[T, 2H]` (it assumes packed fp4 hidden + 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. + """ + + def __call__( + self, + routing_logits: Optional[torch.Tensor], + routing_bias: Optional[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, + num_experts: int, + top_k: int, + num_fused_shared_experts: Optional[int], + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + routing_method_type: int, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None, + act_type: int = 0, + gemm1_clamp_limit: Optional[float] = None, + output: Optional[torch.Tensor] = None, + tune_max_num_tokens: int = 8192, + use_dp: bool = False, + ) -> torch.Tensor: + return torch.ops.trtllm.fp8_block_scale_moe_runner( + routing_logits, + routing_bias, + hidden_states, + hidden_states_scale, + gemm1_weights, + gemm1_weights_scale, + gemm2_weights, + gemm2_weights_scale, + num_experts, + top_k, + num_fused_shared_experts, + n_group, + topk_group, + intermediate_size, + local_expert_offset, + local_num_experts, + routed_scaling_factor, + routing_method_type, + topk_weights=topk_weights, + topk_ids=topk_ids, + act_type=act_type, + gemm1_clamp_limit=gemm1_clamp_limit, + output=output, + tune_max_num_tokens=tune_max_num_tokens, + use_dp=use_dp, + ) + + def reference( + self, + routing_logits: Optional[torch.Tensor], + routing_bias: Optional[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, + num_experts: int, + top_k: int, + num_fused_shared_experts: Optional[int], + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + routing_method_type: int, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None, + act_type: int = 0, + gemm1_clamp_limit: Optional[float] = None, + output: Optional[torch.Tensor] = None, + tune_max_num_tokens: int = 8192, + use_dp: bool = False, + ) -> torch.Tensor: + """The layer over bit-identical operands, fp32 GEMMs, bf16 store. + + Pre-routed only: `routing_logits` must be None. Ignored: `routing_bias`, + `num_experts`, `top_k`, `num_fused_shared_experts`, `n_group`, + `topk_group`, `routed_scaling_factor`, `routing_method_type`, + `act_type`, `output`, `tune_max_num_tokens`, `use_dp` -- each inert on + that path. Returns the `[T, H]` result even when `output` is given. + """ + assert routing_logits is None and topk_ids is not None and topk_weights is not None, ( + "the reference models the pre-routed path only" + ) + num_tokens, hidden = hidden_states.shape + nb = hidden // BLOCK + act_in = ( + hidden_states.float().reshape(num_tokens, nb, BLOCK) + * hidden_states_scale.t().reshape(num_tokens, nb, 1).float() + ).reshape(num_tokens, hidden) + out = torch.zeros(num_tokens, hidden, dtype=torch.float32, device=hidden_states.device) + local = topk_ids.long() - local_expert_offset + live = local[(local >= 0) & (local < local_num_experts)] + with true_fp32_matmul(): + for e in torch.unique(live).tolist(): + tok, slot = (local == e).nonzero(as_tuple=True) + w1 = dequantize_block_weight(gemm1_weights[e], gemm1_weights_scale[e]) + y1 = requantize_1x128(act_in[tok] @ w1.t()) + up, gate = y1[:, :intermediate_size], y1[:, intermediate_size:] + if gemm1_clamp_limit is not None: + gate = gate.clamp(max=gemm1_clamp_limit) + up = up.clamp(-gemm1_clamp_limit, gemm1_clamp_limit) + act = requantize_1x128(torch.nn.functional.silu(gate) * up) + w2 = dequantize_block_weight(gemm2_weights[e], gemm2_weights_scale[e]) + y = act @ w2.t() + out.index_add_(0, tok, y * topk_weights[tok, slot].float().unsqueeze(1)) + return out.to(torch.bfloat16) + + def is_valid( + self, + routing_logits: Optional[torch.Tensor], + routing_bias: Optional[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, + num_experts: int, + top_k: int, + num_fused_shared_experts: Optional[int], + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + routing_method_type: int, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None, + act_type: int = 0, + gemm1_clamp_limit: Optional[float] = None, + output: Optional[torch.Tensor] = None, + tune_max_num_tokens: int = 8192, + use_dp: bool = False, + ) -> None: + """Three inputs the op accepts and answers wrongly, each measured. + + The transposed activation-scale buffer at `T == H/128` is the fourth, + and no metadata check can see it. + """ + # The kernel takes raw data pointers: a strided view is read as dense. + for name, tensor in ( + ("routing_logits", routing_logits), + ("routing_bias", routing_bias), + ("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), + ("topk_weights", topk_weights), + ("topk_ids", topk_ids), + ("output", output), + ): + if tensor is not None: + assert tensor.is_contiguous(), ( + f"{name} must be contiguous; a strided view is read as if dense " + "and silently produces wrong results" + ) + # The Python runner overwrites the activation with SwiGlu. + assert act_type == 0, ( + f"act_type={act_type} is silently ignored -- this op always computes " + "SwiGLU. Use a different runner for a non-gated activation" + ) + # Once topk_ids carry the routing, nothing reads the bias. + if routing_bias is not None: + assert topk_ids is None, ( + "routing_bias is ignored when topk_ids/topk_weights are given " + "(routing has already happened); fold it into the router logits" + ) + + +fp8_block_scale_moe_runner = Fp8BlockScaleMoeRunner() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fused_sigmoid_gate_mul_add.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/moe/fused_sigmoid_gate_mul_add.py new file mode 100644 index 000000000000..f05a5c6e24f1 --- /dev/null +++ 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 @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared-expert combine: final_hidden_states + sigmoid(gate) * shared_expert_output.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# No torch.ops boundary exists around this kernel: the module callable is the +# stable surface, and its body is one Triton launch after argument checks. +from tensorrt_llm._torch.moe.fused_shared_expert import ( + fused_sigmoid_gate_mul_add as _fused_sigmoid_gate_mul_add, +) + + +def _ckpt(tokens: int, **extra) -> dict: + """The checkpoint's MoE combine: hidden 2560 bf16, one gate logit per token, in place.""" + spec = dict( + tokens=tokens, + hidden=2560, + dtype=torch.bfloat16, + gate_dtype=torch.bfloat16, + gate_layout="column", + dest="inplace", + ) + spec.update(extra) + return spec + + +class FusedSigmoidGateMulAdd(OpWrapper): + """`out[t, i] = final[t, i] + sigmoid(gate[t]) * shared[t, i]`, fp32 inside, rounded once. + + Wraps `tensorrt_llm._torch.moe.fused_shared_expert.fused_sigmoid_gate_mul_add` + (one Triton kernel, `_sigmoid_gate_mul_add_kernel`). + + `gate_logits` is one pre-sigmoid scalar per token read as a flat vector of + `num_tokens`: `[T, 1]`, `[T]`, `[1, T]` and a strided column are the same + call, bit for bit, and its dtype is independent of the activations'. + + Destination: with `output=None` the result overwrites `final_hidden_states` + and that tensor is returned; otherwise `output` is written and returned and + both inputs are untouched. Never a fresh tensor. + + Fusion boundary: sigmoid, broadcast, multiply and add. The routed MoE, the + shared expert and the `[hidden, 1]` gate projection are the caller's. The + single rounding is the point: a sigmoid-mul then add rounds the product to + bf16 first and moves the output by up to 0.0312 at [17, 2560] (51.6x the + bf16 tolerance). + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + 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") + ), + 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 = """ + Metadata consumed: none; stateless, no autotuner, no architecture branch. + 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. + bf16 at hidden 2560 is bitwise equal at 1-2 tokens and >= 99.993% bitwise + at 17-8192 tokens, one ulp otherwise -- `tl.sigmoid`'s fast exponential. + Wrong variants land 29,000-152,000x outside the bf16 band: no sigmoid, + multiply instead of add, gate on the routed half, gate dropped, one gate + for the whole batch. + + 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. + """ + + def __call__( + self, + final_hidden_states: torch.Tensor, + gate_logits: torch.Tensor, + shared_expert_output: torch.Tensor, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return _fused_sigmoid_gate_mul_add( + final_hidden_states, + gate_logits, + shared_expert_output, + output, + ) + + def reference( + self, + final_hidden_states: torch.Tensor, + gate_logits: torch.Tensor, + shared_expert_output: torch.Tensor, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """A fresh tensor in `final_hidden_states.dtype`; `output` is ignored. + + Takes the operands as they were before the call: in place, the op + destroys `final_hidden_states`. + """ + s = torch.sigmoid(gate_logits.reshape(-1, 1).float()) + out = final_hidden_states.float() + s * shared_expert_output.float() + return out.to(final_hidden_states.dtype) + + +fused_sigmoid_gate_mul_add = FusedSigmoidGateMulAdd() 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 new file mode 100644 index 000000000000..3e100830050c --- /dev/null +++ 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 @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Renormalized MoE routing: top-k over the raw router logits, then a softmax +over just those k, in one trtllm call.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +def _ckpt(tokens: int, output_dtype: torch.dtype = torch.bfloat16) -> dict: + """The checkpoint's router: 512 experts, top-10, bf16 logits.""" + return dict( + tokens=tokens, + num_experts=512, + topk=10, + logits_dtype=torch.bfloat16, + output_dtype=output_dtype, + dist="randn", + ) + + +class RenormMoeRoutingOp(OpWrapper): + """`norm_topk_prob=True` routing. Returns `(topk_ids, topk_weights)` -- IDS FIRST. + + Per token: `ids` = the `topk` largest `router_logits`, ranked in the + logits' own dtype, by decreasing value, ties to the LOWER expert index + (a stable descending sort -- `torch.topk` breaks ties the other way, and at + 512 bf16 experts ties are common: it disagrees on ~30% of rows); + `weights` = an fp32 softmax over just those `topk` logits, cast to + `output_dtype` afterwards. Equal to HF's softmax-then-top-k-then-divide up + to rounding (5.96e-08 at 512 / top-10). + + `ids` is `[T, topk]` int32; `weights` `[T, topk]` in `output_dtype`, which + defaults to fp32 -- not the logits' dtype. Sibling routing ops return + `(weights, ids)`; both are `[T, topk]`, so a swapped unpack is no shape + error. + + Fusion boundary: routing only. The router GEMM is the caller's, and so is + expert parallelism -- every column of the row is ranked. + """ + + 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="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), + ), + 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 = """ + Metadata consumed: none; stateless, deterministic, runs on the ambient + stream. The launch picks its kernel instance by + `nextPowerOfTwo(num_experts)` and `nextPowerOfTwo(topk)` -- 512 / top-10 is + the (512, 16) instance with a 256-thread block. PDL (`TRTLLM_ENABLE_PDL`) + was measured bitwise inert. + + 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. + + 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. + + `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 + 0.09-0.27 here). `RenormalizeMoeRoutingMethod`'s PyTorch fallback (used past + 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. + + Hard limits + num_experts <= 512 (raises past it) + topk <= 16 (raises past it); > num_experts silent, guarded + T certified 0..8192 (16384 measured once) + """ + + def __call__( + self, + router_logits: torch.Tensor, + topk: int, + output_dtype: Optional[torch.dtype] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops.trtllm.renorm_moe_routing_op(router_logits, topk, output_dtype) + + def reference( + self, + router_logits: torch.Tensor, + topk: int, + output_dtype: Optional[torch.dtype] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Stable descending sort in the logits' dtype, fp32 softmax over the + selected k, then the cast.""" + order = torch.sort(router_logits, dim=-1, descending=True, stable=True).indices[:, :topk] + selected = torch.gather(router_logits, 1, order).float() + e = torch.exp(selected - selected.max(-1, keepdim=True).values) + weights = e / e.sum(-1, keepdim=True) + return order.to(torch.int32), weights.to(output_dtype or torch.float32) + + def is_valid( + self, + router_logits: torch.Tensor, + topk: int, + output_dtype: Optional[torch.dtype] = None, + ) -> None: + # The kernel addresses router_logits from data_ptr() as dense + # [T, num_experts]: a column slice routes on the wider buffer's bytes. + assert router_logits.is_contiguous(), ( + "router_logits must be contiguous; a strided view is read as a dense " + "[num_tokens, num_experts] buffer and silently gives wrong routing" + ) + # Surplus slots come back as duplicated expert 0 at weight 0. + assert topk <= router_logits.shape[-1], ( + f"topk ({topk}) exceeds num_experts ({router_logits.shape[-1]}); the " + "surplus slots come back as duplicates of expert 0 at weight 0" + ) + + +renorm_moe_routing_op = RenormMoeRoutingOp() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/__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/norm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/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 new file mode 100644 index 000000000000..57e25f06609e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/flashinfer_gemma_rmsnorm.py @@ -0,0 +1,139 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gemma-style RMS normalization (weight applied as ``1 + w``) via the flashinfer kernel.""" + +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 + +#: The hyper-connection residual: hc_count 4 streams of hidden 2560, row stride 10240. +_STREAM = dict(row_width=10240, dtype=torch.bfloat16, eps=1e-6) + + +class FlashinferGemmaRmsnorm(OpWrapper): + """`x / sqrt(mean(x^2, -1) + eps) * (1 + weight)`, one op call, new tensor out. + + The reduction, the `1 + w` and the product are all fp32 inside the kernel; + only the store rounds to `x.dtype` -- so `(1 + w)` is not first rounded to + bf16. `x` is 2-D `[T, H]` or 3-D `[B, heads, D]` (a separate per-head + kernel), normalized over the last dim; `x` is not mutated. + + Fusion boundary: normalization and the biased weight only. No residual add + (`flashinfer_gemma_fused_add_rmsnorm` exists for that), no quantization; the + plain-`w` form is `flashinfer_rmsnorm`. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The target's grouped norm is one call per 2560-wide stream of the + # 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), + ), + Cell( + why="prefill 2048 tokens: stream 2 at row stride 10240", + spec=dict(shape=(2048, 2560), col_offset=5120, **_STREAM), + ), + Cell( + 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`, + 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 + 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 -- + `is_valid` refuses them. + + Row stride is free; the strided kernel declares it divisible by the vector + 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. + + Hard limits + T * H contiguous inputs above 2^31 - 1 elements route to the strided + (int64-offset) kernel boundary 2^31 read from the dispatch, not probed + """ + + def __call__(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + return torch.ops.trtllm.flashinfer_gemma_rmsnorm(x, weight, eps) + + def reference(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """fp32 throughout, `(1 + w)` formed in fp32, one rounding at the end.""" + xf = x.float() + normed = xf * torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return (normed * (1.0 + weight.float())).to(x.dtype) + + def is_valid(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> None: + assert x.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2), ( + f"{x.dtype} is accepted by the kernel's dtype table and returns a " + "plausible fp8 result; dequantize before normalizing" + ) + + +flashinfer_gemma_rmsnorm = FlashinferGemmaRmsnorm() 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 new file mode 100644 index 000000000000..886a2e830c40 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/layer_norm.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""LayerNorm over the last dim -- mean subtraction, biased variance, weight and +bias -- in one Triton kernel, through the FLA module callable +`tensorrt_llm._torch.modules.fla.layernorm_gated.layernorm_fn`. + +Tier-2 module callable, not a `torch.ops.trtllm.*` op: the registered norm +family is RMS-norm throughout, and `fused_adaptive_layernorm` refuses every +width but 5120. One call reshapes to 2-D, allocates the output and launches +`_layer_norm_fwd_1pass_kernel` once. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.fla.layernorm_gated import layernorm_fn as _layernorm_fn + +#: The vision tower's norm1 / norm2 / merger.norm: width 1152 with a bias, eps +#: 1e-6. The merger norm runs before the 4-to-1 merge, so it too sees 1152. +WIDTH = 1152 +EPS = 1e-6 + + +class LayerNorm(OpWrapper): + """`(x - mean) / sqrt(var + eps) * weight + bias` per row of the last dim, + variance over `N` (not `N - 1`), new contiguous tensor in `x`'s dtype. + + Every operand is widened to fp32 on load, the reduction and the affine run + in fp32, and only the store rounds. `bias=None` drops the additive term (a + compile-time branch of the kernel, not a zero bias). + + Fusion boundary: the norm and its affine only -- no residual, no gate, no + output quantization. The callable's `z`, `group_size`, `norm_before_gate` + and `is_rms_norm` are fixed at "no gate, one group, LayerNorm"; the gated + and RMS forms are different computations, not options of this entry. + """ + + 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"), + ), + ) + ) + + note = """ + Metadata consumed: none. Inputs are not mutated; two scratch fp32 [M] + tensors (mean, rstd) are allocated per call and discarded. + + 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. + + 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. + + 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. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + eps: float, + ) -> torch.Tensor: + return _layernorm_fn(x, weight, bias, z=None, eps=eps) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + eps: float, + ) -> torch.Tensor: + """fp64 from primitive ops (not `F.layer_norm`), rounded once to `x.dtype`.""" + 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() + if bias is not None: + out = out + bias.double() + return out.to(x.dtype) + + +layer_norm = LayerNorm() 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 new file mode 100644 index 000000000000..dca5f65af9aa --- /dev/null +++ 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 @@ -0,0 +1,186 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gated RMS normalization with the gate read from a token-major 3-D view, +with optional fused static-scale FP8 output quantization.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +# 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 +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") + + +class RmsNormGatedTokenMajor(OpWrapper): + """`rmsnorm(x) * weight * act(z)`, new tensor out, row `r` of `x` gated by `z[r // heads, r % heads]`. + + `x` is `[T * heads, N]`; `z` is `[T, heads, N]` (heads = `z.shape[1]`), + its `(heads, N)` block contiguous per token and its token stride free, so a + column slice of a wider projection is read in place. `act(v)` is + `v * sigmoid(v)` for `"silu"` and `sigmoid(v)` for `"sigmoid"`. All fp32, + rounded once. Norm before gate, plain `w` (not `1 + w`). + + With `fp8_scale` (one fp32 element, the downstream linear's static input + scale) the output is `float8_e4m3fn`: the gated value is rounded to + `x.dtype`, multiplied by `1 / fp8_scale[0]` and saturated at +-448 -- + byte-identical to the unfused norm -> static-quantize pair. + + Fusion boundary: norm, weight, gate and the optional quantize. The caller + owns both projections producing `x` and `z`, and every reshape. + """ + + 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), + ), + Cell( + why="x row stride 2N: only stride(-1) must be 1", + spec=dict(tokens=17, x="strided_rows", **_SILU), + ), + Cell( + why="7 rows: not a multiple of the kernel's 4 rows per program", + spec=dict(_SILU, tokens=7, heads=1), + ), + 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` + (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 + multi-row kernel accepts silently: a weight longer than N (truncated to + `weight[:N]`) or shorter (read out of bounds), and an `fp8_scale` with + 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. + + 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. + + Hard limits + N * itemsize fallback holds a row in 64 KB boundary 65536 B + probed at bf16 N = 32768 (returns) and 40960 (RuntimeError) + """ + + def __call__( + self, + x: torch.Tensor, + z: torch.Tensor, + weight: torch.Tensor, + eps: float, + fp8_scale: Optional[torch.Tensor] = None, + gate_activation: str = "silu", + ) -> torch.Tensor: + return torch.ops.trtllm.rms_norm_gated_token_major( + x, z, weight, eps, fp8_scale, gate_activation + ) + + def reference( + self, + x: torch.Tensor, + z: torch.Tensor, + weight: torch.Tensor, + eps: float, + fp8_scale: Optional[torch.Tensor] = None, + gate_activation: str = "silu", + ) -> torch.Tensor: + """fp32 throughout; row `t * heads + h` pairs with `z[t, h]`. + + With `fp8_scale` the gated value is rounded to `x.dtype` first, then + multiplied by the fp32 reciprocal of the scale, as the kernel does. + """ + num_rows, width = x.shape + xf = x.float() + y = xf * torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) * weight.float() + zf = z.reshape(num_rows, width).float() + gate = torch.sigmoid(zf) + if gate_activation != "sigmoid": + gate = gate * zf + y = y * gate + if fp8_scale is None: + return y.to(x.dtype) + return (y.to(x.dtype).float() * (1.0 / fp8_scale.float())).to(torch.float8_e4m3fn) + + def is_valid( + self, + x: torch.Tensor, + z: torch.Tensor, + weight: torch.Tensor, + eps: float, + fp8_scale: Optional[torch.Tensor] = None, + gate_activation: str = "silu", + ) -> None: + assert weight.shape == (x.shape[-1],), ( + f"weight must be 1-D of length x.shape[-1]={x.shape[-1]}; got " + f"{tuple(weight.shape)} -- the multi-row kernel reads weight[0:N] " + "unmasked, so a longer weight is silently truncated and a shorter " + "one is read out of bounds" + ) + if fp8_scale is not None: + assert fp8_scale.numel() == 1, ( + f"fp8_scale must hold exactly one element; got {fp8_scale.numel()} " + "-- the multi-row kernel loads fp8_scale[0] and ignores the rest" + ) + assert fp8_scale.dtype == torch.float32, ( + f"fp8_scale must be float32; got {fp8_scale.dtype} -- the " + "multi-row kernel widens it and quantizes against the " + "lower-precision scale" + ) + + +rms_norm_gated_token_major = RmsNormGatedTokenMajor() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/__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/ple/__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/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 new file mode 100644 index 000000000000..7cb589e7c5a1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/gated_ngram_injection.py @@ -0,0 +1,219 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""PLE gating: one shared n-gram value injected into every hyper-connection +stream, each scaled by its own key-query gate. + +Tier-3 composed from native torch (two upcasts, one reduction, the signed +square root, one broadcast multiply); it launches no TensorRT-LLM kernel. +Nothing registered under ``torch.ops.trtllm`` computes a dot-product gate or a +signed square root (``fused_sigmoid_mul_inplace`` takes a per-element gate and +could at most stand in for the last multiply). ``Engram.forward`` in +``tensorrt_llm/_torch/modules/engram/engram.py`` computes this gate as four +statements inside a forward that also runs the projections, both norms (plain +``x * w``, not this checkpoint's grouped gemma ``(1 + w)``) and a stateless +short conv -- not an atomic callable, and a different composite. The built-in +Qwen4-Exp implementation has its own, not wrapped here; the semantics follow +the HuggingFace ``Qwen4ExpTextPLELayer``. +""" + +from __future__ import annotations + +import math + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: This checkpoint's PLE geometry and its gate floor (`_PLE_CLAMP_MIN`). +CKPT_HC_COUNT = 4 +CKPT_HIDDEN = 2560 +CKPT_CLAMP_MIN = 1e-6 + + +def gate_ratio(out: torch.Tensor, ref: torch.Tensor) -> float: + """Worst element's distance from `ref` as a fraction of the test's 16-bit + band: `rtol = eps`, `atol = tiny * eps` (one subnormal spacing). Below 1 is + inside.""" + info = torch.finfo(ref.dtype) + ref64 = ref.double() + return ( + ((out.double() - ref64).abs() / (info.tiny * info.eps + info.eps * ref64.abs())) + .max() + .item() + ) + + +class GatedNgramInjection(OpWrapper): + """For token `t`, stream `s`, `D = value.shape[-1]`, with `key_normed` and + `query_normed` packed stream-major as `hc_count` blocks of `D`:: + + dot[t, s] = sum_i key_normed[t, s*D + i] * query_normed[t, s*D + i] + g = dot / sqrt(D) + signed = sign(g) * sqrt(max(|g|, clamp_min)) + out[t, s*D + i] = sigmoid(signed[t, s]) * value[t, i] + + Returns a fresh contiguous `[..., hc_count * D]` in **`value`'s dtype** + (a bf16 key/query with an fp32 value returns fp32). The value is shared: + every stream carries a differently scaled copy of the same vector. + `sign(0) = 0`, so an exactly zero dot gives `sigmoid(0) = 0.5` in every + dtype; `clamp_min` only replaces a small nonzero magnitude. The whole + chain runs in fp32 and rounds once. + + Fusion boundary: the per-stream dot, scale, floor, signed root, sigmoid + and broadcast multiply. The `key_proj` / `value_proj` GEMMs, both grouped + gemma norms producing the operands, the short conv over the result and + the add into the stream are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: kind "ckpt" (default) -- operands with this checkpoint's norm-gamma + # and value statistics: leading, 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. + 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="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 = """ + Metadata consumed: none; every token row and stream independent. Call + site: once per forward, on the PLE layer, positional `clamp_min` = + `_PLE_CLAMP_MIN` (1e-6), `key_normed` / `query_normed` `[T, 10240]` and + `value` `[T, 2560]`, all bf16. + + Preconditions: `key_normed.shape[-1] == hc_count * value.shape[-1]` (the + unflatten raises otherwise, which also catches a wrong `hc_count`). Leading + dims must match exactly (`is_valid`); measured, a single query or value + row broadcasts silently over every key row. + + fp32 rather than the HF reference's bf16 chain: against the unrounded + float64 the fp32 chain sat at 0.498 of a one-epsilon bf16 band, the + literal bf16 chain at 4.693 (4.81 against `reference`), and 31.5% of the + output's bf16 elements differ between the two -- a bitwise comparison + against a bf16 implementation will not match. + + The floor is invisible in bf16: a sub-floor dot moves the sigmoid by ~5e-4 + of its value, measurable in fp32 and below half a bf16 ulp + (`clamp_min` 1e-6 and 0.0 bitwise identical with a bf16 `value`). + + Every near-miss tried (no root, unsigned root, no / misplaced / `1/D` + scale, sigmoid before the root, tanh, stream-minor packing) stays bounded + by `|value|`; they land 105x..1115x outside the band. + + CUDA-graph capture certified, bitwise equal to eager; deterministic. + Cost, not gated: 0.114 ms at 8 tokens, 0.128 at 512, 0.495 at 4096, + 0.956 at 8192 (launch-bound below ~512); transient allocation 3.56x the + operand bytes (the fp32 upcasts). + + Hard limits: not probed (largest operand 8192 x 10240 = 2^26.3 elements). + """ + + def __call__( + self, + key_normed: torch.Tensor, + query_normed: torch.Tensor, + value: torch.Tensor, + hc_count: int, + clamp_min: float = 1e-6, + ) -> torch.Tensor: + activation_dtype = value.dtype + hidden_size = value.shape[-1] + + keys = key_normed.unflatten(-1, (hc_count, hidden_size)).float() + queries = query_normed.unflatten(-1, (hc_count, hidden_size)).float() + + gate = (keys * queries).sum(dim=-1, keepdim=True) / math.sqrt(hidden_size) + gate = gate.abs().clamp_min(clamp_min).sqrt() * gate.sign() + + gated = torch.sigmoid(gate) * value.float().unsqueeze(-2) + return gated.flatten(-2).to(activation_dtype) + + def reference( + self, + key_normed: torch.Tensor, + query_normed: torch.Tensor, + value: torch.Tensor, + hc_count: int, + clamp_min: float = 1e-6, + ) -> torch.Tensor: + """float64 -- the entry already computes in fp32, so an fp32 reference + would be the same arithmetic -- rounded once to `value`'s dtype.""" + hidden = value.shape[-1] + keys = key_normed.double().unflatten(-1, (hc_count, hidden)) + queries = query_normed.double().unflatten(-1, (hc_count, hidden)) + gate = (keys * queries).sum(dim=-1, keepdim=True) / math.sqrt(hidden) + gate = gate.abs().clamp_min(clamp_min).sqrt() * gate.sign() + out = torch.sigmoid(gate) * value.double().unsqueeze(-2) + return out.flatten(-2).to(value.dtype) + + def is_valid( + self, + key_normed: torch.Tensor, + query_normed: torch.Tensor, + value: torch.Tensor, + hc_count: int, + clamp_min: float = 1e-6, + ) -> None: + # Both measured: a [1, width] query beside [64, width] keys, or a + # [1, D] value, returns 64 rows computed from the wrong row. + assert key_normed.shape == query_normed.shape, ( + f"key_normed {tuple(key_normed.shape)} and query_normed " + f"{tuple(query_normed.shape)} must have the same shape; a mismatch " + f"broadcasts silently instead of raising" + ) + assert value.shape[:-1] == key_normed.shape[:-1], ( + f"value {tuple(value.shape)} must carry the same leading dims as " + f"key_normed {tuple(key_normed.shape)}; a mismatch broadcasts silently " + f"instead of raising" + ) + + +gated_ngram_injection = GatedNgramInjection() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/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 new file mode 100644 index 000000000000..9ad4c4d8230a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history.py @@ -0,0 +1,406 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""PLE n-gram id history: a packed token stream becomes one n-gram id window +per token, its out-of-slice predecessors read from -- and its tail published +back to -- the per-slot n-gram context pool. + +Tier-3 composed from native torch integer ops (~20 small launches); exact, so +gated bitwise. Nothing registered under ``torch.ops.trtllm`` assembles an +n-gram id window, and the one kernel-plane callable in the family +(``EngramHashProvider.compute_hashes`` / ``NgramHashMapping._get_ngram_hashes`` +in ``tensorrt_llm/_torch/modules/engram/engram.py``) has no cross-step state: +it left-pads every chunk from a pad id, so a decode step there loses the +sequence's whole context. The built-in Qwen4-Exp implementation keeps its own, +not wrapped here. +""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: This checkpoint's PLE constants: `ngram_size` 3, so a 2-id context row +#: (`Qwen4ExpPLECacheParams.ngram_context_len`). +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. +ENGINE_SLOTS = 461 + + +def _scattered(num_seqs: int) -> tuple[int, ...]: + """`num_seqs` distinct slots over the whole engine pool, both ends in.""" + slots = ((0, ENGINE_SLOTS - 1) + tuple(i * 13 + 4 for i in range(num_seqs)))[:num_seqs] + assert len(set(slots)) == num_seqs and max(slots) < ENGINE_SLOTS + return slots + + +def _cell(why: str, lengths, slots, carries, pool_slots: int, **spec) -> Cell: + return Cell( + why=why, + spec=dict( + lengths=tuple(lengths), + slots=tuple(slots), + carries=tuple(carries), + pool_slots=pool_slots, + **spec, + ), + ) + + +#: The target's own argument dtypes: int32 ids and slots, int64 +#: `query_start_loc_long` (or the int64 `cu_seqlens_from_seq_lens` rebuild). +_CALL_SITE = dict(ids_dtype=torch.int32, index_dtype=torch.int32, qsl_dtype=torch.int64) + + +class NgramTokenHistory(OpWrapper): + """Row `t` of the returned window is `[c[t-n+1], ..., c[t-1], c[t]]` -- + **newest id last**, so column `ngram_size - 1` is the token's own id -- + where `c` is the token's **sequence's** id stream: for sequence `s` over + packed tokens `[begin, end)` with `slot = state_indices[s]`, a predecessor + at shift `k` is `input_ids[t - k]` if `t - begin >= k`, else column + `ngram_size - 1 + (t - begin) - k` of the carry row. The carry row is + `context_pool[slot]` (oldest id first) when `has_initial_states[s]`, and + all `eos_token_id` when not. + + Then, **in place**, each named slot's row becomes the last `ngram_size - 1` + ids of the sequence's history including the new tokens (`window[end - 1, + 1:]`); a zero-token sequence republishes its carry (so a fresh one is + reseeded to EOS). Unnamed slots are bitwise untouched. + + The ids are **raw**: the only EOS written is the fresh-sequence seed, and + an in-stream EOS passes through -- the hash downstream owns the + segmentation, and `[eos, .., eos, id]` is what it treats as a sequence + start. Seeding with 0 would be silently wrong: 0 is a real token. + + Returns a fresh contiguous int64 `[num_tokens, ngram_size]` window; + `context_pool` is the only argument modified. + + Fusion boundary: token-to-sequence attribution, carry read, seed, the + per-shift select and the publish. The hash, the table gather and slot + assignment are the caller's. + """ + + 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]`). + 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", + ), + ) + ) + + note = """ + Metadata consumed: none -- batch layout, pool, slot map and carry flags + are explicit. Call site, once per forward on the PLE layer: + `context_pool` the whole `MambaRole.PLE_NGRAM_CONTEXT` buffer + (`[slots, ngram_size - 1]` int64, allocated by the V2 hybrid cache + manager), `state_indices` the int32 slot vector the GDN mixer uses, + `has_initial_states` the context prefix padded with True for the + generation rows (the engine's own flags cover the context prefix only), + `query_start_loc` int64 -- `query_start_loc_long`, or under MTP verify the + `cu_seqlens_from_seq_lens` rebuild, because the metadata's vector has one + row per generation sequence and would mis-segment silently. + + Preconditions: + - Every slot a real one, `0 <= slot < slots`. The carry read + bounds-checks, so -1 or an index past the pool is a device-side assert + that poisons the CUDA context (observed in its own process). **A -1 + pad sentinel is not supported**; give padded rows real slots. + - Slots distinct: two sequences on one slot is accepted and the publish + is unordered (data; not guarded). + - `query_start_loc[-1] == input_ids.shape[0]`: trailing tokens are + attributed to the last sequence (data; not guarded). + - `context_pool` int64 and `ngram_size - 1` wide, `has_initial_states` + bool, `state_indices` and `input_ids` 1-D, `ngram_size >= 2`, + `num_tokens >= 1`, `num_seqs >= 1`: each refused loudly. + - The two length checks in `is_valid` are the silent ones. + - Token ids are not range-checked (2^62 and -5 pass through verbatim). + + A wrong window is a window of valid ids: newest-first order, a zero + seed, ignored carry flags, packing across sequence boundaries, a rotated + carry read and publishing the head all hash to valid rows. Chunked + prefill (96 tokens split at 41) is bitwise equal to one-shot, window and + pool. At 461 slots the result is bitwise identical to an 11-slot pool + with the same slot content. + + 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). + """ + + def __call__( + self, + input_ids: torch.Tensor, + context_pool: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: torch.Tensor, + query_start_loc: torch.Tensor, + eos_token_id: int, + ngram_size: int, + ) -> torch.Tensor: + ids = input_ids.long() + num_tokens = ids.shape[0] + num_seqs = state_indices.shape[0] + context_len = ngram_size - 1 + + device = ids.device + starts = query_start_loc.long() + slots = state_indices.long() + + # A copy, so the in-place publish cannot disturb rows still being read + # and a zero-token sequence has something to write back. + carried = context_pool.index_select(0, slots) + carried = torch.where( + has_initial_states.reshape(-1, 1), + carried, + carried.new_full((), eos_token_id), + ) + + # Device-side and static-shape, so capturable; `repeat_interleave` on + # device lengths would make the shape data-dependent and sync. + positions = torch.arange(num_tokens, device=device) + seq_id = torch.searchsorted(starts[1:].contiguous(), positions, right=True).clamp_max( + num_seqs - 1 + ) + offset = positions - starts[seq_id] + + # One column per shift, oldest first: the packed vector once the token + # is `shift` into its slice, the carry row (virtual positions + # -context_len .. -1) before that. + columns = [] + for shift in range(context_len, 0, -1): + from_packed = ids.index_select(0, (positions - shift).clamp_min(0)) + pool_column = (context_len + offset - shift).clamp(0, context_len - 1) + from_pool = carried[seq_id, pool_column] + columns.append(torch.where(offset >= shift, from_packed, from_pool)) + columns.append(ids) + window = torch.stack(columns, dim=1) + + # The last token's window minus its oldest id, or the carry when empty. + ends = starts[1:] + last_token = (ends - 1).clamp_min(0) + published = torch.where( + (ends > starts[:-1]).reshape(-1, 1), + window.index_select(0, last_token)[:, 1:], + carried, + ) + context_pool[slots] = published + return window + + def reference( + self, + input_ids: torch.Tensor, + context_pool: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: torch.Tensor, + query_start_loc: torch.Tensor, + eos_token_id: int, + ngram_size: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """`(window, context_pool)` as the call leaves them, from each + sequence's own id stream materialized as a Python list (carry, then + its slice) and read straight out; mutates nothing.""" + context_len = ngram_size - 1 + ids = input_ids.tolist() + pool = context_pool.tolist() + starts = query_start_loc.tolist() + carries = has_initial_states.tolist() + window = [[0] * ngram_size for _ in ids] + published = [list(row) for row in pool] + for s, slot in enumerate(state_indices.tolist()): + carried = list(pool[slot]) if carries[s] else [eos_token_id] * context_len + begin, end = starts[s], starts[s + 1] + stream = carried + ids[begin:end] + for step in range(end - begin): + window[begin + step] = stream[step : step + ngram_size] + published[slot] = stream[len(stream) - context_len :] + device = input_ids.device + return ( + torch.tensor(window, dtype=torch.long, device=device).reshape(len(ids), ngram_size), + torch.tensor(published, dtype=torch.long, device=device).reshape(context_pool.shape), + ) + + def is_valid( + self, + input_ids: torch.Tensor, + context_pool: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: torch.Tensor, + query_start_loc: torch.Tensor, + eos_token_id: int, + ngram_size: int, + ) -> None: + # Both measured silent; the over-long direction of each raises instead. + num_seqs = state_indices.shape[0] + assert has_initial_states.shape[0] == num_seqs, ( + f"has_initial_states must be one flag per sequence; got " + f"{tuple(has_initial_states.shape)} against {num_seqs} sequences -- a " + f"length-1 flag tensor is broadcast over the whole batch rather than " + f"rejected" + ) + assert query_start_loc.shape[0] == num_seqs + 1, ( + f"query_start_loc must be [num_seqs + 1] = [{num_seqs + 1}] cumulative " + f"token counts; got {tuple(query_start_loc.shape)} -- a short one " + f"misattributes tokens to the wrong sequence rather than raising" + ) + + +ngram_token_history = NgramTokenHistory() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/__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/quantization/__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/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 new file mode 100644 index 000000000000..68888e530047 --- /dev/null +++ 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 @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-tensor FP8 dequantization: float8_e4m3fn -> bf16/fp16/fp32, multiplied by one scale.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — loads the extension registering torch.ops.tensorrt_llm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The PLE n-gram table's per-tensor dequantization scale in this checkpoint. +PLE_SCALE = 0.00019931793212890625 + +#: Scale values the whole-domain cells sweep: the checkpoint's, unity, a power +#: of two, an arbitrary one, and the smallest e4m3 subnormal. +_DOMAIN_SCALES = (PLE_SCALE, 1.0, 0.5, 0.037, 2.0**-9) + + +def _ple(tokens: int) -> dict: + """`_ple_block`'s call: 16 gathered 160-wide rows per token, one bf16 scale.""" + return dict(shape=(tokens, 16, 160), dtype=torch.bfloat16, scales=(PLE_SCALE,), source="gather") + + +class DequantizeE4m3PerTensor(OpWrapper): + """`out = to(scales.dtype)(float(input) * float(scales[0]))`, one rounding. + + Wraps `torch.ops.tensorrt_llm.dequantize_e4m3_per_tensor` -- the + `tensorrt_llm` namespace; there is no `trtllm::` alias. + + It MULTIPLIES: pass a checkpoint's `weight_scale` straight, never + reciprocated (the inverse, `static_quantize_e4m3_per_tensor`, divides). The + output dtype follows `scales`, not `input`: a bf16 scale yields bf16 with no + cast. Shape is preserved exactly. `scales` has `input`'s rank with every dim + 1 (`[1, 1, 1]` for the PLE's rank-3 gather). + + Fusion boundary: scale-and-widen only. The table gather before it and the + flatten to `[T, 2560]` after it (a view) are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="_ple_block decode: [1, 16, 160] gathered rows, bf16 scale", spec=_ple(1)), + Cell(why="_ple_block 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, + ), + ), + Cell( + why="input at a non-16-byte-aligned address: same kernel", + spec=dict( + shape=(64, 160), + dtype=torch.bfloat16, + scales=(PLE_SCALE, 1 / 3, 7.3e-3), + source="unaligned", + ), + ), + ) + + note = """ + Metadata consumed: none; stateless. One kernel per call, + `scaleMatrix` at a fixed 1024-block + grid with a grid-stride loop; there is no vectorized variant, so alignment + is not an axis. + + Preconditions, 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. + + 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. + + 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 + + Not certified: `TRTLLM_ENABLE_PDL=1`, `scales` on a different device than + `input`. + """ + + def __call__(self, input: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + return torch.ops.tensorrt_llm.dequantize_e4m3_per_tensor(input, scales) + + def reference(self, input: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """One fp32 multiply, one conversion to the scales' dtype.""" + return (input.float() * scales.float().flatten()[0]).to(scales.dtype) + + +dequantize_e4m3_per_tensor = DequantizeE4m3PerTensor() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/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 new file mode 100644 index 000000000000..b006f4e73adc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/quantization/fp8_quantize_1x128.py @@ -0,0 +1,164 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dynamic 1x128 block-wise FP8 quantization: bf16 -> e4m3 data + transposed fp32 block scales.""" + +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 + +#: Largest finite e4m3 magnitude. +E4M3_MAX = 448.0 + +#: The kernel's amax floor, as the bf16 constant it compares against: +#: bfloat16(1e-10) == 1.000444171950221e-10. +AMAX_FLOOR = float(torch.tensor(1e-10, dtype=torch.bfloat16)) + +#: The two divisions in the kernel are approximate (`--use_fast_math`), each +#: within CUDA's documented 2 ulp; the returned scale is two deep. +SCALE_ULP_GATE = 4 + + +def block_scales(input: torch.Tensor, use_ue8m0: bool) -> tuple[torch.Tensor, torch.Tensor]: + """`(quant_scale, dequant_scale)`, each `[T, ceil(H/128)]` fp32, exactly rounded. + + The amax runs in bf16, as the kernel's running maximum does (exact: |x| of + a bf16 is a bf16), floored at `AMAX_FLOOR`; a short tail block takes the + amax of its own elements. The kernel multiplies by the first before + rounding to e4m3 and returns the second. + """ + t, n = input.shape + nb = (n + 127) // 128 + mag = torch.nn.functional.pad(input.abs(), (0, nb * 128 - n)) + amax = mag.reshape(t, nb, 128).amax(dim=2).float().clamp(min=AMAX_FLOOR) + qs = E4M3_MAX / amax + if not use_ue8m0: + return qs, 1.0 / qs + raw = 1.0 / qs + mant, exp = torch.frexp(raw) + ds = torch.ldexp(torch.ones_like(raw), torch.where(mant == 0.5, exp - 1, exp)) + return 1.0 / ds, ds + + +#: The MTP draft layer's hidden width, quantized with the plain scale the +#: block-scale MoE runner wants. +_CKPT = dict(hidden=2560, use_ue8m0=False) + + +class Fp8Quantize1x128(OpWrapper): + """Quantize a bf16 `[T, H]` matrix to e4m3 with one fp32 scale per 1x128 row block. + + Per token `t` and block `b` of 128 consecutive columns (the last one short + when `H % 128 != 0`): `amax = max(|x[t, block]|, bf16(1e-10))`, + `data = e4m3(x * (448 / amax))` saturating at +-448, `scale[b, t] = + 1 / (448 / amax)`. Returns `(data [T, H] float8_e4m3fn, scales [ceil(H/128), + T] fp32)`. + + Three readings are each a silent wrong answer: `scales` is TRANSPOSED + (`[NB, T]`, the layout `fp8_block_scale_moe_runner` consumes); it + MULTIPLIES the code back (dividing is 6780 relative RMS off); and with + `use_ue8m0=False` it is the plain `amax / 448`, while `use_ue8m0=True` + rounds every scale up to a power of two -- a different quantization of the + same input (99% of scales, 96% of codes change), not a layout choice. + + Fusion boundary: the amax reduction, scaling, rounding and (on the sm_100 + family) the compaction into `[NB, T]` are inside; no GEMM is fused. + """ + + ARCHS = frozenset({Arch.SM_103}) + + 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), + ), + ) + + note = """ + Metadata consumed: none. The returned scale layout is chosen by an + architecture query (`isSM100Family()` in thop/fp8Quantize.cpp): off the + 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. + + Hard limits + T certified 0..8192 at H 2560; above that not probed + H multiple of 16 (raises otherwise); certified 16..7168 + """ + + def __call__( + self, + input: torch.Tensor, + use_ue8m0: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops.trtllm.fp8_quantize_1x128(input, use_ue8m0) + + def reference( + self, + input: torch.Tensor, + use_ue8m0: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + """`(data, scales[NB, T])` from exactly-rounded fp32 divisions.""" + qs, ds = block_scales(input, use_ue8m0) + cols = qs.repeat_interleave(128, dim=1)[:, : input.shape[1]] + data = (input.float() * cols).to(torch.float8_e4m3fn) + return data, ds.t().contiguous() + + +fp8_quantize_1x128 = Fp8Quantize1x128() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/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 new file mode 100644 index 000000000000..d1932d41571f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_fwd.py @@ -0,0 +1,393 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Causal depthwise conv1d over a whole (ragged) batch of sequences, with each +sequence's convolution state read from and published to a paged state pool in +the same call.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: This checkpoint's GDN mixer: `[q|k|v]` is 10240 wide, the conv is width 4, +#: so each pool slot holds a 3-column window. +CKPT_DIM = 10240 +#: The largest conv-pool slot count the cache manager was seen to allocate for +#: this checkpoint (10, 35, 366, 368, 396, 399, 432 and 461 across twelve +#: runtime-value dumps, while `max_num_sequences` never exceeds 32). The target +#: hands the pool over unsliced, so 461 tops the certified interval. +ENGINE_SLOTS = 461 + + +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. + + For sequence `s` over tokens `[b, e)` and `slot = cache_indices[s]` (or `s`): + `xpad = cat([conv_states[slot] if has_initial_state[s] else 0, x[:, b:e]])`, + `out[c, b + t] = act(sum_k weight[c, k] * xpad[c, t + k] + bias[c])` with the + taps accumulated in fp32 and rounded once, and + `conv_states[slot] = xpad[:, L : L + W - 1]` -- the raw input tail, not the + activation, exact in every dtype. + + Fusion boundary: the state read, the convolution, the bias, the SiLU and the + state publish are inside. The projection that produces `x`, and slicing the + result into q/k/v, are the caller's. + + Writes the activation into `out`, or into `x` in place when `out is None`, + and returns that tensor. `conv_states` is advanced in place for every + non-padded sequence and is not returned. + + Two kernels sit behind the op, selected from `x`'s layout: unit stride on + the token axis is channel-major (in place allowed; `out`, if given, must + have `x`'s exact strides); unit stride on the channel axis with a token + stride > 1, `dim > 1` and at least one token is channel-last (`out` + required, channel-last, disjoint from `x`). The two agree bitwise. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec keys: dim, 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". + 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) + ) + ) + + note = """ + Metadata consumed: none. Batch layout (`query_start_loc`, int32), pool + (`conv_states`, `[slots, dim, W - 1]`), slot map (`cache_indices`, int32) + and carry flags (`has_initial_state`, bool) are explicit arguments. + + Call site: the GDN prefill half -- `x` the `[10240, num_prefill_tokens]` + channel-major output of `extract_transpose_prefill_slice`, W 4, no bias, + SiLU on, the layer's whole conv pool unsliced, `pad_slot_id` -1, `out` None. + The 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. + + Preconditions the op checks itself (all raise): dtypes in {bf16, fp16, + fp32}, `conv_states`/`bias` dtype, int32 `query_start_loc`/`cache_indices`, + bool `has_initial_state` of length batch, `out` dtype/shape/layout for the + selected kernel, channel-last in place, `x` rank. + + Not checked, not guarded (content, not metadata, or an illegal access + rather than a wrong number): a `cache_indices` value outside `[0, slots)`; + `has_initial_state[s]` True with `conv_states=None` (reads through null); + `conv_states`' shape (only its dtype is validated). Sequences 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. + + Hard limits + x token/channel strides uint32 stride fields boundary 2^32 elements + not probed. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + conv_states: Optional[torch.Tensor], + query_start_loc: Optional[torch.Tensor], + cache_indices: Optional[torch.Tensor], + has_initial_state: Optional[torch.Tensor], + silu_activation: bool, + pad_slot_id: int, + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + torch.ops.trtllm.causal_conv1d_fwd( + x, + weight, + bias, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + silu_activation, + pad_slot_id, + out, + ) + return x if out is None else out + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + conv_states: Optional[torch.Tensor], + query_start_loc: Optional[torch.Tensor], + cache_indices: Optional[torch.Tensor], + has_initial_state: Optional[torch.Tensor], + silu_activation: bool, + pad_slot_id: int, + out: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """`(activation, conv_states)` as the call leaves them; mutates nothing. + + The activation starts from `out` (or `x`) so a padded sequence's + columns keep whatever was there, as the op leaves them. + """ + width = weight.shape[-1] + act = (x if out is None else out).clone() + states = None if conv_states is None else conv_states.clone() + if query_start_loc is not None: + starts = query_start_loc.tolist() + spans = [(slice(None), slice(starts[s], starts[s + 1])) for s in range(len(starts) - 1)] + else: + spans = [(s,) for s in range(x.shape[0])] + slots = None if cache_indices is None else cache_indices.tolist() + carry = None if has_initial_state is None else has_initial_state.tolist() + for s, where in enumerate(spans): + slot = s if slots is None else slots[s] + if slots is not None and slot == pad_slot_id: + continue + tokens = x[where].float() # [dim, L] + length = tokens.shape[1] + if conv_states is not None and carry is not None and carry[s]: + prefix = conv_states[slot].float() + else: + prefix = tokens.new_zeros(tokens.shape[0], width - 1) + padded = torch.cat([prefix, tokens], dim=1) + acc = torch.zeros_like(tokens) + for k in range(width): + acc = acc + weight[:, k].float().unsqueeze(1) * padded[:, k : k + length] + if bias is not None: + acc = acc + bias.float().unsqueeze(1) + if silu_activation: + acc = acc * torch.sigmoid(acc) + if length: + act[where] = acc.to(act.dtype) + if states is not None: + states[slot] = padded[:, length : length + width - 1].to(states.dtype) + return act, states + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + conv_states: Optional[torch.Tensor], + query_start_loc: Optional[torch.Tensor], + cache_indices: Optional[torch.Tensor], + has_initial_state: Optional[torch.Tensor], + silu_activation: bool, + pad_slot_id: int, + out: Optional[torch.Tensor] = None, + ) -> None: + # Widths outside 2..4 fall off the dispatcher with no launch at all: + # measured at W 1 and W 5, destination and pool bitwise unchanged. + assert weight.shape[-1] in (2, 3, 4), ( + f"width {weight.shape[-1]} has no kernel; only 2, 3 and 4 are dispatched, " + "and any other width is a silent no-op" + ) + # The dispatch reads weight (and bias) as x's type without converting: + # an fp32 weight with bf16 x moved the result by 11.5. + assert weight.dtype == x.dtype, ( + "weight must have x's dtype; the kernel reads it as x's type without converting" + ) + varlen = query_start_loc is not None + c_axis, l_axis = (0, 1) if varlen else (1, 2) + if x.dim() != (2 if varlen else 3): + return # the op rejects a wrong rank itself + channel_last = ( + x.stride(c_axis) == 1 + and x.stride(l_axis) > 1 + and x.shape[c_axis] > 1 + and x.shape[l_axis] > 0 + ) + if not channel_last: + # The channel-major kernel walks the token axis and the state's last + # axis contiguously, ignoring both strides: token stride 2 landed + # 16.3 off; a state last-dim stride 2 landed 7.9 off (5.4 in the + # published state). The channel-last kernel honours every stride. + assert x.shape[l_axis] == 0 or x.stride(l_axis) == 1, ( + "the channel-major kernel loads the token axis contiguously; " + "x must have stride 1 along it" + ) + assert conv_states is None or conv_states.stride(-1) == 1, ( + "the channel-major kernel ignores conv_states' last-dim stride; it must be 1" + ) + + +causal_conv1d_fwd = CausalConv1dFwd() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/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 new file mode 100644 index 000000000000..bbeb9ca56599 --- /dev/null +++ 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 @@ -0,0 +1,387 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Multi-token causal depthwise conv1d step (Triton), with optional per-step +window snapshots for speculative verification. + +Wraps ``tensorrt_llm._torch.modules.mamba.causal_conv1d_triton.causal_conv1d_update`` +-- one Triton launch (``_causal_conv1d_update_kernel``) after allocating the +output. There is no torch-op boundary around this kernel. It is *not* +``torch.ops.trtllm.causal_conv1d_update``: that op writes `x` in place, has a +ring mode, takes `state_len` from the pool, has no snapshots, and multiplies in +fp32 -- the two are not substitutable. +""" + +from __future__ import annotations + +from typing import Optional, Union + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.mamba.causal_conv1d_triton import ( + causal_conv1d_update as _causal_conv1d_triton_update, +) + +#: 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. +ENGINE_SLOTS = 461 +#: The engine's `max_num_sequences`: the snapshot buffer's row count. +ENGINE_SNAPSHOT_ROWS = 32 + + +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,) + + +class CausalConv1dTritonUpdate(OpWrapper): + """Convolve each row's new tokens against its cached window, out of place, + and advance the window; optionally snapshot the window after every token. + + Row `b`, slot `s = conv_state_indices[b]` (or `b`), `off = 0` (or + `num_accepted_tokens[b] - 1`): `window = [conv_state[s, :, off : off+W-1], + x[b]]`, `out[b, c, i] = act(bias[c] + sum_j fp32(weight[c, j] * window[i+j]))` + -- each product rounded to the operand dtype, accumulated in fp32, SiLU as + `acc / (1 + exp(-acc))`, rounded once. The state is then rewritten from + column 0: the last `W - 1` window columns (or, with `num_accepted_tokens`, + `W - 1 + seqlen - 1` columns rebased past the offset). Columns beyond that + are untouched -- `state_len` comes from `W`, never from the pool. + + With `intermediate_conv_window`, snapshot `i` of row `b` -- + `[intermediate_state_indices[b], i, :, :W-1]` -- is the window after token + `i`; after acceptance of `k` tokens the cache is snapshot `k - 1`. + + Fusion boundary: window read, conv, bias, SiLU, state write and snapshots + are inside. The projection, the q/k/v split and the post-acceptance + promotion of a snapshot into the pool are the caller's. + + Returns a fresh tensor with `x`'s shape (2-D in, 2-D out); `x` is read-only. + A row whose slot equals `pad_slot_id` is not written at all -- its result + rows are allocator garbage -- and its pool slot and snapshot rows stay + bitwise unchanged. `conv_state` and `intermediate_conv_window` are written + in place and not returned. + """ + + 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"), + ), + ) + + 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 + ), + ), + ) + ) + + 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 + 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. + + 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. + + Accepted and inert: `launch_dependent_kernels` (bitwise identical either + way). Accepted but not exposed: `cache_seqlens` (ignored by the kernel -- + there is no ring mode), the eagle-tree mask trio, `validate_data`, + `metadata`, the `_block_n`/`_num_warps`/`_num_stages` overrides. A + `conv_state` dtype different from `x` is accepted and uncertified. + + Hard limits + none probed. + """ + + def __call__( + self, + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + activation: Union[bool, str, None], + conv_state_indices: Optional[torch.Tensor], + pad_slot_id: int, + num_accepted_tokens: Optional[torch.Tensor] = None, + intermediate_conv_window: Optional[torch.Tensor] = None, + intermediate_state_indices: Optional[torch.Tensor] = None, + launch_dependent_kernels: bool = False, + ) -> torch.Tensor: + return _causal_conv1d_triton_update( + x, + conv_state, + weight, + bias=bias, + activation=activation, + conv_state_indices=conv_state_indices, + num_accepted_tokens=num_accepted_tokens, + intermediate_conv_window=intermediate_conv_window, + intermediate_state_indices=intermediate_state_indices, + pad_slot_id=pad_slot_id, + launch_dependent_kernels=launch_dependent_kernels, + ) + + def reference( + self, + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + activation: Union[bool, str, None], + conv_state_indices: Optional[torch.Tensor], + pad_slot_id: int, + num_accepted_tokens: Optional[torch.Tensor] = None, + intermediate_conv_window: Optional[torch.Tensor] = None, + intermediate_state_indices: Optional[torch.Tensor] = None, + launch_dependent_kernels: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """`(out, conv_state, intermediate_conv_window)` as the call leaves + them; mutates nothing. `launch_dependent_kernels` is ignored. + + Rows the op leaves unwritten (padded) are NaN in `out`; the test + skips them. Products are formed in the operand dtype, as the kernel + forms them -- multiplying in fp32 instead misses the default band on + about a third of the elements. + """ + x3 = x.unsqueeze(-1) if x.dim() == 2 else x + batch, dim, seqlen = x3.shape + width = weight.shape[1] + silu = activation in (True, "silu", "swish") + state_len = width - 1 + (seqlen - 1 if num_accepted_tokens is not None else 0) + out = torch.full(x3.shape, float("nan"), dtype=x.dtype, device=x.device) + new_state = conv_state.clone() + inter = None if intermediate_conv_window is None else intermediate_conv_window.clone() + slots = None if conv_state_indices is None else conv_state_indices.tolist() + accepted = None if num_accepted_tokens is None else num_accepted_tokens.tolist() + rows = None if intermediate_state_indices is None else intermediate_state_indices.tolist() + for b in range(batch): + slot = b if slots is None else slots[b] + if slots is not None and slot == pad_slot_id: + continue + off = 0 if accepted is None else accepted[b] - 1 + state = conv_state[slot] + full = torch.cat([state[:, off : off + width - 1], x3[b]], dim=1) + acc = torch.zeros(dim, seqlen, dtype=torch.float32, device=x.device) + for j in range(width): + acc = acc + (weight[:, j].unsqueeze(1) * full[:, j : j + seqlen]).float() + if bias is not None: + acc = acc + bias.float().unsqueeze(1) + if silu: + acc = acc / (1.0 + torch.exp(-acc)) + out[b] = acc.to(x.dtype) + if inter is not None: + row = b if rows is None else rows[b] + for i in range(seqlen): + inter[row, i, :dim, : width - 1] = full[:, i + 1 : i + width] + if accepted is None: + new_state[slot, :, : width - 1] = full[:, seqlen : seqlen + width - 1] + else: + new_state[slot, :, :state_len] = torch.cat( + [state[:, off + 1 : off + width - 1], x3[b]], dim=1 + ) + return (out.squeeze(-1) if x.dim() == 2 else out), new_state, inter + + def is_valid( + self, + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + activation: Union[bool, str, None], + conv_state_indices: Optional[torch.Tensor], + pad_slot_id: int, + num_accepted_tokens: Optional[torch.Tensor] = None, + intermediate_conv_window: Optional[torch.Tensor] = None, + intermediate_state_indices: Optional[torch.Tensor] = None, + launch_dependent_kernels: bool = False, + ) -> None: + dim = x.shape[1] + seqlen = 1 if x.dim() == 2 else x.shape[2] + width = weight.shape[1] + state_len = width - 1 + (seqlen - 1 if num_accepted_tokens is not None else 0) + # The tap loop is specialised for W 2..4: W 5 runs and returned an + # answer up to 47.2 off (mean 1.89). W 1 fails to compile. + assert 2 <= width <= 4, ( + f"weight width must be 2, 3 or 4; got {width} -- width 5 is accepted " + "by the kernel and silently computes a wrong answer" + ) + # Reads and writes `state_len` columns unmasked against the pool: a + # 2-column pool at W 4 wrote one element into the next slot. + assert conv_state.shape[2] >= state_len, ( + f"conv_state must hold at least {state_len} columns (width={width}, " + f"seqlen={seqlen}, num_accepted_tokens=" + f"{'given' if num_accepted_tokens is not None else 'None'}); got " + f"{conv_state.shape[2]} -- a shorter pool is accepted and writes past " + "the slot into the next one" + ) + # The channel axis is masked against x's dim only: a half-width pool + # changed 384 elements of unnamed slots; a half-height weight was read + # past its end. + assert conv_state.shape[1] >= dim and weight.shape[0] >= dim, ( + f"conv_state and weight must cover all {dim} channels of x; got " + f"{conv_state.shape[1]} and {weight.shape[0]} -- a narrower pool or " + "weight is accepted and reads/writes out of bounds" + ) + if intermediate_conv_window is not None: + # No bound on the step axis: a 2-step buffer at seqlen 4 wrote + # into the next row. + assert ( + intermediate_conv_window.shape[1] >= seqlen + and intermediate_conv_window.shape[2] >= dim + and intermediate_conv_window.shape[3] >= width - 1 + ), ( + f"intermediate_conv_window must be at least [rows, {seqlen}, {dim}, " + f"{width - 1}]; got {tuple(intermediate_conv_window.shape)} -- an " + "undersized buffer is accepted and writes past the row it was given" + ) + + +causal_conv1d_triton_update = CausalConv1dTritonUpdate() 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 new file mode 100644 index 000000000000..d397c03bb956 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/causal_conv1d_update.py @@ -0,0 +1,280 @@ +# 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.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: This checkpoint's GDN `[q|k|v]` width; the conv is width 4 over a 3-column pool. +CKPT_DIM = 10240 +#: Largest conv-pool slot count the cache manager was seen to allocate for this +#: checkpoint (10 .. 461 over twelve runtime-value dumps, batch never above 32). +#: The decode step hands the whole pool over unsliced. +ENGINE_SLOTS = 461 + + +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,) + + +class CausalConv1dUpdate(OpWrapper): + """One causal depthwise conv step per batch row against that row's cached + window, then the window's advance -- `x` and the pool both written in place. + + Row `b`, slot `s = conv_state_indices[b]` (or `b`), new tokens `i < seqlen`: + `history = [conv_state[s, c, :], x[b, c, :]]`, the window is its last + `W - 1 + seqlen` columns, `x[b, c, i] = act(bias[c] + sum_j weight[c, j] * + window[i + j])` in fp32 rounded once, and `conv_state[s] = history[-state_len:]` + with `state_len = conv_state.shape[2]`. With `cache_seqlens` the pool is a + ring instead: position `p` lives at `p % state_len`, `cache_seqlens[b]` is + the first new token's absolute position, and nothing shifts. + + Fusion boundary: window read, convolution, bias, SiLU and state write are + inside. The projection producing `x`, slot allocation and the q/k/v split + are the caller's. + + Returns `x` -- the same object, written in place. A 2-D `[batch, dim]` `x` + is handed to the op as an `unsqueeze(-1)` view (metadata only; the op + writes through it). `conv_state` is advanced in place and not returned. + """ + + 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. + 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="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 + `cache_seqlens`, a batch-sized pool when indices are absent. So no guard. + + Not guarded because it is content, not metadata: two rows naming one slot + race -- the slot ended up matching neither row's expected state in 10 of + 10 runs while both outputs matched. Slot uniqueness is the caller's. + 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. + + Hard limits + none probed. + """ + + def __call__( + self, + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + silu_activation: bool, + cache_seqlens: Optional[torch.Tensor], + conv_state_indices: Optional[torch.Tensor], + pad_slot_id: int, + ) -> torch.Tensor: + torch.ops.trtllm.causal_conv1d_update( + x.unsqueeze(-1) if x.dim() == 2 else x, + conv_state, + weight, + bias, + silu_activation, + cache_seqlens, + conv_state_indices, + pad_slot_id, + ) + return x + + def reference( + self, + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + silu_activation: bool, + cache_seqlens: Optional[torch.Tensor], + conv_state_indices: Optional[torch.Tensor], + pad_slot_id: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """`(x, conv_state)` as the call leaves them; mutates nothing.""" + x3 = x.unsqueeze(-1) if x.dim() == 2 else x + batch, dim, seqlen = x3.shape + width = weight.shape[1] + state_len = conv_state.shape[2] + out = x3.clone() + new_state = conv_state.clone() + slots = None if conv_state_indices is None else conv_state_indices.tolist() + starts = None if cache_seqlens is None else cache_seqlens.tolist() + for b in range(batch): + slot = b if slots is None else slots[b] + if slots is not None and slot == pad_slot_id: + continue + state = conv_state[slot].float() + tokens = x3[b].float() + if starts is None: + history = torch.cat([state, tokens], dim=1) + window = history[:, state_len - (width - 1) :] + new_state[slot] = history[:, -state_len:].to(conv_state.dtype) + else: + start = starts[b] % state_len + read = [(start - (width - 1) + i) % state_len for i in range(width - 1)] + window = torch.cat([state[:, read], tokens], dim=1) + advanced = state.clone() + for i in range(seqlen): + advanced[:, (start + i) % state_len] = tokens[:, i] + new_state[slot] = advanced.to(conv_state.dtype) + acc = torch.zeros(dim, seqlen, dtype=torch.float32, device=x.device) + for j in range(width): + acc = acc + weight[:, j].float().unsqueeze(1) * window[:, j : j + seqlen] + if bias is not None: + acc = acc + bias.float().unsqueeze(1) + if silu_activation: + acc = acc * torch.sigmoid(acc) + out[b] = acc.to(x.dtype) + return (out.squeeze(-1) if x.dim() == 2 else out), new_state + + +causal_conv1d_update = CausalConv1dUpdate() 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 new file mode 100644 index 000000000000..703afea315b0 --- /dev/null +++ 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 @@ -0,0 +1,310 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Chunked gated delta rule: the Gated DeltaNet prefill core. + +Wraps the in-tree module callable +``tensorrt_llm._torch.modules.fla.chunk.chunk_gated_delta_rule`` -- the +**vendored Triton** implementation (tier-2 sourcing: the gated delta rule has +no ``torch.ops.trtllm`` surface). Its interior is a fixed sequence of Triton +steps (gate cumsum, scaled K K^T, triangular solve, WY recompute, chunked state +recurrence, output epilogue) no caller invokes on its own: one fused op. + +The GDN mixer resolves between this and a FlashInfer-backed sibling of the same +name (``tensorrt_llm._torch.modules.fla.flashinfer_chunk``) through +``gdn_mixer._resolve_chunk_gated_delta_rule``, which picks FlashInfer on sm_103 +unless ``TLLM_USE_FLASHINFER_GDN_PREFILL=0``. This entry names the Triton one +directly, so neither the resolver nor the variable changes what it runs. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.fla.chunk import chunk_gated_delta_rule as _chunk_gated_delta_rule + +#: This checkpoint's Gated DeltaNet geometry at tp1. +NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM = 16, 48, 128, 128 +#: 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. +ENGINE_SLOTS = 461 + + +def _engine_slots(num_seqs: int) -> tuple[int, ...]: + """`num_seqs` distinct slots of the engine pool, slot 0 and the last live.""" + return ((0, ENGINE_SLOTS - 1) + tuple(i * 13 + 4 for i in range(num_seqs)))[:num_seqs] + + +class ChunkGatedDeltaRule(OpWrapper): + """Run the gated delta rule over whole sequences, returning `(o, final_state)`. + + Per token `t` and value head, `h` the `[V, K]` state, `q_h`/`k_h` the key + heads broadcast by `repeat_interleave` (value head `i` reads key head + `i // (H // Hg)`), fp32 accumulation: + + h *= exp(g[t]); u = (v[t] - h k_h[t]) * beta[t]; h += u k_h[t]^T + o[t] = h (scale * q_h[t]) + + computed 64 tokens at a time over a WY representation rounded to + `q.dtype` -- the same recurrence, a different sum. `scale=None` is + `K ** -0.5`; `g` is in log space. `use_qk_l2norm_in_kernel` normalizes q/k + as `x / sqrt(sum(x^2) + 1e-6)` first. + + With `cu_seqlens` the batch is one packed row (`q.shape[0] == 1`) split at + those offsets; without it each row is a sequence. Initial state: zeros + (`initial_state=None`); `initial_state[n]`, read-only; or, with + `initial_state_indices` and `inplace_indexed_state_update=True`, + `initial_state[indices[n]]` with the final state **written back into that + slot**. `output_final_state` only decides whether a fresh fp32 + `[N, H, V, K]` final state is also returned (on the indexed path, a copy + gathered after the write-back). + + Fusion boundary: gate cumsum, optional q/k norm, scale, recurrence, state + read and write, output epilogue. The conv, q/k/v/g/beta preparation, slot + allocation and the gated norm on `o` are the caller's. + + Returns `(o, final_state)` unchanged: `o` `[B, T, H, V]` in `q.dtype` + (`output` itself when given), `final_state` or `None`. The pool write is in + place and not returned. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: seqlens (packed with cu_seqlens) 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). + 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="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 = """ + Metadata consumed: none. + + Call site (`_gdn_block`): `(q, k, v, g, beta, None, ssm_states, si_all, + True, False, qsl_long, False, None)` after `fused_gdn_post_conv` (q/k + already normalized) -- the layer's whole 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. + + 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. + + Loud, from the callable: fp32 or mixed-dtype q/k/v, a 2-D `beta`, + `q.shape[0] != 1` with `cu_seqlens`, a sequence-count mismatch of indices + or initial states, indices without `inplace_indexed_state_update`. + + Not checked: slots must not alias (not probed); q/k must be unit-norm, + by the caller or the in-kernel flag (raw q/k with the flag off return + finite values ~1e25). `head_first` is pinned False. + + Sequence lengths are unconstrained: 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. + + Not certified: the FlashInfer sibling, fp16 (accepted), bf16 `g`/`beta`/ + pool, `K = 256` (accepted, unmeasured), empty `T` or zero-length spans. + + Hard limits + K state kernel asserts K <= 256 (read, not probed) + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: Optional[float] = None, + initial_state: Optional[torch.Tensor] = None, + initial_state_indices: Optional[torch.Tensor] = None, + inplace_indexed_state_update: bool = False, + output_final_state: bool = False, + cu_seqlens: Optional[torch.Tensor] = None, + use_qk_l2norm_in_kernel: bool = False, + output: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + return _chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + initial_state_indices=initial_state_indices, + inplace_indexed_state_update=inplace_indexed_state_update, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + head_first=False, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + output=output, + ) + + def reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: Optional[float] = None, + initial_state: Optional[torch.Tensor] = None, + initial_state_indices: Optional[torch.Tensor] = None, + inplace_indexed_state_update: bool = False, + output_final_state: bool = False, + cu_seqlens: Optional[torch.Tensor] = None, + use_qk_l2norm_in_kernel: bool = False, + output: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + """`(o, final_state, initial_state as the call leaves it)`. + + fp64 token-at-a-time recurrence, rounded once: `o` to `q.dtype`, the + states to fp32 / the pool's dtype. Mutates nothing. `output` is ignored + (it only chooses where `o` lands). + """ + dt = torch.float64 + batch, num_tokens, num_k_heads, head_k = q.shape + num_v_heads, head_v = v.shape[2], v.shape[3] + repeats = num_v_heads // num_k_heads + scale = head_k**-0.5 if scale is None else scale + if cu_seqlens is None: + spans = [(b, 0, num_tokens) for b in range(batch)] + else: + bounds = cu_seqlens.tolist() + spans = [(0, bounds[n], bounds[n + 1]) for n in range(len(bounds) - 1)] + slots = None if initial_state_indices is None else initial_state_indices.tolist() + out = torch.zeros(batch, num_tokens, num_v_heads, head_v, dtype=dt, device=q.device) + finals = [] + for n, (row, bos, eos) in enumerate(spans): + if initial_state is None: + state = torch.zeros(num_v_heads, head_v, head_k, dtype=dt, device=q.device) + else: + state = initial_state[n if slots is None else slots[n]].to(dt) + for t in range(bos, eos): + q_t, k_t = q[row, t].to(dt), k[row, t].to(dt) + if use_qk_l2norm_in_kernel: + q_t = q_t / torch.sqrt(q_t.pow(2).sum(-1, keepdim=True) + 1e-6) + k_t = k_t / torch.sqrt(k_t.pow(2).sum(-1, keepdim=True) + 1e-6) + q_t = q_t.repeat_interleave(repeats, dim=0) * scale + k_t = k_t.repeat_interleave(repeats, dim=0) + state = state * g[row, t].to(dt).exp()[:, None, None] + u = (v[row, t].to(dt) - (state * k_t[:, None, :]).sum(-1)) * beta[row, t].to(dt)[ + :, None + ] + state = state + u[:, :, None] * k_t[:, None, :] + out[row, t] = (state * q_t[:, None, :]).sum(-1) + finals.append(state) + final_state = torch.stack(finals).float() if output_final_state else None + pool = None if initial_state is None else initial_state.clone() + if slots is not None and inplace_indexed_state_update: + for slot, state in zip(slots, finals): + pool[slot] = state.to(pool.dtype) + return out.to(q.dtype), final_state, pool + + def is_valid( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: Optional[float] = None, + initial_state: Optional[torch.Tensor] = None, + initial_state_indices: Optional[torch.Tensor] = None, + inplace_indexed_state_update: bool = False, + output_final_state: bool = False, + cu_seqlens: Optional[torch.Tensor] = None, + use_qk_l2norm_in_kernel: bool = False, + output: Optional[torch.Tensor] = None, + ) -> None: + if output is None: + return + # Measured: a head-slice view (strides (1179648, 12288, 128, 1)) was + # accepted and came back 9.9e-2 off where correct is ~4e-4; an fp32 + # buffer was filled correctly but the returned bf16 tensor no longer + # aliased it. A token-window slice that reports contiguous is fine. + assert output.is_contiguous(), ( + f"output must be contiguous; got strides {tuple(output.stride())} -- the " + "epilogue addresses it through strides derived from its shape" + ) + assert output.dtype == q.dtype, ( + f"output must have q's dtype {q.dtype}; got {output.dtype} -- the result " + "is cast to q.dtype and stops aliasing the buffer" + ) + + +chunk_gated_delta_rule = ChunkGatedDeltaRule() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/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 new file mode 100644 index 000000000000..6a10c7d409cd --- /dev/null +++ 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 @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Prepend-zero cumulative sum of a per-request query-row count: the +``cu_seqlens`` segmentation vector every varlen recurrent op takes. + +Tier-3 composed from native torch (``torch.zeros`` + ``torch.cumsum``, three +launches). No registered op or ``modules/`` callable produces this vector: the +recurrent metadata's ``query_start_loc*`` is built inline by +``Mamba2Metadata.prepare`` and segments the generation tail one row per +sequence, which is short under speculative decoding. +""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Integer dtypes accepted on either side. ``torch.cumsum`` silently accepts +#: 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). +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`. + + The scan accumulates in `dtype`, not in `seq_lens`'s dtype. Fusion boundary: + the leading zero, the scan and the dtype. Choosing which length vector to + scan, and keeping the packed activation exactly `out[-1]` rows long, are the + caller's. `seq_lens` is read-only. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: lengths, in_dtype, out_dtype; `layout` "contiguous" (default), + # "prefix" (the first len(lengths) of a longer buffer -- the call site's + # `split(seq_lens_cuda, [batch, ...])[0]`), or "strided" (a [::2] view). + 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="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="empty batch returns [0]", + spec=dict(lengths=(), in_dtype=torch.int32, out_dtype=torch.int32), + ), + ) + + 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 = """ + Metadata consumed: none. The call site reads `attn_metadata.seq_lens_cuda` + prefix, which still holds the true per-request query-row count during the + trunk's forward (the spec worker rewrites it only afterwards). Not the + recurrent metadata's `query_start_loc_long`: that one has one row per + generation sequence and is short under MTP; the two agree exactly when + every generation request contributes one row. + + Preconditions, none checked (data, not metadata -- a check would sync): + entries non-negative (a negative length was silently wrong in three + consumer probes and a device-side assert in the fourth); a padded row is + length 0, never -1; `out[-1]` equals the packed activation's row count. + + `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. + + Hard limits + out[-1] accumulated in `dtype` boundary 2^31 - 1 for int32 + probed at [2e9, 2e9]: int64 gives [0, 2e9, 4e9]; int32 wraps silently + to [0, 2000000000, -294967296]. Not guarded (data). + """ + + def __call__(self, seq_lens: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + out = torch.zeros(seq_lens.shape[0] + 1, dtype=dtype, device=seq_lens.device) + torch.cumsum(seq_lens, dim=0, dtype=dtype, out=out[1:]) + return out + + def reference(self, seq_lens: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """A scalar Python prefix sum; no torch op computes it.""" + total, out = 0, [0] + for length in seq_lens.tolist(): + total += int(length) + out.append(total) + return torch.tensor(out, dtype=dtype, device=seq_lens.device) + + def is_valid(self, seq_lens: torch.Tensor, dtype: torch.dtype) -> None: + # Measured: a [1, n] row vector returned a length-2 [0, n_0]; an [n, 2] + # input a length-(n + 1) vector of nonsense (only a resize warning); a + # float32 [7.9, 1.2, 13.6] was truncated to [0, 7, 8, 21]; a bool input + # was counted as nonzero rows; a float `dtype` returned a float vector. + assert seq_lens.dim() == 1, ( + f"seq_lens must be a 1-D per-request row count, got shape {tuple(seq_lens.shape)}" + ) + assert seq_lens.dtype in _INTEGER_DTYPES, ( + f"seq_lens must have an integer dtype, got {seq_lens.dtype}" + ) + assert dtype in _INTEGER_DTYPES, f"dtype must be an integer dtype, got {dtype}" + + +cu_seqlens_from_seq_lens = CuSeqlensFromSeqLens() 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 new file mode 100644 index 000000000000..5bf4c9257981 --- /dev/null +++ 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 @@ -0,0 +1,386 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dilated depthwise short convolution over a packed, ragged serving batch, with +each sequence's convolution state read from and published to a paged pool. + +Composed from native torch ops (sourcing tier 3): no TensorRT-LLM kernel takes +a dilation. The ladder, as measured on this machine in this paged form: + +* tier 1 -- no registered op schema mentions `dilation`; the three convolutions + (`causal_conv1d_fwd`, `causal_conv1d_update`, `mamba_conv1d`) are the + non-dilated family. `causal_conv1d_fwd` fits the batch layout exactly, and a + zero-interleaved width-10 weight (taps at 0, 3, 6, 9) over the very + `[slots, channels, 9]` pool **returns normally and computes nothing** -- + activation and pool bitwise unchanged -- because its dispatcher has no branch + past width 4 and the width check sits on the update entry point only. + `causal_conv1d_update` raises on the same weight. +* tier 2 -- the vendored Triton `causal_conv1d_fn` does not compile at width + 10; engram's `ShortConv` is dilated but rectangular and stateless. +* tier 3 -- `F.conv1d(dilation=, groups=)` over a packed history plus index + arithmetic. This entry. + +Not taken: `d` phase-separated non-dilated convolutions through +`causal_conv1d_fwd` (three engine calls plus gathers -- not one entry), and a +rectangular `[num_seqs, channels, state_len + max_tokens]` buffer (5.4 GB at 32 +x 8192 tokens). A fused dilation-aware paged kernel is the faster path. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: `hc_count * hidden_size = 4 * 2560` channels, kernel 4, dilation `ngram_size` +#: 3, hence the 9-column state the mamba cache manager allocates per slot. +CKPT_CHANNELS = 10240 +CKPT_DILATION = 3 +#: 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. +ENGINE_SLOTS = 461 + + +def _scattered(num_seqs: int) -> tuple[int, ...]: + """`num_seqs` distinct slots over the whole engine pool, both ends in.""" + slots = ((0, ENGINE_SLOTS - 1) + tuple(i * 13 + 4 for i in range(num_seqs)))[:num_seqs] + assert len(set(slots)) == num_seqs and max(slots) < ENGINE_SLOTS + return slots + + +def _ckpt(why: str, lengths: tuple[int, ...], **spec) -> Cell: + return Cell(why=why, spec=dict(channels=CKPT_CHANNELS, lengths=lengths, **spec)) + + +class DilatedShortConvPaged(OpWrapper): + """For sequence `s` over columns `[b, e)` of the packed channel-major `x`, + `slot = cache_indices[s]` (or `s`), `state_len = (kernel_size - 1) * + dilation`: `window = [conv_states[slot] if has_initial_state[s] else 0, + x[:, b:e]]`, `out[c, b + t] = silu(round(sum_k weight[c, 0, k] * + window[c, t + k * dilation] + bias[c]))` -- taps in fp32, rounded by the + convolution, then the SiLU reads that back and rounds again -- and + `conv_states[slot] = window[:, L : L + state_len]`, the raw input tail. + + The highest tap index multiplies the token being emitted. A padded + sequence (`cache_indices[s] == pad_slot_id`) publishes nothing and still + gets output columns -- its tokens against a zero carry, defined and + meaningless. Two sequences naming one slot: the last live one's state + survives (defined, not a race). + + Fusion boundary: state read, zero fill for a fresh sequence, the dilated + depthwise conv, bias, SiLU and the state publish are inside. The + pre-conv norm producing `x`, slot allocation and reset, and the residual + add are the caller's. + + Returns a fresh contiguous `[channels, num_tokens]` activation; `x`, + `weight` and `bias` are read-only. `conv_states` is advanced in place and + not returned. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec keys: 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. + 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 + ), + ), + ) + ) + + note = """ + Metadata consumed: none -- the same explicit argument set as + `causal_conv1d_fwd` (packed batch, pool, int32/int64 slot map, bool carry). + + Call site: the PLE short conv on one layer -- `x` a bare `.t()` of the + `[T, 10240]` normed activation, weight `[10240, 1, 4]` bf16, no bias, + dilation 3 (state 9), SiLU on, the whole `[slots, 10240, 9]` pool, int32 + slots, `query_start_loc_long` (int64, rebuilt from `seq_lens` under MTP), + carry flags covering the whole batch, `pad_slot_id` -1. Certified at 8, 32 + and 461 slots; at 461 bitwise identical to 8 for the same slot content. + + 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. + + CUDA-graph capture is certified: 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. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: Optional[torch.Tensor], + has_initial_state: Optional[torch.Tensor], + dilation: int, + silu_activation: bool, + pad_slot_id: int, + ) -> torch.Tensor: + channels, num_tokens = x.shape + kernel_size = weight.shape[-1] + state_len = (kernel_size - 1) * dilation + num_slots = conv_states.shape[0] + num_seqs = query_start_loc.numel() - 1 + device = x.device + + # Device-resident arithmetic only: no `.item()`, no boolean indexing, + # no data-dependent shape -- which is what makes the call capturable. + starts = query_start_loc[:-1].to(torch.int64) + ends = query_start_loc[1:].to(torch.int64) + rows = torch.arange(num_seqs, device=device) + columns = torch.arange(num_tokens, device=device) + taps = torch.arange(state_len, device=device) + + # The first sequence whose end is strictly past the column; skips + # zero-token sequences by construction. + seq_of_column = torch.bucketize(columns, ends.contiguous(), right=True) + + if cache_indices is None: + slots = rows + live = torch.ones(num_seqs, dtype=torch.bool, device=device) + else: + live = cache_indices != pad_slot_id + slots = torch.where(live, cache_indices.to(torch.int64), torch.zeros_like(rows)) + carries = ( + torch.zeros_like(live) if has_initial_state is None else (has_initial_state & live) + ) + + # Each sequence's carry placed right ahead of its own tokens, so one + # convolution over the buffer never reads across a sequence boundary. + history = x.new_zeros((channels, num_tokens + num_seqs * state_len)) + history.index_copy_(1, columns + state_len * (seq_of_column + 1), x) + + previous = conv_states.index_select(0, slots) + history.index_copy_( + 1, + ((starts + rows * state_len).unsqueeze(1) + taps).reshape(-1), + torch.where(carries.reshape(num_seqs, 1, 1), previous, torch.zeros_like(previous)) + .permute(1, 0, 2) + .reshape(channels, -1), + ) + + convolved = F.conv1d( + history, + weight, + bias=bias, + stride=1, + padding=0, + dilation=dilation, + groups=channels, + ) + if silu_activation: + convolved = F.silu(convolved) + out = convolved.index_select(1, columns + state_len * seq_of_column) + + published = ( + history.index_select(1, ((ends + rows * state_len).unsqueeze(1) + taps).reshape(-1)) + .view(channels, num_seqs, state_len) + .permute(1, 0, 2) + ) + + # Skip padded rows without a host sync: every row writes, but writes + # the bytes of the highest-numbered live row owning that slot (or the + # slot's own contents when no live row does), so duplicate + # destinations carry identical bytes and landing order cannot matter. + writer = torch.full((num_slots,), -1, dtype=torch.int64, device=device) + writer.scatter_reduce_( + 0, + slots, + torch.where(live, rows, torch.full_like(rows, -1)), + reduce="amax", + include_self=True, + ) + owner = writer.index_select(0, slots) + conv_states.index_copy_( + 0, + slots, + torch.where( + (owner >= 0).reshape(num_seqs, 1, 1), + published.index_select(0, owner.clamp_min(0)), + previous, + ), + ) + return out + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: Optional[torch.Tensor], + has_initial_state: Optional[torch.Tensor], + dilation: int, + silu_activation: bool, + pad_slot_id: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """`(out, conv_states)` as the call leaves them, one sequence at a time + as an explicit fp32 tap-sum (no convolution op); mutates nothing. + + Rounds where the entry's two calls round: once after the taps, again + after the SiLU. Publishes in sequence order, so a duplicated slot ends + with the last live sequence's state. + """ + channels = x.shape[0] + kernel_size = weight.shape[-1] + state_len = (kernel_size - 1) * dilation + starts = query_start_loc.tolist() + slots = None if cache_indices is None else cache_indices.tolist() + carry = None if has_initial_state is None else has_initial_state.tolist() + w = weight.float() + out = torch.zeros(channels, x.shape[1], dtype=x.dtype, device=x.device) + pool = conv_states.clone() + for s in range(len(starts) - 1): + begin, end = starts[s], starts[s + 1] + length = end - begin + slot = s if slots is None else slots[s] + live = slots is None or slot != pad_slot_id + if live and carry is not None and carry[s]: + prefix = conv_states[slot] + else: + prefix = x.new_zeros(channels, state_len) + window = torch.cat([prefix, x[:, begin:end]], dim=1) + xf = window.float() + acc = torch.zeros(channels, length, dtype=torch.float32, device=x.device) + for k in range(kernel_size): + acc = acc + w[:, 0, k].reshape(-1, 1) * xf[:, k * dilation : k * dilation + length] + if bias is not None: + acc = acc + bias.float().reshape(-1, 1) + y = acc.to(x.dtype) + if silu_activation: + y = F.silu(y.float()).to(x.dtype) + out[:, begin:end] = y + if live: + pool[slot] = window[:, length : length + state_len] + return out, pool + + +dilated_short_conv_paged = DilatedShortConvPaged() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/extract_transpose_prefill_slice.py new file mode 100644 index 000000000000..b2f09209e898 --- /dev/null +++ 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 @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Channel-major extraction of a prefill row/column slice, the operand a causal +conv1d over a token-major projection needs. + +Wraps the in-tree Triton callable +``tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops.extract_transpose_prefill_slice`` +(one kernel launch, ``_extract_transpose_prefill_kernel``). A module-plane +callable rather than a ``torch.ops.trtllm.*`` op: there is no op boundary +around this kernel. +""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops import ( + extract_transpose_prefill_slice as _extract_transpose_prefill_slice, +) + +#: The GDN mixer's packed `[q|k|v]` width for this checkpoint; the target hands +#: the whole row to this op. +GDN_MIXED_QKV_WIDTH = 10240 + + +class ExtractTransposePrefillSlice(OpWrapper): + """`out[c, t] = src[t, start_col + c]` for `t < num_prefill_tokens`, + `c < width`: `src[:n, a:a+w].transpose(0, 1).contiguous()` in one pass. + + Fusion boundary: the row crop, the column crop and the transpose. No + convolution, no conv state, no channel split -- the caller hands the result + to `causal_conv1d_fwd`. Returns a fresh contiguous `[width, + num_prefill_tokens]` tensor in `src`'s dtype; `src` is read-only. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # `num_rows`/`num_cols` build `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). + CELLS: tuple[Cell, ...] = ( + tuple( + Cell( + why=f"GDN prefill, whole 10240-wide row, {n} tokens", + spec=dict( + num_rows=n, + num_cols=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 + ), + ), + ) + ) + + note = """ + Metadata consumed: none. Stateless. + + Preconditions: `src` 2-D on CUDA with `src.stride(1) == 1`; the row stride + is free. A non-unit column stride (`src.t()`, `src[:, ::2]`) is refused + loudly by the callable's own bare `assert`, before any launch. + + The kernel masks its loads only against the requested `num_prefill_tokens` + and `width`, never against `src`'s extent: 128 rows of a 64-row `src`, or + columns `[200, 328)` of a 256-column one, were accepted and returned a + full tensor read from out-of-bounds memory. `is_valid` bounds both. + + Hard limits + width * num_prefill_tokens int32 destination offset boundary 2^31 + probed at width 16384: n 131000 (2.146e9) bitwise exact; n 140000 + (2.294e9) raises `CUDA error: an illegal memory access` (loud, so not + a guard). The source offset is widened to int64 in-kernel: a source + offset of 2.62e9 was bitwise exact. At width 10240 the destination + boundary is n > 209715, within reach of a large `max_num_tokens`. + """ + + def __call__( + self, src: torch.Tensor, num_prefill_tokens: int, start_col: int, width: int + ) -> torch.Tensor: + return _extract_transpose_prefill_slice(src, num_prefill_tokens, start_col, width) + + def reference( + self, src: torch.Tensor, num_prefill_tokens: int, start_col: int, width: int + ) -> torch.Tensor: + return src[:num_prefill_tokens, start_col : start_col + width].transpose(0, 1).contiguous() + + def is_valid( + self, src: torch.Tensor, num_prefill_tokens: int, start_col: int, width: int + ) -> None: + # A 3-D src with a unit dim-1 stride passes the callable's own assert + # and is read as 2-D through stride(0). + assert src.dim() == 2, f"src must be 2D, got {src.dim()}D" + assert 0 <= num_prefill_tokens <= src.shape[0], ( + f"num_prefill_tokens {num_prefill_tokens} outside src rows [0, {src.shape[0]}]; " + "the kernel reads past src rather than failing" + ) + assert 0 <= start_col and 0 <= width and start_col + width <= src.shape[1], ( + f"column range [{start_col}, {start_col + width}) outside src columns " + f"[0, {src.shape[1]}); the kernel reads past src rather than failing" + ) + + +extract_transpose_prefill_slice = ExtractTransposePrefillSlice() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_gating.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_gating.py new file mode 100644 index 000000000000..e81f6f95b7b8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_gating.py @@ -0,0 +1,221 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gated-DeltaNet decay gate: g = -exp(A_log) * softplus(a + dt_bias), per token +and per value head, in one Triton kernel. + +Wraps the in-tree module callable +``tensorrt_llm._torch.modules.mamba.gdn_mixer.fused_gdn_gating`` (tier-2 +sourcing: no op boundary exists). Its body allocates the fp32 output and +launches ``fused_gdn_gating_kernel`` once. +""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.mamba.gdn_mixer import fused_gdn_gating as _fused_gdn_gating + +#: This checkpoint's GDN value-head count at tp1, and the `[b | a]` width the +#: stock mixer slices `a` out of. +NUM_V_HEADS = 48 +BA_WIDTH = 2 * NUM_V_HEADS + +#: Read off this checkpoint's 48 layers of `linear_attn.A_log` / +#: `linear_attn.dt_bias` (bf16, 48 each): exp(A_log) spans [3.839e-03, 1.394e+02], +#: and dt_bias reaching 19.25 against the default threshold of 20.0 makes the +#: softplus's linear branch a live path, not a corner case. +A_LOG_RANGE = (-5.5625, 4.9375) +DT_BIAS_RANGE = (-8.9375, 19.25) + +#: Relative error of the SASS approximate transcendentals (`ex2.approx.f32`, +#: `lg2.approx.f32`) Triton lowers `tl.exp` / `tl.log` to. +APPROX_ULP = 2.0**-22 + + +class FusedGdnGating(OpWrapper): + """`g[t, h] = -exp(A_log[h]) * softplus(a[t, h] + dt_bias[h])`, computed in + fp32 whatever the operand dtypes, into a fresh contiguous fp32 + `[batch, num_heads]` tensor. `batch, num_heads = a.shape`. + + `softplus(x) = log(1 + exp(beta * x)) / beta` while `beta * x <= threshold`, + else `x` -- the unstabilised form: below `x = -24 ln 2 = -16.6355` the fp32 + sum drops `exp(x)` and `g` is exactly 0 (absolute error <= 6e-8 times + exp(A_log), relative error 1.0). Harmless for a log-domain decay. + + Fusion boundary: bias add, softplus with its threshold branch, `exp(A_log)` + and the product are inside. The projection producing `a`, the + recurrence's other gate `beta = sigmoid(b)`, and the recurrence are not. + No input is modified. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: rows (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. + 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)), + ) + ) + + note = """ + Metadata consumed: none. Stateless and per (token, head). + + Call site (`_gdn_verify`): `(A_log, a_d, dt_bias, 1.0, 20.0)` with `a_d` + the `[n * p, 48]` bf16 decode rows of the `a` projection (all of it on a + pure verify step) and fp32 `A_log` / `dt_bias`; `g` is then reshaped to + `[n, p, 48]` for `fused_recurrent_gated_delta_rule_update`. + + 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 + view moved g by 6.031e+01 and a stride-2 dt_bias by 3.832e-01; 40-element + A_log / dt_bias against 48 heads moved it by 1.009e+02 / 1.010e+02 (the + per-head load is masked against `a.shape[1]` only); an `a` with head + stride 2 by 1.153e+00. + + Deterministic: two launches agree bitwise. + + Hard limits + none probed. + """ + + def __call__( + self, + A_log: torch.Tensor, + a: torch.Tensor, + dt_bias: torch.Tensor, + beta: float = 1.0, + threshold: float = 20.0, + ) -> torch.Tensor: + return _fused_gdn_gating(A_log, a, dt_bias, beta, threshold) + + def reference( + self, + A_log: torch.Tensor, + a: torch.Tensor, + dt_bias: torch.Tensor, + beta: float = 1.0, + threshold: float = 20.0, + ) -> torch.Tensor: + """fp64 with the stabilised `log1p(exp(x))`, rounded once to fp32. + + The stabilised form makes the reference the accurate side; the + kernel's unstabilised tail differs inside the test's absolute floor. + """ + num_heads = a.shape[-1] + x = a.double() + dt_bias[:num_heads].double() + softplus = torch.where(beta * x <= threshold, torch.log1p(torch.exp(beta * x)) / beta, x) + return (-torch.exp(A_log[:num_heads].double()) * softplus).float() + + def is_valid( + self, + A_log: torch.Tensor, + a: torch.Tensor, + dt_bias: torch.Tensor, + beta: float = 1.0, + threshold: float = 20.0, + ) -> None: + # `shape[-1]` / `stride(-1)` so a non-2-D `a` still reaches the op's own + # ValueError rather than being preempted here. + num_heads = a.shape[-1] + assert a.stride(-1) == 1, ( + f"`a`'s head axis must be unit-stride; got strides {a.stride()} -- the kernel " + "reads a row as `row * a.stride(0) + head`, so another column stride returns " + "a finite wrong g" + ) + assert A_log.stride(0) == 1 and dt_bias.stride(0) == 1, ( + f"A_log and dt_bias must be contiguous 1-D; got strides {A_log.stride()} and " + f"{dt_bias.stride()} -- both are read as `ptr + head`" + ) + assert A_log.numel() >= num_heads and dt_bias.numel() >= num_heads, ( + f"A_log and dt_bias must hold at least a.shape[-1]={num_heads} elements; got " + f"{A_log.numel()} and {dt_bias.numel()} -- the per-head load is masked against " + "a.shape[-1] only, so a short tensor is read out of bounds" + ) + + +fused_gdn_gating = FusedGdnGating() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_gdn_post_conv.py new file mode 100644 index 000000000000..b97bbdec86f4 --- /dev/null +++ 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 @@ -0,0 +1,262 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gated-DeltaNet post-convolution preparation: the q/k/v split of the causal +conv output with Q/K L2 normalization, plus the decay gate and beta. + +Wraps the in-tree module callable +``tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops.fused_gdn_post_conv`` +(tier-2 sourcing: no op boundary exists). One Triton launch, +``_fused_gdn_post_conv_kernel``, after allocating its five outputs. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops import ( + fused_gdn_post_conv as _fused_gdn_post_conv, +) + +#: This checkpoint's linear-attention geometry at tp1. +NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM = 16, 128, 48, 128 +QKV_DIM = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240 + +#: Read off this checkpoint's `linear_attn.A_log` / `linear_attn.dt_bias` +#: (bf16, 48 per layer): A_log spans [-5.5625, 3.5625] and dt_bias +#: [-8.9375, 19.25]. dt_bias reaching 19.25 makes the softplus threshold branch +#: a live path, not a corner case. +A_LOG_RANGE = (-5.5625, 3.5625) +DT_BIAS_RANGE = (-8.9375, 19.25) +SOFTPLUS_THRESHOLD = 20.0 + + +class FusedGdnPostConv(OpWrapper): + """Split a causal-conv output into q/k/v and build the GDN gates. + + The token axis is the `num_prefill_tokens` columns of the dim-major + `prefill` `[>= qkv_dim, n_p]`, then the rows of the token-major `decode` + `[n_d, >= qkv_dim]` (`None` for a prefill-only batch), with + `qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim` packed + `[q | k | v]`. Per token and head, in fp32: + + q, k = x / sqrt(sum(x^2) + l2_norm_eps) (per head) + v = copied bit-exact + g = -exp(A_log) * softplus(a + dt_bias) (fp32 always) + beta = sigmoid(b) (fp32, rounded once to beta_dtype) + + `softplus(x)` is the unstabilised `log(1 + exp(x))` below + `softplus_threshold` and `x` above it. + + Fusion boundary: split, Q/K norm, `g` and `beta` are inside. The conv is + not (`prefill`/`decode` are already its output), and the recurrence that + consumes these must not normalize q/k again (`use_qk_l2norm_in_kernel=False`). + + Returns `(q, k, v, g, beta)`, five fresh contiguous tensors: + `[1, T, num_k_heads, head_k_dim]` x2 and `[1, T, num_v_heads, head_v_dim]` + in `prefill.dtype`, then `[1, T, num_v_heads]` fp32 and `beta_dtype`. + No input is modified. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: num_prefill, num_decode; 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. + 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", + 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), + ), + ) + ) + + note = """ + Metadata consumed: none. Stateless and per token; it never needs sequence + boundaries. + + Call sites (`_gdn_block`): prefill-only `(mq_t, None, a, b, ...)`, mixed + `(mq_p_t, mq_d, a, b, ...)` and the MTP-verify context half + `(mq_p_t, None, a_p, b_p, ...)`, all with `l2_norm_eps=1e-6`, + `softplus_threshold=20.0` and the fp32 `beta_dtype` default; `prefill` is + `extract_transpose_prefill_slice`'s contiguous output, `a`/`b` the + `[T, 48]` projections. The stock mixer's verify call passes + `beta_dtype=b.dtype`. + + 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 + case) and are read through their own strides. `head_k_dim`/`head_v_dim` + are rounded up to a power of two and masked. + + The op raises on none of the metadata violations probed; `is_valid` + guards the four the kernel reads out of bounds or through the wrong + addresses, measured: stride-2 A_log moved g by 2.8e+02, stride-2 dt_bias by + 3.2e+02; A_log, dt_bias, a, b short by rows/elements moved g by 4.9e+12, + 2.9e+02, 3.2e+37 and made beta NaN. + + The softplus is unstabilised: for very negative `x` the fp32 `1 + exp(x)` + loses `exp(x)` and `g` rounds to exactly 0. Absolute error stays ~1 fp32 + ulp of the log argument times `exp(A_log)` (4.8e-07 measured over + x in [-9, 5]); relative error in that tail reaches 1.0. Inconsequential + for a log-domain decay. `exp(x)` is evaluated on the linear side too + (`tl.where` is a select); harmless below fp32 overflow at x ~ 88, which no + checkpoint `a + dt_bias` reaches. + + Not certified: fp16/fp32 activations, other head geometries. + + Hard limits + none probed. + """ + + def __call__( + self, + prefill: torch.Tensor, + decode: Optional[torch.Tensor], + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + num_k_heads: int, + head_k_dim: int, + num_v_heads: int, + head_v_dim: int, + l2_norm_eps: float = 1e-6, + softplus_threshold: float = 20.0, + beta_dtype: torch.dtype = torch.float32, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return _fused_gdn_post_conv( + prefill, + decode, + a, + b, + A_log, + dt_bias, + num_k_heads, + head_k_dim, + num_v_heads, + head_v_dim, + l2_norm_eps, + softplus_threshold, + beta_dtype, + ) + + def reference( + self, + prefill: torch.Tensor, + decode: Optional[torch.Tensor], + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + num_k_heads: int, + head_k_dim: int, + num_v_heads: int, + head_v_dim: int, + l2_norm_eps: float = 1e-6, + softplus_threshold: float = 20.0, + beta_dtype: torch.dtype = torch.float32, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """q/k normalized in fp32 and rounded once; `g` in fp64 with the + stabilised `log1p(exp(x))` (the kernel's unstabilised form differs + inside the test's floor, not around it), rounded to fp32; `beta` is + the fp32 sigmoid rounded once to `beta_dtype`.""" + blocks = [prefill.t()[:, : 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim]] + if decode is not None: + blocks.append(decode[:, : blocks[0].shape[1]]) + mixed = torch.cat(blocks, dim=0) + n = mixed.shape[0] + kd = num_k_heads * head_k_dim + + def l2norm(x): + x = x.reshape(n, num_k_heads, head_k_dim).float() + return (x / torch.sqrt(x.pow(2).sum(-1, keepdim=True) + l2_norm_eps)).to(prefill.dtype) + + q, k = l2norm(mixed[:, :kd]), l2norm(mixed[:, kd : 2 * kd]) + v = mixed[:, 2 * kd :].reshape(n, num_v_heads, head_v_dim).contiguous() + x = a[:n, :num_v_heads].double() + dt_bias[:num_v_heads].double() + softplus = torch.where(x <= softplus_threshold, torch.log1p(torch.exp(x)), x) + g = (-torch.exp(A_log[:num_v_heads].double()) * softplus).float() + beta = torch.sigmoid(b[:n, :num_v_heads].float()).to(beta_dtype) + return tuple(t.unsqueeze(0) for t in (q, k, v, g, beta)) + + def is_valid( + self, + prefill: torch.Tensor, + decode: Optional[torch.Tensor], + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + num_k_heads: int, + head_k_dim: int, + num_v_heads: int, + head_v_dim: int, + l2_norm_eps: float = 1e-6, + softplus_threshold: float = 20.0, + beta_dtype: torch.dtype = torch.float32, + ) -> None: + # The kernel indexes A_log/dt_bias as `ptr + head` with no stride and + # no mask, and does not mask a/b's token axis against their extent. + num_tokens = prefill.shape[1] + (0 if decode is None else decode.shape[0]) + assert A_log.stride(0) == 1 and dt_bias.stride(0) == 1, ( + f"A_log and dt_bias must be contiguous 1-D; got strides {A_log.stride()} " + f"and {dt_bias.stride()} -- read as `ptr + head`, a strided view gives a " + "finite wrong g" + ) + assert A_log.numel() >= num_v_heads and dt_bias.numel() >= num_v_heads, ( + f"A_log and dt_bias must hold at least num_v_heads={num_v_heads} elements; " + f"got {A_log.numel()} and {dt_bias.numel()} -- the load is unmasked" + ) + assert a.shape[0] >= num_tokens and b.shape[0] >= num_tokens, ( + f"a and b must have at least num_tokens={num_tokens} rows (prefill + decode); " + f"got {a.shape[0]} and {b.shape[0]} -- the token axis is unmasked" + ) + + +fused_gdn_post_conv = FusedGdnPostConv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_update.py new file mode 100644 index 000000000000..a290b44717ab --- /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_update.py @@ -0,0 +1,367 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gated DeltaNet recurrent core over a slot pool, several tokens per sequence: +the speculative verify step's recurrence. + +Wraps the in-tree module callable +``tensorrt_llm._torch.modules.fla.fused_recurrent.fused_recurrent_gated_delta_rule_update`` +(tier-2 sourcing: the gated delta rule has no op boundary). One launch of the +vendored Triton kernel ``fused_recurrent_gated_delta_rule_update_fwd_kernel`` -- +``triton.heuristics`` + ``triton.jit``, no autotuner. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.fla.fused_recurrent import ( + fused_recurrent_gated_delta_rule_update as _fused_recurrent_gated_delta_rule_update, +) + +#: This checkpoint's GDN geometry at tp1. +NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM = 16, 48, 128, 128 + +#: Rows of the engine's per-step buffer (`max_num_sequences`), as the +#: runtime-value dumps record `intermediate_ssm_states`: [32, p, 48, 128, 128]. +ENGINE_VERIFY_ROWS = 32 +#: 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. +SLOTS = 8 + + +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 + h *= exp(g); u = (v - h^T k) * beta; h += k u^T; o = h^T (scale * q) + + in fp32 throughout. `h` starts at `initial_state_source[idx]` (`[S, HV, V, + K]`, K innermost), or zero when `idx < 0` or there is no pool. `scale=None` + is `K ** -0.5`. `beta` with `v`'s rank is read per channel. + + Two independent sinks for `h`: + `disable_state_update=False` stores the final state back into slot `idx`; + `intermediate_states_buffer` (`[S', cache_steps, HV, V, K]`) receives the + state after every token at `[idx, step]`. The verify call is the second + alone: the pool must not move until the sampler has accepted a prefix. + + Fusion boundary: L2 norm, scale, the indexed gather and both scatters are + inside. The gates are *not*: `g = -exp(A_log) * softplus(a + dt_bias)` and + `beta = sigmoid(b)` arrive evaluated. The conv and q/k/v split are the + caller's too. + + Returns the `[N, P, HV, V]` output (`[1, T, HV, V]` with `cu_seqlens`) in + `output`'s dtype, else `q`'s; it *is* `output` when one is given. The pool + and the per-step buffer are written in place and not returned. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: mode ("verify": disable_state_update + buffer with cache_steps = p + # | "write_back" | "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" + # 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) + ) + + ( + _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), + ) + ) + + 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. + + Preconditions not checkable from metadata: `HV % H == 0` (q/k heads are + repeat_interleaved, not tiled); indices in `[0, S)` and `[0, S')` or + negative (the device check is `idx >= 0` only; out of range is an + out-of-bounds access, not probed); distinct indices when writing back -- + measured at two sequences on one slot: outputs exact, the slot a + per-(v-tile, head) mixture 0.83 .. 0.89 from a sequential application; in the + varlen form `cache_steps >= max(seq_len)`; q/k unit-norm-conditioned (pass + `use_qk_l2norm_in_kernel=True` or normalize upstream). `g` is log-space. + + Raise loudly, so not guarded: `cu_seqlens` with `q.shape[0] > 1` and an + index count other than `len(cu_seqlens) - 1` (ValueError); a `None` pool + with the write-back on (triton CompilationError). + + `is_valid` guards what is accepted and silently wrong, each measured: + a strided `output` (the one tensor not made contiguous) is written with + packed arithmetic -- 7.04 off and the neighbouring half overwritten; + `output` with `disable_output_calculation=True` is never written and the + return is a [1, 1, 1, 1] scratch; a strided pool or per-step buffer is + silently replaced by a contiguous copy, so the writes land in the copy and + the caller's tensor comes back bitwise unchanged; `cache_steps` below the + token count (or None, i.e. 0) overwrites the next row's steps -- 0.850 off + at cache_steps 2, P 4; `cache_steps` other than the buffer's step axis + misplaces every row past 0 against the buffer's own indexing. + + Deterministic; the first call in a fresh process is bitwise the second. + + 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. + + Hard limits + N * HV gridDim.z boundary 65535 probed at N=1300 (62400): runs; + N=1400 (67200): RuntimeError "invalid argument", context survives + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state_source: Optional[torch.Tensor], + initial_state_indices: torch.Tensor, + scale: Optional[float] = None, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: Optional[torch.Tensor] = None, + disable_state_update: bool = False, + disable_output_calculation: bool = False, + intermediate_states_buffer: Optional[torch.Tensor] = None, + cache_steps: Optional[int] = None, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return _fused_recurrent_gated_delta_rule_update( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state_source=initial_state_source, + initial_state_indices=initial_state_indices, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + disable_state_update=disable_state_update, + disable_output_calculation=disable_output_calculation, + intermediate_states_buffer=intermediate_states_buffer, + cache_steps=cache_steps, + output=output, + ) + + def reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state_source: Optional[torch.Tensor], + initial_state_indices: torch.Tensor, + scale: Optional[float] = None, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: Optional[torch.Tensor] = None, + disable_state_update: bool = False, + disable_output_calculation: bool = False, + intermediate_states_buffer: Optional[torch.Tensor] = None, + cache_steps: Optional[int] = None, + output: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + """`(out, initial_state_source, intermediate_states_buffer)` as the call + leaves them; mutates nothing. Token-sequential, fp32 throughout, each + store rounded once. `disable_output_calculation` and `cache_steps` are + ignored (the output is always computed; the buffer is indexed by its + own step axis).""" + H, K = k.shape[2], k.shape[3] + HV, V = v.shape[2], v.shape[3] + rep = HV // H + scale = K**-0.5 if scale is None else scale + if cu_seqlens is None: + spans = [(b, 0, q.shape[1]) for b in range(q.shape[0])] + else: + cu = cu_seqlens.tolist() + spans = [(0, cu[i], cu[i + 1]) for i in range(len(cu) - 1)] + out = torch.zeros(v.shape, dtype=torch.float32, device=v.device) + pool = None if initial_state_source is None else initial_state_source.clone() + cache = None if intermediate_states_buffer is None else intermediate_states_buffer.clone() + per_channel = beta.dim() == v.dim() + for i, (b, start, end) in enumerate(spans): + slot = int(initial_state_indices[i]) + if slot >= 0 and initial_state_source is not None: + h = initial_state_source[slot].float().transpose(1, 2) # [HV, K, V] + else: + h = torch.zeros(HV, K, V, dtype=torch.float32, device=v.device) + for step, t in enumerate(range(start, end)): + q_t = q[b, t].float().repeat_interleave(rep, dim=0) # [HV, K] + k_t = k[b, t].float().repeat_interleave(rep, dim=0) + if use_qk_l2norm_in_kernel: + q_t = q_t / (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 * scale + h = h * torch.exp(g[b, t].float())[:, None, None] + u = v[b, t].float() - (h * k_t[:, :, None]).sum(1) # [HV, V] + u = u * (beta[b, t].float() if per_channel else beta[b, t].float()[:, None]) + h = h + k_t[:, :, None] * u[:, None, :] + out[b, t] = (h * q_t[:, :, None]).sum(1) + if cache is not None and slot >= 0: + cache[slot, step] = h.transpose(1, 2).to(cache.dtype) + if pool is not None and not disable_state_update and slot >= 0: + pool[slot] = h.transpose(1, 2).to(pool.dtype) + return out.to(q.dtype if output is None else output.dtype), pool, cache + + def is_valid( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state_source: Optional[torch.Tensor], + initial_state_indices: torch.Tensor, + scale: Optional[float] = None, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: Optional[torch.Tensor] = None, + disable_state_update: bool = False, + disable_output_calculation: bool = False, + intermediate_states_buffer: Optional[torch.Tensor] = None, + cache_steps: Optional[int] = None, + output: Optional[torch.Tensor] = None, + ) -> None: + assert not (output is not None and disable_output_calculation), ( + "output and disable_output_calculation=True are mutually exclusive: with the " + "output calculation off the buffer is never written and the return value is a " + "[1, 1, 1, 1] scratch tensor" + ) + if output is not None: + assert output.is_contiguous(), ( + f"output must be contiguous; got strides {output.stride()} for shape " + f"{tuple(output.shape)} -- the kernel writes it with packed [B, T, HV, V] " + "arithmetic, misplacing values and clobbering whatever lies between" + ) + if initial_state_source is not None and not disable_state_update: + assert initial_state_source.is_contiguous(), ( + f"initial_state_source must be contiguous when disable_state_update is False; " + f"got strides {initial_state_source.stride()} -- a strided pool is copied " + "before the call and the final state is written into the copy" + ) + if intermediate_states_buffer is not None: + assert intermediate_states_buffer.is_contiguous(), ( + f"intermediate_states_buffer must be contiguous; got strides " + f"{intermediate_states_buffer.stride()} -- a strided buffer is copied before " + "the call and every per-step state is written into the copy" + ) + if cu_seqlens is None: + assert cache_steps is not None and cache_steps >= q.shape[1], ( + f"cache_steps must be at least the per-sequence token count when an " + f"intermediate_states_buffer is given; got cache_steps={cache_steps} for " + f"{q.shape[1]} tokens per sequence -- a smaller stride makes each " + "sequence's tail steps overwrite the next row's" + ) + assert cache_steps == intermediate_states_buffer.shape[1], ( + f"cache_steps={cache_steps} must equal intermediate_states_buffer's step axis " + f"({intermediate_states_buffer.shape[1]}) -- the kernel strides rows by " + "cache_steps, so any other value misplaces every row past 0" + ) + + +fused_recurrent_gated_delta_rule_update = FusedRecurrentGatedDeltaRuleUpdate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_sigmoid_gating_delta_rule_update.py new file mode 100644 index 000000000000..12dd485c4bfe --- /dev/null +++ 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 @@ -0,0 +1,343 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gated DeltaNet recurrent core with the gating fused in, over a slot pool: +the GDN decode step. + +Wraps the in-tree module callable +``tensorrt_llm._torch.modules.fla.fused_sigmoid_gating_recurrent.fused_sigmoid_gating_delta_rule_update`` +(tier-2 sourcing: the gated delta rule has no ``torch.ops.trtllm`` surface). +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.fla.fused_sigmoid_gating_recurrent import ( + fused_sigmoid_gating_delta_rule_update as _fused_sigmoid_gating_delta_rule_update, +) + +#: 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. +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] + + +class FusedSigmoidGatingDeltaRuleUpdate(OpWrapper): + """Advance the gated delta rule over every sequence, gates computed in-kernel. + + Sequence `n` (tokens `cu_seqlens[n]:cu_seqlens[n+1]`, or row `n` without + `cu_seqlens`), v-head `hv` reading k/q head `hv // (HV // H)`, state `h` + `[K, V]` loaded from `initial_state_source[initial_state_indices[n]]` + (stored `[V, K]`, K innermost), per token, all in fp32: + + g = -exp(A_log) * softplus_b(a + dt_bias) beta = sigmoid(b) + q, k = q / (|q| + 1e-6), k / (|k| + 1e-6) (only with use_qk_l2norm_in_kernel) + h *= exp(g); u = (v - h^T k) * beta; h += k u^T; o = h^T (scale * q) + + `softplus_b(x) = log(1 + exp(beta_s x)) / beta_s` below `softplus_threshold` + (in `beta_s x`), `x` above; `scale=None` is `K ** -0.5`, bitwise. The final + `h` is written back to the same slot. A negative index means no slot: zero + initial state, nothing written. `initial_state_source=None`: zero state, + nothing written. + + Fusion boundary: both gates, the optional q/k norm, scale, recurrence, + indexed gather and scatter are inside. The conv, the q/k/v split (plain + views) and the gated RMS norm on `o` are the caller's. + + Returns `o` `[B, T, HV, V]` in `output`'s dtype (it *is* `output` when + given), else `q`'s. The pool is advanced in place and not returned. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: seq_lens, slots (pool rows; 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). + 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="cu_seqlens None: batched [N, T, ...] layout, N=3, T=4", + spec=dict(seq_lens=(4, 4, 4), slots=SLOTS, layout="batched"), + ), + ) + ) + + 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. + + 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. + + Preconditions not checked by anything: + - `HV % H == 0`. + - Pool trailing axes `(V, K)`, K innermost. The callable asserts the + strides `(V*K, K, 1)`, but at K == V a transposed `[S, HV, K, V]` pool + has the same strides: accepted, measured 2.15e-2 off (vs 3.8e-6 + correct). The slot stride is free. + - Indices distinct (a shared slot races on write-back: outputs fine, the + slot holds one writer's state) and in `[0, S)` or negative + (out-of-range is an unchecked device access -- `tl.device_assert` is + compiled out -- not probed). + - q/k unit-norm conditioned: in-kernel norm on, or normalized by the + caller. Raw multi-token q/k with the norm off drifted the fp32 state + 1.7e-2 relative from an fp32 sequential reference: domain, not defect. + - `output` contiguous, `[B, T, HV, V]`: guarded by `is_valid`. Its dtype + wins (an fp32 `output` with bf16 q returns fp32). + + Deterministic: repeated calls are bitwise identical. Grid + `(1, ceil(V/32), min(N*HV, 65535))`, grid-striding past the cap (read, + not measured: N > 1365 at this geometry). + + Not certified: the FlashInfer branch, K or V other than 128, fp16/fp32 + operands, `softplus_beta`/`softplus_threshold` other than 1.0/20.0, + `B > 1` with `cu_seqlens`. + + Hard limits + none probed. + """ + + def __call__( + self, + A_log: torch.Tensor, + a: torch.Tensor, + dt_bias: torch.Tensor, + softplus_beta: float, + softplus_threshold: float, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + b: torch.Tensor, + initial_state_source: Optional[torch.Tensor], + initial_state_indices: torch.Tensor, + scale: Optional[float] = None, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return _fused_sigmoid_gating_delta_rule_update( + A_log=A_log, + a=a, + dt_bias=dt_bias, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + q=q, + k=k, + v=v, + b=b, + initial_state_source=initial_state_source, + initial_state_indices=initial_state_indices, + scale=scale, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + cu_seqlens=cu_seqlens, + output=output, + ) + + def reference( + self, + A_log: torch.Tensor, + a: torch.Tensor, + dt_bias: torch.Tensor, + softplus_beta: float, + softplus_threshold: float, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + b: torch.Tensor, + initial_state_source: Optional[torch.Tensor], + initial_state_indices: torch.Tensor, + scale: Optional[float] = None, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """`(o, pool_after)`: sequential over tokens, fp32 throughout like the + kernel, `o` rounded once to the result dtype and the state once to the + pool's. `output` contributes only its dtype; nothing is mutated.""" + batch, total, num_k_heads, k_dim = q.shape + hv, v_dim = v.shape[2], v.shape[3] + repeat = hv // num_k_heads + scale = k_dim**-0.5 if scale is None else scale + if cu_seqlens is None: + spans = [(n, 0, total) for n in range(batch)] + else: + bounds = cu_seqlens.tolist() + spans = [(0, bounds[n], bounds[n + 1]) for n in range(len(bounds) - 1)] + a3, b3 = (a, b) if a.dim() == 3 else (a.unsqueeze(0), b.unsqueeze(0)) + out = torch.zeros(batch, total, hv, v_dim, device=q.device, dtype=torch.float32) + pool = None if initial_state_source is None else initial_state_source.clone() + neg_exp_a = -torch.exp(A_log.float()) + slots = initial_state_indices.tolist() + for n, (row, bos, eos) in enumerate(spans): + slot = slots[n] + if pool is not None and slot >= 0: + state = initial_state_source[slot].float().transpose(1, 2) + else: + state = torch.zeros(hv, k_dim, v_dim, device=q.device, dtype=torch.float32) + for t in range(bos, eos): + q_t = q[row, t].float().repeat_interleave(repeat, dim=0) + k_t = k[row, t].float().repeat_interleave(repeat, dim=0) + x = a3[row, t].float() + dt_bias.float() + bx = softplus_beta * x + softplus = torch.where( + bx <= softplus_threshold, torch.log(1.0 + torch.exp(bx)) / softplus_beta, x + ) + beta = torch.sigmoid(b3[row, t].float()) + if use_qk_l2norm_in_kernel: + q_t = q_t / (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) + state = state * torch.exp(neg_exp_a * softplus)[:, None, None] + u = (v[row, t].float() - (state * k_t[:, :, None]).sum(1)) * beta[:, None] + state = state + k_t[:, :, None] * u[:, None, :] + out[row, t] = (state * (q_t * scale)[:, :, None]).sum(1) + if pool is not None and slot >= 0: + pool[slot] = state.transpose(1, 2).to(pool.dtype) + return out.to(q.dtype if output is None else output.dtype), pool + + def is_valid( + self, + A_log: torch.Tensor, + a: torch.Tensor, + dt_bias: torch.Tensor, + softplus_beta: float, + softplus_threshold: float, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + b: torch.Tensor, + initial_state_source: Optional[torch.Tensor], + initial_state_indices: torch.Tensor, + scale: Optional[float] = None, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + ) -> None: + # Measured on the Triton branch: a last-dim slice of a doubly wide + # buffer and a token-strided view were both accepted, came back 5.3e-2 + # and 4.8e-2 off at an output absmax of 4.1e-2, and overwrote elements + # outside the view. + assert output is None or output.is_contiguous(), ( + f"output must be contiguous; got strides {output.stride()} for shape " + f"{tuple(output.shape)} -- the kernel writes it with packed " + "[NK, B, T, HV, V] arithmetic, misplacing values and clobbering the gaps" + ) + + +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/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 new file mode 100644 index 000000000000..175b7f695dfd --- /dev/null +++ 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 @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gated-DeltaNet decode-phase q/k/v packing: the plain split of the causal +conv-update output into per-head q, k and v, with no normalization and no gates. + +Wraps the in-tree module callable +``tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops.pack_gdn_decode_qkv`` +(tier-2 sourcing: no op boundary exists). One Triton launch, +``_pack_gdn_decode_qkv_kernel``, after allocating the three outputs. +""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops import ( + pack_gdn_decode_qkv as _pack_gdn_decode_qkv, +) + +#: This checkpoint's linear-attention geometry at tp1. +NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM = 16, 128, 48, 128 +CKPT_HEADS = (NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM) +QKV_DIM = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240 + + +class PackGdnDecodeQkv(OpWrapper): + """Cut a token-major `[num_tokens, >= qkv_dim]` conv output, packed + `[q | k | v]` along features with + `qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim`, into + + q, k [1, num_tokens, num_k_heads, head_k_dim] + v [1, num_tokens, num_v_heads, head_v_dim] + + Every element is a bit-exact copy: no Q/K L2 norm, no gate, no dtype + change. Columns at or past `qkv_dim` are dropped. + + Fusion boundary: only the split and per-head reshape are inside. The conv + (and its activation) is the caller's, and so is the Q/K normalization -- + the recurrence consuming these must run `use_qk_l2norm_in_kernel=True`. + That is the one way to misuse this op and still get plausible output. + + Row-wise, with no notion of a sequence: a verify step of `B` sequences x + `P` tokens is passed as `B * P` rows, the token axis keeps the input row + order, and `q.reshape(B, P, ...)` is a view, not a re-gather. + + Returns three fresh contiguous allocations in `mixed_qkv.dtype`; + `mixed_qkv` is read through both of its strides and is not modified. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: num_tokens, or (num_decodes, draft) for a verify batch; 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"), + ) + ) + + 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)),) + ) + + note = """ + Metadata consumed: none. Stateless; sequence boundaries and draft grouping + are carried by the caller's row order alone. + + Call site (`_gdn_verify`): `mixed_qkv` the `[n * p, 10240]` token-major + conv output (contiguous: the conv result carries its input's strides, so + transposing back is a view), heads `16, 128, 48, 128`, bf16. + + Raise loudly, so not guarded: `head_k_dim` and `head_v_dim` rounding up to + different powers of two (`128 / 64`: triton `CompilationError` + "Mismatched type for dim_mask", before any launch); a 1-D input + (`IndexError` reading `stride(1)`). + + The feature axis is never masked against the input's own extent; the + token axis is. `is_valid` guards that: `[32, 9216]` at qkv_dim 10240 was + accepted and returned finite values (83.3% of `v` bit-equal to the block + it should have read, the rest up to 6.25 off); over-declaring + `num_v_heads` likewise; a 3-D `[4, 2, 10240]` input faulted with an + illegal memory access that poisons the process. + + Hard limits + none probed. + """ + + def __call__( + self, + mixed_qkv: torch.Tensor, + num_k_heads: int, + head_k_dim: int, + num_v_heads: int, + head_v_dim: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return _pack_gdn_decode_qkv(mixed_qkv, num_k_heads, head_k_dim, num_v_heads, head_v_dim) + + def reference( + self, + mixed_qkv: torch.Tensor, + num_k_heads: int, + head_k_dim: int, + num_v_heads: int, + head_v_dim: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """`[q | k | v]` cut with `narrow` and reshaped per head; views where torch allows.""" + n = mixed_qkv.shape[0] + qk_width, v_width = num_k_heads * head_k_dim, num_v_heads * head_v_dim + q = mixed_qkv.narrow(1, 0, qk_width).reshape(n, num_k_heads, head_k_dim) + k = mixed_qkv.narrow(1, qk_width, qk_width).reshape(n, num_k_heads, head_k_dim) + v = mixed_qkv.narrow(1, 2 * qk_width, v_width).reshape(n, num_v_heads, head_v_dim) + return q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0) + + def is_valid( + self, + mixed_qkv: torch.Tensor, + num_k_heads: int, + head_k_dim: int, + num_v_heads: int, + head_v_dim: int, + ) -> None: + # `shape[1]` rather than a rank check: a 1-D input still reaches the + # op's own IndexError, and a 3-D one (whose shape[1] is unrelated) is + # refused here instead of faulting. + qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim + assert mixed_qkv.shape[1] >= qkv_dim, ( + f"mixed_qkv must be at least qkv_dim={qkv_dim} wide " + f"(2 * {num_k_heads} * {head_k_dim} + {num_v_heads} * {head_v_dim}); got " + f"shape {tuple(mixed_qkv.shape)} -- the feature axis is not masked against " + "the input's extent, so a short row is read out of bounds" + ) + + +pack_gdn_decode_qkv = PackGdnDecodeQkv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/reset_recurrent_state_rows.py new file mode 100644 index 000000000000..5416d7bc8808 --- /dev/null +++ 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 @@ -0,0 +1,245 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Zero the recurrent (and optional convolution) pool rows of fresh requests. + +Wraps the in-tree module callable +``tensorrt_llm._torch.modules.mamba.recurrent_state_cache.reset_recurrent_state_rows`` +(tier-2 sourcing: no ``torch.ops.trtllm`` op exists). One Triton launch, +``_reset_recurrent_state_rows_kernel``; the rest of the callable is shape and +stride validation. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.mamba import recurrent_state_cache as _recurrent_state_cache + +#: This checkpoint's per-layer pool rows: a [48, 128, 128] recurrent state and a +#: [conv_dim, d_conv - 1] = [10240, 3] convolution window. +REC_SHAPE = (48, 128, 128) +CONV_SHAPE = (10240, 3) +#: 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. +ENGINE_SLOTS = 461 + + +def _engine_case(num_seqs: int, flip: bool) -> dict: + """`num_seqs` distinct slots over the engine pool, slot 0 and 460 live; + `flip` moves which of the two is the one reset.""" + slots = (0, ENGINE_SLOTS - 1) + tuple(i * 13 + 4 for i in range(num_seqs - 2)) + flags = tuple((i % 2 == 1) != flip for i in range(num_seqs)) + return dict(slots=ENGINE_SLOTS, indices=slots, flags=flags) + + +class ResetRecurrentStateRows(OpWrapper): + """For each request `i` with `slot = state_indices[i]`: when + `has_initial_states[i] == 0` and `0 <= slot < recurrent_states.shape[0]`, + fill row `slot` of `recurrent_states` -- and of `conv_states`, when given -- + with +0.0. Nothing is read from either pool. + + Untouched, bitwise: every unnamed slot, every named slot whose flag is + nonzero (the row a resumed or chunked request keeps), every negative or + out-of-range slot (skipped, not clamped), padding between rows, the whole + conv pool when it is `None`, both index tensors. A slot named twice is + zeroed if any of its entries asks. + + Fusion boundary: the call decides nothing. Computing `has_initial_states`, + assigning slots, and slicing both vectors to the context prefix are the + caller's. + + Returns `(recurrent_states, conv_states)` -- the pools it was handed, + written in place (`conv_states` is `None` when not passed). + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: slots (pool rows), indices, 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. + 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"), + ), + 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. + + Loud, from the callable itself (ValueError): an empty recurrent pool, conv + and recurrent slot counts differing (the in-range test uses the recurrent + pool's), flag/index shape mismatch, non-contiguous or overlapping rows in + either pool, any CPU tensor. + + Silent, so `is_valid` guards it: both index tensors are read as flat + contiguous buffers of `state_indices.shape[0]` elements. + + Not guarded (data, not metadata): an index dtype too narrow for its slot + wraps -- an int8 130 arrives as -126 and is skipped as invalid. int32 is + what the runtime uses. + + Launch: grid `(num_requests, cdiv(max(rec_row, conv_row), 256))`, the two + row widths masked independently; ordered on the current stream. + + Hard limits + slot * stride(0) int64 row base probed at 460 * 786432 = 3.6e8 exact + within-row offset int32 boundary 2^31 elements per row not probed + (this checkpoint's largest row is 786432) + """ + + def __call__( + self, + recurrent_states: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: torch.Tensor, + conv_states: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + _recurrent_state_cache.reset_recurrent_state_rows( + recurrent_states, + state_indices, + has_initial_states, + conv_states, + ) + return recurrent_states, conv_states + + def reference( + self, + recurrent_states: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: torch.Tensor, + conv_states: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Both pools as the call leaves them: clone, zero the rows the rule names.""" + slots = recurrent_states.shape[0] + rec = recurrent_states.clone() + conv = None if conv_states is None else conv_states.clone() + for slot, flag in zip(state_indices.tolist(), has_initial_states.tolist()): + if flag or not 0 <= slot < slots: + continue + rec[slot] = 0 + if conv is not None: + conv[slot] = 0 + return rec, conv + + def is_valid( + self, + recurrent_states: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: torch.Tensor, + conv_states: Optional[torch.Tensor] = None, + ) -> None: + # Measured: a stride-2 view of a 6-element index buffer reset rows + # [0, 1] instead of [0, 1, 2] (read through the base pointer); a + # stride-2 flag view reset row 1 instead of nothing; a [2, 3] + # index/flag pair considered only its first 2 of 6 requests. + assert state_indices.dim() == 1 and state_indices.is_contiguous(), ( + "state_indices must be a contiguous 1-D tensor" + ) + assert has_initial_states.dim() == 1 and has_initial_states.is_contiguous(), ( + "has_initial_states must be a contiguous 1-D tensor" + ) + + +reset_recurrent_state_rows = ResetRecurrentStateRows() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__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 new file mode 100644 index 000000000000..585c2a3b9a40 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/patch_embed_conv3d.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Vision patch embedding: one 3-D convolution per pixel patch. + +Tier-3 composed: one `torch.nn.functional.conv3d` call between two views that +launch nothing. No registered op is a 3-D convolution, and the CuTe DSL +implicit-GEMM conv3d kernels in `cute_dsl_kernels/blackwell/conv/` refuse this +checkpoint's three input channels (bf16 needs `in_channels % 8 == 0`; +`can_implement` swallows the refusal and returns False). +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +#: This checkpoint: in_channels 3, kernel (temporal_patch_size 2, patch_size +#: 16, 16), out_channels 1152 -- weight [1152, 3, 2, 16, 16], a 1536-pixel row. +IN_CHANNELS = 3 +KERNEL = (2, 16, 16) +EMBED_DIM = 1152 + + +class PatchEmbedConv3d(OpWrapper): + """`out[n, k] = bias[k] + sum_{c,t,i,j} pixel[n, c, t, i, j] * weight[k, c, t, i, j]`. + + One whole `[in_channels, kt, kh, kw]` patch per row of `pixel_values`, flat + `[N, C*kt*kh*kw]` (what the image processor emits, channel-slowest) or + explicit `[N, C, kt, kh, kw]`; the kernel extent is read from `weight`. + Cross-correlation (no kernel flip), like `nn.Conv3d`. Returns a fresh + contiguous `[N, out_channels]` in the input dtype. + + Fusion boundary: the patch projection and its bias. Patchification, the + position embedding added after it, and the row order (rows pass through in + the order they arrive -- the merge-block order the merger needs) are the + caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: num_patches, dtype; `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"), + ), + ) + ) + + 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. + + fp32 runs in TF32 unless the caller sets `torch.backends.cudnn.allow_tf32 = + False` (global state the caller owns): measured 5.84e-4 from an exact fp32 + matmul with it on (99.99% of outputs differ), 3.58e-6 with it off. bf16 and + fp16 are unaffected. + + At this shape the op is a GEMM: one output position per row, and the + `[1152, 1536]` reshape of the filter is a view of the same bytes. Not timed. + CUDA-graph capturable; the output is a fresh allocation. + + Not certified: float8 operands; patch counts past 4096. + """ + + def __call__( + self, + pixel_values: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + out_channels, in_channels, kt, kh, kw = weight.shape + patches = pixel_values.view(-1, in_channels, kt, kh, kw) + return F.conv3d(patches, weight, bias).view(-1, out_channels) + + def reference( + self, + pixel_values: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """True-fp32 matmul of the flattened rows against the flattened filter.""" + rows = pixel_values.reshape(pixel_values.shape[0], -1).float() + filters = weight.reshape(weight.shape[0], -1).float() + with true_fp32_matmul(): + out = rows @ filters.t() + if bias is not None: + out = out + bias.float() + return out.to(pixel_values.dtype) + + def is_valid( + self, + pixel_values: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> None: + # The view reinterprets anything whose element count divides: a whole + # [1, 3, 2, 64, 64] volume becomes 16 "patches" -- a [16, 1152] result + # of the right shape from the wrong pixels, 2.52 from the strided conv + # the caller meant -- and [N, 3072] rows become 2N, without an error. + _, in_channels, kt, kh, kw = weight.shape + patch_numel = in_channels * kt * kh * kw + if pixel_values.dim() == 2: + one_patch_per_row = pixel_values.shape[1] == patch_numel + elif pixel_values.dim() == 5: + one_patch_per_row = tuple(pixel_values.shape[1:]) == (in_channels, kt, kh, kw) + else: + one_patch_per_row = False + assert one_patch_per_row, ( + f"each row must be exactly one patch: pixel_values must be " + f"[num_patches, {patch_numel}] or " + f"[num_patches, {in_channels}, {kt}, {kh}, {kw}] for this weight; got " + f"{tuple(pixel_values.shape)}" + ) + + +patch_embed_conv3d = PatchEmbedConv3d() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/pos_embed_interpolate.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/pos_embed_interpolate.py new file mode 100644 index 000000000000..ba26e3250f11 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/vision/pos_embed_interpolate.py @@ -0,0 +1,276 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Vision position-embedding bilinear interpolation onto each image's patch +grid, emitted in spatial-merge order. + +Tier-3 composed: index and corner-weight arithmetic on `[h]`, `[w]` and +`[h, w]` tensors, then one `torch.nn.functional.embedding_bag` for the whole +batch. No registered op interpolates anything, and `modules/` has no vision +position-embedding callable. The exact-math sibling, +`_triton_pos_embed_interpolate` in `models/modeling_qwen3vl.py`, is a private +symbol of another model's modeling file; the other in-tree resamplers are +bicubic and/or `align_corners=False`, a different function. +""" + +from __future__ import annotations + +import math +from typing import Sequence + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: This checkpoint: num_position_embeddings 2304 (a 48 x 48 grid), hidden 1152, +#: spatial_merge_size 2. +NUM_GRID = 48 +HIDDEN = 1152 +MERGE = 2 + +#: 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.""" + tail = x.shape[2:] + return ( + x.reshape(height // merge, merge, width // merge, merge, *tail) + .permute(0, 2, 1, 3, *range(4, 4 + len(tail))) + .reshape(height * width, *tail) + ) + + +class PosEmbedInterpolate(OpWrapper): + """Resample a square row-major `[N*N, hidden]` table onto each `(t, h, w)` + patch grid; return `[sum(t*h*w), hidden]` in the table's dtype. + + Corner-aligned bilinear: `pos = linspace(0, N - 1, n)` per axis, `lo = + floor(pos)`, `hi = min(lo + 1, N - 1)`; `h` indexes table rows (stride `N`), + `w` columns. Rows come out in spatial-merge order -- whole `merge x merge` + blocks row-major, row-major inside a block -- which is what the patch + merger's 4-to-1 reshape consumes; raster order is silently wrong. `t` + repeats the frame's rows (no temporal axis); images concatenate in order. + The four-term sum accumulates in fp32 and rounds once. + + Fusion boundary: the position embeddings alone. Adding them to the patch + embeddings, the 2-D rope and the attention segmentation are the caller's. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # spec: grid (tuple of (t, h, w)), merge, dtype, rows (table rows), hidden. + CELLS: tuple[Cell, ...] = ( + 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="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) + ) + ) + + note = """ + Metadata consumed: none. `grid_thw` is host ints (a device tensor works but + `int()` syncs per element). The table is not mutated; the result is fresh. + + Loud failures, left to torch: an `h` or `w` not divisible by + `spatial_merge_size` (or a merge larger than the grid) raises "shape ... is + invalid for input of size ..."; a non-2-D table raises "weight has to be a + 2D Tensor". An empty grid side returns `[0, hidden]` rather than raising. + + At `h = w = N` the samples are exactly `arange(N)` and the result is + bitwise the table in merge order. The bf16 result is within one bf16 + half-ulp of the fp64 value (0.0147 at 62x90); a bf16 multiply-add chain + (a literal HF transcription) lands up to one bf16 ulp (0.031) from this. + + CUDA graphs: capturable, but the grid is a host argument, so it is frozen + at capture. Cost is ~225 us of launch-bound index construction per image, + flat in token count; the gather is the only part that scales. + + Not certified: a CPU table (it runs and returns on CPU). + """ + + def __call__( + self, + pos_embed_weight: torch.Tensor, + grid_thw: Sequence[Sequence[int]], + spatial_merge_size: int, + ) -> torch.Tensor: + num_grid_per_side = math.isqrt(pos_embed_weight.shape[0]) + device = pos_embed_weight.device + merge = spatial_merge_size + idx_parts = [] + weight_parts = [] + + for frames, height, width in grid_thw: + frames, height, width = int(frames), int(height), int(width) + + # Corner-aligned; non-negative positions, so the int64 cast floors. + h_pos = torch.linspace( + 0, num_grid_per_side - 1, height, device=device, dtype=torch.float32 + ) + w_pos = torch.linspace( + 0, num_grid_per_side - 1, width, device=device, dtype=torch.float32 + ) + h_lo = h_pos.to(torch.int64) + w_lo = w_pos.to(torch.int64) + h_hi = (h_lo + 1).clamp(max=num_grid_per_side - 1) + w_hi = (w_lo + 1).clamp(max=num_grid_per_side - 1) + dh = (h_pos - h_lo).unsqueeze(1) + dw = (w_pos - w_lo).unsqueeze(0) + + row_lo = (h_lo * num_grid_per_side).unsqueeze(1) + row_hi = (h_hi * num_grid_per_side).unsqueeze(1) + col_lo = w_lo.unsqueeze(0) + col_hi = w_hi.unsqueeze(0) + + # [h, w, 4]: corners in (lo,lo) (lo,hi) (hi,lo) (hi,hi) order. + idx = torch.stack( + (row_lo + col_lo, row_lo + col_hi, row_hi + col_lo, row_hi + col_hi), dim=-1 + ) + weights = torch.stack( + ((1 - dh) * (1 - dw), (1 - dh) * dw, dh * (1 - dw), dh * dw), dim=-1 + ) + + # Spatial-merge reorder; `reshape` is loud when a side does not divide. + merged_h, merged_w = height // merge, width // merge + idx = ( + idx.reshape(merged_h, merge, merged_w, merge, 4) + .permute(0, 2, 1, 3, 4) + .reshape(height * width, 4) + ) + weights = ( + weights.reshape(merged_h, merge, merged_w, merge, 4) + .permute(0, 2, 1, 3, 4) + .reshape(height * width, 4) + ) + + if frames > 1: + idx = idx.repeat(frames, 1) + weights = weights.repeat(frames, 1) + idx_parts.append(idx) + weight_parts.append(weights) + + idx = idx_parts[0] if len(idx_parts) == 1 else torch.cat(idx_parts) + weights = weight_parts[0] if len(weight_parts) == 1 else torch.cat(weight_parts) + + # One launch: each row of `idx` is a 4-element bag scaled by `weights` + # and summed, fp32 table and fp32 weights. + table = pos_embed_weight.float() + out = F.embedding_bag(idx, table, mode="sum", per_sample_weights=weights) + return out.to(pos_embed_weight.dtype) + + def reference( + self, + pos_embed_weight: torch.Tensor, + grid_thw: Sequence[Sequence[int]], + spatial_merge_size: int, + ) -> torch.Tensor: + """The explicit four-corner weighted sum in fp64, rounded once. + + Shares the entry's sample positions (fp32 `linspace`) but not its + gather or accumulation; the test cross-checks the structure separately + against `F.interpolate`. + """ + num_grid = math.isqrt(pos_embed_weight.shape[0]) + hidden = pos_embed_weight.shape[1] + grid = pos_embed_weight.double().view(num_grid, num_grid, hidden) + device = pos_embed_weight.device + outs = [] + for frames, height, width in grid_thw: + frames, height, width = int(frames), int(height), int(width) + h_pos = torch.linspace(0, num_grid - 1, height, device=device) + w_pos = torch.linspace(0, num_grid - 1, width, device=device) + h_lo, w_lo = h_pos.to(torch.int64), w_pos.to(torch.int64) + h_hi = (h_lo + 1).clamp(max=num_grid - 1) + w_hi = (w_lo + 1).clamp(max=num_grid - 1) + d_h = (h_pos - h_lo).double().unsqueeze(1) + d_w = (w_pos - w_lo).double().unsqueeze(0) + raster = ( + grid[h_lo][:, w_lo] * ((1 - d_h) * (1 - d_w))[..., None] + + grid[h_lo][:, w_hi] * ((1 - d_h) * d_w)[..., None] + + grid[h_hi][:, w_lo] * (d_h * (1 - d_w))[..., None] + + grid[h_hi][:, w_hi] * (d_h * d_w)[..., None] + ) + rows = _merge_order(raster, height, width, spatial_merge_size) + outs.append(rows.repeat(frames, 1) if frames > 1 else rows) + out = torch.cat(outs) if len(outs) > 1 else outs[0] + return out.to(pos_embed_weight.dtype) + + def is_valid( + self, + pos_embed_weight: torch.Tensor, + grid_thw: Sequence[Sequence[int]], + spatial_merge_size: int, + ) -> None: + # With 2300 rows isqrt truncates to 47, every gather index stays in + # range, and the result is a correctly shaped resample of the wrong grid. + num_rows = pos_embed_weight.shape[0] + side = math.isqrt(num_rows) + assert side * side == num_rows, ( + f"pos_embed_weight must hold a square grid; got {num_rows} rows, " + f"which is not a perfect square" + ) + + +pos_embed_interpolate = PosEmbedInterpolate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py new file mode 100644 index 000000000000..e839617d3405 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/weights.py @@ -0,0 +1,1149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weight table and loader: Qwen3.8-Flash-Next-NVFP4 / sm_103 / tp1. + +One `W` per weight role (shape, checkpoint source, dtype, load-time transform); +`Qwen4ExpTp1Weights` unrolls the table into declaration, manifest and load. tp1 +means no per-rank slicing. Storage is HF `[out, in]`; GEMM views, grouped-norm +slices, fp32 gating params and NVFP4 scalars are derived in `modeling.py` after load. + +Transforms: + +* 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); +* n-gram table: 128 row chunks concatenated into one `[rows, 160]` buffer (the hash + emits global row indices). + +The n-gram table is not in the table: its height is a checkpoint fact, so `load` +registers it with `torch.empty` and adds its rows to the manifest; +`post_load_weights` reconciles the height. Residency (device or host) is set by +`TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD` (`_host_resident_empty`); the manifest is the +same either way. Offloading frees ~47.7 GiB, which the KV cache manager turns into +more slots; several entries certify at most 461 slots, so raising +`kv_cache_fraction` with the offload on needs those cells re-checked. + +Not loaded: `mtp.*` when there is no `speculative_config`, 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 +`weight_scale_inv` cast to fp32. `lm_head` is the shell's parameter. + +`load(model, weights)` fills every declared parameter once and asserts coverage +both ways: every key is consumed or in the non-load set. +""" + +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any, Callable + +import torch +from cuda.bindings import runtime as cudart +from torch import nn + +from tensorrt_llm._utils import CUASSERT, prefer_pinned + +#: `layers=ALL_LAYERS` is the common case and the default. A sentinel rather +#: than a lambda so the table reads as data, and rather than `True` so the +#: field has one type of answer: which layers. +ALL_LAYERS = object() + + +@dataclass(frozen=True) +class W: + """One weight role: what it is, where its bytes come from, what happens on the way. + + `shape` takes the dimension bundle and returns the declared shape; every + shape in these models is derived from configuration rather than constant. + + `src` says where the bytes come from, in one of three forms: + + * a checkpoint key template, with `{p}` standing for the layer prefix; + * a sequence of `(template, slice_fn)` pairs, when one parameter is + assembled from several checkpoint tensors side by side -- `slice_fn` + takes the dimension bundle and returns an index into the destination; + * a callable `(d, i) -> [(key, index), ...]`, taking the dimension bundle + and the layer index. The first two forms are special cases of it; this + one exists for sources that are not a fixed template -- a per-expert + row, or a prefix that differs between layer families. A `key` may itself + be a tuple, meaning the transform is handed those tensors together. + + `layers` says which layers carry this weight: `ALL_LAYERS` for every one, + `None` for a weight that is not per-layer at all, or a callable + `d -> iterable[int]` when only some layers have it. + + `transform` runs on the source tensor before the copy, on the + destination's device. `None` means the checkpoint tensor is already the + declared shape. + """ + + name: str + shape: Callable[[Any], tuple[int, ...]] + src: str | tuple[tuple[str, Callable], ...] | Callable + dtype: torch.dtype | None = None + transform: Callable | None = None + layers: object | None = ALL_LAYERS + + def layer_indices(self, d: SimpleNamespace) -> list[int] | None: + """The layers this weight exists on, or None if it is not per-layer.""" + if self.layers is None: + return None + if self.layers is ALL_LAYERS: + return list(range(d.num_layers)) + return list(self.layers(d)) + + def sources(self, d: SimpleNamespace, i: int | None, prefix: str) -> list[tuple[Any, Any]]: + """`[(checkpoint key or key tuple, index into the param | None), ...]`.""" + if callable(self.src): + return list(self.src(d, i)) + if isinstance(self.src, str): + return [(self.src.format(p=prefix), None)] + return [(template.format(p=prefix), slice_fn(d)) for template, slice_fn in self.src] + + +# ----------------------------------------------------------- host residency -- +# +# Allocates the one PLE n-gram table. The table stays a `cuda` tensor of the same +# shape/dtype. Pinned, device-mapped pages (not a host gather) because the gather's +# row indices are produced on device inside the captured decode graph. + +#: `cudaMemoryType` values; the core asserts on them after load. +CUDA_MEMORY_TYPE_HOST = 1 +CUDA_MEMORY_TYPE_DEVICE = 2 + + +def memory_kind(ptr: int) -> int: + """The driver's own answer to "what kind of memory is this pointer". + + Returns a `cudaMemoryType`. This is the residency question asked of the + process rather than of `nvidia-smi`, and it is the only one a + working-but-resident implementation cannot answer the way an offloaded one + does. + """ + (attributes,) = CUASSERT(cudart.cudaPointerGetAttributes(ptr)) + return int(attributes.type) + + +def _host_resident_empty(rows: int, cols: int, dtype, device) -> tuple[torch.Tensor, torch.Tensor]: + """`[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 + caller **must** keep the second alive for as long as the first: the tensor + is a view through a device address and owns nothing. + """ + itemsize = torch.empty(0, dtype=dtype).element_size() + assert itemsize == 1, ( + f"the host-resident allocation reinterprets a `|u1` buffer, so its " + f"element must be one byte; {dtype} is {itemsize}" + ) + assert device.type == "cuda", f"host offload needs a cuda device to map into; got {device}" + + # 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)) + + # No public torch API builds a tensor over a pointer the caching allocator + # did not hand out; `__cuda_array_interface__` is the numba/cupy path. + # `|u1` because the protocol has no float8 typestr; reinterpreted below. + interface = SimpleNamespace( + __cuda_array_interface__={ + "shape": (rows, cols), + "typestr": "|u1", + "data": (int(device_pointer), False), + "version": 3, + "strides": None, + } + ) + table = torch.as_tensor(interface, device=device).view(dtype) + # `torch.as_tensor` copies when it decides the buffer is not already on the + # requested device; a copy here would be a silent 48 GiB device allocation. + assert table.data_ptr() == int(device_pointer), ( + "torch.as_tensor copied the host-mapped buffer instead of viewing it, " + "so the table would be device-resident after all" + ) + assert table.shape == (rows, cols) and table.dtype == dtype, (table.shape, table.dtype) + return table, host + + +# ------------------------------------------------------------------ the table -- + +_LM = "model.language_model" +_VISUAL = "model.visual" +_MTP = "mtp" +_MTP_LAYER = "mtp.layers.0" + +# 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 + +# The FP8 block-scale family's block edge, on both axes of the MTP experts' +# weight scales and along the activation row. +FP8_BLOCK = 128 + +# Parameter-key suffixes after which to release transform scratch. The expert +# relayouts allocate per-expert scratch 512 times per parameter and the table +# fill stages one shard at a time; without this, peak load memory is the whole +# model rather than one parameter. +_RELEASE_AFTER = ("_fc2_s", "_fc1_s", "ple_table", "mtp_fc1_w", "mtp_fc2_w") + + +def _block32_perm(rows: int, device) -> torch.Tensor: + """Gather index of the 32-row block shuffle: inside each aligned block of + 32 rows, source row `4u + v` (0<=u<8, 0<=v<4) lands at destination row + `8v + u`.""" + assert rows % 32 == 0, rows + src = torch.arange(32) + dst = (src % 4) * 8 + src // 4 + idx = torch.empty(32, dtype=torch.long) + idx[dst] = src + blocks = rows // 32 + return (idx.repeat(blocks) + torch.arange(blocks).repeat_interleave(32) * 32).to(device) + + +def _interleave_perm(rows: int, device) -> torch.Tensor: + """Gather index that interleaves the `[up | gate]` halves of a `2*I` row + stack into up0, gate0, up1, gate1, ...""" + p = torch.empty(rows, dtype=torch.long) + p[0::2] = torch.arange(0, rows // 2) + p[1::2] = torch.arange(rows // 2, rows) + return p.to(device) + + +def _fc1_perm(rows: int, device) -> torch.Tensor: + """Interleave then 32-row block shuffle, composed into one gather.""" + return _interleave_perm(rows, device)[_block32_perm(rows, device)] + + +def _swizzle(x: torch.Tensor) -> torch.Tensor: + """128x4 block-scale swizzle of an `[M, C]` uint8 scale matrix, returned + flat (`M % 128 == 0` and `C % 4 == 0` hold at this geometry, so the result + is exactly `M * C` bytes).""" + return torch.ops.trtllm.block_scale_interleave(x.contiguous()) + + +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() + + +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() + + +def _expert_fc1_s(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: + """(up, gate) e4m3 scales -> one expert's FC1 scale operand.""" + 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)) + 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).""" + 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.""" + s = t.view(torch.uint8) + rows, cols = s.shape + s = torch.index_select(s, 0, _block32_perm(rows, s.device)) + return _swizzle(s).reshape(rows, cols) + + +def _mtp_fc1_w(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: + """(up, gate) e4m3 codes -> one draft expert's FC1 operand, `[up ; gate]` on + the row axis and nothing else. The FP8 block-scale runner reads the + checkpoint's plain row-major layout.""" + up, gate = t + return torch.cat([up, gate], dim=0).contiguous() + + +def _mtp_fc1_s(t: tuple[torch.Tensor, ...], core) -> torch.Tensor: + """(up, gate) 128x128 block scales -> one draft expert's FC1 scale operand: + the same `[up ; gate]` concatenation, cast bf16 -> fp32.""" + up, gate = t + return torch.cat([up, gate], dim=0).to(torch.float32).contiguous() + + +def _mtp_fc2_s(t: torch.Tensor, core) -> torch.Tensor: + """`down_proj.weight_scale_inv` -> the FC2 scale operand: a bf16 -> fp32 + cast, no relayout.""" + return t.to(torch.float32).contiguous() + + +def _materialize(entry) -> torch.Tensor: + """Realize one checkpoint entry: `[:]` on a lazy safetensors slice, but a + 0-dim tensor (every NVFP4 per-tensor scalar on this checkpoint) rejects that + index.""" + return entry if getattr(entry, "ndim", 1) == 0 else entry[:] + + +# --- which layers carry a weight --- + + +def _attn_layers(d): + """The full-attention layers (`i % 4 == 3` on this checkpoint).""" + return [i for i, a in enumerate(d.is_attn) if a] + + +def _gdn_layers(d): + """The Gated-DeltaNet (linear-attention) layers.""" + return [i for i, a in enumerate(d.is_attn) if not a] + + +def _vision_blocks(d): + """The vision tower's blocks. Keyed `l{b}_vis_*`: block `b`, not decoder layer `b`.""" + return range(d.v_depth) + + +def _mtp_only(d): + """The MTP module, at its own layer index `num_hidden_layers`, only when it is live.""" + return [d.num_layers] if d.mtp_enabled else [] + + +# --- where the bytes come from --- + + +def _ple(suffix): + """A PLE weight. The PLE layer's index is config, so the key is computed.""" + return lambda d, i: [(f"{_LM}.layers.{d.ple_layer}.ple.{suffix}", None)] + + +def _ngram_prefix(ple_layer: int) -> str: + """The quantized n-gram table submodule: the 128 shards and the one + per-tensor scale that dequantizes them both hang off this, one level below + `ple_embedding` itself.""" + return f"{_LM}.layers.{ple_layer}.ple.ple_embedding.ngram_embedding" + + +def _visual(suffix): + """A vision-tower weight that is not per-block.""" + return lambda d, i: [(f"{_VISUAL}.{suffix}", None)] + + +def _vision_block(suffix): + """One vision block's weight.""" + return lambda d, i: [(f"{_VISUAL}.blocks.{i}.{suffix}", None)] + + +def _mtp_key(key): + """One MTP-module tensor, at its own checkpoint key (the module's prefix is not the + trunk's layer prefix).""" + return lambda d, i: [(key, None)] + + +def _experts(prefix, *suffixes): + """One source per routed expert, written at index `e`; a pair of suffixes hands the + transform both tensors together.""" + + def build(d, i): + base = prefix(i) + out = [] + for e in range(d.num_experts): + keys = tuple(f"{base}.experts.{e}.{s}" for s in suffixes) + out.append((keys if len(keys) > 1 else keys[0], (e,))) + return out + + return build + + +def _trunk_mlp(i): + return f"{_LM}.layers.{i}.mlp" + + +def _mtp_mlp(i): + return f"{_MTP_LAYER}.mlp" + + +#: The six per-tensor NVFP4 scalars every routed expert carries, as +#: (parameter suffix, checkpoint suffix). The checkpoint stores modelopt +#: reciprocals; the core folds them into the runner's scalars after load. +_EXPERT_SCALARS = ( + ("e_isc1", "gate_proj.input_scale"), + ("e_isc1_up", "up_proj.input_scale"), + ("e_ws2_1", "gate_proj.weight_scale_2"), + ("e_ws2_1_up", "up_proj.weight_scale_2"), + ("e_isc2", "down_proj.input_scale"), + ("e_ws2_2", "down_proj.weight_scale_2"), +) + + +#: (parameter tag, checkpoint module) of the two per-layer hyper-connection gates. +_HC_MODULES = (("ahc", "attn_hyper_connection"), ("mhc", "mlp_hyper_connection")) + + +def _hyper_connection(tag, src, layers): + """The four weights of one gated-residual hyper-connection; `src(name)` names the source.""" + return ( + W(f"{tag}_norm", shape=lambda d: (d.stream,), src=src("hc_norm"), layers=layers), + W( + f"{tag}_down", + shape=lambda d: (d.hc_lowrank, d.stream), + src=src("input_mix_weight_down"), + layers=layers, + ), + W( + f"{tag}_up", + shape=lambda d: (d.stream, d.hc_lowrank), + src=src("input_mix_weight_up"), + layers=layers, + ), + W( + f"{tag}_inject", + shape=lambda d: (d.hc_count, d.stream), + src=src("block_inject_weight"), + layers=layers, + ), + ) + + +def _mixer(tag, prefix, layers): + """A stack-final hyper-connection mixer: a layer's gate minus the injection weight.""" + src = _mtp_key if layers is not None else (lambda key: key) + return ( + W( + f"{tag}_norm", + shape=lambda d: (d.stream,), + src=src(f"{prefix}.hc_norm.weight"), + layers=layers, + ), + W( + f"{tag}_down", + shape=lambda d: (d.hc_lowrank, d.stream), + src=src(f"{prefix}.input_mix_weight_down.weight"), + layers=layers, + ), + W( + f"{tag}_up", + shape=lambda d: (d.stream, d.hc_lowrank), + src=src(f"{prefix}.input_mix_weight_up.weight"), + layers=layers, + ), + ) + + +def _gated_attention(tag, src, layers): + """The gated GQA block's six weights. `q_proj` is twice as wide as q: q and the output + gate interleave per head.""" + return ( + W( + f"{tag}q_w", + shape=lambda d: (2 * d.q_width, d.hidden), + src=src("q_proj.weight"), + layers=layers, + ), + W( + f"{tag}k_w", + shape=lambda d: (d.kv_width, d.hidden), + src=src("k_proj.weight"), + layers=layers, + ), + W( + f"{tag}v_w", + shape=lambda d: (d.kv_width, d.hidden), + src=src("v_proj.weight"), + layers=layers, + ), + W( + f"{tag}o_w", + shape=lambda d: (d.hidden, d.q_width), + src=src("o_proj.weight"), + layers=layers, + ), + W(f"{tag}q_norm", shape=lambda d: (d.head_dim,), src=src("q_norm.weight"), layers=layers), + W(f"{tag}k_norm", shape=lambda d: (d.head_dim,), src=src("k_norm.weight"), layers=layers), + ) + + +def _shared_expert(tag, prefix, layers): + """The bf16 shared expert beside the routed ones, and its per-token gate.""" + return ( + W( + f"{tag}sh_gu", + shape=lambda d: (2 * d.shared_inter, d.hidden), + src=lambda d, i: [ + ( + ( + f"{prefix(i)}.shared_expert.gate_proj.weight", + f"{prefix(i)}.shared_expert.up_proj.weight", + ), + None, + ) + ], + transform=_cat_rows, + layers=layers, + ), + W( + f"{tag}sh_dn", + shape=lambda d: (d.hidden, d.shared_inter), + src=lambda d, i: [(f"{prefix(i)}.shared_expert.down_proj.weight", None)], + layers=layers, + ), + 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, + ), + ) + + +u8, f8, f32, i64 = torch.uint8, torch.float8_e4m3fn, torch.float32, torch.int64 + + +class Qwen4ExpTp1Weights: + """Qwen3.8-Flash-Next-NVFP4 / sm_103 / tp1. + + The table, and the passes over it: `declare` allocates every weight it names, + `manifest` maps checkpoint keys onto them, `load` fills them. `declare` and + `manifest` share one "which parameters exist" decision (`_targets`), so they + cannot disagree about a weight. + """ + + WEIGHTS: tuple[W, ...] = ( + # --- the two per-layer hyper-connection gates: every trunk layer --- + *( + w + for tag, module in _HC_MODULES + for w in _hyper_connection( + tag, lambda name, m=module: f"{{p}}.{m}.{name}.weight", ALL_LAYERS + ) + ), + # --- the gated GQA layers and their QSA index branch --- + *_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, + ), + # --- 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", + layers=_gdn_layers, + ), + W( + "A_log", + shape=lambda d: (d.num_v_heads,), + src="{p}.linear_attn.A_log", + layers=_gdn_layers, + ), + W( + "dt_bias", + shape=lambda d: (d.num_v_heads,), + src="{p}.linear_attn.dt_bias", + layers=_gdn_layers, + ), + # Stored [conv_dim, 1, width]; the conv entries want [conv_dim, width], which is + # a zero-copy view the targets take once, after load. + W( + "conv_w", + shape=lambda d: (d.conv_dim, 1, d.conv_kernel), + src="{p}.linear_attn.conv1d.weight", + layers=_gdn_layers, + ), + W( + "gnorm", + shape=lambda d: (d.head_v_dim,), + src="{p}.linear_attn.norm.weight", + layers=_gdn_layers, + ), + W( + "out_w", + shape=lambda d: (d.hidden, d.value_dim), + src="{p}.linear_attn.out_proj.weight", + layers=_gdn_layers, + ), + # --- MoE on every layer: NVFP4 routed experts, a bf16 router and shared expert --- + W("router", shape=lambda d: (d.num_experts, d.hidden), src="{p}.mlp.gate.weight"), + W( + "fc1_w", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // 2), + dtype=u8, + src=_experts(_trunk_mlp, "up_proj.weight", "gate_proj.weight"), + transform=_expert_fc1_w, + ), + W( + "fc1_s", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // SF_VEC), + dtype=u8, + src=_experts(_trunk_mlp, "up_proj.weight_scale", "gate_proj.weight_scale"), + transform=_expert_fc1_s, + ), + W( + "fc2_w", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // 2), + dtype=u8, + src=_experts(_trunk_mlp, "down_proj.weight"), + transform=_expert_fc2_w, + ), + W( + "fc2_s", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // SF_VEC), + dtype=u8, + src=_experts(_trunk_mlp, "down_proj.weight_scale"), + transform=_expert_fc2_s, + ), + *( + W(name, shape=lambda d: (d.num_experts,), dtype=f32, src=_experts(_trunk_mlp, suffix)) + for name, suffix in _EXPERT_SCALARS + ), + *_shared_expert("", _trunk_mlp, ALL_LAYERS), + # --- not per-layer: the stack-final mixer, which is also the final norm --- + *_mixer("fhc", f"{_LM}.hyper_connection_mixer", None), + W( + "embed", + shape=lambda d: (d.vocab, d.hidden), + src=f"{_LM}.embed_tokens.weight", + layers=None, + ), + # --- PLE, on its one decoder layer (the n-gram table itself is declared at load) --- + W( + "ple_key_w", + shape=lambda d: (d.stream, d.hidden), + src=_ple("key_proj.weight"), + layers=None, + ), + W( + "ple_val_w", + shape=lambda d: (d.ple_embed_dim, d.hidden), + src=_ple("value_proj.weight"), + layers=None, + ), + W("ple_nk", shape=lambda d: (d.stream,), src=_ple("norm_key.weight"), layers=None), + W("ple_nq", shape=lambda d: (d.stream,), src=_ple("norm_query.weight"), layers=None), + W("ple_nc", shape=lambda d: (d.stream,), src=_ple("norm_conv.weight"), layers=None), + W( + "ple_conv_w", + shape=lambda d: (d.stream, 1, d.ple_conv_kernel), + src=_ple("conv1d.weight"), + layers=None, + ), + W( + "ple_mult", + shape=lambda d: (d.ngram_size,), + dtype=i64, + src=_ple("ple_embedding.layer_multipliers"), + layers=None, + ), + W( + "ple_vocab", + shape=lambda d: (d.ngram_heads,), + dtype=i64, + src=_ple("ple_embedding.ngram_heads_vocab_sizes"), + layers=None, + ), + W( + "ple_offs", + shape=lambda d: (d.ngram_heads,), + dtype=i64, + src=_ple("ple_embedding.ngram_heads_offsets"), + layers=None, + ), + # The table's per-tensor scale sits with the **table**, one level below the three + # metadata tensors above: the shards and this scale are the quantized + # `ngram_embedding` submodule, while the multipliers, offsets and vocabulary sizes + # describe the id space and stay on `ple_embedding` itself. The two groups also + # live in different safetensors files, which is the other way to notice the split. + W( + "ple_tscale", + shape=lambda d: (1,), + src=lambda d, i: [(f"{_ngram_prefix(d.ple_layer)}.weight_scale", None)], + layers=None, + ), + # --- the vision tower: all bf16, one-to-one, no transform --- + W( + "vis_patch_w", + shape=lambda d: (d.v_hidden, d.v_channels, d.v_tpatch, d.v_patch, d.v_patch), + src=_visual("patch_embed.proj.weight"), + layers=None, + ), + W( + "vis_patch_b", + shape=lambda d: (d.v_hidden,), + src=_visual("patch_embed.proj.bias"), + layers=None, + ), + W( + "vis_pos", + shape=lambda d: (d.v_pos_rows, d.v_hidden), + src=_visual("pos_embed.weight"), + layers=None, + ), + # LayerNorm **with a bias**, unlike every norm in the language tower. + W( + "vis_n1_w", + shape=lambda d: (d.v_hidden,), + src=_vision_block("norm1.weight"), + layers=_vision_blocks, + ), + W( + "vis_n1_b", + shape=lambda d: (d.v_hidden,), + src=_vision_block("norm1.bias"), + layers=_vision_blocks, + ), + # One fused `[q | k | v]` projection, MHA, with a bias. + W( + "vis_qkv_w", + shape=lambda d: (3 * d.v_hidden, d.v_hidden), + src=_vision_block("attn.qkv.weight"), + layers=_vision_blocks, + ), + W( + "vis_qkv_b", + shape=lambda d: (3 * d.v_hidden,), + src=_vision_block("attn.qkv.bias"), + layers=_vision_blocks, + ), + W( + "vis_proj_w", + shape=lambda d: (d.v_hidden, d.v_hidden), + src=_vision_block("attn.proj.weight"), + layers=_vision_blocks, + ), + W( + "vis_proj_b", + shape=lambda d: (d.v_hidden,), + src=_vision_block("attn.proj.bias"), + layers=_vision_blocks, + ), + W( + "vis_n2_w", + shape=lambda d: (d.v_hidden,), + src=_vision_block("norm2.weight"), + layers=_vision_blocks, + ), + W( + "vis_n2_b", + shape=lambda d: (d.v_hidden,), + src=_vision_block("norm2.bias"), + layers=_vision_blocks, + ), + # A **non-gated** MLP: nothing here is split in half, unlike every SwiGLU in the + # language tower. + W( + "vis_fc1_w", + shape=lambda d: (d.v_inter, d.v_hidden), + src=_vision_block("mlp.linear_fc1.weight"), + layers=_vision_blocks, + ), + W( + "vis_fc1_b", + shape=lambda d: (d.v_inter,), + src=_vision_block("mlp.linear_fc1.bias"), + layers=_vision_blocks, + ), + W( + "vis_fc2_w", + shape=lambda d: (d.v_hidden, d.v_inter), + src=_vision_block("mlp.linear_fc2.weight"), + layers=_vision_blocks, + ), + W( + "vis_fc2_b", + shape=lambda d: (d.v_hidden,), + src=_vision_block("mlp.linear_fc2.bias"), + layers=_vision_blocks, + ), + # The patch merger. Its norm is applied **before** the 4-to-1 shuffle, so it is + # `hidden` wide and not `merge_dim` wide. + W( + "vis_m_n_w", + shape=lambda d: (d.v_hidden,), + src=_visual("merger.norm.weight"), + layers=None, + ), + W("vis_m_n_b", shape=lambda d: (d.v_hidden,), src=_visual("merger.norm.bias"), layers=None), + W( + "vis_m_fc1_w", + shape=lambda d: (d.v_merge_dim, d.v_merge_dim), + src=_visual("merger.linear_fc1.weight"), + layers=None, + ), + W( + "vis_m_fc1_b", + shape=lambda d: (d.v_merge_dim,), + src=_visual("merger.linear_fc1.bias"), + layers=None, + ), + W( + "vis_m_fc2_w", + shape=lambda d: (d.hidden, d.v_merge_dim), + src=_visual("merger.linear_fc2.weight"), + layers=None, + ), + W( + "vis_m_fc2_b", + shape=lambda d: (d.hidden,), + src=_visual("merger.linear_fc2.bias"), + layers=None, + ), + # --- the MTP draft module, at layer index num_hidden_layers, only when live --- + # + # Structurally one more decoder layer of this architecture with a two-branch + # front end; every shape equals a trunk full-attention layer's except the front + # end and the expert quantization (FP8 128x128 block scales, not NVFP4). + W( + "mtp_pre_e_norm", + shape=lambda d: (d.hidden,), + src=_mtp_key(f"{_MTP}.pre_fc_norm_embedding.weight"), + layers=_mtp_only, + ), + # `[hc_count * hidden]`: the module's input is the trunk's four-stream residual. + W( + "mtp_pre_h_norm", + shape=lambda d: (d.stream,), + src=_mtp_key(f"{_MTP}.pre_fc_norm_hidden.weight"), + layers=_mtp_only, + ), + W( + "mtp_fc_e", + shape=lambda d: (d.hidden, d.hidden), + src=_mtp_key(f"{_MTP}.fc_embedding.weight"), + layers=_mtp_only, + ), + W( + "mtp_fc_h", + shape=lambda d: (d.hidden, d.hidden), + src=_mtp_key(f"{_MTP}.fc_hidden.weight"), + layers=_mtp_only, + ), + *( + w + for tag, module in _HC_MODULES + for w in _hyper_connection( + f"mtp_{tag}", + lambda name, m=module: _mtp_key(f"{_MTP_LAYER}.{m}.{name}.weight"), + _mtp_only, + ) + ), + *_gated_attention("mtp_", lambda s: _mtp_key(f"{_MTP_LAYER}.self_attn.{s}"), _mtp_only), + W( + "mtp_router", + shape=lambda d: (d.num_experts, d.hidden), + src=_mtp_key(f"{_MTP_LAYER}.mlp.gate.weight"), + layers=_mtp_only, + ), + # e4m3 codes stored row-major with one fp32 scale per 128x128 weight block, the + # FC1 halves concatenated `[up ; gate]` on the row axis. + W( + "mtp_fc1_w", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden), + dtype=f8, + src=_experts(_mtp_mlp, "up_proj.weight", "gate_proj.weight"), + transform=_mtp_fc1_w, + layers=_mtp_only, + ), + W( + "mtp_fc1_s", + shape=lambda d: ( + d.num_experts, + (2 * d.moe_inter) // FP8_BLOCK, + d.hidden // FP8_BLOCK, + ), + dtype=f32, + src=_experts(_mtp_mlp, "up_proj.weight_scale_inv", "gate_proj.weight_scale_inv"), + transform=_mtp_fc1_s, + layers=_mtp_only, + ), + # FC2 needs no relayout at all: the checkpoint's `[hidden, I]` e4m3 block is + # exactly the operand. + W( + "mtp_fc2_w", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter), + dtype=f8, + src=_experts(_mtp_mlp, "down_proj.weight"), + layers=_mtp_only, + ), + W( + "mtp_fc2_s", + shape=lambda d: (d.num_experts, d.hidden // FP8_BLOCK, d.moe_inter // FP8_BLOCK), + dtype=f32, + src=_experts(_mtp_mlp, "down_proj.weight_scale_inv"), + transform=_mtp_fc2_s, + layers=_mtp_only, + ), + *_shared_expert("mtp_", _mtp_mlp, _mtp_only), + # The module's own stack-final mixer, which is its `shared_head`. + *_mixer("mtp_fhc", f"{_MTP}.hyper_connection_mixer", _mtp_only), + ) + + def dims(self, core) -> SimpleNamespace: + """Shapes are read off the normalized text config and the composite's vision + config, which the core holds as `vision_config` before it declares anything.""" + cfg = core.model_config.pretrained_config + vcfg = core.vision_config + hidden = cfg.hidden_size + index_heads, index_kv_heads = cfg.indexer_n_heads, cfg.indexer_kv_heads + num_k_heads, num_v_heads = cfg.linear_num_key_heads, cfg.linear_num_value_heads + key_dim = num_k_heads * cfg.linear_key_head_dim + value_dim = num_v_heads * cfg.linear_value_head_dim + ngram_heads = cfg.heads_per_ngram * (cfg.ngram_size - 1) + v_hidden = vcfg.hidden_size + return SimpleNamespace( + num_layers=cfg.num_hidden_layers, + is_attn=[t == "full_attention" for t in cfg.layer_types], + hidden=hidden, + hc_count=cfg.hc_count, + hc_lowrank=cfg.hc_lowrank, + stream=cfg.hc_count * hidden, + head_dim=cfg.head_dim, + q_width=cfg.num_attention_heads * cfg.head_dim, + kv_width=cfg.num_key_value_heads * cfg.head_dim, + index_head_dim=cfg.indexer_head_dim, + index_width=(index_heads + index_kv_heads) * cfg.indexer_head_dim, + 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, + conv_kernel=cfg.linear_conv_kernel_dim, + num_experts=cfg.num_experts, + moe_inter=cfg.moe_intermediate_size, + shared_inter=cfg.shared_expert_intermediate_size, + vocab=cfg.vocab_size, + ple_layer=list(cfg.ple_layer_ids)[0] - 1, # one-based in the config + ple_embed_dim=cfg.ple_embed_dim, + ple_conv_kernel=cfg.ple_conv_kernel_size, + ngram_size=cfg.ngram_size, + ngram_heads=ngram_heads, + v_depth=vcfg.depth, + v_hidden=v_hidden, + v_inter=vcfg.intermediate_size, + v_channels=vcfg.in_channels, + v_tpatch=vcfg.temporal_patch_size, + v_patch=vcfg.patch_size, + v_pos_rows=vcfg.num_position_embeddings, + v_merge_dim=v_hidden * vcfg.spatial_merge_size * vcfg.spatial_merge_size, + mtp_enabled=getattr(core.model_config, "spec_config", None) is not None, + 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.""" + 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 {k for k in weights if k.startswith(f"{_MTP}.")} + + def declare(self, d: SimpleNamespace) -> nn.ParameterDict: + """Allocate every weight the table declares. + + Meta-init intercepts `torch.empty` here -- real CUDA storage arrives + when the engine materializes the registry. + """ + w = nn.ParameterDict() + for entry in self.WEIGHTS: + dtype = entry.dtype or d.dtype + for key, _, _ in self._targets(entry, d): + w[key] = nn.Parameter( + torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False + ) + return w + + def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None, str]]: + """`[(param key, layer index | None, checkpoint layer prefix), ...]`. + + One place decides which parameters an entry declares, so `declare` and + `manifest` cannot disagree about whether a weight exists. + """ + layers = entry.layer_indices(d) + if layers is None: + return [(entry.name, None, "")] + return [(f"l{i}_{entry.name}", i, f"{_LM}.layers.{i}") for i in layers] + + def manifest(self, core) -> dict[str, list[tuple[Any, Any, Callable | None]]]: + """param key -> [(checkpoint key, index into the param | None, transform | None)].""" + d = self.dims(core) + rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {} + for entry in self.WEIGHTS: + for key, i, prefix in self._targets(entry, d): + rows[key] = [ + (ckpt, index, entry.transform) for ckpt, index in entry.sources(d, i, prefix) + ] + return rows + + def load(self, model, weights) -> None: + """Fill every declared parameter. The n-gram table is declared here first, + multi-key rows feed their transform several tensors at once, the per-tensor + scalars are 0-dim and reject `[:]`, and coverage is asserted in both + directions.""" + core = model.model + # The one load-time declaration, before the manifest is built: the table's + # height comes from the shards, and the manifest addresses it by row range. + shard_keys = _declare_ngram_table(core, weights) + manifest = self.manifest(core) + manifest["ple_table"] = [ + (key, (slice(s * core.ple_rows_per_shard, (s + 1) * core.ple_rows_per_shard),), None) + for s, key in enumerate(shard_keys) + ] + consumed: set = set() + + def fill(param: nn.Parameter, ckpt_key, index, transform) -> None: + keys = ckpt_key if isinstance(ckpt_key, tuple) else (ckpt_key,) + for key in keys: + assert key in weights, f"checkpoint key missing: {key}" + dst = param.data if index is None else param.data[index] + # Materialize where the destination lives: the expert relayouts are row + # gathers and the scale swizzle is a kernel, both on the device. + src_all = tuple( + _materialize(weights[key]).to(dst.device, non_blocking=True) for key in keys + ) + if transform is not None: + src = transform(src_all if isinstance(ckpt_key, tuple) else src_all[0], core) + else: + assert len(src_all) == 1, "a multi-key row needs a transform" + src = src_all[0] + assert dst.shape == src.shape, (ckpt_key, tuple(dst.shape), tuple(src.shape)) + assert src.dtype == dst.dtype, (ckpt_key, src.dtype, dst.dtype) + dst.copy_(src, non_blocking=True) + consumed.update(keys) + + assert set(manifest.keys()) == set(core.w.keys()), ( + "manifest/parameter drift", + set(manifest.keys()) ^ set(core.w.keys()), + ) + for param_key, sources in manifest.items(): + for ckpt_key, index, transform in sources: + fill(core.w[param_key], ckpt_key, index, transform) + if param_key.endswith(_RELEASE_AFTER): + torch.cuda.empty_cache() + + # Shell-registered exception: the base class owns lm_head (untied), and this + # checkpoint leaves it bf16, so it is a single dense weight. Asserted rather than + # adapted — a shape change is a different logits path. + assert tuple(model.lm_head.weight.shape) == (core.vocab, core.hidden), ( + f"the shell built lm_head {tuple(model.lm_head.weight.shape)} instead " + f"of {(core.vocab, core.hidden)}" + ) + fill(model.lm_head.weight, "lm_head.weight", None, None) + # Parameter-side coverage, the other direction: name every registered parameter + # outside the target's own ParameterDict instead of assuming lm_head is the only + # one. + shell = {n for n, _ in model.named_parameters() if not n.startswith("model.w.")} + assert shell == {"lm_head.weight"}, f"unfed shell parameters: {sorted(shell)}" + + torch.cuda.synchronize() + leftover = set(weights.keys()) - consumed + expected = self.expected_unconsumed(core, weights) + what = ( + "the draft layer's three QSA indexer tensors" if core.mtp_enabled else "the MTP module" + ) + assert leftover == expected, ( + f"checkpoint coverage: leftover keys are not exactly {what}; " + f"unexpected {sorted(leftover - expected)[:8]}, " + f"missed {sorted(expected - leftover)[:8]}" + ) + + +MODEL_WEIGHTS = Qwen4ExpTp1Weights() + + +def _declare_ngram_table(core, weights) -> list[str]: + """Register the n-gram table with the height its shards concatenate to, and return + the shard keys in concatenation order. + + The shards are **contiguous row chunks**, not an interleaving: global row `r` lives in + shard `r // rows_per_shard` at local row `r % rows_per_shard`. That is what makes the + flat buffer indexable by the hash's global rows. + + **This is the whole of PLE host offload.** `core.ple_host_offload`, resolved at + construction from `TRTLLM_MODELING_V2_PLE_HOST_OFFLOAD`, picks between an ordinary + device allocation and host memory the device addresses. Nothing else changes: same + shape, same dtype, same `ParameterDict` key, same rows filling it shard by shard, same + gather in the forward. An offloaded tensor is **loaded, not skipped** — it stays + fully covered in both directions. What the manifest cannot record is *where* it + landed, so the core asserts that separately against the driver. + """ + prefix = _ngram_prefix(core.ple_layer) + keys = [f"{prefix}.shard_{s}.weight" for s in range(core.ple_shards)] + for key in keys: + assert key in weights, f"checkpoint key missing: {key}" + shapes = {tuple(weights[key].shape) for key in keys} + assert len(shapes) == 1, ( + f"the n-gram table's {len(keys)} shards do not share one shape: " + f"{sorted(shapes)}; they are concatenated as equal contiguous row " + f"chunks and an uneven one would shift every row above it" + ) + rows_per_shard, head_dim = shapes.pop() + assert head_dim == core.ple_head_dim, ( + f"the n-gram table is {head_dim} wide but {core.ngram_heads} heads must " + f"tile ple_embed_dim {core.ple_embed_dim}, i.e. {core.ple_head_dim}" + ) + dtype = weights[keys[0]].dtype + assert dtype == core.ple_table_dtype, ( + f"the n-gram table is stored {dtype}; this target declares " + f"{core.ple_table_dtype} and does not convert it" + ) + device = core.w["embed"].device + rows = core.ple_shards * rows_per_shard + if core.ple_host_offload: + table, core._ple_table_host = _host_resident_empty(rows, head_dim, dtype, device) + else: + table = torch.empty(rows, head_dim, dtype=dtype, device=device) + core._ple_table_host = None + core.w["ple_table"] = nn.Parameter(table, requires_grad=False) + # `nn.Parameter` is `_make_subclass`, which shares storage rather than copying — but + # "rather than copying" is the entire feature here, so it is asserted instead of + # assumed. + assert core.w["ple_table"].data_ptr() == table.data_ptr(), ( + "wrapping the n-gram table in nn.Parameter moved its storage" + ) + core.ple_rows_per_shard = rows_per_shard + return keys diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py new file mode 100644 index 000000000000..b915c58b4e92 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/routing.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Routing for Qwen4ExpForConditionalGeneration. + +`route` compares the facts that fix which code runs against the one deployment the +target was written for, and raises `NoTarget` stating them otherwise (`auto` falls +back to the built-in implementation; `require` raises). The checkpoint's published +architecture name is kept by `load_pretrained_config`, so `_router_index` keys on it. +All criteria are read off `text_config` / `vision_config`; the composite parent's +mirrored fields are unreliable. +""" + +from __future__ import annotations + +from typing import Any, Optional, Tuple + +from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig + +from ..._router_index import NoTarget + +# (layers, hidden, experts, attention heads, linear value heads, vocab). Shape-sniffing +# also matches fine-tunes; the accuracy gate pins identity. +_TEXT_SHAPE = (48, 2560, 512, 24, 48, 248320) + +# (depth, hidden, heads, merge, out_hidden, num_pos, deepstack_indexes, act). +# `deepstack_visual_indexes == ()` is the key member: a deepstack sibling has a merger +# list this target lacks. +_VISION_SHAPE = (27, 1152, 16, 2, 2560, 2304, (), "gelu_pytorch_tanh") + + +def route( + pretrained_config: Any, + mapping: Any, + compute_capability: Tuple[int, int], + spec_config: Optional[Any], +) -> type: + """Return the target class for this deployment, or raise ``NoTarget`` saying why none fits. + + The target is imported where it is returned, so declining costs no import + of modeling code. + """ + text = getattr(pretrained_config, "text_config", None) + vision = getattr(pretrained_config, "vision_config", None) + # Absent means this is not the normalized Qwen4-Exp composite at all, and + # nothing below would be readable. + if text is None or vision is None: + raise NoTarget( + f"{type(pretrained_config).__name__} carries no text_config/vision_config; " + "the qwen4_exp target reads both halves of the normalized composite" + ) + + shape = ( + text.num_hidden_layers, + text.hidden_size, + text.num_experts, + text.num_attention_heads, + text.linear_num_value_heads, + text.vocab_size, + ) + vshape = ( + vision.depth, + vision.hidden_size, + vision.num_heads, + vision.spatial_merge_size, + vision.out_hidden_size, + vision.num_position_embeddings, + tuple(getattr(vision, "deepstack_visual_indexes", None) or ()), + vision.hidden_act, + ) + + # 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. + 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 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 + ) + + # Qwen3.8-Flash-Next-NVFP4 on one GB300 (sm103), MTP-Eagle or no speculation. + if ( + compute_capability == (10, 3) + and shape == _TEXT_SHAPE + and vshape == _VISION_SHAPE + and mapping.world_size == 1 + and mtp_eagle_or_none + ): + # One line on purpose: the claims suite reads the routed targets off + # `from ..modeling import ` lines. + # fmt: off + # isort: off + from .qwen3_8_flash_next_nvfp4__sm_103__tp1.modeling import ModelingV2Qwen38FlashNextNvfp4Sm103Tp1 + # isort: on + # fmt: on + + return ModelingV2Qwen38FlashNextNvfp4Sm103Tp1 + + raise NoTarget( + f"compute capability {compute_capability}, text shape {shape}, vision shape " + f"{vshape}, {mapping.world_size} ranks, speculative decoding " + f"{spec_config.spec_dec_mode.name if spec_config is not None else 'off'}" + f"{f' max_draft_len={spec_config.max_draft_len}' if spec_config is not None else ''}; " + "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" + ) diff --git a/tests/integration/defs/.test_durations b/tests/integration/defs/.test_durations index fd942361b96a..aa549002e5ac 100644 --- a/tests/integration/defs/.test_durations +++ b/tests/integration/defs/.test_durations @@ -478,10 +478,10 @@ "accuracy/test_llm_api_pytorch_multimodal.py::TestNanoV3Omni::test_auto_dtype[nvfp4]": 484.5171071428572, "accuracy/test_llm_api_pytorch_multimodal.py::TestQwen3VL::test_auto_dtype[forced_chunked_prefill]": 164.90103846153846, "accuracy/test_llm_api_pytorch_multimodal.py::TestQwen3_8_Flash_Next_VL::test_nvfp4_1gpu_mtp3_trtllm_ple_offload": 760.0499642857143, - "accuracy/test_modeling_v2_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_gsm8k_identity_vs_mtp3": 417.92940000000004, - "accuracy/test_modeling_v2_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_mtp3_acceptance[modeling_v2]": 292.89186666666666, - "accuracy/test_modeling_v2_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_mtp3_acceptance[stock]": 600.2221333333333, - "accuracy/test_modeling_v2_gpt_oss.py::TestModelingV2GptOss120bSm103Tp1::test_gsm8k": 210.3445, + "accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_gsm8k_identity_vs_mtp3": 417.92940000000004, + "accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_mtp3_acceptance[modeling_v2]": 292.89186666666666, + "accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_mtp3_acceptance[stock]": 600.2221333333333, + "accuracy/modeling_v2/test_gpt_oss.py::TestModelingV2GptOss120bSm103Tp1::test_gsm8k": 210.3445, "cpp/test_multi_gpu.py::test_cache_transceiver[2proc-ucx_kvcache-90]": 856.4715789473685, "cpp/test_multi_gpu.py::test_fused_gemm_allreduce[4proc-90]": 21.290479166666668, "cpp/test_multi_gpu.py::test_mpi_utils[90]": 897.56684375, diff --git a/tests/integration/defs/accuracy/modeling_v2/__init__.py b/tests/integration/defs/accuracy/modeling_v2/__init__.py new file mode 100644 index 000000000000..3159bfe65645 --- /dev/null +++ b/tests/integration/defs/accuracy/modeling_v2/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v3.py b/tests/integration/defs/accuracy/modeling_v2/test_deepseek_v3.py similarity index 99% rename from tests/integration/defs/accuracy/test_modeling_v2_deepseek_v3.py rename to tests/integration/defs/accuracy/modeling_v2/test_deepseek_v3.py index a2ef52bc341c..838ab46151b8 100644 --- a/tests/integration/defs/accuracy/test_modeling_v2_deepseek_v3.py +++ b/tests/integration/defs/accuracy/modeling_v2/test_deepseek_v3.py @@ -38,8 +38,8 @@ from tensorrt_llm._utils import get_sm_version from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, MTPDecodingConfig -from ..conftest import llm_models_root -from .accuracy_core import ( +from ...conftest import llm_models_root +from ..accuracy_core import ( GSM8K, LlmapiAccuracyTestHarness, assert_acceptance_length, diff --git a/tests/integration/defs/accuracy/test_modeling_v2_gpt_oss.py b/tests/integration/defs/accuracy/modeling_v2/test_gpt_oss.py similarity index 97% rename from tests/integration/defs/accuracy/test_modeling_v2_gpt_oss.py rename to tests/integration/defs/accuracy/modeling_v2/test_gpt_oss.py index 2b033d506a9c..80e0b1d02007 100644 --- a/tests/integration/defs/accuracy/test_modeling_v2_gpt_oss.py +++ b/tests/integration/defs/accuracy/modeling_v2/test_gpt_oss.py @@ -33,8 +33,8 @@ from tensorrt_llm._torch._experimental.modeling_v2 import MODELING_V2_VALIDATE_ENV from tensorrt_llm._utils import get_sm_version -from ..conftest import llm_models_root -from .accuracy_core import GSM8K, LlmapiAccuracyTestHarness +from ...conftest import llm_models_root +from ..accuracy_core import GSM8K, LlmapiAccuracyTestHarness # The targets assert their own SM at construction: certification is per GPU # architecture, and a receipt from another one says nothing here. diff --git a/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py b/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py new file mode 100644 index 000000000000..f7764bd3fa60 --- /dev/null +++ b/tests/integration/defs/accuracy/modeling_v2/test_qwen4_exp.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Whole-model gates for the qwen4_exp modeling_v2 target. + +Every test passes ``modeling_v2="require"``: under ``"auto"`` a configuration +that missed the target's routing would fall back to the built-in model, pass, +and report the built-in's numbers as the target's. +""" + +import pytest + +from tensorrt_llm import LLM +from tensorrt_llm._torch._experimental.modeling_v2 import MODELING_V2_VALIDATE_ENV +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.llmapi import ( + CudaGraphConfig, + KvCacheConfig, + MoeConfig, + MTPDecodingConfig, + SamplingParams, +) + +from ...conftest import llm_models_root +from ..accuracy_core import ( + GSM8K, + MMMU, + LlmapiAccuracyTestHarness, + assert_acceptance_length, + compute_acceptance_length, +) + +# The target is certified on sm_103 only. +skip_not_sm103 = pytest.mark.skipif( + 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"} + + +def _require_target(monkeypatch: pytest.MonkeyPatch, **env: str) -> dict: + """``LLM(...)`` kwargs that build the target, with ``env`` on the ranks. + + ``LLM`` applies ``env_overrides`` to this process and never restores it, so + ``monkeypatch`` sets it too, keeping it from leaking into later tests. + """ + for key, value in env.items(): + monkeypatch.setenv(key, value) + return dict(modeling_v2="require", env_overrides=dict(env)) + + +class TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1(LlmapiAccuracyTestHarness): + """qwen3.8-flash-next-nvfp4 / sm_103 / tp1: text, MTP and multimodal.""" + + # References are keyed by model name plus the quantization the engine + # resolves for this checkpoint (MIXED_PRECISION). + MODEL_NAME = "Qwen/Qwen3.8-Flash-Next" + MODEL_PATH = f"{llm_models_root()}/Qwen3.8-Flash-Next-NVFP4" + + # The gsm8k anchors were measured under this protocol. `enable_thinking: + # false` is required: with thinking on, the answer does not finish inside + # GSM8K_MAX_OUTPUT_LEN. + EXTRA_EVALUATOR_KWARGS = dict( + apply_chat_template=True, + fewshot_as_multiturn=True, + chat_template_kwargs=dict(enable_thinking=False), + ) + + # The anchor's `max_tokens`; the class default (256) truncates this checkpoint's answers. + GSM8K_MAX_OUTPUT_LEN = 512 + + # `mamba_ssm_cache_dtype` is float32 because the checkpoint declares + # `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. + KV_CACHE = KvCacheConfig( + free_gpu_memory_fraction=0.6, + enable_block_reuse=True, + mamba_ssm_cache_dtype="float32", + ) + + # Required, not a tuning choice: the MTP draft layer's FP8 block-scale + # experts are not supported by the default CUTLASS MoE on sm_103. Set for + # every case so they share one engine configuration. + MOE = MoeConfig(backend="TRTLLM") + + ENGINE = dict( + max_seq_len=8192, + max_batch_size=32, + enable_chunked_prefill=True, + kv_cache_config=KV_CACHE, + cuda_graph_config=CudaGraphConfig(), + moe_config=MOE, + ) + + # One-model MTP-Eagle with the checkpoint's single MTP module. + MTP1 = MTPDecodingConfig(max_draft_len=1) + + # Acceptance-length anchor, keyed by target and variant. + ACCEPTANCE_KEY = "ModelingV2Qwen38FlashNextNvfp4Sm103Tp1::mtp1" + + @skip_not_sm103 + def test_gsm8k(self, mocker, monkeypatch): + """The text model without speculation, against the gsm8k anchor.""" + mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN) + + with LLM( + self.MODEL_PATH, + **self.ENGINE, + **_require_target(monkeypatch, **_VALIDATE), + ) as llm: + task = GSM8K(self.MODEL_NAME) + 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) + + with LLM( + self.MODEL_PATH, + speculative_config=self.MTP1, + enable_iter_perf_stats=True, + **self.ENGINE, + **_require_target(monkeypatch, **_GDN_LEGACY, **_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) + + @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. + + Block reuse is off here (on in the cases above) because the vision path + requires it. + """ + engine = dict(self.ENGINE) + engine["kv_cache_config"] = KvCacheConfig( + free_gpu_memory_fraction=0.6, + enable_block_reuse=False, + mamba_ssm_cache_dtype="float32", + ) + + # 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|>", + ) + + with LLM( + self.MODEL_PATH, + speculative_config=self.MTP1, + **engine, + **_require_target(monkeypatch, **_GDN_LEGACY, **_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)), + ) diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml index dab70b6e5949..8d658362ebc8 100644 --- a/tests/integration/defs/accuracy/references/acceptance_length.yaml +++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml @@ -122,3 +122,34 @@ ModelingV2DeepseekR10528Nvfp4Sm103Dep4::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 ModelingV2DeepseekR10528Nvfp4Sm103Dep4::mtp3 entry above. +# +# 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. +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. + 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 + # benchmark, which is a flaky test rather than a gate. 1.5 is the midpoint of + # 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 + # measured, and tighten min_al then. + min_al: 1.5 diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 5ec4a2aabce6..0a29a32feca5 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -134,7 +134,13 @@ l0_b300: - 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. - unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm TIMEOUT (30) - - accuracy/test_modeling_v2_gpt_oss.py::TestModelingV2GptOss120bSm103Tp1::test_gsm8k + - accuracy/modeling_v2/test_gpt_oss.py::TestModelingV2GptOss120bSm103Tp1::test_gsm8k + # 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. + - 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 - condition: ranges: system_gpu_count: diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml index c335bf4ac973..1ae313f8fdad 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml @@ -17,9 +17,9 @@ l0_gb300_multi_gpus: tests: # One ordered sequence in one job: the cases share module-global rank state. - unittest/_torch/modeling_v2/comm - - accuracy/test_modeling_v2_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_gsm8k_identity_vs_mtp3 - - accuracy/test_modeling_v2_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_mtp3_acceptance[modeling_v2] - - accuracy/test_modeling_v2_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_mtp3_acceptance[stock] + - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_gsm8k_identity_vs_mtp3 + - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_mtp3_acceptance[modeling_v2] + - accuracy/modeling_v2/test_deepseek_v3.py::TestModelingV2DeepseekR10528Nvfp4Sm103Dep4::test_mtp3_acceptance[stock] # Covers tests/unittest/_torch/attention/. Two sub-trees moved in here from elsewhere # under tests/unittest/_torch/ and have no entry of their own on any list, so this entry # is what picks them up: diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py new file mode 100644 index 000000000000..f7a304db86bd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_fused_qk_norm_rope.py @@ -0,0 +1,288 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the fused_qk_norm_rope catalog entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( # noqa: E501 + fused_qk_norm_rope as op, +) + +__extra_import_path__ = [".."] +from _validating import ( # noqa: E402,F401 — needs the path declared above + ULP, + assert_within_ulp, + certifying, + validating, +) + +assert torch.cuda.is_available(), "fused_qk_norm_rope requires a CUDA device" + +# Spec keys that shape the test's inputs or gate rather than the op call. +_NOT_OP_ARGS = ("tokens", "max_position", "gemma_prefill_band") + + +#: The element gate of a Gemma-norm cell at prefill token counts; see `_compare`. +GEMMA_PREFILL_ATOL = 0.125 + + +def _compare(out: torch.Tensor, ref: torch.Tensor, *, element_atol: float | None = None) -> None: + """A rotation is a two-term combination, and two terms can cancel. + + `x1 * cos - x2 * sin` goes to zero wherever the two products nearly + match, and there torch's per-element rtol is meaningless: at 2048 tokens + roughly 2e-4 of the outputs land close enough that a 7.3e-4 absolute + difference reads as a relative difference of 23. The decode and 37-token + cells pass under torch's default band only because they are too small to + sample one. + + So: 2 ulp element-wise against the row scale, 1 of relative RMS -- tight, + because one rotation is the shortest chain in this catalog. Over 16 + seeds the gpt-oss cells' worst is 0.0156 absolute, one output rounding + step at magnitude 2-4: 1.3 ulp of the row scale, 0.04 RMS. + + The qwen4_exp cells sit closer, for two reasons that are not defects: + `1 + w` with `w` near 1 doubles the q/k magnitude to 8-16, where one + bf16 ulp is 0.0625, and the kernel's fast sincos loses phase as + `position * freq` grows, so positions up to 32768 rotate a large pair a + little off. Below a hundred tokens, 16 seeds per cell, that is 0.03125 + max abs and 1.5 ulp of the row scale (0.21 RMS). At 8192 tokens a few + elements per tensor reach 0.0625 -- one ulp of the rotated pair -- and + 1.98 ulp of the row scale (6 seeds per cell), which a 2-ulp gate passes + only by luck of where the row's largest element falls. A cell that + names `gemma_prefill_band` (GEMMA_PREFILL_ATOL = 0.125, one ulp above that + ceiling) gates each element at `GEMMA_PREFILL_ATOL + 1.6e-2 |ref|` instead + and keeps the RMS gate (0.13 measured). Every wrong mRoPE wiring lands + 11-19 away, ~100x outside it. + """ + if element_atol is None: + assert_within_ulp(out, ref, element_ulp=2.0, rms_ulp=1.0) + return + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + o, r = out.float(), ref.float() + torch.testing.assert_close(o, r, rtol=1.6e-2, atol=element_atol) + rel_rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt().clamp_min(1e-9)).item() + assert rel_rms <= ULP[ref.dtype], f"relative RMS {rel_rms:.3e} exceeds 1 ulp" + + +def _seed(text: str) -> int: + """A seed that is the same in every process; `hash()` of a str is not.""" + return zlib.crc32(text.encode()) + + +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. + + Rows that agree make every section wiring compute the same rotation, so a + cell that drew them would certify nothing about the split. + """ + 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]) + assert (rows[0] != rows[1]).all() and (rows[0] != rows[2]).all() + assert (rows[1] != rows[2]).all() + return rows.to(torch.int32).contiguous() + + +def _build(spec, seed): + """Materialize one cell as (qkv, op kwargs, compare kwargs), with qkv + separate because it is the buffer the op writes and the caller has to keep + a copy of. + + Norm weights are drawn near 1, which under the Gemma form puts the q/k + heads at their largest magnitude. Positions are `0..T-1` unless the spec + names a `max_position`, in which case they are uniform below it -- three + distinct rows under mRoPE, one flat row otherwise. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + tokens = spec["tokens"] + heads = spec["num_heads_q"] + spec["num_heads_k"] + spec["num_heads_v"] + width = heads * spec["head_dim"] + qkv = torch.randn(tokens, width, generator=g, device="cuda").to(torch.bfloat16) + q_w, k_w = ( + (torch.randn(spec["head_dim"], generator=g, device="cuda") * 0.1 + 1.0).to(torch.bfloat16) + for _ in range(2) + ) + max_position = spec.get("max_position") + if max_position is None: + position_ids = torch.arange(tokens, dtype=torch.int32, device="cuda") + elif spec.get("use_mrope", False): + position_ids = _distinct_rows(tokens, max_position, g) + else: + position_ids = torch.randint( + 0, max_position, (tokens,), generator=g, device="cuda", dtype=torch.int32 + ) + kwargs = {k: v for k, v in spec.items() if k not in _NOT_OP_ARGS} + kwargs.update(eps=1e-6, q_weight=q_w, k_weight=k_w, position_ids=position_ids) + compare_kwargs = {"element_atol": GEMMA_PREFILL_ATOL} if spec.get("gemma_prefill_band") else {} + return qkv, kwargs, compare_kwargs + + +def _deviation(out: torch.Tensor, ref: torch.Tensor) -> str: + """Max abs, worst element against the row scale and relative RMS, in bf16 ulp.""" + 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" + + +@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) + qkv, kwargs, compare_kwargs = _build(cell.spec, seed=_seed(cell.why)) + before = qkv.clone() + ref = op.reference(qkv, **kwargs) + out = op(qkv, **kwargs) + assert out is qkv, "the op writes in place" + print(f"\n {cell.why[:60]}: {_deviation(out, ref)}") + _compare(out, ref, **compare_kwargs) + # The v heads pass through bit-exactly; `_compare` alone would allow them + # an ulp. + v_start = (kwargs["num_heads_q"] + kwargs["num_heads_k"]) * kwargs["head_dim"] + assert torch.equal(out[:, v_start:], before[:, v_start:]), "the v heads were written" + + +def test_the_v_heads_are_left_alone() -> None: + """The op's boundary, and nothing about the output shape shows it. + + A kernel that rotated the v heads too would still return a buffer of the + right size, and every cell compares against a reference that makes the same + claim -- so this pins it against the input instead. + """ + spec = dict(op.CELLS[1].spec) + qkv, kwargs, _ = _build(spec, seed=11) + before = qkv.clone() + op(qkv, **kwargs) + v_start = (kwargs["num_heads_q"] + kwargs["num_heads_k"]) * kwargs["head_dim"] + assert torch.equal(qkv[:, v_start:], before[:, v_start:]) + assert not torch.equal(qkv[:, :v_start], before[:, :v_start]), "q and k were not touched" + + +def _qwen_spec(num_heads_q: int, tokens: int) -> dict: + """The certified qwen4_exp mRoPE cell at this geometry, at `tokens` rows.""" + matches = [ + c.spec + for c in op.CELLS + if c.spec["num_heads_q"] == num_heads_q and c.spec.get("use_mrope", False) + ] + assert matches, f"no certified mRoPE cell with {num_heads_q} query heads" + spec = dict(matches[0]) + spec.pop("gemma_prefill_band", None) + spec["tokens"] = tokens + return spec + + +# The two geometries qwen4_exp runs the mRoPE form at: the main attention rope +# (24 query heads over 256) and the QSA index plane's query rope (4 over 128). +_QWEN_GEOMETRIES = pytest.mark.parametrize("num_heads_q", [24, 4], ids=["main", "index"]) + + +@_QWEN_GEOMETRIES +def test_mrope_section_split_is_11_11_10(num_heads_q: int) -> None: + """The [11, 11, 10] split, measured against the wrong wirings as well as the right one. + + At rotary_dim 64 there are 32 frequency pairs and the checkpoint's + mrope_section sums to exactly 32: right only if 11 pairs land on row 0, 11 + on row 1 and 10 on row 2. Every wrong wiring is invisible whenever the + three position rows agree -- a text batch, a single image -- so this + compares against them explicitly. The rule is over frequency pairs, so it + is independent of head_dim and the head counts; running it at both + geometries makes that a measurement rather than an argument. + """ + qkv0, kwargs, _ = _build(_qwen_spec(num_heads_q, 19), seed=10 + num_heads_q) + pos = kwargs["position_ids"] + + def call(**over): + qkv = qkv0.clone() + with validating(op): + return op(qkv, **{**kwargs, **over}) + + def reference(**over): + return op.reference(qkv0, **{**kwargs, **over}) + + got = call() + _compare(got, reference()) + + # The correct wiring sits at ~0.03 max abs; measured, every wrong one lands + # 11-19 away (~350-600x), so an absolute 1.0 is ~30x above the right + # answer and an order of magnitude below every wrong one. + for what, wrong in ( + ( + "row 0 for every pair (mRoPE inert)", + reference(position_ids=pos[0].contiguous(), use_mrope=False), + ), + ("rows 1 and 2 exchanged", reference(position_ids=pos[[0, 2, 1]].contiguous())), + ): + far = (got.float() - wrong.float()).abs().max().item() + print(f"\n {what}: max abs {far:.6f}") + assert far > 1.0, f"{what} is indistinguishable from the correct wiring ({far})" + with pytest.raises(AssertionError): + _compare(got, wrong) + + # mrope_section1 is load-bearing here: 3*11 = 33 keeps pair 31 on row 1, + # 3*10 = 30 pushes it back to row 0. + assert not torch.equal(call(mrope_section1=10), got) + # mrope_section2 saturates instead: the largest j with j % 3 == 2 below 32 + # is 29, so every mrope_section2 >= 10 selects the same 10 pairs. + assert torch.equal(call(mrope_section2=10), got) + assert torch.equal(call(mrope_section2=12), got) + + +@_QWEN_GEOMETRIES +def test_mrope_reduces_to_plain_rope_when_rows_agree(num_heads_q: int) -> None: + """Text tokens put the same position in all three rows; that must be the plain path. + + Bitwise, which is what lets the index plane run row 0 alone on the plain + form and lets a text batch stay where its accuracy gates measured it. + """ + qkv0, kwargs, _ = _build(_qwen_spec(num_heads_q, 5), seed=16 + num_heads_q) + flat = kwargs["position_ids"][0].contiguous() + + def call(**over): + qkv = qkv0.clone() + with validating(op): + return op(qkv, **{**kwargs, **over}) + + three_equal_rows = flat.expand(3, flat.shape[0]).contiguous() + assert torch.equal( + call(position_ids=three_equal_rows), + call(position_ids=flat, use_mrope=False, mrope_section1=0, mrope_section2=0), + ) + + +def _ok(**over): + spec = dict(op.CELLS[0].spec) + spec.update(over) + return _build(spec, seed=3)[:2] + + +def test_guard_refuses_a_width_that_disagrees_with_the_head_counts() -> None: + """Read as whatever the counts say; the heads come out shifted.""" + qkv, kwargs = _ok() + with pytest.raises(AssertionError, match="last dim must be"): + with validating(op): + op(qkv[:, :-64], **kwargs) + + +def test_guard_refuses_a_rotary_dim_past_the_head() -> None: + qkv, kwargs = _ok() + kwargs["rotary_dim"] = 128 + with pytest.raises(AssertionError, match="exceeds head_dim"): + with validating(op): + op(qkv, **kwargs) + + +def test_guard_refuses_a_position_id_count_that_is_not_the_token_count() -> None: + qkv, kwargs = _ok() + kwargs["position_ids"] = torch.arange(3, dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="one position per token"): + with validating(op): + op(qkv, **kwargs) 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 0ca8bfec2de9..e999aa9396cf 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 @@ -164,6 +164,32 @@ page. Wrong batch states (the flag left off, context_lengths carrying the full KV length or 0, a replaced in-window cached page) are gated far outside the tolerance band. +10. GQA 24/4/256 (Qwen3.8-27B) over an fp8-e4m3 pool at s = 1.0: which + operands each context path reads, settled bitwise (packed: bf16 q rows, + never the pool; paged: e4m3, so the flag is not output-inert over fp8), + the kv-scale semantics off 1.0 (append within one e4m3 code point, + read-side dequant against a dequant-skipping rival), the quant_mode bits a + mixed-precision checkpoint passes (384/1152/4224/4480 bit-identical to 128, + 4352 the control), and a torn-append positive control. +11. predicted_tokens_per_seq > 1 in the standard configuration: the + 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. +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 + -- the engine's two-pool hybrid arrangement at all twelve driven rows, a + 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 + 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). +The cells at the end drive the entry's own `reference` and `_compare` over the +configurations the targets run. """ import math @@ -188,9 +214,10 @@ from tensorrt_llm.mapping import Mapping __extra_import_path__ = [".."] -from _validating import ( # noqa: E402 — needs the path declared above +from _validating import ( # noqa: E402,F401 — needs the path declared above ULP, assert_within_ulp, + certifying, validating, ) @@ -274,7 +301,12 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: Scored in bf16 ulp whatever the output dtype: on the fp8 paths the operands are e4m3 and the accumulation is bf16-limited, so fp32's ulp - would measure the container. + would measure the container. The standard configuration's fp8 paths + also round the softmax probabilities to e4m3, which this band does not + absorb (5-6 ulp RMS measured); `reference(e4m3_probabilities=True)` + models it instead. Measured on sm_103: the qwen4_exp, encoder and fp8 + context cells use at most 0.25 of the element gate and 0.16 of the RMS + one; the fp8 200-token decode, whose KV range spans tiles, 0.58 of both. """ assert_within_ulp(out, ref, element_ulp=8.0, rms_ulp=4.0, ulp=ULP[torch.bfloat16]) @@ -392,6 +424,9 @@ def call_op( ctx_lens_override: Optional[List[int]] = None, max_seq_len_override: Optional[int] = None, use_paged_context_fmha: bool = False, + predicted_tokens_per_seq: int = 1, + alias_pages: bool = False, + op_overrides: Optional[Dict] = None, ) -> torch.Tensor: """One thop_attention call over explicitly constructed batch state. @@ -405,6 +440,15 @@ def call_op( count rather than the registered prompt length. use_paged_context_fmha=True selects the paged-context execution path (see the paged-context section). + predicted_tokens_per_seq > 1 makes each generation sequence contribute + that many q rows (one speculative-decoding step); seq_lens then carries + P for a generation row, which is also its new-token count, so the page + bookkeeping and the K/V mirror need no special case. + alias_pages points every entry of a sequence's offsets row at its first + physical page -- the torn-append arming sequence, used only as a control + that the pool comparison can see a racing append. + op_overrides replaces named arguments of the op call itself, for cases + that sweep an argument this helper does not otherwise expose. """ ns = len(request_ids) kv_lens = [] # cached + new, per sequence @@ -418,8 +462,9 @@ def call_op( # K page offset of page p is 2*p, V is 2*p + 1 (single-layer pool). row = self.block_offsets[0, len(kv_lens) - 1] for j, p in enumerate(self.pages[rid]): - row[0, j] = 2 * p - row[1, j] = 2 * p + 1 + page = self.pages[rid][0] if alias_pages else p + row[0, j] = 2 * page + row[1, j] = 2 * page + 1 if seq_lens_override is not None: kv_lens = list(seq_lens_override) if ctx_lens_override is not None: @@ -437,94 +482,96 @@ def call_op( dtype=qkv.dtype, device=qkv.device, ) + op_kwargs = dict( + q=qkv, + k=None, # packed QKV rides inside q + v=None, + output=output, + output_sf=None, + workspace_=self.workspace, + sequence_length=torch.tensor(kv_lens, dtype=torch.int32, device="cuda"), + host_past_key_value_lengths=torch.tensor(kv_lens, dtype=torch.int32), + host_total_kv_lens=torch.tensor([total_ctx_kv, total_gen_kv], dtype=torch.int32), + context_lengths=torch.tensor(ctx_lens, dtype=torch.int32, device="cuda"), + host_context_lengths=torch.tensor(ctx_lens, dtype=torch.int32), + host_request_types=torch.tensor(req_types, dtype=torch.int32), + max_context_q_len_override=None, + kv_cache_block_offsets=self.block_offsets, + host_kv_cache_pool_pointers=self.pool_pointers, + host_kv_cache_pool_mapping=self.pool_mapping, + cache_indirection=None, + kv_scale_orig_quant=self.kv_scale_orig_quant, + kv_scale_quant_orig=self.kv_scale_quant_orig, + out_scale=None, + rotary_inv_freq=None, + rotary_cos_sin=None, + latent_cache=None, + q_pe=None, + block_ids_per_seq=None, + attention_sinks=attention_sinks, + is_fused_qkv=True, + update_kv_cache=True, + predicted_tokens_per_seq=predicted_tokens_per_seq, + local_layer_idx=0, + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_size=self.head_dim, + tokens_per_block=self.tokens_per_block, + max_num_requests=self.max_batch, + max_context_length=self.max_seq_len, + max_seq_len=( + self.max_seq_len if max_seq_len_override is None else max_seq_len_override + ), + attention_window_size=self.attention_window_size, + beam_width=1, + mask_type=mask_type, + quant_mode=self.quant_mode, + q_scaling=1.0, + position_embedding_type=0, # no in-kernel RoPE + 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=use_paged_context_fmha, + attention_input_type=0, # mixed + 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=num_ctx_tokens, + ) + op_kwargs.update(op_overrides or {}) with validating(op): - op( - q=qkv, - k=None, # packed QKV rides inside q - v=None, - output=output, - output_sf=None, - workspace_=self.workspace, - sequence_length=torch.tensor(kv_lens, dtype=torch.int32, device="cuda"), - host_past_key_value_lengths=torch.tensor(kv_lens, dtype=torch.int32), - host_total_kv_lens=torch.tensor([total_ctx_kv, total_gen_kv], dtype=torch.int32), - context_lengths=torch.tensor(ctx_lens, dtype=torch.int32, device="cuda"), - host_context_lengths=torch.tensor(ctx_lens, dtype=torch.int32), - host_request_types=torch.tensor(req_types, dtype=torch.int32), - max_context_q_len_override=None, - kv_cache_block_offsets=self.block_offsets, - host_kv_cache_pool_pointers=self.pool_pointers, - host_kv_cache_pool_mapping=self.pool_mapping, - cache_indirection=None, - kv_scale_orig_quant=self.kv_scale_orig_quant, - kv_scale_quant_orig=self.kv_scale_quant_orig, - out_scale=None, - rotary_inv_freq=None, - rotary_cos_sin=None, - latent_cache=None, - q_pe=None, - block_ids_per_seq=None, - attention_sinks=attention_sinks, - is_fused_qkv=True, - update_kv_cache=True, - predicted_tokens_per_seq=1, - local_layer_idx=0, - num_heads=self.num_heads, - num_kv_heads=self.num_kv_heads, - head_size=self.head_dim, - tokens_per_block=self.tokens_per_block, - max_num_requests=self.max_batch, - max_context_length=self.max_seq_len, - max_seq_len=( - self.max_seq_len if max_seq_len_override is None else max_seq_len_override - ), - attention_window_size=self.attention_window_size, - beam_width=1, - mask_type=mask_type, - quant_mode=self.quant_mode, - q_scaling=1.0, - position_embedding_type=0, # no in-kernel RoPE - 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=use_paged_context_fmha, - attention_input_type=0, # mixed - 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=num_ctx_tokens, - ) + op(**op_kwargs) torch.cuda.synchronize() # Record the K/V slices just appended to the cache, per request. @@ -553,12 +600,17 @@ def reference( q_transform=None, kv_transform=None, window: Optional[int] = None, + full_block: bool = False, ) -> torch.Tensor: """fp32 SDPA over each sequence's full K/V history (must be called after call_op so the history includes this call's K/V). q_transform / kv_transform, when given, map the bf16 q / K/V history first (e.g. the e4m3 round-trip a quantized cache imposes). window, when given, - additionally drops keys older than the sliding window.""" + additionally drops keys older than the sliding window. + full_block=True is the rival model for a predicted_tokens_per_seq > 1 + call -- every q row of a sequence sees its whole cached+new range, i.e. + a draft row also sees its later siblings -- and exists only to be gated + against.""" q_w = self.num_heads * self.head_dim outs = [] start = 0 @@ -578,7 +630,7 @@ def reference( ) for i in range(sl): lo = 0 if window is None else max(0, cached + i + 1 - window) - mask[i, lo : cached + i + 1] = True + mask[i, lo : skv if full_block else cached + i + 1] = True else: # padding: every query attends to every key assert window is None, "sliding window certified for causal only" mask = None @@ -858,6 +910,7 @@ def call_op( pool_pointers: Optional[torch.Tensor] = None, block_offsets: Optional[torch.Tensor] = None, local_layer_idx: Optional[int] = None, + predicted_tokens_per_seq: int = 1, ) -> torch.Tensor: """One thop_attention call for one layer of the shared pool. @@ -881,6 +934,10 @@ def call_op( that coincide in a single-pool layout: layer_idx is this env's own layer (its K/V history and buffer view), while local_layer_idx is the row of the composed mapping. Across pools they differ. + + predicted_tokens_per_seq > 1 makes each generation sequence contribute + that many token-major q rows; seq_lens then carries P for a generation + row, which is also its new-token count. """ ns = len(request_ids) kv_lens = [self.cached_len(layer_idx, rid) + new for rid, new in zip(request_ids, seq_lens)] @@ -936,7 +993,7 @@ def call_op( attention_sinks=None, is_fused_qkv=True, update_kv_cache=True, - predicted_tokens_per_seq=1, + predicted_tokens_per_seq=predicted_tokens_per_seq, local_layer_idx=local_layer_idx, num_heads=self.num_heads, num_kv_heads=self.num_kv_heads, @@ -1063,14 +1120,16 @@ def expected_cache(self, kv: torch.Tensor) -> torch.Tensor: assert self.kv_scale_orig_quant is not None return (kv.float() * self.kv_scale_orig_quant).to(pool_dtype) - def check_caches(self, request_ids: List[int]) -> None: + def check_caches(self, request_ids: List[int], layers: Optional[List[int]] = None) -> None: """Every layer's slabs must hold exactly the expected rows, bitwise (the bf16 K/V fed to that layer, or their scaled e4m3 casts for an fp8 pool), at the pages the manager assigned (offsets are - layer-agnostic and count slabs: one page = num_layers * 2 slabs).""" + layer-agnostic and count slabs: one page = num_layers * 2 slabs). + layers restricts the check to the layers that were driven: an undriven + layer has no history to compare against.""" tpb = self.tokens_per_block slabs_per_page = self.num_layers * 2 - for layer_idx in range(self.num_layers): + for layer_idx in range(self.num_layers) if layers is None else layers: view = self.layer_views[layer_idx] for s, rid in enumerate(request_ids): k_seq = self.expected_cache(torch.cat(self.k_history[layer_idx, rid])) @@ -5308,6 +5367,1448 @@ def test_mla_rejects_null_q_lora_rank() -> None: raise AssertionError("q_lora_rank=None was accepted on the MLA path") +# ─── GQA 24/x/256 and the encoder-only mode: shared helpers ──────────── + + +def _history(env, rid: int, layer: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]: + """The K/V one request's cache holds (for one layer of a shared pool).""" + if layer is None: + return torch.cat(env.k_history[rid]), torch.cat(env.v_history[rid]) + return torch.cat(env.k_history[layer, rid]), torch.cat(env.v_history[layer, rid]) + + +def _op_reference( + env, + qkv: torch.Tensor, + seq_lens: List[int], + request_ids: List[int], + causal: bool = True, + e4m3: bool = False, + layer: Optional[int] = None, +) -> torch.Tensor: + """`op.reference` over the history the env mirrored for these requests. + + 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 -- + every draft row seeing its later siblings. e4m3=True is the standard + configuration's fp8 arithmetic: q, K, V and the softmax probabilities all + e4m3 operands. + """ + 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)) + return op.reference( + qkv[:, : hq * d], + torch.cat(ks).reshape(-1, hkv * d), + torch.cat(vs).reshape(-1, hkv * d), + list(seq_lens), + [k.shape[0] for k in ks], + num_heads=hq, + num_kv_heads=hkv, + qk_head_dim=d, + v_head_dim=d, + softmax_scale=1.0 / math.sqrt(d), + causal=causal, + e4m3_inputs=e4m3, + e4m3_probabilities=e4m3, + ) + + +def _packed_reference( + qkv: torch.Tensor, + seg_lens: List[int], + num_heads: int, + num_kv_heads: int, + head_dim: int, + causal: bool, +) -> torch.Tensor: + """`op.reference` over each segment's own packed rows and nothing else: + what the no-KV-cache mode computes, and the rival a cached-prefix call + must be separated from (its prefix ignored).""" + q_w, kv_w = num_heads * head_dim, num_kv_heads * head_dim + return op.reference( + qkv[:, :q_w], + qkv[:, q_w : q_w + kv_w], + qkv[:, q_w + kv_w :], + list(seg_lens), + list(seg_lens), + num_heads=num_heads, + num_kv_heads=num_kv_heads, + qk_head_dim=head_dim, + v_head_dim=head_dim, + softmax_scale=1.0 / math.sqrt(head_dim), + causal=causal, + ) + + +def _std_onehot_qkv( + env: _PagedAttnEnv, n_tokens: int, positions: List[int], t0: int +) -> torch.Tensor: + """Packed QKV that turns the op into an attention-weight read-out: K is all + zero, so every unmasked logit is 0 and the softmax is uniform over exactly + the attended set, and V[t] = e_{t - t0}, so output column d of a query row + IS the weight that row gives key t0 + d. 1.0 and 0.0 are both exact in + e4m3, so the probe survives the quantized pool unchanged.""" + q_w = env.num_heads * env.head_dim + kv_w = env.num_kv_heads * env.head_dim + qkv = torch.zeros(n_tokens, q_w + 2 * kv_w, dtype=torch.bfloat16, device="cuda") + qkv[:, :q_w] = torch.randn(n_tokens, q_w, dtype=torch.bfloat16, device="cuda") + v = torch.zeros(n_tokens, env.num_kv_heads, env.head_dim, dtype=torch.bfloat16, device="cuda") + for i, t in enumerate(positions): + d = t - t0 + if 0 <= d < env.head_dim: + v[i, :, d] = 1.0 + qkv[:, q_w + kv_w :] = v.reshape(n_tokens, kv_w) + return qkv + + +def _assert_draft_row_weights( + out: torch.Tensor, + env: _PagedAttnEnv, + row: int, + last_key: int, + label: str, + weight_rtol: float, +) -> None: + """Under the one-hot probe, output column d of `row` is the weight it gives + key d. A bottom-right-aligned causal block puts the last attended key at + `last_key`: every column above it must be an exact zero (a later sibling + contributing anything at all is the failure this exists to catch) and every + column at or below it carries the uniform 1 / (last_key + 1).""" + w = out.view(-1, env.num_heads, env.head_dim)[row].float() + assert bool((w[:, last_key + 1 :] == 0).all()), ( + f"{label}: keys above {last_key} carry weight {w[:, last_key + 1 :].abs().max().item():.4g}" + ) + torch.testing.assert_close( + w[:, : last_key + 1], + torch.full_like(w[:, : last_key + 1], 1.0 / (last_key + 1)), + rtol=weight_rtol, + atol=0.0, + msg=lambda m: f"{label}: attended weights are not uniform\n{m}", + ) + + +# ─── Standard configuration over an fp8-e4m3 pool at GQA 24/4/256 ────── +# +# The Qwen3.8-27B full-attention layer -- 24 query / 4 kv heads, head_size 256, +# tokens_per_block 32, causal, no sinks or window -- served with an fp8 KV +# cache at s = 1.0. Its cells gate packed and paged context, decode, +# P = 2..4 and a mixed step against `op.reference`; what is here is what a +# cell cannot carry: the operand dtype each context path reads (settled +# bitwise), the scale semantics off s = 1.0, the rivals each claim is +# separated from, and the quant_mode bits a mixed-precision checkpoint hands +# the op. + +STD_FP8_HQ, STD_FP8_HKV, STD_FP8_D = 24, 4, 256 +STD_FP8_TPB = 32 + +# How far a rival model has to sit outside a band to count as separated. +FP8_STD_MIN_SEPARATION = 3.0 + + +def _e4m3_round_trip(scale: float = 1.0): + """The round trip an fp8 KV pool imposes on whatever passes through it: + quantize by kv_scale_orig_quant (= 1/s), dequantize by kv_scale_quant_orig + (= s). Returns a callable usable as reference()'s q_transform/kv_transform.""" + + def _apply(x: torch.Tensor) -> torch.Tensor: + return (x.float() * (1.0 / scale)).to(torch.float8_e4m3fn).float() * scale + + return _apply + + +def _e4m3_code_distance(got: torch.Tensor, expected: torch.Tensor) -> Tuple[int, int]: + """(bytes that differ, largest distance in e4m3 code points). + + e4m3's byte encoding is sign | 4-bit exponent | 3-bit mantissa, so within + one sign the magnitude code is monotone in the value and adjacent codes are + exactly one ulp apart -- which makes "at most one ulp" a bit-level claim + rather than a tolerance. + """ + g = got.contiguous().view(torch.uint8).to(torch.int16) + e = expected.contiguous().view(torch.uint8).to(torch.int16) + assert torch.equal(g >> 7, e >> 7), "sign bits differ -- not an ulp comparison" + return int((g != e).sum()), int(((g & 0x7F) - (e & 0x7F)).abs().max()) + + +def _std_fp8_env(**kwargs) -> _PagedAttnEnv: + """The 24/4/256 cell over an fp8-e4m3 pool at the production kv scale.""" + kwargs.setdefault("num_heads", STD_FP8_HQ) + kwargs.setdefault("num_kv_heads", STD_FP8_HKV) + kwargs.setdefault("head_dim", STD_FP8_D) + kwargs.setdefault("tokens_per_block", STD_FP8_TPB) + kwargs.setdefault("num_blocks", 128) + kwargs.setdefault("max_blocks_per_seq", 16) + kwargs.setdefault("pool_dtype", torch.float8_e4m3fn) + kwargs.setdefault("quant_mode", QUANT_MODE_FP8_KV_CACHE) + kwargs.setdefault("kv_scaling_factor", 1.0) + return _PagedAttnEnv(**kwargs) + + +def test_fp8_std_packed_context_does_not_read_the_pool() -> None: + """The packed context path at 24/4/256 reads the bf16 q rows, not the fp8 + pool: inside `_compare` of the bf16-KV reference, while the + "context sees the quantized cache" model -- K/V round-tripped through + e4m3 -- sits 26.0x outside the bf16 band (sm_103).""" + torch.manual_seed(400) + env = _std_fp8_env() + env.add_request(0, 48) # crosses a 32-token page boundary + env.add_request(1, 17) + qkv = env.random_qkv(65) + out = env.call_op(qkv, [48, 17], 2, [0, 1], MASK_CAUSAL) + _compare(out, _op_reference(env, qkv, [48, 17], [0, 1])) + _assert_far_outside_band( + out, + env.reference(qkv, [48, 17], [0, 1], [0, 0], MASK_CAUSAL, kv_transform=_e4m3_round_trip()), + FP8_STD_MIN_SEPARATION, + "packed context FMHA at 24/4/256 read the fp8 pool", + ATOL, + RTOL, + ) + env.check_cache(0) + env.check_cache(1) + env.check_unwritten_pool_zero() + + +def test_fp8_std_kv_scale_semantics() -> None: + """The two kv_scale_* tensors at 24/4/256, swept off the production 1.0. + + The append mirror is exact arithmetic -- e4m3(k / s), fp32 multiply then + round-to-nearest cast -- so it is asserted at the bit level: bit-exact at + the power-of-2 scale, at most one e4m3 code point off at 1.5 (first + measured: 0.6% of the bytes, every one exactly one ulp). + + The read side is what needs a scale at all, and s != 1.0 is where it + becomes observable: a model that keeps the write-side quantization but + skips the read-side dequant -- the forgotten kv_scale_quant_orig -- sits + 15.1x (s = 2.0) and 10.8x (s = 1.5) outside the decode band on sm_103. At + s = 1.0 the two coincide, which is why the sweep exists. + + The correct model is gated two ways. At s = 2.0, `_compare` against + `op.reference` over e4m3(q/s)*s, the stored pool rows times s, and e4m3 + probabilities -- measured bitwise-exact on sm_103. At s = 1.5 no model of + the kernel's probability rounding tried fits that band (448, 448/s, 448*s, + 448/s**2, fp32: 4.2 ulp RMS at best), so the fp8 decode band of the + q/K/V round-trip model gates it: a non-power-of-2 scale is outside every + cell. + """ + for scale, max_inexact_bytes in ((2.0, 0), (1.5, 1000)): + torch.manual_seed(410) + env = _std_fp8_env(kv_scaling_factor=scale) + env.add_request(0, 48) + qkv = env.random_qkv(48) + env.call_op(qkv, [48], 1, [0], MASK_CAUSAL) + got_k, got_v, exp_k, exp_v = env.cache_pages_content(0) + for name, got, exp in (("K", got_k, exp_k), ("V", got_v, exp_v)): + differing, ulps = _e4m3_code_distance(got, exp) + assert differing <= max_inexact_bytes, ( + f"s={scale}: {differing} {name} bytes differ from the " + f"e4m3(k/s) mirror, allowed {max_inexact_bytes}" + ) + assert ulps <= 1, f"s={scale}: {name} append is {ulps} e4m3 ulps off the mirror" + cached = [env.cached_len(0)] + qkv_d = env.random_qkv(1) + out_d = env.call_op(qkv_d, [1], 0, [0], MASK_CAUSAL) + if scale == 2.0: + # What the decode attended: q through the same scaled round trip, + # and the pool rows exactly as stored, times s. + got_k, got_v, _, _ = env.cache_pages_content(0) + q_w = env.num_heads * env.head_dim + kv_w = env.num_kv_heads * env.head_dim + q_hat = _e4m3_round_trip(scale)(qkv_d[:, :q_w]).to(torch.bfloat16) + ref = op.reference( + q_hat, + (got_k.float() * scale).reshape(-1, kv_w), + (got_v.float() * scale).reshape(-1, kv_w), + [1], + [got_k.shape[0]], + num_heads=env.num_heads, + num_kv_heads=env.num_kv_heads, + qk_head_dim=env.head_dim, + v_head_dim=env.head_dim, + softmax_scale=1.0 / math.sqrt(env.head_dim), + e4m3_probabilities=True, + ) + _compare(out_d, ref) + else: + torch.testing.assert_close( + out_d, + env.reference( + qkv_d, + [1], + [0], + cached, + MASK_CAUSAL, + q_transform=_e4m3_round_trip(scale), + kv_transform=_e4m3_round_trip(scale), + ), + rtol=FP8_DECODE_RTOL, + atol=FP8_DECODE_ATOL, + ) + _assert_far_outside_band( + out_d, + env.reference( + qkv_d, + [1], + [0], + cached, + MASK_CAUSAL, + q_transform=_e4m3_round_trip(scale), + # write-side quantization kept, read-side dequant dropped + kv_transform=lambda x, s=scale: (x.float() / s).to(torch.float8_e4m3fn).float(), + ), + FP8_STD_MIN_SEPARATION, + f"s={scale}: decode ignored kv_scale_quant_orig", + FP8_DECODE_ATOL, + FP8_DECODE_RTOL, + ) + + +def _std_peaked_qkv(env: _PagedAttnEnv, n_tokens: int, hot: int) -> Tuple[torch.Tensor, ...]: + """Packed QKV whose softmax collapses onto one key. + + Every K row is zero except row `hot`, which is a large constant vector, and + q is that same vector -- so every query row that can see `hot` gives it a + weight of 1 to within fp32, and its output IS that row's V. The V operand's + dtype is then readable bit for bit out of `output`. + """ + q_w = env.num_heads * env.head_dim + kv_w = env.num_kv_heads * env.head_dim + qkv = torch.zeros(n_tokens, q_w + 2 * kv_w, dtype=torch.bfloat16, device="cuda") + key = torch.full((env.head_dim,), 6.0, dtype=torch.bfloat16, device="cuda") + qkv[:, :q_w] = key.repeat(env.num_heads) + k = torch.zeros(n_tokens, env.num_kv_heads, env.head_dim, dtype=torch.bfloat16, device="cuda") + k[hot] = key + qkv[:, q_w : q_w + kv_w] = k.reshape(n_tokens, kv_w) + v = torch.randn(n_tokens, env.num_kv_heads, env.head_dim, dtype=torch.bfloat16, device="cuda") + qkv[:, q_w + kv_w :] = v.reshape(n_tokens, kv_w) + return qkv, v + + +def test_fp8_std_paged_context_operands_are_e4m3() -> None: + """What the paged-context FMHA sees over an fp8 pool, settled bitwise. + + The peaked-softmax readout returns exactly e4m3(V) on the paged path and + exactly the untouched bf16 V on the packed one, in the same process, from + the same inputs; the two candidate answers are asserted more than 0.05 + apart, so the comparison is not vacuous. + + Equality is elementwise float equality rather than a byte compare: a V + element that underflows to e4m3's negative zero comes back +0.0, because + the kernel's fp32 accumulator starts at zero and 0.0 + (-0.0) is +0.0. + """ + groups = STD_FP8_HQ // STD_FP8_HKV + for use_paged, want_e4m3 in ((True, True), (False, False)): + torch.manual_seed(420) + env = _std_fp8_env() + env.add_request(0, 48) + qkv, v = _std_peaked_qkv(env, 48, hot=10) + out = env.call_op(qkv, [48], 1, [0], MASK_CAUSAL, use_paged_context_fmha=use_paged) + row = out.view(48, STD_FP8_HQ, STD_FP8_D)[47] # sees the hot key causally + for h in range(STD_FP8_HQ): + v_head = v[10, h // groups] + wanted = _e4m3_round_trip()(v_head).to(torch.bfloat16) if want_e4m3 else v_head + assert torch.equal(row[h].float(), wanted.float()), ( + f"use_paged_context_fmha={use_paged}: head {h}'s V operand is not " + f"{'e4m3' if want_e4m3 else 'bf16'} (max abs diff " + f"{(row[h].float() - wanted.float()).abs().max().item():.4g})" + ) + gap = ( + (v[10].float() - _e4m3_round_trip()(v[10]).to(torch.bfloat16).float()) + .abs() + .max() + .item() + ) + assert gap > 0.05, f"bf16 V and e4m3(V) differ by only {gap:.4g}" + + +def test_fp8_std_paged_context_flag_is_not_inert_over_an_fp8_pool() -> None: + """With nothing cached the flag is bitwise inert over a bf16 pool -- and is + not over an fp8 one. + + The pool content is bitwise identical either way (the append does not care + about the flag), but over fp8 the paged output reads K/V through e4m3: the + packed output stays inside `_compare` of the bf16-KV reference, the paged + one sits 26.2x outside the bf16 band around it (sm_103). + """ + for pool_dtype, quant_mode, scale, expect_same in ( + (torch.bfloat16, 0, None, True), + (torch.float8_e4m3fn, QUANT_MODE_FP8_KV_CACHE, 1.0, False), + ): + runs = [] + for use_paged in (False, True): + torch.manual_seed(430) + env = _std_fp8_env( + pool_dtype=pool_dtype, quant_mode=quant_mode, kv_scaling_factor=scale + ) + env.add_request(0, 48) + env.add_request(1, 17) + qkv = env.random_qkv(65) + out = env.call_op( + qkv, [48, 17], 2, [0, 1], MASK_CAUSAL, use_paged_context_fmha=use_paged + ) + runs.append((out, env.kv_cache.clone(), _op_reference(env, qkv, [48, 17], [0, 1]))) + packed, paged = runs + assert _bitwise_equal(packed[1], paged[1]), ( + f"pool_dtype={pool_dtype}: the append moved with use_paged_context_fmha" + ) + same = _bitwise_equal(packed[0], paged[0]) + assert same == expect_same, ( + f"pool_dtype={pool_dtype}: packed and paged context outputs " + f"{'differ' if not same else 'agree'}, expected the opposite" + ) + _compare(packed[0], packed[2]) + if not expect_same: + _assert_far_outside_band( + paged[0], + paged[2], + FP8_STD_MIN_SEPARATION, + "paged-context output over an fp8 pool matched the bf16-KV model", + ATOL, + RTOL, + ) + + +def test_fp8_std_paged_context_bf16_kv_rival() -> None: + """The paged-context cells over the fp8 pool are separated from the bf16-KV + model the bf16-pool paged-context surface is certified on. + + At 24/4/256 and 32/8/128, a fresh 48-row context call and a 33-row one over + that 48-token prefix are each inside `_compare` of the e4m3-operand + reference, and outside the *bf16* band around the bf16-KV one by 22.0x / + 11.9x and 22.9x / 12.0x on sm_103. The decode that follows is gated too, + and the append stays bit-exact against the e4m3 mirror. + """ + for i, (num_heads, num_kv_heads, head_dim) in enumerate( + [(STD_FP8_HQ, STD_FP8_HKV, STD_FP8_D), (32, 8, 128)] + ): + torch.manual_seed(440 + i) + env = _std_fp8_env(num_heads=num_heads, num_kv_heads=num_kv_heads, head_dim=head_dim) + env.add_request(0, 96) # served by two context calls: 48 + 33, then decode + for new_tokens in (48, 33): + cached = [env.cached_len(0)] + qkv = env.random_qkv(new_tokens) + out = _paged_ctx_call(env, qkv, [new_tokens], 1, [0]) + _compare(out, _op_reference(env, qkv, [new_tokens], [0], e4m3=True)) + _assert_far_outside_band( + out, + env.reference(qkv, [new_tokens], [0], cached, MASK_CAUSAL), + FP8_STD_MIN_SEPARATION, + f"{num_heads}/{num_kv_heads}/{head_dim} paged context over " + f"{cached[0]} cached tokens matched the bf16-KV model", + ATOL, + RTOL, + ) + env.check_cache(0) + qkv_d = env.random_qkv(1) + out_d = env.call_op(qkv_d, [1], 0, [0], MASK_CAUSAL, use_paged_context_fmha=True) + _compare(out_d, _op_reference(env, qkv_d, [1], [0], e4m3=True)) + env.check_cache(0) + env.check_unwritten_pool_zero() + + +def test_fp8_std_quant_mode_extra_bits() -> None: + """A mixed-precision checkpoint served with an fp8 KV cache does not hand + the op a bare 128, and the standard configuration reads only the KV-cache + bit. + + Bit-identical to 128 across a context prefill, a decode step, a P = 2 + generation call and the whole pool: 384 (| FP8_QDQ), 1152 + (| FP8_1x128_128x128), 4224 (NVFP4 | FP8_KV_CACHE) and 4480 (NVFP4 | + FP8_QDQ | FP8_KV_CACHE). 4352 (NVFP4 | FP8_QDQ, the same extra bits without + the KV-cache one) is the control: it must differ everywhere the pool is + read while leaving the packed context output bitwise unchanged. Bits nobody + tried stay unknown, not inert. + """ + runs: Dict[int, Tuple[torch.Tensor, ...]] = {} + for quant_mode in (128, 384, 1152, 4224, 4480, 4352): + torch.manual_seed(480) + env = _std_fp8_env() + env.quant_mode = quant_mode + env.add_request(0, 48) + env.add_request(1, 17) + out_ctx = env.call_op(env.random_qkv(65), [48, 17], 2, [0, 1], MASK_CAUSAL) + out_dec = env.call_op(env.random_qkv(2), [1, 1], 0, [0, 1], MASK_CAUSAL) + out_p2 = env.call_op( + env.random_qkv(4), [2, 2], 0, [0, 1], MASK_CAUSAL, predicted_tokens_per_seq=2 + ) + runs[quant_mode] = (out_ctx, out_dec, out_p2, env.kv_cache.clone()) + base = runs[128] + for quant_mode in (384, 1152, 4224, 4480): + for what, got, want in zip(("context", "decode", "P=2", "pool"), runs[quant_mode], base): + assert _bitwise_equal(got, want), ( + f"quant_mode {quant_mode} is not bit-identical to 128: {what}" + ) + control = runs[4352] + assert _bitwise_equal(control[0], base[0]), ( + "quant_mode 4352 changed the packed context output, which never reads the pool" + ) + for what, got, want in zip(("decode", "P=2", "pool"), control[1:], base[1:]): + assert not _bitwise_equal(got, want), ( + f"quant_mode 4352 (no FP8-KV bit) is bit-identical to 128 on {what}: " + f"the comparison cannot see a quant_mode change at all" + ) + + +def test_fp8_std_append_race_control() -> None: + """Positive control, not a certification. + + Every append assertion here rests on comparing pool bytes, so the + comparison has to be able to see this entry's documented intermittent + defect: a call whose new tokens do not map to distinct physical slots races + and leaves different bytes on identical repeats. Armed at 24/4/256 over the + fp8 pool by aliasing all three pages of a 96-token prefill onto one, four + identical seeded repeats should not all agree (a race need not show, so + agreement skips rather than fails) -- while the certified geometry + reproduces bitwise. + """ + + def repeat_pools(alias: bool) -> List[torch.Tensor]: + pools = [] + for _ in range(4): + torch.manual_seed(490) + env = _std_fp8_env() + env.add_request(0, 96) + env.call_op(env.random_qkv(96), [96], 1, [0], MASK_CAUSAL, alias_pages=alias) + pools.append(env.kv_cache.clone()) + return pools + + certified = repeat_pools(alias=False) + for i in range(1, 4): + assert _bitwise_equal(certified[i], certified[0]), ( + f"the certified 24/4/256 fp8 append is not reproducible: repeat {i}" + ) + armed = repeat_pools(alias=True) + if all(_bitwise_equal(armed[i], armed[0]) for i in range(1, 4)): + pytest.skip("the armed append happened to reproduce bitwise in all four repeats") + + +# ─── predicted_tokens_per_seq > 1 in the standard configuration ──────── +# +# One speculative-decoding step: each generation sequence contributes P q rows +# (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. + +# 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 +# tolerance is the bf16 one-hot bound over a bf16 pool; over fp8 the softmax +# probabilities are an e4m3 MMA operand and carry its 2**-4 relative error -- +# the zeros stay exactly zero, which is the part that pins the mask. +_DRAFT_POOLS = [ + ("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), +] + + +def _draft_env(num_kv_heads: int, pool_dtype, quant_mode: int, scale) -> _PagedAttnEnv: + return _std_fp8_env( + num_kv_heads=num_kv_heads, + pool_dtype=pool_dtype, + quant_mode=quant_mode, + kv_scaling_factor=scale, + ) + + +@pytest.mark.parametrize("pool", _DRAFT_POOLS, ids=[c[0] for c in _DRAFT_POOLS]) +def test_std_predicted_tokens_mask_is_bottom_right_causal(pool) -> None: + """The within-block mask, read out rather than inferred. + + Cache rows carrying a one-hot V turn the output into the attention-weight + vector. Draft row t of a sequence with L cached tokens attends keys + [0, L + t] uniformly and its later siblings' columns come back bitwise + zero, at P = 2, 3, 4. Rows are token-major within a sequence, all P tokens + are appended, and mask_type 0 and 1 are bitwise identical. + """ + _, hkv, pool_dtype, quant_mode, scale, weight_rtol = pool + for p in (2, 3, 4): + torch.manual_seed(450 + p) + env = _draft_env(hkv, pool_dtype, quant_mode, scale) + cached = 40 + env.add_request(0, cached) + env.add_request(1, cached) + prefill = torch.cat([_std_onehot_qkv(env, cached, list(range(cached)), 0)] * 2) + env.call_op(prefill, [cached, cached], 2, [0, 1], MASK_CAUSAL) + gen = torch.cat([_std_onehot_qkv(env, p, [cached + t for t in range(p)], 0)] * 2) + out = env.call_op(gen, [p, p], 0, [0, 1], MASK_CAUSAL, predicted_tokens_per_seq=p) + for n in range(2 * p): + _assert_draft_row_weights( + out, + env, + n, + cached + (n % p), + f"{pool[0]} P={p} row {n} (sequence {n // p}, draft {n % p})", + weight_rtol, + ) + assert env.cached_len(0) == cached + p and env.cached_len(1) == cached + p + env.check_cache(0) + env.check_cache(1) + env.check_unwritten_pool_zero() + # mask_type is not what produces the within-block mask. + gen2 = torch.cat([_std_onehot_qkv(env, p, [cached + p + t for t in range(p)], 0)] * 2) + out_causal = env.call_op( + gen2, [p, p], 0, [0, 1], MASK_CAUSAL, predicted_tokens_per_seq=p, record=False + ) + out_padding = env.call_op(gen2, [p, p], 0, [0, 1], MASK_PADDING, predicted_tokens_per_seq=p) + assert _bitwise_equal(out_causal, out_padding), ( + f"{pool[0]} P={p}: mask_type 0 and 1 differ in the standard configuration" + ) + + +@pytest.mark.parametrize("pool", _DRAFT_POOLS, ids=[c[0] for c in _DRAFT_POOLS]) +def test_std_predicted_tokens_full_block_rival(pool) -> None: + """Realistic inputs at P = 2, 3, 4 (and a mixed step at P = 3): inside + `_compare` of the bottom-right causal reference, while the full-block + model -- every draft row seeing its later siblings -- sits outside the + band by, at P = 2, 3, 4 and the mixed step on sm_103: 45x, 89x, 113x, 75x + (bf16 24/2/256); 61x, 59x, 118x, 104x (bf16 24/4/256); 4.8x, 6.2x, 11.3x, + 10.0x (fp8 24/4/256, in the fp8 decode band). At P = 1 the two coincide.""" + name, hkv, pool_dtype, quant_mode, scale, _ = pool + fp8 = pool_dtype == torch.float8_e4m3fn + atol, rtol = (FP8_DECODE_ATOL, FP8_DECODE_RTOL) if fp8 else (ATOL, RTOL) + for p in (2, 3, 4): + torch.manual_seed(460 + p) + env = _draft_env(hkv, pool_dtype, quant_mode, scale) + env.add_request(0, 70) + env.add_request(1, 35) + _paged_ctx_call(env, env.random_qkv(105), [70, 35], 2, [0, 1]) + gen = env.random_qkv(2 * p) + out = env.call_op( + gen, + [p, p], + 0, + [0, 1], + MASK_CAUSAL, + predicted_tokens_per_seq=p, + use_paged_context_fmha=True, + ) + _compare(out, _op_reference(env, gen, [p, p], [0, 1], e4m3=fp8)) + _assert_far_outside_band( + out, + _op_reference(env, gen, [p, p], [0, 1], causal=False, e4m3=fp8), + FP8_STD_MIN_SEPARATION, + f"{name} P={p}: a draft row saw its later siblings", + atol, + rtol, + ) + + # A context row ahead of two draft blocks of different histories. + torch.manual_seed(470) + p = 3 + env = _draft_env(hkv, pool_dtype, quant_mode, scale) + env.add_request(0, 31) + env.add_request(1, 64) + env.add_request(2, 20) + _paged_ctx_call(env, env.random_qkv(95), [31, 64], 2, [0, 1]) + mixed = env.random_qkv(20 + 2 * p) + out = _paged_ctx_call(env, mixed, [20, p, p], 1, [2, 0, 1], predicted_tokens_per_seq=p) + _compare(out, _op_reference(env, mixed, [20, p, p], [2, 0, 1], e4m3=fp8)) + _assert_far_outside_band( + out[20:], + _op_reference(env, mixed, [20, p, p], [2, 0, 1], causal=False, e4m3=fp8)[20:], + FP8_STD_MIN_SEPARATION, + f"{name} mixed step: a draft row saw its later siblings", + atol, + rtol, + ) + + +@pytest.mark.parametrize("pool", _DRAFT_POOLS, ids=[c[0] for c in _DRAFT_POOLS]) +def test_std_predicted_tokens_page_crossing_and_consecutive_steps(pool) -> None: + """A draft block that starts mid-page and crosses a page boundary (31 + cached, P = 3, pages of 32), read out one-hot; then three consecutive + P = 4 steps over a growing history, so the block lands at several offsets, + each inside `_compare`.""" + name, hkv, pool_dtype, quant_mode, scale, weight_rtol = pool + fp8 = pool_dtype == torch.float8_e4m3fn + torch.manual_seed(471) + p = 3 + env = _draft_env(hkv, pool_dtype, quant_mode, scale) + cached = 31 + env.add_request(0, cached) + env.call_op(_std_onehot_qkv(env, cached, list(range(cached)), 0), [cached], 1, [0], MASK_CAUSAL) + out = env.call_op( + _std_onehot_qkv(env, p, [cached + t for t in range(p)], 0), + [p], + 0, + [0], + MASK_CAUSAL, + predicted_tokens_per_seq=p, + ) + for n in range(p): + _assert_draft_row_weights( + out, env, n, cached + n, f"{name} page-crossing draft row {n}", weight_rtol + ) + env.check_cache(0) + env.check_unwritten_pool_zero() + + torch.manual_seed(472) + env = _draft_env(hkv, pool_dtype, quant_mode, scale) + env.add_request(0, 100) + env.call_op(env.random_qkv(100), [100], 1, [0], MASK_CAUSAL) + for _ in range(3): + gen = env.random_qkv(4) + out = env.call_op(gen, [4], 0, [0], MASK_CAUSAL, predicted_tokens_per_seq=4) + _compare(out, _op_reference(env, gen, [4], [0], e4m3=fp8)) + env.check_cache(0) + env.check_unwritten_pool_zero() + + +# ─── qwen4_exp's full-attention cell: GQA 24/2/256 over a bf16 pool ──── +# +# 24 query / 2 kv heads / head_size 256, tokens_per_block 32, causal, no sinks +# or window, the paged-context path on, P up to 4. 2 kv heads put 12 query heads +# in a group where 24/4 puts 6, which moves the compiled decode kernel's q-tile +# (VarSeqQ16 against VarSeqQ8), so 24/2/256 is certified as a triple of its own. +# Every rival here is a reference model computed on the host, never a second +# deliberately wrong call. + +QWEN_HQ, QWEN_HKV, QWEN_D = 24, 2, 256 +QWEN_TPB = 32 + +# How far a rival model has to sit outside the bf16 band. Measured on sm_103: +# 45x-118x (a draft row seeing its later siblings), 183x (a context call's +# cached prefix dropped), 17.8x-582x (an encoder's frame segmentation dropped). +QWEN_MIN_SEPARATION = 3.0 + + +def _qwen_bf16_env(**kwargs) -> _PagedAttnEnv: + """The checkpoint's full-attention cell over a bf16 paged pool.""" + kwargs.setdefault("num_heads", QWEN_HQ) + kwargs.setdefault("num_kv_heads", QWEN_HKV) + kwargs.setdefault("head_dim", QWEN_D) + kwargs.setdefault("tokens_per_block", QWEN_TPB) + kwargs.setdefault("num_blocks", 128) + kwargs.setdefault("max_blocks_per_seq", 16) + return _PagedAttnEnv(**kwargs) + + +def test_bf16_gqa24_2_cached_prefix_rival() -> None: + """A 96-token request served by two paged-context calls, 48 then 33 over + the 48-token prefix: each inside `_compare`, and the second far outside + the band around the prefix-ignored model -- every row seeing only this + call's own 33 tokens, which is what a call that never reached the pool + computes (183x on sm_103). On the fresh call the two coincide.""" + torch.manual_seed(523) + env = _qwen_bf16_env() + env.add_request(0, 96) + for new_tokens in (48, 33): + cached = env.cached_len(0) + qkv = env.random_qkv(new_tokens) + out = _paged_ctx_call(env, qkv, [new_tokens], 1, [0]) + _compare(out, _op_reference(env, qkv, [new_tokens], [0])) + if cached: + _assert_far_outside_band( + out, + _packed_reference(qkv, [new_tokens], QWEN_HQ, QWEN_HKV, QWEN_D, causal=True), + QWEN_MIN_SEPARATION, + f"the {cached}-token cached prefix never reached the FMHA", + ATOL, + RTOL, + ) + env.check_cache(0) + env.check_unwritten_pool_zero() + + +def test_bf16_gqa24_2_sparse_group_is_inert() -> None: + """The sparse / skip-softmax arguments, at and off their inert values. + + In the standard configuration nothing in the group moves a byte: each of + nine settings -- the four index/offset tensors filled with a real + 8-key-per-row selection, a non-zero sparse_attn_indices_block_size, a + num_sparse_topk, sparse_attn_kv_lens, and the two skip-softmax thresholds + with and without a stats buffer -- returns output bitwise identical to the + all-inert call. So a caller passing something here is not silently getting + sparse attention either. + + A negative result needs a control: the same byte comparison over arguments + the op does read (q_scaling 2.0, mask_type 0) must fail. + """ + torch.manual_seed(522) + env = _qwen_bf16_env() + env.add_request(0, 40) + qkv = env.random_qkv(40) + + def again(mask_type: int = MASK_CAUSAL, **overrides) -> torch.Tensor: + # record=False throughout: the history mirror is what would turn a + # repeat into a cached-prefix call, and these stay one fresh prefill. + return env.call_op( + qkv, + [40], + 1, + [0], + mask_type, + record=False, + use_paged_context_fmha=True, + op_overrides=overrides, + ) + + base = again() + assert _bitwise_equal(again(), base), ( + "the identical call is not reproducible, so no byte comparison below means anything" + ) + + keep = 8 # a selection that could not reproduce full causal attention + attn_idx = torch.tensor([list(range(keep))] * 40, dtype=torch.int32, device="cuda").reshape(-1) + attn_off = torch.tensor([keep * i for i in range(41)], dtype=torch.int32, device="cuda") + off_inert = { + "sparse_attn_indices + sparse_attn_offsets": dict( + sparse_attn_indices=attn_idx, sparse_attn_offsets=attn_off + ), + "... + sparse_attn_indices_block_size": dict( + sparse_attn_indices=attn_idx, + sparse_attn_offsets=attn_off, + sparse_attn_indices_block_size=1, + ), + "... + num_sparse_topk": dict( + sparse_attn_indices=attn_idx, + sparse_attn_offsets=attn_off, + sparse_attn_indices_block_size=1, + num_sparse_topk=keep, + ), + "sparse_kv_indices + sparse_kv_offsets": dict( + sparse_kv_indices=attn_idx, sparse_kv_offsets=attn_off + ), + "num_sparse_topk alone": dict(num_sparse_topk=8), + "sparse_attn_indices_block_size alone": dict(sparse_attn_indices_block_size=64), + "sparse_attn_kv_lens": dict( + sparse_attn_kv_lens=torch.zeros(1, dtype=torch.int32, device="cuda") + ), + "skip_softmax thresholds": dict( + skip_softmax_threshold_scale_factor_prefill=0.9, + skip_softmax_threshold_scale_factor_decode=0.9, + ), + "skip_softmax thresholds + stat": dict( + skip_softmax_threshold_scale_factor_prefill=0.9, + skip_softmax_threshold_scale_factor_decode=0.9, + skip_softmax_stat=torch.zeros(40, QWEN_HQ, 2, dtype=torch.float32, device="cuda"), + ), + } + for label, overrides in off_inert.items(): + assert _bitwise_equal(again(**overrides), base), ( + f"{label} changed the output: the sparse / skip-softmax group is not inert here" + ) + assert not _bitwise_equal(again(q_scaling=2.0), base), ( + "q_scaling 2.0 is bitwise identical to 1.0: this case's byte comparison is blind" + ) + assert not _bitwise_equal(again(mask_type=MASK_PADDING), base), ( + "mask_type 0 is bitwise identical to 1 on a 40-token context call: " + "this case's byte comparison is blind" + ) + + +# ─── local_layer_idx over qwen4_exp's 48-row mapping ─────────────────── +# +# The 48-layer tower carries a full-attention layer every fourth position, and +# the target passes the layer's `layer_offsets` entry, so it drives rows 3, 7, +# ..., 47. Two arrangements, because the mapping's two columns move +# independently and only one of them is the row index: +# +# * the engine's own -- a 48-row mapping whose twelve attention rows sit on +# pool 0 at layer-in-pool 0..11 and whose other 36 rows name a second pool, +# so the row index (to 47) and the layer-in-pool column (to 11) differ; +# * a real 48-layer KVCacheManager, whose identity mapping drives both to 47. +# +# Both assert the pair a single-layer pool cannot: the named layer's slabs +# moved, AND every sibling layer's slabs are bitwise untouched. + +QWEN_NUM_LAYERS = 48 +QWEN_ATTN_LAYERS = list(range(3, QWEN_NUM_LAYERS, 4)) + + +def _qwen_layer_pools( + attn_blocks: int = 48, other_blocks: int = 32, max_seq_len: int = 256, max_batch: int = 4 +) -> Tuple[_MultiLayerPagedAttnEnv, _MultiLayerPagedAttnEnv]: + """The two pools of the hybrid mapping, at the full-attention cell. + + The second holds the 36 rows this cell never drives; it exists so those + 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. + """ + kw = dict( + num_heads=QWEN_HQ, + num_kv_heads=QWEN_HKV, + head_dim=QWEN_D, + tokens_per_block=QWEN_TPB, + max_seq_len=max_seq_len, + max_batch=max_batch, + ) + attn = _MultiLayerPagedAttnEnv(num_layers=len(QWEN_ATTN_LAYERS), num_blocks=attn_blocks, **kw) + other = _MultiLayerPagedAttnEnv( + num_layers=QWEN_NUM_LAYERS - len(QWEN_ATTN_LAYERS), num_blocks=other_blocks, **kw + ) + return attn, other + + +def _compose_engine_mapping( + attn: _MultiLayerPagedAttnEnv, other: _MultiLayerPagedAttnEnv +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Pool pointers, the 48-row layer->pool mapping and the per-pool block + offsets of the hybrid arrangement. Row i is what local_layer_idx=i selects: + rows 3, 7, ..., 47 are pool 0 at layer-in-pool 0..11, every other row is + pool 1. Recompose after each refresh_offsets: offsets are concatenated by + value.""" + pool_pointers = torch.zeros(2, 2, dtype=torch.int64, device="cpu") + pool_pointers[0] = attn.mgr.kv_cache_pool_pointers[0] + pool_pointers[1] = other.mgr.kv_cache_pool_pointers[0] + attn_rows = set(QWEN_ATTN_LAYERS) + rows, in_attn, in_other = [], 0, 0 + for i in range(QWEN_NUM_LAYERS): + if i in attn_rows: + rows.append([0, in_attn]) + in_attn += 1 + else: + rows.append([1, in_other]) + in_other += 1 + mapping = torch.tensor(rows, dtype=torch.int32, device="cpu") + offsets = torch.cat([attn.block_offsets, other.block_offsets], dim=0) + return pool_pointers, mapping, offsets + + +def _assert_only_this_layer_moved( + env: _MultiLayerPagedAttnEnv, + snapshot: List[torch.Tensor], + layer: int, + label: str, +) -> None: + """This layer's slabs moved, and every other layer's slabs are bitwise as + they were. The complement's size is asserted too, so the sibling loop + cannot pass by being empty.""" + assert not _bitwise_equal(snapshot[layer], env.layer_views[layer]), ( + f"{label}: the named layer's slabs did not move at all" + ) + siblings = [m for m in range(env.num_layers) if m != layer] + assert len(siblings) == env.num_layers - 1 and siblings, f"{label}: empty complement" + for m in siblings: + assert _bitwise_equal(snapshot[m], env.layer_views[m]), ( + f"{label}: the call wrote layer {m}'s slabs as well as its own" + ) + + +def test_bf16_gqa24_2_layer_index_engine_48_row_mapping() -> None: + """24/2/256 at the twelve local_layer_idx values the target drives. + + Each of rows 3, 7, ..., 47 runs a fresh context chunk, a context chunk + over its own cached prefix on the paged-context path, and generation at + P = 1 and 2. Each call is inside `_compare` of its own layer's history, + so a read landing in a sibling's slabs cannot pass; each append moves its + own layer's slabs and leaves the other eleven bitwise untouched; and the + 36-layer second pool comes back bitwise as it started. + """ + torch.manual_seed(530) + attn, other = _qwen_layer_pools() + other_before = [v.clone() for v in other.layer_views] + _, mapping, _ = _compose_engine_mapping(attn, other) + assert mapping.shape == (QWEN_NUM_LAYERS, 2) + assert [mapping[row].tolist() for row in QWEN_ATTN_LAYERS] == [ + [0, k] for k in range(len(QWEN_ATTN_LAYERS)) + ], "the twelve driven rows are pool 0, layer-in-pool 0..11" + + attn.add_request(0, 96) # the manager allocates the whole prompt up front + attn.add_request(1, 40) + attn.refresh_offsets([0, 1], num_contexts=2) + ptrs, mapping, offsets = _compose_engine_mapping(attn, other) + assert int(ptrs[0, 0]) != int(ptrs[1, 0]), "the two pools must have distinct bases" + + for chunk in ([64, 25], [32, 15]): + for layer, row in enumerate(QWEN_ATTN_LAYERS): + qkv = attn.random_qkv(sum(chunk)) + snapshot = [v.clone() for v in attn.layer_views] + out = attn.call_op( + layer, + qkv, + chunk, + 2, + [0, 1], + pool_mapping=mapping, + pool_pointers=ptrs, + block_offsets=offsets, + local_layer_idx=row, + use_paged_context_fmha=True, + ctx_lens=chunk, + ) + _compare(out, _op_reference(attn, qkv, chunk, [0, 1], layer=layer)) + _assert_only_this_layer_moved( + attn, snapshot, layer, f"context chunk {chunk} at local_layer_idx {row}" + ) + attn.check_caches([0, 1]) + + for p in (1, 2): + for _ in range(p): + attn.add_decode_token(0) + attn.add_decode_token(1) + attn.refresh_offsets([0, 1], num_contexts=0) + ptrs, mapping, offsets = _compose_engine_mapping(attn, other) + for layer, row in enumerate(QWEN_ATTN_LAYERS): + gen = attn.random_qkv(2 * p) + snapshot = [v.clone() for v in attn.layer_views] + out = attn.call_op( + layer, + gen, + [p, p], + 0, + [0, 1], + pool_mapping=mapping, + pool_pointers=ptrs, + block_offsets=offsets, + local_layer_idx=row, + use_paged_context_fmha=True, + predicted_tokens_per_seq=p, + ) + _compare(out, _op_reference(attn, gen, [p, p], [0, 1], layer=layer)) + _assert_only_this_layer_moved( + attn, snapshot, layer, f"P={p} decode at local_layer_idx {row}" + ) + attn.check_caches([0, 1]) + + for before, after in zip(other_before, other.layer_views): + assert _bitwise_equal(before, after), ( + "a call on a pool-0 row wrote the pool the mapping's other 36 rows name" + ) + + +def test_bf16_gqa24_2_layer_index_identity_48_layer_pool() -> None: + """The same question with the mapping's other column moved: a real 48-layer + KVCacheManager's identity rows take local_layer_idx and the layer-in-pool + column to 47 together, so the pool-base shift runs 47 * kv_factor slabs + deep. Rows 0, 3, 23 and 47 over context, cached-prefix context and P = 2 + generation, each inside `_compare` of its own history, with the + moved/untouched pair over all 48 layers.""" + torch.manual_seed(532) + env = _MultiLayerPagedAttnEnv( + num_layers=QWEN_NUM_LAYERS, + num_heads=QWEN_HQ, + num_kv_heads=QWEN_HKV, + head_dim=QWEN_D, + tokens_per_block=QWEN_TPB, + num_blocks=24, + max_seq_len=128, + ) + assert env.pool_mapping.tolist() == [[0, i] for i in range(QWEN_NUM_LAYERS)], ( + "a real 48-layer manager produces identity mapping rows" + ) + driven = [0, 3, 23, 47] + + env.add_request(0, 64) + env.refresh_offsets([0], num_contexts=1) + for chunk in ([40], [24]): + for layer in driven: + qkv = env.random_qkv(sum(chunk)) + snapshot = [v.clone() for v in env.layer_views] + out = env.call_op( + layer, qkv, chunk, 1, [0], use_paged_context_fmha=True, ctx_lens=chunk + ) + _compare(out, _op_reference(env, qkv, chunk, [0], layer=layer)) + _assert_only_this_layer_moved( + env, snapshot, layer, f"context chunk {chunk} at identity row {layer}" + ) + env.check_caches([0], layers=driven) + + p = 2 + for _ in range(p): + env.add_decode_token(0) + env.refresh_offsets([0], num_contexts=0) + for layer in driven: + gen = env.random_qkv(p) + snapshot = [v.clone() for v in env.layer_views] + out = env.call_op( + layer, gen, [p], 0, [0], use_paged_context_fmha=True, predicted_tokens_per_seq=p + ) + _compare(out, _op_reference(env, gen, [p], [0], layer=layer)) + _assert_only_this_layer_moved(env, snapshot, layer, f"P={p} decode at identity row {layer}") + env.check_caches([0], layers=driven) + + +def test_bf16_gqa24_2_layer_index_ab_single_layer_vs_row_47() -> None: + """A/B: the same call, once in a single-layer pool and once at row 47. + + Identical operands and identical layer content either side, and every + other op argument equal (max_seq_len, max_context_length, + attention_window_size, max_num_requests, tokens_per_block), so the only + variable is where the call's slice lives -- a one-layer pool at base, or + layer-in-pool 11 of a twelve-layer pool selected by mapping row 47 of 48. + The outputs are bitwise equal on both the paged-context prefill and the + decode that reads the history back. + + The control is the degeneracy itself: the same decode sent to row 43 + (layer-in-pool 10), whose slabs hold no history, must not come back equal. + """ + torch.manual_seed(531) + single = _qwen_bf16_env() + single.add_request(0, 64) + qkv = single.random_qkv(64) + out_single = _paged_ctx_call(single, qkv, [64], 1, [0]) + qkv_d = single.random_qkv(1) + out_single_d = _paged_ctx_call(single, qkv_d, [1], 0, [0]) + single.check_cache(0) + + attn, other = _qwen_layer_pools(attn_blocks=64, other_blocks=64, max_seq_len=512) + assert attn.max_seq_len == single.max_seq_len + deep_layer, deep_row = len(QWEN_ATTN_LAYERS) - 1, QWEN_ATTN_LAYERS[-1] + assert deep_row == 47 + attn.add_request(0, 64) + attn.refresh_offsets([0], num_contexts=1) + ptrs, mapping, offsets = _compose_engine_mapping(attn, other) + tables = dict(pool_mapping=mapping, pool_pointers=ptrs, block_offsets=offsets) + out_multi = attn.call_op( + deep_layer, + qkv, + [64], + 1, + [0], + local_layer_idx=deep_row, + use_paged_context_fmha=True, + ctx_lens=[64], + **tables, + ) + attn.add_decode_token(0) + attn.refresh_offsets([0], num_contexts=0) + ptrs, mapping, offsets = _compose_engine_mapping(attn, other) + tables = dict(pool_mapping=mapping, pool_pointers=ptrs, block_offsets=offsets) + out_multi_d = attn.call_op( + deep_layer, + qkv_d, + [1], + 0, + [0], + local_layer_idx=deep_row, + use_paged_context_fmha=True, + **tables, + ) + attn.check_caches([0], layers=[deep_layer]) + + for what, a, b in ( + ("paged-context prefill", out_single, out_multi), + ("decode", out_single_d, out_multi_d), + ): + assert _bitwise_equal(a, b), ( + f"{what}: a single-layer pool and layer-in-pool 11 of a twelve-layer " + f"pool at mapping row 47 disagree (max|delta| " + f"{(a.float() - b.float()).abs().max().item():.6g})" + ) + + wrong_row = QWEN_ATTN_LAYERS[-2] + wrong = attn.call_op( + deep_layer - 1, + qkv_d, + [1], + 0, + [0], + local_layer_idx=wrong_row, + use_paged_context_fmha=True, + record=False, + **tables, + ) + assert not _bitwise_equal(wrong, out_single_d), ( + f"a decode routed to mapping row {wrong_row}, whose slabs hold no history, " + f"is bitwise identical to the correct one: this A/B cannot see a wrong layer index" + ) + + +# ─── Encoder-only: the op with no KV cache at all ────────────────────── +# +# A vision tower has no KV cache: one bidirectional pass over each frame's +# patch tokens. The op is told so by receiving no cache addressing at all -- +# kv_cache_block_offsets, host_kv_cache_pool_pointers and +# host_kv_cache_pool_mapping all None -- with update_kv_cache still True. +# Frame segmentation is the batch: one context sequence per frame, and +# mask_type 0 makes each frame's tokens attend to exactly their own frame. + +# The segment shapes the multimodal runtime drives, read off the runtime-value +# dump of the gate that printed `mmmu_val average accuracy: 58.33`: 170 +# distinct segment sizes from 256 to 12544, totals to 22144 rows, 1, 2, 4 and +# 5 segments per call. These are its four largest calls plus the two +# multi-segment counts; only a call's total, count and maximum are in the dump, +# so the splits realize those triples with lengths a 2x2-merged grid produces. +_ENCODER_SERVING_SEGMENTS: List[List[int]] = [ + [5580], + [8160], + [728, 728, 460, 288], + [480, 480, 400, 352, 180], + [8000, 1008], + [12544, 9600], +] + +# What `workspace_` grows to across that sequence, from empty: the op sizes it +# in place per call and never shrinks it, and the dump's largest recorded size +# is this figure. +_ENCODER_SERVING_WORKSPACE = 33732608 + + +def _encoder_call( + qkv: torch.Tensor, + seg_lens: List[int], + num_heads: int, + num_kv_heads: int, + head_dim: int, + mask_type: int, + workspace: torch.Tensor, + **overrides, +) -> torch.Tensor: + """One context-phase call with every KV-cache argument absent, batch state + as the vision tower builds it (`max_num_requests` = the segment count).""" + ns = len(seg_lens) + total = sum(seg_lens) + output = torch.empty(total, num_heads * head_dim, dtype=qkv.dtype, device="cuda") + op_kwargs = dict( + q=qkv, + k=None, + v=None, + output=output, + output_sf=None, + workspace_=workspace, + sequence_length=torch.tensor(seg_lens, dtype=torch.int32, device="cuda"), + host_past_key_value_lengths=torch.tensor(seg_lens, dtype=torch.int32), + host_total_kv_lens=torch.tensor([total, 0], dtype=torch.int32), + context_lengths=torch.tensor(seg_lens, dtype=torch.int32, device="cuda"), + host_context_lengths=torch.tensor(seg_lens, dtype=torch.int32), + host_request_types=torch.tensor([0] * ns, dtype=torch.int32), + max_context_q_len_override=None, + # The three whose absence selects the no-KV-cache mode. + kv_cache_block_offsets=None, + host_kv_cache_pool_pointers=None, + host_kv_cache_pool_mapping=None, + cache_indirection=None, + kv_scale_orig_quant=None, + kv_scale_quant_orig=None, + out_scale=None, + 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=1, + local_layer_idx=0, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_size=head_dim, + tokens_per_block=None, + max_num_requests=ns, + max_context_length=max(seg_lens), + max_seq_len=max(seg_lens), + attention_window_size=max(seg_lens), + beam_width=1, + mask_type=mask_type, + quant_mode=0, + 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=False, + 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=ns, + num_ctx_tokens=total, + ) + op_kwargs.update(overrides) + with validating(op): + op(**op_kwargs) + torch.cuda.synchronize() + return output + + +def _merged_rival( + qkv: torch.Tensor, + num_heads: int, + head_dim: int, + mask_type: int, + q_chunk: int = 1024, +) -> torch.Tensor: + """The answer with the frame segmentation dropped -- every row attending + to every row of the call, as one sequence -- in fp32 SDPA, query rows in + blocks of `q_chunk` so a 22144-row call does not ask for one 16 x 22144 x + 22144 score matrix. A rival model only: the chunking changes no row's + arithmetic.""" + n = qkv.shape[0] + w = num_heads * head_dim + q = qkv[:, :w].view(n, num_heads, head_dim).transpose(0, 1).float() + k = qkv[:, w : 2 * w].view(n, num_heads, head_dim).transpose(0, 1).float() + v = qkv[:, 2 * w :].view(n, num_heads, head_dim).transpose(0, 1).float() + mask = None + if mask_type == MASK_CAUSAL: + mask = torch.ones(n, n, dtype=torch.bool, device=qkv.device).tril() + parts = [ + F.scaled_dot_product_attention( + q[:, first : first + q_chunk], + k, + v, + attn_mask=None if mask is None else mask[first : first + q_chunk], + ) + for first in range(0, n, q_chunk) + ] + return torch.cat(parts, dim=1).transpose(0, 1).reshape(n, -1).to(qkv.dtype) + + +def test_encoder_no_kv_cache_segmentation_is_the_batch() -> None: + """The frame segmentation is the whole of the correctness here and nothing + else in the call would report its absence, so every multi-segment shape is + gated against the merged-segment model -- the same rows run as one + sequence, what a caller who drops `num_contexts` gets. + + Three frames of 27, 9 and 41 rows at head_size 72 and 128 under both masks, + each inside `_compare`, the rival 315x / 495x (d72 padding / causal) and + 390x / 582x (d128) outside the band on sm_103; then the mmmu dump's + multi-segment calls at 72 under padding, where the cells hold the accuracy + and this holds the rival: 107x, 114x, 67x and 17.8x for four, five, two + (9008 rows) and two (22144 rows) frames. + """ + for head_dim in (72, 128): + torch.manual_seed(520) + workspace = torch.empty(0, dtype=torch.int8, device="cuda") + seg_lens = [27, 9, 41] + qkv = torch.randn(sum(seg_lens), 48 * head_dim, dtype=torch.bfloat16, device="cuda") + for mask_type in (MASK_PADDING, MASK_CAUSAL): + out = _encoder_call(qkv, seg_lens, 16, 16, head_dim, mask_type, workspace) + _compare( + out, + _packed_reference(qkv, seg_lens, 16, 16, head_dim, mask_type == MASK_CAUSAL), + ) + _assert_far_outside_band( + out, + _merged_rival(qkv, 16, head_dim, mask_type), + QWEN_MIN_SEPARATION, + f"D={head_dim} mask={mask_type}: the frame segmentation was ignored", + ATOL, + RTOL, + ) + workspace = torch.empty(0, dtype=torch.int8, device="cuda") + for seg_lens in _ENCODER_SERVING_SEGMENTS: + if len(seg_lens) == 1: + continue + torch.manual_seed(900) + qkv = torch.randn(sum(seg_lens), 48 * 72, dtype=torch.bfloat16, device="cuda") + out = _encoder_call(qkv, seg_lens, 16, 16, 72, MASK_PADDING, workspace) + _assert_far_outside_band( + out, + _merged_rival(qkv, 16, 72, MASK_PADDING), + QWEN_MIN_SEPARATION, + f"{seg_lens}: the frame segmentation was ignored", + ATOL, + RTOL, + ) + del qkv, out + torch.cuda.empty_cache() + + +def test_encoder_no_kv_cache_cache_arguments_are_inert() -> None: + """Arguments that address a cache cannot matter when there is none: + tokens_per_block 32 (against None) and use_paged_context_fmha True (there + is no pool for the context FMHA to read) return bitwise the same output.""" + for head_dim in (72, 128): + torch.manual_seed(521) + workspace = torch.empty(0, dtype=torch.int8, device="cuda") + seg_lens = [27, 9, 41] + qkv = torch.randn(sum(seg_lens), 48 * head_dim, dtype=torch.bfloat16, device="cuda") + plain = _encoder_call(qkv, seg_lens, 16, 16, head_dim, MASK_PADDING, workspace) + for label, overrides in ( + ("tokens_per_block=32", dict(tokens_per_block=32)), + ("use_paged_context_fmha=True", dict(use_paged_context_fmha=True)), + ): + got = _encoder_call( + qkv, seg_lens, 16, 16, head_dim, MASK_PADDING, workspace, **overrides + ) + assert _bitwise_equal(got, plain), ( + f"D={head_dim}: {label} changed a call that owns no KV cache" + ) + + +def test_encoder_no_kv_cache_workspace_high_water() -> None: + """`workspace_` over the mmmu serving sequence, from empty: sized in place + per call, never shrunk, and topping out at the largest size the runtime + dump records -- which a tower that preallocates it must cover.""" + workspace = torch.empty(0, dtype=torch.int8, device="cuda") + sizes = [] + for seg_lens in _ENCODER_SERVING_SEGMENTS: + torch.manual_seed(900) + qkv = torch.randn(sum(seg_lens), 48 * 72, dtype=torch.bfloat16, device="cuda") + _encoder_call(qkv, seg_lens, 16, 16, 72, MASK_PADDING, workspace) + sizes.append(workspace.numel()) + assert sizes == sorted(sizes), f"workspace_ did not grow monotonically: {sizes}" + assert sizes[-1] == _ENCODER_SERVING_WORKSPACE, ( + f"workspace_ grew to {sizes[-1]} bytes over the serving shapes, expected " + f"{_ENCODER_SERVING_WORKSPACE}; the sequence was {sizes}" + ) + + +def test_encoder_no_kv_cache_boundaries() -> None: + """The two things that are not how to ask for the no-KV-cache mode. + + `update_kv_cache=False` is rejected outright, so the mode is selected by + the absence of the three addressing tensors and nothing else. And + head_size 72 exists only in the context phase: a context call over a paged + pool computes it correctly (a fallback kernel takes it), while a + generation-phase call raises. Both raises are host-side RuntimeErrors, not + a process abort -- a correct call after each one is still right. + """ + torch.manual_seed(530) + workspace = torch.empty(0, dtype=torch.int8, device="cuda") + seg_lens = [20] + qkv = torch.randn(20, 48 * 72, dtype=torch.bfloat16, device="cuda") + with pytest.raises(RuntimeError, match="KV cache update cannot be disabled now"): + _encoder_call(qkv, seg_lens, 16, 16, 72, MASK_PADDING, workspace, update_kv_cache=False) + + env = _PagedAttnEnv( + num_heads=16, + num_kv_heads=16, + head_dim=72, + tokens_per_block=32, + num_blocks=64, + max_blocks_per_seq=8, + ) + env.add_request(0, 40) + ctx = env.random_qkv(40) + _compare(env.call_op(ctx, [40], 1, [0], MASK_CAUSAL), _op_reference(env, ctx, [40], [0])) + env.check_cache(0) + with pytest.raises(RuntimeError, match="unsupported head_size 72"): + env.call_op(env.random_qkv(1), [1], 0, [0], MASK_CAUSAL, record=False) + + _compare( + _encoder_call(qkv, seg_lens, 16, 16, 72, MASK_PADDING, workspace), + _packed_reference(qkv, seg_lens, 16, 16, 72, causal=False), + ) + + # ─── The entry's own cell list ────────────────────────────────────────── @@ -5413,8 +6914,176 @@ def _cell_mla_context(cell) -> None: _compare(out, ref) +def _seed(cell) -> None: + """A cell's own seed, stable across processes (str hash is not).""" + torch.manual_seed(zlib.crc32(cell.why.encode())) + + +def _cell_paged(cell, arm) -> None: + """A standard-configuration cell over a single-layer paged pool. + + 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 + 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) + env = _PagedAttnEnv( + num_heads=spec["num_heads"], + num_kv_heads=spec["num_kv_heads"], + head_dim=spec["head_size"], + tokens_per_block=spec["tokens_per_block"], + num_blocks=spec["num_blocks"], + max_batch=spec["max_num_requests"], + 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, + ) + cached, new, nc = spec["cached"], spec["new"], spec["num_contexts"] + p, paged = spec["predicted_tokens_per_seq"], spec["use_paged_context_fmha"] + assert all(n == p for n in new[nc:]), "a generation row carries P q rows" + assert paged or not any(cached[:nc]), "a cached prefix needs the paged-context path" + rids = list(range(len(new))) + for rid in rids: + env.add_request(rid, cached[rid] + (new[rid] if rid < nc else 0)) + prior = [rid for rid in rids if cached[rid]] + if prior: + hist = [cached[rid] for rid in prior] + env.call_op( + env.random_qkv(sum(hist)), + hist, + len(prior), + prior, + MASK_CAUSAL, + ctx_lens_override=hist, + use_paged_context_fmha=paged, + ) + qkv = env.random_qkv(sum(new)) + out = env.call_op( + qkv, + list(new), + nc, + rids, + spec["mask_type"], + ctx_lens_override=[new[i] if i < nc else env.prompt_lens[i] for i in rids], + use_paged_context_fmha=paged, + predicted_tokens_per_seq=p, + ) + # 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) + causal = spec["mask_type"] == MASK_CAUSAL + _compare(out, _op_reference(env, qkv, list(new), rids, causal=causal, e4m3=e4m3)) + for rid in rids: + env.check_cache(rid) + env.check_unwritten_pool_zero() + + +def _cell_layered(cell, arm) -> None: + """A qwen4_exp cell at one row of a 48-row mapping. + + `hybrid48` is the engine's arrangement (rows 3, 7, ..., 47 on pool 0 at + layer-in-pool 0..11, the other 36 rows on pool 1); `identity48` a real + 48-layer manager's single pool. The history goes in through the same row + by an uncertified context call. Beside `_compare`: the call moved its own + layer's slabs and no sibling's, and the second pool stayed untouched. + """ + spec = cell.spec + arm(op) + _seed(cell) + 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" + if spec["pool_layout"] == "hybrid48": + env, other = _qwen_layer_pools(max_batch=spec["max_num_requests"]) + layer = QWEN_ATTN_LAYERS.index(row) + else: + env = _MultiLayerPagedAttnEnv( + num_layers=QWEN_NUM_LAYERS, + num_heads=spec["num_heads"], + num_kv_heads=spec["num_kv_heads"], + head_dim=spec["head_size"], + tokens_per_block=spec["tokens_per_block"], + num_blocks=24, + max_batch=spec["max_num_requests"], + max_seq_len=128, + ) + other, layer = None, row + + def tables() -> Dict: + if other is None: + return {} + ptrs, mapping, offsets = _compose_engine_mapping(env, other) + return dict(pool_pointers=ptrs, pool_mapping=mapping, block_offsets=offsets) + + other_before = [] if other is None else [v.clone() for v in other.layer_views] + rids = list(range(len(new))) + for rid in rids: + env.add_request(rid, cached[rid] + (new[rid] if rid < nc else 0)) + env.refresh_offsets(rids, num_contexts=len(rids)) + hist = list(cached) + env.call_op( + layer, + env.random_qkv(sum(hist)), + hist, + len(rids), + rids, + local_layer_idx=row, + use_paged_context_fmha=True, + ctx_lens=hist, + **tables(), + ) + if nc == 0: + for rid in rids: + for _ in range(p): + env.add_decode_token(rid) + env.refresh_offsets(rids, num_contexts=nc) + qkv = env.random_qkv(sum(new)) + snapshot = [v.clone() for v in env.layer_views] + out = env.call_op( + layer, + qkv, + list(new), + nc, + rids, + local_layer_idx=row, + use_paged_context_fmha=True, + ctx_lens=[new[i] if i < nc else env.prompt_lens[i] for i in rids], + predicted_tokens_per_seq=p, + **tables(), + ) + _compare(out, _op_reference(env, qkv, list(new), rids, layer=layer)) + 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): + assert _bitwise_equal(before, after), "a pool-0 row wrote the second pool" + + +def _cell_encoder(cell, arm) -> None: + """The no-KV-cache mode: one context sequence per frame, no cache + addressing, gated against `op.reference` over each frame's own rows.""" + spec = cell.spec + arm(op) + _seed(cell) + seg_lens = list(spec["seg_lens"]) + h, hkv, d = spec["num_heads"], spec["num_kv_heads"], spec["head_size"] + qkv = torch.randn(sum(seg_lens), (h + 2 * hkv) * d, dtype=torch.bfloat16, device="cuda") + workspace = torch.empty(0, dtype=torch.int8, device="cuda") + out = _encoder_call(qkv, seg_lens, h, hkv, d, spec["mask_type"], workspace) + causal = spec["mask_type"] == MASK_CAUSAL + _compare(out, _packed_reference(qkv, seg_lens, h, hkv, d, causal=causal)) + del qkv, out + torch.cuda.empty_cache() + + @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell) -> None: +def test_certified_cells(cell, certifying) -> None: # noqa: F811 """Drive every configuration the entry claims. The MLA generation cell is driven by the tests above rather than here: its @@ -5423,9 +7092,20 @@ def test_certified_cells(cell) -> None: it is a configuration a shipped target runs and `is_valid` is checked against it; what this file does not do is compare it to `reference`, and saying so is better than a cell that quietly proves less than its siblings. + + The cells that state their batch as `cached`/`new` (or `seg_lens`, with no + KV cache) arm `certifying`, so their run is also proof that the entry's + own `__call__` and `reference` were the ones exercised. """ - kind = cell.spec["input_type"] - if kind == STANDARD: + spec = cell.spec + kind = spec["input_type"] + if spec.get("kv_cache") == "none": + _cell_encoder(cell, certifying) + elif "pool_layout" in spec: + _cell_layered(cell, certifying) + elif "cached" in spec: + _cell_paged(cell, certifying) + elif kind == STANDARD: _cell_standard(cell) elif kind == MLA_CONTEXT: _cell_mla_context(cell) diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_nvfp4_gemm.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_nvfp4_gemm.py index 4e6a0579e963..8f7d3f6b8ce0 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_nvfp4_gemm.py +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_nvfp4_gemm.py @@ -8,7 +8,7 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.nvfp4_gemm import nvfp4_gemm as op -from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v3.r1_0528_nvfp4__sm_103__dep4.ops.quantization.fp4_quantize import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( # noqa: E501 fp4_quantize, ) diff --git a/tests/unittest/_torch/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/attention/test_modeling_v2_fused_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/attention/test_modeling_v2_fused_qk_norm_rope.py deleted file mode 100644 index e13ebb8e4456..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_103__tp1/attention/test_modeling_v2_fused_qk_norm_rope.py +++ /dev/null @@ -1,110 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU test for the fused_qk_norm_rope catalog entry.""" - -import zlib - -import pytest -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.models.gpt_oss.gpt_oss_120b__sm_103__tp1.ops.attention.fused_qk_norm_rope import ( # noqa: E501 - fused_qk_norm_rope as op, -) - -__extra_import_path__ = ["../../../.."] -from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above - -assert torch.cuda.is_available(), "fused_qk_norm_rope requires a CUDA device" - - -def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: - """A rotation is a two-term combination, and two terms can cancel. - - `x1 * cos - x2 * sin` goes to zero wherever the two products nearly - match, and there torch's per-element rtol is meaningless: at 2048 tokens - roughly 2e-4 of the outputs land close enough that a 7.3e-4 absolute - difference reads as a relative difference of 23. The decode and 37-token - cells pass under torch's default band only because they are too small to - sample one. - - So: 2 ulp element-wise against the row scale, 1 of relative RMS -- tight, - because one rotation is the shortest chain in this catalog. Measured - worst case across these cells is 7.3e-4 absolute at a row scale of ~4, - about 0.05 ulp. - """ - assert_within_ulp(out, ref, element_ulp=2.0, rms_ulp=1.0) - - -def _build(spec, seed): - """Materialize one cell as (kwargs, qkv), with qkv separate because it is - the buffer the op writes and the caller has to keep a copy of.""" - g = torch.Generator(device="cuda").manual_seed(seed) - tokens = spec["tokens"] - heads = spec["num_heads_q"] + spec["num_heads_k"] + spec["num_heads_v"] - width = heads * spec["head_dim"] - qkv = torch.randn(tokens, width, generator=g, device="cuda").to(torch.bfloat16) - weight = torch.randn(spec["head_dim"], generator=g, device="cuda").to(torch.bfloat16) - kwargs = {k: v for k, v in spec.items() if k != "tokens"} - kwargs.update( - eps=1e-6, - q_weight=weight, - k_weight=weight, - position_ids=torch.arange(tokens, dtype=torch.int32, device="cuda"), - ) - return qkv, kwargs - - -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) -def test_certified_cells(cell) -> None: - qkv, kwargs = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) - ref = op.reference(qkv, **kwargs) - with validating(op): - out = op(qkv, **kwargs) - assert out is qkv, "the op writes in place" - _compare(qkv, ref) - - -def test_the_v_heads_are_left_alone() -> None: - """The op's boundary, and nothing about the output shape shows it. - - A kernel that rotated the v heads too would still return a buffer of the - right size, and every cell compares against a reference that makes the same - claim -- so this pins it against the input instead. - """ - spec = dict(op.CELLS[1].spec) - qkv, kwargs = _build(spec, seed=11) - before = qkv.clone() - op(qkv, **kwargs) - v_start = (kwargs["num_heads_q"] + kwargs["num_heads_k"]) * kwargs["head_dim"] - assert torch.equal(qkv[:, v_start:], before[:, v_start:]) - assert not torch.equal(qkv[:, :v_start], before[:, :v_start]), "q and k were not touched" - - -def _ok(**over): - spec = dict(op.CELLS[0].spec) - spec.update(over) - return _build(spec, seed=3) - - -def test_guard_refuses_a_width_that_disagrees_with_the_head_counts() -> None: - """Read as whatever the counts say; the heads come out shifted.""" - qkv, kwargs = _ok() - with pytest.raises(AssertionError, match="last dim must be"): - with validating(op): - op(qkv[:, :-64], **kwargs) - - -def test_guard_refuses_a_rotary_dim_past_the_head() -> None: - qkv, kwargs = _ok() - kwargs["rotary_dim"] = 128 - with pytest.raises(AssertionError, match="exceeds head_dim"): - with validating(op): - op(qkv, **kwargs) - - -def test_guard_refuses_a_position_id_count_that_is_not_the_token_count() -> None: - qkv, kwargs = _ok() - kwargs["position_ids"] = torch.arange(3, dtype=torch.int32, device="cuda") - with pytest.raises(AssertionError, match="one position per token"): - with validating(op): - op(qkv, **kwargs) 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 new file mode 100644 index 000000000000..973a12dc0501 --- /dev/null +++ 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 @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the fused_sigmoid_mul_inplace entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.activation.fused_sigmoid_mul_inplace import ( # noqa: E501 + fused_sigmoid_mul_inplace as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "fused_sigmoid_mul_inplace requires a CUDA device" + + +def _seed(text: str) -> int: + return zlib.crc32(text.encode()) + + +def _build(spec, seed): + """`(buffer, out, region, 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. + """ + 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 + + gate_layout = spec["gate"] + if gate_layout == "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 + + +@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() + 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`.""" + 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] + + 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 new file mode 100644 index 000000000000..ab0b602f1189 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_erf.py @@ -0,0 +1,231 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp gelu_erf entry. + +The entry's identity is *which* GELU it is, and the wrong answer -- the tanh +GELU this same tower's block MLPs use -- rounds to within about one bf16 ulp of +it. The discrimination test is the point of this file; only fp32 tells the two +apart. +""" + +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.activation.gelu_erf import ( # noqa: E501 + gelu_erf as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "gelu_erf requires a CUDA device" + +MERGER_WIDTH = 4608 +TOWER_HIDDEN = 1152 +#: model.visual.merger, measured from the checkpoint: linear_fc1.weight std +#: 0.010927, linear_fc1.bias std 0.002407, norm.weight mean 1.174581 / std +#: 0.140536. The real tensors give an activand at std 0.875 / absmax 4.23. +FC1_WEIGHT_STD = 0.010927 +FC1_BIAS_STD = 0.002407 +NORM_GAMMA_MEAN = 1.174581 +NORM_GAMMA_STD = 0.140536 + +SQRT_2 = math.sqrt(2.0) +SQRT_2_OVER_PI = math.sqrt(2.0 / math.pi) + +#: torch.testing's untouched per-dtype defaults, restated only so the +#: discrimination test can say how far outside that gate each variant lands. +_RTOL = {torch.bfloat16: 1.6e-2, torch.float16: 1e-3, torch.float32: 1.3e-6} +_ATOL = 1e-5 + + +def _seed(text: str) -> int: + return zlib.crc32(text.encode()) + + +def _uniform(shape, lo, hi, dtype, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + x = torch.empty(shape, device="cuda", dtype=torch.float32).uniform_(lo, hi, generator=g) + return x.to(dtype) + + +def _build(spec, seed): + shape, dtype = spec["shape"], spec["dtype"] + if spec.get("dist", "uniform") == "merger": + g = torch.Generator(device="cuda").manual_seed(seed) + rows = shape[0] + gamma = NORM_GAMMA_MEAN + NORM_GAMMA_STD * torch.randn( + TOWER_HIDDEN, generator=g, device="cuda" + ) + hidden = torch.randn(4 * rows, TOWER_HIDDEN, generator=g, device="cuda") + grouped = (hidden * gamma.to(dtype).float()).reshape(rows, MERGER_WIDTH) + weight = FC1_WEIGHT_STD * torch.randn( + MERGER_WIDTH, MERGER_WIDTH, generator=g, device="cuda" + ) + 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) + + +@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 = _build(cell.spec, seed=_seed(cell.why)) + 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 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)) + + +def _erf64(x): + x64 = x.double() + return 0.5 * x64 * (1.0 + torch.erf(x64 / SQRT_2)) + + +def _tanh64(x): + """The *other* GELU -- `gelu_pytorch_tanh`, this tower's block-MLP activation.""" + x64 = x.double() + return 0.5 * x64 * (1.0 + torch.tanh(SQRT_2_OVER_PI * (x64 + 0.044715 * x64 * x64 * x64))) + + +def test_identity_is_the_erf_form_and_not_tanh() -> None: + """Measured in fp32, where the question has an answer; and shown to have + none in bf16, so a bf16 match cannot quietly be read as confirmation.""" + x32 = _uniform((512, MERGER_WIDTH), -4.0, 4.0, torch.float32, 24) + out = op(x32).double() + ref = _erf64(x32) + correct = (out - ref).abs().max().item() + gate = _ATOL + _RTOL[torch.float32] * ref.abs().max().item() + + # The entry's own gate refuses the tanh form in fp32. + with pytest.raises(AssertionError): + torch.testing.assert_close(out.float(), _tanh64(x32).float()) + + x64 = x32.double() + variants = { + "tanh gelu (gelu_pytorch_tanh)": _tanh64(x32), + "sigmoid approx x*sigmoid(1.702x)": x64 * torch.sigmoid(1.702 * x64), + "silu": x64 * torch.sigmoid(x64), + "relu": torch.clamp(x64, min=0.0), + } + print(f"\n fp32: gate {gate:.3e}; correct-path deviation {correct:.3e}") + for name, wrong in variants.items(): + dev = (out - wrong).abs().max().item() + print( + f" {name:36s} dev {dev:.6e} ({dev / correct:9.1f}x correct, {dev / gate:7.1f}x gate)" + ) + assert dev > 100.0 * correct, (name, dev, correct) + assert dev > 10.0 * gate, (name, dev, gate) + + # erf's scale truncated from 1/sqrt(2) to 1/1.414 -- a plausible + # transcription slip, weak by construction, so a lower multiple. + perturbed = 0.5 * x64 * (1.0 + torch.erf(x64 / 1.414)) + dev = (out - perturbed).abs().max().item() + print(f" {'erf scale 1/sqrt(2) -> 1/1.414':36s} dev {dev:.6e} ({dev / gate:7.1f}x gate)") + assert dev > 20.0 * correct, (dev, correct) + assert dev > 1.5 * gate, (dev, gate) + + # bf16 cannot see the difference. + xb = x32.to(torch.bfloat16) + outb = op(xb).double() + correct_b = (outb - _erf64(xb)).abs().max().item() + tanh_b = (outb - _tanh64(xb)).abs().max().item() + 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 new file mode 100644 index 000000000000..aa52e550fce0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/activation/test_modeling_v2_gelu_tanh.py @@ -0,0 +1,211 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp gelu_tanh entry. + +The entry's identity is *which* GELU it is, so the discrimination test is the +point of this file: in bf16 the tanh and erf GELUs round to within about one +ulp of each other, and only fp32 tells them apart. +""" + +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.activation.gelu_tanh import ( # noqa: E501 + gelu_tanh as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "gelu_tanh requires a CUDA device" + +INTERMEDIATE_SIZE = 4304 +HIDDEN_SIZE = 1152 +#: model.visual.blocks.*.mlp.linear_fc1, measured from the checkpoint: weight +#: std 0.0197 (block 0), bias std 1.676 / absmax 9.44 -- the bias dominates. +FC1_WEIGHT_STD = 0.0197 +FC1_BIAS_STD = 1.676 + +SQRT_2_OVER_PI = math.sqrt(2.0 / math.pi) + +#: torch.testing's untouched per-dtype defaults, restated only so the +#: discrimination test can say how far outside that gate each variant lands. +_RTOL = {torch.bfloat16: 1.6e-2, torch.float16: 1e-3, torch.float32: 1.3e-6} +_ATOL = 1e-5 + + +def _seed(text: str) -> int: + return zlib.crc32(text.encode()) + + +def _uniform(shape, lo, hi, dtype, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + x = torch.empty(shape, device="cuda", dtype=torch.float32).uniform_(lo, hi, generator=g) + return x.to(dtype) + + +def _build(spec, seed): + shape, dtype = spec["shape"], spec["dtype"] + if spec.get("dist", "uniform") == "fc1": + g = torch.Generator(device="cuda").manual_seed(seed) + rows = shape[0] + hidden = torch.randn(rows, HIDDEN_SIZE, generator=g, device="cuda").to(dtype) + weight = ( + FC1_WEIGHT_STD * torch.randn(shape[-1], HIDDEN_SIZE, generator=g, device="cuda") + ).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) + + +@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 = _build(cell.spec, seed=_seed(cell.why)) + 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 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)) + + +def _tanh64(x): + x64 = x.double() + return 0.5 * x64 * (1.0 + torch.tanh(SQRT_2_OVER_PI * (x64 + 0.044715 * x64 * x64 * x64))) + + +def _erf64(x): + """The *other* GELU -- exact, erf-based; `torch.nn.GELU()` and `moe_gelu`.""" + x64 = x.double() + return 0.5 * x64 * (1.0 + torch.erf(x64 / math.sqrt(2.0))) + + +def test_identity_is_the_tanh_form_and_not_erf() -> None: + """Measured in fp32, where the question has an answer; and shown to have + none in bf16, so a bf16 match cannot quietly be read as confirmation.""" + x32 = _uniform((512, INTERMEDIATE_SIZE), -4.0, 4.0, torch.float32, 24) + out = op(x32).double() + ref = _tanh64(x32) + correct = (out - ref).abs().max().item() + gate = _ATOL + _RTOL[torch.float32] * ref.abs().max().item() + + # The entry's own gate refuses the erf form in fp32. + with pytest.raises(AssertionError): + torch.testing.assert_close(out.float(), _erf64(x32).float()) + + x64 = x32.double() + variants = { + "erf gelu (nn.GELU, and moe_gelu)": _erf64(x32), + "sigmoid approx x*sigmoid(1.702x)": x64 * torch.sigmoid(1.702 * x64), + "silu": x64 * torch.sigmoid(x64), + "relu": torch.clamp(x64, min=0.0), + } + print(f"\n fp32: gate {gate:.3e}; correct-path deviation {correct:.3e}") + for name, wrong in variants.items(): + dev = (out - wrong).abs().max().item() + print( + f" {name:36s} dev {dev:.6e} ({dev / correct:9.1f}x correct, {dev / gate:7.1f}x gate)" + ) + assert dev > 100.0 * correct, (name, dev, correct) + assert dev > 10.0 * gate, (name, dev, gate) + + # A perturbation of this same formula: weak by construction (bounded by the + # coefficient change times the tanh derivative), so a lower multiple. + perturbed = 0.5 * x64 * (1.0 + torch.tanh(SQRT_2_OVER_PI * (x64 + 0.045 * x64**3))) + dev = (out - perturbed).abs().max().item() + print(f" {'coefficient 0.044715 -> 0.045':36s} dev {dev:.6e} ({dev / gate:7.1f}x gate)") + assert dev > 100.0 * correct, (dev, correct) + assert dev > 5.0 * gate, (dev, gate) + + # bf16 cannot see the difference: the erf variant's max deviation is within + # a small multiple of the correct path's own. + xb = x32.to(torch.bfloat16) + outb = op(xb).double() + correct_b = (outb - _tanh64(xb)).abs().max().item() + erf_b = (outb - _erf64(xb)).abs().max().item() + 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/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 new file mode 100644 index 000000000000..3c09c66fc05f --- /dev/null +++ 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 @@ -0,0 +1,310 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_expand_block_indices catalog entry. + +Selections are built the way `qsa_topk_select` emits them: a random subset of +the complete visible groups, prefix-packed, then `-1`. Every gate is bitwise. +""" + +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_expand_block_indices import ( # noqa: E501 + qsa_expand_block_indices as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "qsa_expand_block_indices requires a CUDA device" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Bitwise: these are KV-cache token indices, and one wrong entry reads + a different key. There is no rounding for a tolerance to absorb.""" + assert out.dtype == ref.dtype == torch.int32, (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, TOPK = 4, 2048 +BLOCK_TOPK = TOPK // CR +FINAL_TOPK = TOPK + CR - 1 +ENGINE_MAX_SEQ_LEN = 8192 + + +def _serving_rows(requests: int, context_lens: tuple) -> tuple[list[int], list[int]]: + """Per-row positions and lengths of one packed step. + + Chunk 0 starts at 0, so its rows span every visible-group count up to the + budget; later chunks are chunked-prefill continuations. Each remaining + request is an MTP generation request with two rows at consecutive positions. + """ + positions: list[int] = [] + lengths: list[int] = [] + for chunk, length in enumerate(context_lens): + start = 0 if chunk == 0 else (907 * chunk) % (ENGINE_MAX_SEQ_LEN - length + 1) + positions.extend(range(start, start + length)) + lengths.extend([start + length] * length) + for gen in range(requests - len(context_lens)): + position = 100 + (499 * gen) % 8000 + positions.extend((position, position + 1)) + lengths.extend((position + 2, position + 2)) + assert max(lengths, default=0) <= ENGINE_MAX_SEQ_LEN + return positions, lengths + + +def _rows(batch) -> tuple[list[int], list[int]]: + kind = batch[0] + if kind == "sweep": + positions = [4096 + 13 * i for i in range(batch[1])] + elif kind == "serving": + return _serving_rows(batch[1], batch[2]) + else: + positions = list(batch[1]) + return positions, [p + 1 for p in positions] + + +def _packed_selection(positions: torch.Tensor, compress_ratio: int, token_topk: int, seed: int): + """A legal Top-K result: distinct complete visible groups, then -1 padding.""" + block_topk = token_topk // compress_ratio + rows = positions.numel() + if rows == 0: + return torch.empty((0, block_topk), dtype=torch.int32, device=DEV) + available = (positions.long() + 1) // compress_ratio + width = max(block_topk, int(available.max())) + gen = torch.Generator(device=DEV).manual_seed(seed) + scores = torch.rand(rows, width, generator=gen, device=DEV) + groups = torch.arange(width, device=DEV) + scores = torch.where(groups[None, :] < available[:, None], scores, torch.inf) + chosen = scores.argsort(dim=1)[:, :block_topk] + taken = available.clamp(max=block_topk) + chosen = torch.where(torch.arange(block_topk, device=DEV)[None, :] < taken[:, None], chosen, -1) + return chosen.to(torch.int32) + + +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) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args, geometry = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + out = op(*args, **geometry) + rows = args[0].shape[0] + final = geometry["token_topk"] + geometry["compress_ratio"] - 1 + assert out.dtype == torch.int32 and out.is_cuda and out.is_contiguous() + assert tuple(out.shape) == (rows, final) + _compare(out, op.reference(*args, **geometry)) + if cell.spec["batch"][0] == "serving": + positions, lengths = args[1], args[2] + assert bool((lengths.long() > positions.long() + 1).any()), ( + "a serving cell must drive sequence_lengths past position + 1" + ) + + +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) + lengths = (positions + 1).to(torch.int32) + blocks = _packed_selection(positions, CR, TOPK, seed=31) + + def call(): + return op(blocks, positions, lengths, compress_ratio=CR, token_topk=TOPK) + + eager = call().clone() + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(3): + call() + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = call() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(captured, eager) + + blocks.copy_(_packed_selection(positions, CR, TOPK, seed=32)) + graph.replay() + torch.cuda.synchronize() + _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.""" + 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) 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 new file mode 100644 index 000000000000..b2ce01b5cbb0 --- /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_index_scores.py @@ -0,0 +1,553 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_paged_index_scores catalog entry. + +The paged state is real: a `[pages, tokens_per_block, 1, dim]` index cache +sized from the pages the table maps, and an int32 block table whose rows are +request *slots* -- more of them than the batch has requests, assigned by a +seeded permutation -- with shuffled physical pages and `-1` past each +sequence on the serving layouts, which is what the engine's table carries. + +The call has two kernels behind it. `_watch_launches` swaps the two Triton +entry points for recording proxies that still launch the real kernel, so every +cell asserts the branch it names. +""" + +import math + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_index_scores import ( # noqa: E501 + qsa_paged_index_scores as op, +) +from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "qsa_paged_index_scores requires a CUDA device" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """torch's default fp32 band on every specified column; `-inf` exactly. + + bf16 and fp16 products are exact in the kernel's fp32 accumulator, so + only summation order separates it from the reference: measured at most + 2.9e-6 over every cell (2.4e-6 at 8191 rows x 2080 columns, scores to + 11.6), inside the default rtol 1.3e-6 / atol 1e-5. Wrong variants (no + relu, mean over heads, first group member as key, unscaled, + 1/head_dim scale, block table ignored) 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. + """ + assert out.dtype == ref.dtype == torch.float32, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + specified = ~torch.isnan(ref) + torch.testing.assert_close(out[specified], ref[specified]) + + +DEV = torch.device("cuda") +#: The callable's own threshold: at or below it every call is row-wise. +DECODE_ROW_LIMIT = 256 +#: Index heads must divide this into more than one row for the query tile. +DOT_TILE_WIDTH = 16 +CAPTURE_PAGES = 260 + + +class _LaunchRecorder: + """Transparent proxy over a Triton entry point: records the grid, then launches.""" + + def __init__(self, inner, name, log): + self._inner, self._name, self._log = inner, name, log + + def __getitem__(self, grid): + self._log.append((self._name, tuple(grid))) + return self._inner[grid] + + def __getattr__(self, item): + return getattr(self._inner, item) + + +class _watch_launches: + """Record which of the entry's two Triton kernels a call launches.""" + + _NAMES = { + "_qsa_paged_index_scores_kernel": "row_wise", + "_qsa_paged_index_scores_qtile_kernel": "query_tile", + } + + def __enter__(self): + self.log = [] + self._saved = {name: getattr(_qsa_kernels, name) for name in self._NAMES} + for name, tag in self._NAMES.items(): + setattr(_qsa_kernels, name, _LaunchRecorder(self._saved[name], tag, self.log)) + return self + + def __exit__(self, *exc): + for name, kernel in self._saved.items(): + setattr(_qsa_kernels, name, kernel) + return False + + @property + def branch(self): + assert len(self.log) == 1, f"expected exactly one kernel launch, saw {self.log}" + return self.log[0][0] + + @property + def grid(self): + assert len(self.log) == 1, f"expected exactly one kernel launch, saw {self.log}" + return self.log[0][1] + + +def _case(rows, requests, **overrides): + """A narrow-table spec for the non-cell tests, same shape as the entry's own.""" + spec = dict( + layout="random", + rows=rows, + requests=requests, + slots=requests + 4, + pages=8, + heads=4, + head_dim=128, + tokens_per_block=32, + compress_ratio=4, + dtype=torch.bfloat16, + positions_dtype=torch.int32, + request_dtype=torch.int32, + q_layout="packed", + only_visible_blocks=False, + context_rows=False, + ) + spec.update(overrides) + return spec + + +def _rows_and_requests(spec, gen): + """Per-row positions, per-row request ordinal, and each request's mapped pages.""" + layout, tpb, pages = spec["layout"], spec["tokens_per_block"], spec["pages"] + + def pages_for(seq_len): + return math.ceil(seq_len / tpb) + + if layout in ("random", "shuffled", "past_table", "arange"): + 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) + 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 + 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]) + req = torch.cat([torch.full((n,), i, device=DEV) for i, (_, n) in enumerate(segments)]) + mapped = [pages_for(s + n) for s, n in segments] + assert max(mapped) <= pages + return positions, req, mapped + + +def _build(spec, seed): + """Real paged state for one spec, keyed by `__call__`'s parameter names.""" + gen = torch.Generator(device=DEV).manual_seed(seed) + heads, head_dim, dtype = spec["heads"], spec["head_dim"], spec["dtype"] + positions, req, mapped = _rows_and_requests(spec, gen) + rows, slots = positions.numel(), spec["slots"] + assert len(mapped) <= slots + + slot_of = torch.randperm(slots, generator=torch.Generator().manual_seed(seed))[: len(mapped)] + pool_pages = sum(mapped) + 3 + page_ids = torch.randperm(pool_pages, generator=gen, device=DEV).to(torch.int32) + block_table = torch.full((slots, spec["pages"]), -1, dtype=torch.int32, device=DEV) + cursor = 0 + for slot, count in zip(slot_of.tolist(), mapped): + block_table[slot, :count] = page_ids[cursor : cursor + count] + cursor += count + index_cache = torch.randn( + pool_pages, spec["tokens_per_block"], 1, head_dim, generator=gen, device=DEV + ).to(dtype) + + if spec["q_layout"] == "scratch_slice": + # [T, index heads | 1 k | 1 v, D], q is the first `heads` slots. + scratch = torch.randn(rows, heads + 2, head_dim, generator=gen, device=DEV).to(dtype) + q = scratch[:, :heads] + 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 + assert not q.is_contiguous() + + return dict( + q=q, + index_cache=index_cache, + block_table=block_table, + query_positions=positions.to(spec["positions_dtype"]), + request_indices=slot_of.to(DEV)[req].to(spec["request_dtype"]), + tokens_per_block=spec["tokens_per_block"], + compress_ratio=spec["compress_ratio"], + only_visible_blocks=spec["only_visible_blocks"], + context_rows=spec["context_rows"], + ) + + +def _branch_rule(rows, spec): + heads = spec["heads"] + packs = DOT_TILE_WIDTH % heads == 0 and DOT_TILE_WIDTH // heads > 1 + tiled = spec["context_rows"] and rows > DECODE_ROW_LIMIT and packs + return "query_tile" if tiled else "row_wise" + + +def _columns(args): + return args["block_table"].shape[1] * args["tokens_per_block"] // args["compress_ratio"] + + +def _visible(args): + blocks = torch.arange(_columns(args), device=DEV) + bound = (args["query_positions"].long() + 1) // args["compress_ratio"] + return blocks[None, :] < bound[:, None] + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 + certifying(op) + args = _build(cell.spec, seed=sum(map(ord, cell.why))) + rows = args["q"].shape[0] + assert cell.spec["branch"] == _branch_rule(rows, cell.spec), "the cell names the wrong kernel" + with _watch_launches() as watcher: + out = op(**args) + torch.cuda.synchronize() + assert watcher.branch == cell.spec["branch"] + assert out.shape == (rows, _columns(args)) and out.dtype == torch.float32 + _compare(out, op.reference(**args)) + + +@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}) + 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() + graph = torch.cuda.CUDAGraph() + try: + with torch.cuda.graph(graph): + captured = op(**args) + 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) + + +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): + 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 new file mode 100644 index 000000000000..a8c3b5b8fd53 --- /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_index_store.py @@ -0,0 +1,387 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_paged_index_store catalog entry. + +The paged state is real: an index-key plane and a position plane of `pages` +pages, and the engine's `[slots, width]` table -- more slots than requests, +each live row's pages drawn from a seeded permutation of the pool, `-1` past +its allocation, rows past the batch holding a previous step's plausible page +ids. 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. +""" + +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_index_store import ( # noqa: E501 + qsa_paged_index_store as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "qsa_paged_index_store requires a CUDA device" + + +def _compare(out, ref) -> None: + """Bitwise over both whole planes, as raw bits. + + A store is a move with one conversion that `reference` performs too, so + there is nothing to spend a tolerance on; comparing every slot rather + than the written ones is the stray-write check. + """ + for name, got, want in zip(("index_cache", "position_cache"), out, ref): + assert got.dtype == want.dtype, (name, got.dtype, want.dtype) + assert got.shape == want.shape, (name, got.shape, want.shape) + torch.testing.assert_close( + _bits(got), _bits(want), rtol=0, atol=0, msg=lambda m, n=name: f"{n}: {m}" + ) + + +DEV = "cuda" +KEY_FILL = -7.0 +POS_FILL = -7 +ENGINE_MAX_NUM_TOKENS = 8192 + + +def _bits(tensor: torch.Tensor) -> torch.Tensor: + as_int = {1: torch.uint8, 2: torch.int16, 4: torch.int32}[tensor.element_size()] + return tensor.contiguous().view(as_int) + + +def _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) + tpb, hkv, head_dim, axes = (spec[k] for k in ("tokens_per_block", "hkv", "head_dim", "axes")) + slots, width, pages = spec["slots"], spec["width"], spec["pages"] + 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)] + rows = len(positions) + + table = torch.randint(0, pages, (slots, width), generator=gen, dtype=torch.int32) + order = torch.randperm(pages, generator=gen).to(torch.int32) + used = 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 + assert used <= pages and len(segments) < slots, "pool or table too small for the batch" + 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() + + 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() + + 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() + 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) + args = dict( + token_k=token_k, + position_coordinates=coords, + index_cache=index_cache, + position_cache=position_cache, + request_indices=request_indices, + logical_positions=logical_positions, + block_table=table, + tokens_per_block=tpb, + ) + return args, extras + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args, 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) + ) + } + 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" + + +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 test_cuda_graph_capture_and_replay() -> None: + """Captured once, replayed against new positions written into the same buffer.""" + args = _small(seed=52) + positions = args["logical_positions"] + replayed = positions.clone() + positions.zero_() + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(3): + op(**args) + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(**args) + + positions.copy_(replayed) + args["index_cache"].fill_(KEY_FILL) + args["position_cache"].fill_(POS_FILL) + ref = op.reference(**args) + graph.replay() + torch.cuda.synchronize() + _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.""" + 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) 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 new file mode 100644 index 000000000000..2a1fe4cc9440 --- /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.py @@ -0,0 +1,269 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_paged_kv_store catalog entry. + +The paged state is real: one `[pages, 2, Hkv, tpb, D]` HND pool whose K and V +role views are the caches, and the engine's `[slots, width]` table -- a row +per request slot, more slots than requests, each live row's pages drawn from a +seeded permutation of the pool, `-1` past its allocation, and the rows past +the batch holding a previous step's plausible page ids. Pool sizes are the old +receipts'; the pool is pre-filled with a sentinel so a stray write shows. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_kv_store import ( # noqa: E501 + qsa_paged_kv_store as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "qsa_paged_kv_store requires a CUDA device" + + +def _compare(out, ref) -> None: + """Bitwise over both role planes, compared as raw bits. + + A store is a move with one conversion, and `reference` performs the + same conversion, so there is no rounding to spend a tolerance on. In + the target both planes together are the whole pool allocation, so this + is also the stray-write check. + """ + for name, got, want in zip(("k_cache", "v_cache"), out, ref): + assert got.dtype == want.dtype, (name, got.dtype, want.dtype) + assert got.shape == want.shape, (name, got.shape, want.shape) + torch.testing.assert_close( + _bits(got), _bits(want), rtol=0, atol=0, msg=lambda m, n=name: f"{n}: {m}" + ) + + +DEV = "cuda" +FILL = -2.0 + + +def _bits(tensor: torch.Tensor) -> torch.Tensor: + as_int = {1: torch.uint8, 2: torch.int16, 4: torch.int32}[tensor.element_size()] + return tensor.contiguous().view(as_int) + + +def _hole(spec, table, 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 _build(spec, seed: int): + """Real pool, table and rows for one cell; returns the op's kwargs and the pool.""" + gen = torch.Generator(device="cpu").manual_seed(seed) + tpb, hkv, head_dim = spec["tokens_per_block"], spec["hkv"], spec["head_dim"] + slots, width, pages = spec["slots"], spec["width"], spec["pages"] + segments = spec["segments"] + requests = [r for r, (_, n) in enumerate(segments) for _ in range(n)] + positions = [p for start, n in segments for p in range(start, start + n)] + + table = torch.randint(0, pages, (slots, width), generator=gen, dtype=torch.int32) + order = torch.randperm(pages, generator=gen).to(torch.int32) + used, allocated_cols = 0, [] + for r, (start, n) in enumerate(segments): + cols = (start + n - 1) // tpb + 1 + table[r] = -1 + table[r, :cols] = order[used : used + cols] + used += cols + allocated_cols.append(cols) + assert used <= pages and len(segments) < slots, "pool or table too small for the batch" + if spec["hole"] is not None: + request, position = _hole(spec, table, len(segments), 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 + ).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 + + 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"] + ) + args = dict( + k=k, + v=v, + k_cache=pool[:, 0], + v_cache=pool[:, 1], + request_indices=torch.tensor(requests, dtype=spec["request_dtype"], device=DEV), + logical_positions=torch.tensor(positions, dtype=spec["position_dtype"], device=DEV), + block_table=table, + tokens_per_block=tpb, + ) + return args, pool + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args, pool = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + before = pool.clone() + ref = op.reference(**args) + out = op(**args) + assert out[0] is args["k_cache"] and out[1] is args["v_cache"] + _compare(out, ref) + wrote = not torch.equal(_bits(pool), _bits(before)) + assert wrote == bool(cell.spec["segments"]), ( + "a cell with rows must write, an empty one must not" + ) + + +def _decode_args(seed: int, **overrides): + spec = op.CELLS[0].spec + 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 + + +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) + positions = args["logical_positions"] + saved = positions.clone() + positions.zero_() + + def call(): + op(**args) + + for _ in range(3): # compile the Triton kernel before capture + call() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + call() + + pool.fill_(FILL) + positions.copy_(saved) + ref = op.reference(**args) + graph.replay() + torch.cuda.synchronize() + _compare((args["k_cache"], args["v_cache"]), ref) + + +def test_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. + + The kernel receives strides for k, v, both caches and the table, but reads + the two metadata vectors as flat unit-stride buffers. A column of a 2-wide + buffer is read through its base pointer: accepted, and rows land elsewhere. + """ + args, pool = _decode_args(seed=10) + requests, positions = args["request_indices"], args["logical_positions"] + 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) + + 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" + ) + assert not torch.equal(_bits(out[0]), _bits(ref[0])), "the reference skips that row" 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 new file mode 100644 index 000000000000..f7d69fa1df94 --- /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.py @@ -0,0 +1,490 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_paged_sparse_gqa catalog entry. + +The paged state is real: one `[pages, 2, kv, tpb, D]` HND pool sized from the +pages its requests map, and the engine's slot table -- a row per request slot, +more slots than requests, assigned by a seeded permutation, `-1` past each +sequence. Selection rows are built the way QSA's top-k + expansion lay them +out: complete groups expanded to tokens, then the trailing partial group, then +`-1`, all packed at the front. + +The call has two kernels behind it (and a merge). The `launches` fixture +records which ones a call indexed with a grid, so every cell asserts its +branch. +""" + +import math +import re + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_paged_sparse_gqa import ( # noqa: E501 + qsa_paged_sparse_gqa as op, +) +from tensorrt_llm._torch.attention.backends.sparse.qsa import kernels as _qsa_kernels + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "qsa_paged_sparse_gqa requires a CUDA device" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """4 ulp of each (row, head)'s magnitude, 2 ulp relative RMS, in q's dtype. + + The kernel rounds its online-softmax probabilities to the cache dtype + before P.V, so the error tracks the magnitude of the values averaged, + not of the output: against an fp32 reference it measured 1.2e-3 on rows + averaging a thousand tokens (outputs 0.14 to 0.4) and 1.2e-2 on rows + averaging a handful (outputs to 5). Scaled by each head row's own + magnitude that is one band. Against this reference (fp32, rounded once) + every cell measured at most 2.0 element ulp -- a single final-rounding + flip, since bf16's step is 2 of these units at the bottom of a binade -- + and at most 0.62 ulp RMS: 2x and 3.2x margin. Six wrong references (no + scale, interleaved heads, padding attended, identity pages, K/V + swapped, page/slot transposed) are each more than 0.1 off and rejected. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + if out.numel() == 0: + return + assert_within_ulp(out, ref, element_ulp=4.0, rms_ulp=2.0) + + +DEV = "cuda" +ENGINE_MAX_SEQ_LEN = 8192 + +_KERNELS = { + "_qsa_paged_sparse_gqa_kernel": "prefill", + "_qsa_paged_sparse_gqa_splitk_kernel": "splitk", + "_qsa_merge_splitk_kernel": "merge", +} + + +class _LaunchRecorder: + """Records a Triton kernel's grid indexing, then delegates to the kernel.""" + + def __init__(self, inner, name, log): + self._inner, self._name, self._log = inner, name, log + + def __getitem__(self, grid): + self._log.append((self._name, tuple(grid))) + return self._inner[grid] + + +@pytest.fixture +def launches(monkeypatch): + """Which branch kernels the next call launches, in order, with grids.""" + log: list = [] + for name, tag in _KERNELS.items(): + inner = getattr(_qsa_kernels, name) + monkeypatch.setattr(_qsa_kernels, name, _LaunchRecorder(inner, tag, log)) + return log + + +# -------------------------------------------------------------------------- +# building real state from a cell spec +# -------------------------------------------------------------------------- + + +def _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. + + A random `min(available, block_topk)`-subset of the complete visible + groups, each expanded to its `compress_ratio` tokens, then the incomplete + trailing group, then `-1` -- packed at the front, as the producer emits. + """ + block_topk = token_topk // compress_ratio + final_topk = token_topk + compress_ratio - 1 + gen = torch.Generator(device=DEV).manual_seed(seed) + rows = positions.numel() + visible = positions.long() + 1 + available = visible // compress_ratio + taken = available.clamp(max=block_topk) + width = max(block_topk, int(available.max().item()) if rows else 0) + + scores = torch.rand(rows, width, generator=gen, device=DEV) + group_ids = torch.arange(width, device=DEV)[None, :] + scores = torch.where(group_ids < available[:, None], scores, torch.inf) + chosen = scores.argsort(dim=1)[:, :block_topk] + + column = torch.arange(final_topk, device=DEV)[None, :] + group_of_column = (column // compress_ratio).clamp(max=block_topk - 1).expand(rows, -1) + body = torch.gather(chosen, 1, group_of_column) * compress_ratio + column % compress_ratio + in_body = column < (taken * compress_ratio)[:, None] + tail_column = column - (taken * compress_ratio)[:, None] + tail = (available * compress_ratio)[:, None] + tail_column + in_tail = ~in_body & (tail_column >= 0) & (tail < visible[:, None]) + selected = torch.where(in_body, body, torch.where(in_tail, tail, torch.full_like(body, -1))) + return selected.to(torch.int32) + + +def _build(spec, seed): + """Real paged state for one spec, keyed by `__call__`'s parameter names.""" + gen = torch.Generator(device=DEV).manual_seed(seed) + tpb, kv, d, q_heads = ( + spec["tokens_per_block"], + spec["kv_heads"], + spec["head_dim"], + spec["q_heads"], + ) + segments, generation = _segments(spec) + 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] + 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] + 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) + + 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 + + args = dict( + q=q, + k_cache=k_cache, + v_cache=v_cache, + block_table=block_table, + selected_tokens=selected, + request_indices=request_indices, + tokens_per_block=tpb, + softmax_scale=d**-0.5, + ) + if spec["prefill_pair"]: + args.update(query_positions=positions, compress_ratio=spec["compress_ratio"]) + return args, generation + + +def _bound(positions, spec): + """The prefill kernel's per-row scan bound, from the positions alone.""" + cr = spec["compress_ratio"] + visible = positions.long() + 1 + return (visible // cr).clamp(max=spec["token_topk"] // cr) * cr + visible % cr + + +def _is_prefix_packed(table): + valid = table >= 0 + columns = torch.arange(table.shape[1], device=table.device)[None, :] + return (valid == (columns < valid.sum(1)[:, None])).all(1) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying, launches) -> None: # noqa: F811 + certifying(op) + spec = cell.spec + args, _ = _build(spec, seed=sum(map(ord, cell.why))) + rows = args["q"].shape[0] + if spec["prefill_pair"] and not spec.get("holes"): + # The precondition the op cannot check: front-packed, and exactly as + # many valid entries as the prefill kernel will scan. + selected = args["selected_tokens"] + assert bool(_is_prefix_packed(selected).all()) + valid = (selected >= 0).sum(1) + assert torch.equal(valid, _bound(args["query_positions"], spec)) + + out = op(**args) + assert [name for name, _ in launches] == list(spec["launches"]) + assert launches[0][1][:2] == (rows, spec["kv_heads"]) + assert out.shape == (rows, spec["q_heads"], spec["head_dim"]) + assert out.dtype == args["q"].dtype and out.is_contiguous() + assert not bool(out.isnan().any()) + _compare(out, op.reference(**args)) + + +# -------------------------------------------------------------------------- +# the prefill kernel's bound +# -------------------------------------------------------------------------- + +#: The first cell is this checkpoint's own decode geometry; the probes below +#: start from it. +CKPT = op.CELLS[0].spec +assert (CKPT["q_heads"], CKPT["kv_heads"], CKPT["head_dim"]) == (24, 2, 256) + + +def _with(**overrides): + return {**CKPT, **overrides} + + +def _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) + 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) + args["selected_tokens"] = args["selected_tokens"][:, :64].contiguous() + mutate, message = _REJECTIONS[case] + 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_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 new file mode 100644 index 000000000000..05b1b5a03d10 --- /dev/null +++ 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 @@ -0,0 +1,374 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_prefill_compress catalog entry. + +The paged state is real: an index-key pool of random raw keys, a position +plane holding `(t, h, w)` for every allocated position, and the engine's +`[slots, width]` table -- more slots than requests, each live row's pages drawn +from a seeded permutation of the pool, `-1` past its allocation, rows past the +batch holding a previous step's plausible page ids. The cos|sin table is built +the way the target builds it at load. + +A padding row at logical position -1 is deliberately absent: it is an illegal +memory access that poisons the process for every test after it. +""" + +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_prefill_compress import ( # noqa: E501 + qsa_prefill_compress as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +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. + """ + 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) + + +DEV = "cuda" +ROTARY_DIM = 64 +THETA = 1.0e7 +MAX_POSITIONS = 16384 + + +def _cos_sin() -> torch.Tensor: + """The target's `[max_positions, 2 * pairs]` float32 cos|sin table.""" + exponent = torch.arange(0, ROTARY_DIM, 2, dtype=torch.float64, device=DEV) + inv_freq = (1.0 / (THETA ** (exponent / ROTARY_DIM))).to(torch.float32) + positions = torch.arange(MAX_POSITIONS, dtype=torch.float32, device=DEV) + angles = positions[:, None] * inv_freq[None, :] + return torch.cat([torch.cos(angles), torch.sin(angles)], dim=1).contiguous() + + +def _build(spec, seed: int): + """Real caches, table and rows for one cell; returns the op's kwargs and extras.""" + gen = torch.Generator(device="cpu").manual_seed(seed) + tpb, head_dim = spec["tokens_per_block"], 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) + index_cache = torch.randn(pages, tpb, 1, head_dim, generator=gen).to(torch.bfloat16) + position_cache = torch.zeros(pages, tpb, 3, dtype=torch.int32) + used = 0 + for r, (start, n) in enumerate(segments): + cols = (start + n - 1) // tpb + 1 + table[r] = -1 + 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) + 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()) + + args = dict( + logical_positions=logical.to(DEV), + request_indices=request.to(DEV), + block_table=table, + index_cache=index_cache, + position_cache=position_cache, + k_norm_weight=k_norm_weight.to(DEV), + cos_sin=_cos_sin(), + eps=spec["eps"], + tokens_per_block=tpb, + compress_ratio=spec["compress_ratio"], + mrope_section=spec["mrope_section"], + ) + return args, extras + + +def _anchors(args) -> int: + """Distinct anchor slots the rows name.""" + logical = args["logical_positions"].long() + keyed = args["request_indices"].long() * (1 << 32) + logical + return int(torch.unique(keyed[(logical + 1) % args["compress_ratio"] == 0]).numel()) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + args, extras = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + before = args["index_cache"].clone() + ref = op.reference(**args) + out = op(**args) + assert out is args["index_cache"] + _compare(out, ref) + 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" + + +def _scene(seed: int, **overrides): + """Two requests of 64 tokens from 0 under mRoPE, 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) + + +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) + cache = args["index_cache"] + before = cache.clone() + ref = op.reference(**args) + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(3): + cache.copy_(before) + op(**args) + torch.cuda.current_stream().wait_stream(warmup) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + cache.copy_(before) + with torch.cuda.graph(graph): + op(**args) + for _ in range(2): + cache.copy_(before) + graph.replay() + torch.cuda.synchronize() + _compare(cache, ref) + + +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" + + +def _strided(tensor: torch.Tensor, dim: int) -> torch.Tensor: + """The same values as a stride-2 view along `dim` of a zero-filled buffer.""" + shape = list(tensor.shape) + shape[dim] *= 2 + wide = torch.zeros(shape, dtype=tensor.dtype, device=tensor.device) + index = [slice(None)] * tensor.dim() + index[dim] = slice(None, None, 2) + wide[tuple(index)] = 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. + """ + 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() + + 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}" + + args["index_cache"].copy_(before) + with pytest.raises(AssertionError, match=match): + 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_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 new file mode 100644 index 000000000000..11421ec14395 --- /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_select.py @@ -0,0 +1,233 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qsa_topk_select catalog entry. + +The op selects; it computes nothing. `_compare` therefore checks the output's +structure exactly and the *multiset of selected values* bit for bit, never the +permutation, which the op leaves unspecified. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_select import ( # noqa: E501 + TopkExpectation, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_topk_select import ( # noqa: E501 + qsa_topk_select as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "qsa_topk_select requires a CUDA device" + + +def _compare(out: torch.Tensor, ref: TopkExpectation) -> None: + """Structure exactly, then the selected values bit for bit. + + A selection does no arithmetic, so nothing is tolerated. Values rather + than offsets are compared so that a tied k-th score -- common here: the + relu cells are half zeros -- is a pass rather than a coin flip. + """ + assert out.dtype == torch.int32, out.dtype + assert out.shape == ref.values.shape, (out.shape, ref.values.shape) + index_topk = out.shape[1] + valid = out >= 0 + assert not bool((valid[:, 1:] & ~valid[:, :-1]).any()), "a valid offset follows a -1 pad" + assert bool((out >= -1).all()), "padding must be exactly -1" + assert torch.equal(valid.sum(1), ref.window.clamp(max=index_topk)), ( + "each row must select exactly min(index_topk, window) offsets" + ) + assert bool((out < ref.window[:, None]).all()), "an offset reaches past row_end" + descending = torch.sort(out, dim=1, descending=True).values + assert not bool(((descending[:, :-1] == descending[:, 1:]) & (descending[:, 1:] >= 0)).any()), ( + "duplicate offset within a row" + ) + absolute = (out.long() + ref.row_starts[:, None]).clamp(0, ref.logits.shape[1] - 1) + picked = torch.gather(ref.logits, 1, absolute).masked_fill(~valid, float("-inf")) + got = torch.sort(picked, dim=1, descending=True).values + assert torch.equal(got, ref.values), "selected values are not the window's top-k" + + +DEV = "cuda" +COMPRESS_RATIO = 4 +BLOCK_TOPK = 512 +#: Pre-fill for every output buffer, so a slot the op failed to write shows up +#: as an offset past `row_end` in `_compare`. +SENTINEL = 12345 + + +def _logits(kind: str, rows: int, cols: int, gen: torch.Generator) -> torch.Tensor: + x = torch.randn(rows, cols, generator=gen, device=DEV) + # QSA's scores are relu(q.k) summed over heads: non-negative with ties at 0. + return torch.relu(x) * 3.0 if kind == "relu" else x + + +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) + 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) + + if spec["tail"] == "nan": + columns = torch.arange(cols, device=DEV)[None, :] + outside = (columns < row_starts[:, None]) | (columns >= row_ends[:, None]) + logits.masked_fill_(outside, float("nan")) + indices = torch.full((rows, index_topk), SENTINEL, dtype=torch.int32, device=DEV) + return logits, row_starts, row_ends, indices, index_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 + certifying(op) + args = _build(cell.spec, seed=sum(map(ord, cell.why))) + out = op(*args) + assert out is args[3], "the op writes `indices` and returns it" + _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.""" + gen = torch.Generator(device=DEV).manual_seed(8) + rows, cols = 8, 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) + indices = torch.empty(rows, BLOCK_TOPK, dtype=torch.int32, device=DEV) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + op(logits, row_starts, row_ends, indices, BLOCK_TOPK) + torch.cuda.current_stream().wait_stream(warmup) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(logits, row_starts, row_ends, indices, BLOCK_TOPK) + + logits.copy_(_logits("relu", rows, cols, gen)) + row_ends.fill_(1500) + indices.fill_(SENTINEL) + graph.replay() + torch.cuda.synchronize() + _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 + logits = torch.randn(rows, cols, device=DEV) + row_starts = torch.zeros(rows, dtype=torch.int32, device=DEV) + row_ends = torch.full((rows,), cols, dtype=torch.int32, device=DEV) + return logits, row_starts, row_ends, top_k + + +def test_guard_refuses_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 + guard catches it at the call that forgot to pass the block budget.""" + logits, row_starts, row_ends, _ = _layout_case() + indices = torch.empty(logits.shape[0], BLOCK_TOPK, dtype=torch.int32, device=DEV) + 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_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 new file mode 100644 index 000000000000..8843c837443b --- /dev/null +++ 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 @@ -0,0 +1,326 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp vision_rope_2d entry. + +`op.reference` shares no code with the entry: positions from nested loops, +angles in fp64 from `theta ** (-2j / rope_dim)`, rotation as a complex +multiply. 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. +""" + +import zlib +from typing import List, Tuple + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.vision_rope_2d import ( # noqa: E501 + vision_rope_2d as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "vision_rope_2d requires a CUDA device" + + +def _compare( + out: Tuple[torch.Tensor, torch.Tensor], + ref: Tuple[torch.Tensor, torch.Tensor], +) -> None: + """Default band on each of `(q_out, k_out)`; nothing is loosened.""" + for got, want in zip(out, ref, strict=True): + torch.testing.assert_close(got, want) + + +HEADS = 16 +HEAD_DIM = 72 +MERGE = 2 +THETA = 10000.0 + + +def _num_tokens(grid) -> int: + return sum(t * h * w for t, h, w in grid) + + +def _operands(tokens, heads, heads_k, head_dim, dtype, seed, layout="separate"): + g = torch.Generator(device="cuda").manual_seed(seed) + if layout == "packed_qkv": + assert heads == heads_k + qkv = torch.randn(tokens, 3 * heads, head_dim, generator=g, device="cuda").to(dtype) + q, k, _ = torch.split(qkv, [heads] * 3, 1) + assert not q.is_contiguous() and q.stride(-1) == 1 + return q, k + q = torch.randn(tokens, heads, head_dim, generator=g, device="cuda").to(dtype) + k = torch.randn(tokens, heads_k, head_dim, generator=g, device="cuda").to(dtype) + return q, k + + +def _build(spec, seed): + grid = spec["grid"] + heads = spec.get("heads", HEADS) + q, k = _operands( + _num_tokens(grid), + heads, + spec.get("heads_k", heads), + spec.get("head_dim", HEAD_DIM), + spec["dtype"], + seed, + spec.get("layout", "separate"), + ) + grid_thw = torch.tensor( + [list(row) for row in grid], + dtype=spec.get("grid_dtype", torch.int32), + device=spec.get("grid_device", "cpu"), + ) + return q, k, grid_thw, MERGE, spec.get("theta", THETA) + + +def _max_deviation(lhs, rhs) -> float: + return (lhs.float() - rhs.float()).abs().max().item() + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + q, k, grid_thw, merge, theta = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + q_before, k_before = q.clone(), k.clone() + q_out, k_out = op(q, k, grid_thw, merge, theta) + assert q_out.shape == q.shape and q_out.dtype == q.dtype + assert k_out.shape == k.shape and k_out.dtype == k.dtype + assert torch.equal(q, q_before) and torch.equal(k, k_before), "inputs must not be mutated" + assert q_out.data_ptr() != q.data_ptr() and k_out.data_ptr() != k.data_ptr() + q_ref, k_ref = op.reference(q, k, grid_thw, merge, theta) + print( + f"[{cell.why}] q_dev={_max_deviation(q_out, q_ref):.3e} " + f"k_dev={_max_deviation(k_out, k_ref):.3e}" + ) + _compare((q_out, k_out), (q_ref, k_ref)) + + +# --- 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: + grid = [(1, 16, 24), (1, 8, 12)] + q, k = _operands(_num_tokens(grid), HEADS, HEADS, HEAD_DIM, torch.bfloat16, seed=29) + host = torch.tensor(grid, dtype=torch.int32) + q_host, k_host = op(q, k, host, MERGE, THETA) + for other in (host.to("cuda"), host.to(torch.int64), host.to(torch.int64).to("cuda")): + q_other, k_other = op(q, k, other, MERGE, THETA) + assert torch.equal(q_host, q_other) and torch.equal(k_host, k_other) + + +def test_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) + q, k = _operands(16, HEADS, HEADS, HEAD_DIM, torch.bfloat16, seed=5) + _, k_out = op(q, k[:1], grid_thw, MERGE, THETA) + assert k_out.shape == (16, HEADS, HEAD_DIM) + + with pytest.raises(AssertionError, match="must agree on tokens and head_dim"): + with validating(op): + op(q, k[:1], grid_thw, MERGE, THETA) + + +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 new file mode 100644 index 000000000000..e88fb04f3fb5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/embedding/test_modeling_v2_ngram_hash.py @@ -0,0 +1,277 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""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. +""" + +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_CONTEXT_LEN, + CKPT_EOS_TOKEN_ID, + CKPT_HEAD_OFFSETS, + CKPT_HEAD_VOCAB_SIZES, + CKPT_MULTIPLIERS, + CKPT_VOCAB_SIZE, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.embedding.ngram_hash import ( # noqa: E501 + ngram_hash as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "ngram_hash requires a CUDA device" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Bitwise: exact integer arithmetic, no tolerance to loosen.""" + assert out.dtype == ref.dtype == torch.int64, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + assert out.is_contiguous() + assert torch.equal(out, ref), f"{int((out != ref).sum())} indices differ" + + +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 + + +def _constants(): + return ( + torch.tensor(CKPT_MULTIPLIERS, dtype=torch.long, device=DEV), + torch.tensor(CKPT_HEAD_VOCAB_SIZES, dtype=torch.long, device=DEV), + torch.tensor(CKPT_HEAD_OFFSETS, dtype=torch.long, device=DEV), + ) + + +def _tokens(batch, length, seed, eos_positions=(), lead_eos=True, eos_density=0.0): + g = torch.Generator().manual_seed(seed) + ids = torch.randint(0, CKPT_VOCAB_SIZE, (batch, length), generator=g, dtype=torch.long) + if eos_density > 0.0: + ids[torch.rand(batch, length, generator=g) < eos_density] = EOS + if lead_eos: + ids[:, : min(CKPT_CONTEXT_LEN, length)] = EOS + for position in eos_positions: + ids[:, position] = EOS + return ids.to(DEV) + + +def _hash(ids, num_query_tokens): + return op(ids, *_constants(), EOS, num_query_tokens) + + +@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 + ids = _tokens( + spec["batch"], + spec["length"], + seed=zlib.crc32(cell.why.encode()), + 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) + assert out.shape == (spec["batch"], spec["num_query_tokens"], NGRAM_HEADS) + assert out.data_ptr() != ids.data_ptr() + assert torch.equal(ids, before), "the history is read-only" + _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 + fresh = torch.tensor([[EOS] * CKPT_CONTEXT_LEN + [token]], device=DEV) + after_eos = torch.tensor([[11, 22, 33, EOS, token]], device=DEV) + no_eos = torch.tensor([[11, 22, 33, 44, token]], device=DEV) + assert torch.equal(_hash(fresh, 1), _hash(after_eos, 1)) + assert not torch.equal(_hash(fresh, 1), _hash(no_eos, 1)) + + +def test_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() + narrow = multipliers.to(torch.int32) + unguarded = op(ids, narrow, sizes, offsets, EOS, 4) + assert not torch.equal(unguarded, _hash(ids, 4)), "the hazard this guard covers is gone" + assert int(unguarded.min()) >= 0 and int(unguarded.max()) < PADDED_TABLE_ROWS + with pytest.raises(AssertionError, match="int64"): + with validating(op): + op(ids, narrow, sizes, offsets, EOS, 4) + + +def test_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() + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(3): + op(ids, *constants, EOS, 1) + torch.cuda.current_stream().wait_stream(warmup) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(ids, *constants, EOS, 1) + + fresh = _tokens(8, CKPT_CONTEXT_LEN + 1, seed=54, lead_eos=False, eos_density=0.2) + ids.copy_(fresh) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(captured, op.reference(fresh, *constants, EOS, 1)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/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 new file mode 100644 index 000000000000..4fa6444fc13d --- /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_gated_residual.py @@ -0,0 +1,249 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_gated_residual entry. + +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. +""" + +import zlib +from typing import Optional, Tuple + +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, +) +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 + qwen4_exp_gated_residual 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_gated_residual requires a CUDA device" + + +def _compare( + out: Tuple[torch.Tensor, Optional[torch.Tensor]], + ref: Tuple[torch.Tensor, Optional[torch.Tensor]], +) -> None: + """Each result within `rtol = eps`, `atol = eps * max|ref|` of the + activation dtype (`gate_ratio <= 1`). + + The absolute term is not torch's 1e-5: `mixed` is a mean of signed + products, so its near-zero elements are cancellations whose error is + set by the whole stream's rounding, and a relative gate on them gates + nothing. The relative term is half torch's bf16 default. Measured on + sm_103 at 512 tokens: the correct result at 0.341 (`mixed`) and 0.404 + (`inject`) of this band (0.198 / 0.210 against the unrounded float64, + the rest is the reference's own rounding), ten wrong formulas at + 32x..192x. + """ + (mixed, inject), (ref_mixed, ref_inject) = out, ref + assert (inject is None) == (ref_inject is None), "inject must be None exactly when ref's is" + pairs = [(mixed, ref_mixed)] + ([] if inject is None else [(inject, ref_inject)]) + for got, want in pairs: + assert got.dtype == want.dtype, (got.dtype, want.dtype) + assert got.shape == want.shape, (got.shape, want.shape) + assert got.is_contiguous() + ratio = gate_ratio(got, want) + assert ratio <= 1.0, f"{ratio:.4g}x the band" + + +DEV = "cuda" +HC_COUNT = 4 +HIDDEN = 2560 +WIDTH = HC_COUNT * HIDDEN +LOWRANK = 320 +RMS_NORM_EPS = 1e-6 + +# Element statistics of this checkpoint's layer-0 attn_hyper_connection tensors. +# The normed stream is scaled by (1 + gamma), so gamma's tail -- |6.625| -- is +# what puts it at ~+-30 rather than ~+-4, and the error floor scales with it. +CKPT_DOWN_RMS = 0.0296 +CKPT_UP_RMS = 0.0681 +CKPT_INJECT_RMS = 0.0107 +CKPT_GAMMA_RMS = 0.477 +CKPT_GAMMA_ABSMAX = 6.625 + + +def _make_weights(seed, dtype=torch.bfloat16, hc_count=HC_COUNT): + g = torch.Generator(device=DEV).manual_seed(seed) + width = hc_count * HIDDEN + + def normal(*shape, scale): + return (torch.randn(*shape, generator=g, device=DEV) * scale).to(dtype) + + down = normal(LOWRANK, width, scale=CKPT_DOWN_RMS) + up = normal(width, LOWRANK, scale=CKPT_UP_RMS) + inject = normal(hc_count, width, scale=CKPT_INJECT_RMS) + gamma = normal(width, scale=CKPT_GAMMA_RMS) + # A Gaussian does not reproduce the real gamma's ~14-sigma tail; plant it. + gamma[0] = CKPT_GAMMA_ABSMAX + gamma[width // 2] = -CKPT_GAMMA_ABSMAX + return down, up, inject, gamma + + +def _grouped_gemma_rmsnorm(x, gamma, hc_count): + """The hc_norm the entry's input comes from -- here only to give the + operand its real distribution, not under test.""" + hidden = x.shape[-1] // hc_count + xf = x.float().reshape(*x.shape[:-1], hc_count, hidden) + xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + RMS_NORM_EPS) + return (xf.flatten(-2) * (1.0 + gamma.float())).to(x.dtype) + + +def _build(leading, seed, hc_count=HC_COUNT, dtype=torch.bfloat16, inject=True): + down, up, inject_w, gamma = _make_weights(seed, dtype, hc_count) + g = torch.Generator(device=DEV).manual_seed(seed + 1) + raw = torch.randn(*leading, hc_count * HIDDEN, generator=g, device=DEV).to(dtype) + normed = _grouped_gemma_rmsnorm(raw, gamma, hc_count) + return raw, (normed, down, up, inject_w if inject else None, hc_count) + + +@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 + hc_count = spec.get("hc_count", HC_COUNT) + _, args = _build( + spec["leading"], + seed=zlib.crc32(cell.why.encode()), + hc_count=hc_count, + dtype=spec.get("dtype", torch.bfloat16), + inject=spec.get("inject", True), + ) + before = args[0].clone() + mixed, inject = op(*args) + assert mixed.shape == (*spec["leading"], HIDDEN) + if spec.get("inject", True): + assert inject.shape == (*spec["leading"], hc_count) + assert torch.equal(args[0], before), "the normed stream is read-only" + _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.""" + _, (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: + _, (static_in, down, up, inject_w, c) = _build((8,), seed=8) + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(3): + op(static_in, down, up, inject_w, c) + torch.cuda.current_stream().wait_stream(warmup) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_mixed, graph_inject = op(static_in, down, up, inject_w, c) + + _, (fresh, *_) = _build((8,), seed=18) + static_in.copy_(fresh) + graph.replay() + torch.cuda.synchronize() + eager_mixed, eager_inject = op(fresh, down, up, inject_w, c) + assert torch.equal(graph_mixed, eager_mixed) + assert torch.equal(graph_inject, eager_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_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 new file mode 100644 index 000000000000..65853efa4a69 --- /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_post_inject.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp_post_inject entry. + +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. +""" + +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_post_inject import ( # noqa: E501 + qwen4_exp_post_inject 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_post_inject requires a CUDA device" + + +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.""" + 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) + + +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): + """Operands in the ranges the model produces: `inject` is 2*sigmoid(.), so [0, 2].""" + g = torch.Generator(device=DEV).manual_seed(seed) + hyper_input = torch.randn(*leading, hc_count * hidden, generator=g, device=DEV).to(dtype) + block_out = torch.randn(*leading, hidden, generator=g, device=DEV).to(dtype) + inject = (2.0 * torch.sigmoid(torch.randn(*leading, hc_count, generator=g, device=DEV))).to( + dtype + ) + return hyper_input, block_out, inject + + +def _build(spec, seed): + hc_count, hidden = spec["hc_count"], spec["hidden"] + dtype = spec.get("dtype", torch.bfloat16) + 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": + buf = torch.randn(hyper_input.shape[0], 2 * width, device=DEV).to(dtype) + buf[:, :width] = hyper_input + hyper_input = buf[:, :width] + if layout != "contiguous": + assert not hyper_input.is_contiguous() and hyper_input.stride(-1) == 1 + return hyper_input, block_out, 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 — pytest resolves the fixture by this name + certifying(op) + hyper_input, block_out, inject = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + before = hyper_input.clone() + out = op(hyper_input, block_out, inject) + assert out.shape == hyper_input.shape and out.dtype == hyper_input.dtype + assert out.data_ptr() != hyper_input.data_ptr(), "the result is a fresh tensor" + assert torch.equal(hyper_input, before), "the residual it was given is read-only" + _compare(out, op.reference(hyper_input, block_out, inject)) + + +def test_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.""" + 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 + + +def test_cuda_graph_capture_and_replay() -> None: + hyper_input, block_out, inject = _operands((64,), seed=14) + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(3): + op(hyper_input, block_out, inject) + torch.cuda.current_stream().wait_stream(warmup) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(hyper_input, block_out, inject) + + fresh, _, _ = _operands((64,), seed=15) + hyper_input.copy_(fresh) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(captured, op.reference(hyper_input, block_out, inject)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/moe/test_modeling_v2_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 new file mode 100644 index 000000000000..dc92b8dd9c6b --- /dev/null +++ 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 @@ -0,0 +1,512 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the fp8_block_scale_moe_runner catalog entry. + +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. +""" + +import zlib + +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, +) + +assert torch.cuda.is_available(), "fp8_block_scale_moe_runner requires a CUDA device" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """16 ulp element-wise, 0.75 ulp RMS, both in bf16 ulp of each row's max. + + Both statistics normalize each element by its reference row's largest + magnitude: `d = (out - ref) / (rowmax * 2^-8)`, gated on `max|d|` and + on `sqrt(mean(d^2))` -- the RMS of the row-normalized deviation, not + `assert_within_ulp`'s ratio of norms, which reads higher on the same + data (a row's max sits several sigma above its RMS). Kernel and + reference consume bit-identical operands and model both + requantizations, so they differ by accumulation order and by FC1 sums + that land on an e4m3 rounding boundary. + + Measured worst over the cells: 6.65 element-wise (T=4096; the max over + T*H samples grows with T -- 0.91 at T=1) and 0.306 RMS (H=I=512), the + RMS sitting at 0.17-0.31 everywhere; earlier draws reached 10.03 and + 0.342. Wrong references at the checkpoint geometry, RMS: gate/up swap + 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. + """ + assert out.dtype == ref.dtype == torch.bfloat16, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + o, r = out.float(), ref.float() + rowmax = r.abs().amax(dim=-1, keepdim=True).clamp(min=1e-6) + d = (o - r) / (rowmax * 2.0**-8) + elt, rms = d.abs().max().item(), d.pow(2).mean().sqrt().item() + assert elt <= 16.0, f"{elt:.2f} ulp element-wise > 16" + assert rms <= 0.75, f"{rms:.3f} ulp RMS > 0.75" + + +DEV = "cuda" +ULP = 2.0**-8 +CKPT = dict(hidden=2560, inter=640, num_experts=512, top_k=10) + +# ── operands ────────────────────────────────────────────────────────────── + +_WEIGHTS: dict = {} + + +def _weights(hidden: int, inter: int, experts: int, seed: int): + """`(w1, w1s, w2, w2s)` drawn directly in e4m3 / fp32, so every value is exact. + + Cached: the checkpoint stack is 2.5 GB and every checkpoint-geometry cell + reads the same one. Drawn 32 experts at a time to cap the fp32 transient. + """ + key = (hidden, inter, experts, seed) + if key not in _WEIGHTS: + g = torch.Generator(device=DEV).manual_seed(seed) + w1 = torch.empty(experts, 2 * inter, hidden, dtype=torch.float8_e4m3fn, device=DEV) + w2 = torch.empty(experts, hidden, inter, dtype=torch.float8_e4m3fn, device=DEV) + for lo in range(0, experts, 32): + hi = min(lo + 32, experts) + w1[lo:hi] = torch.randn(hi - lo, 2 * inter, hidden, device=DEV, generator=g) + w2[lo:hi] = torch.randn(hi - lo, hidden, inter, device=DEV, generator=g) + w1s = 0.01 + 0.02 * torch.rand( + experts, 2 * inter // BLOCK, hidden // BLOCK, device=DEV, generator=g + ) + w2s = 0.01 + 0.02 * torch.rand( + experts, hidden // BLOCK, inter // BLOCK, device=DEV, generator=g + ) + _WEIGHTS[key] = (w1, w1s, w2, w2s) + return _WEIGHTS[key] + + +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, +): + """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) + 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. + hs_scale = 0.5 + torch.rand(hidden // BLOCK, tokens, device=DEV, generator=g) + 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(), + 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, + 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, + ) + + +def _ulp(out: torch.Tensor, ref: torch.Tensor): + """(max, RMS) deviation in bf16 ulp of the row max -- `_compare`'s statistics.""" + o, r = out.float(), ref.float() + d = (o - r) / (r.abs().amax(dim=-1, keepdim=True).clamp(min=1e-6) * ULP) + 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()) + + +# ── the certified cells ─────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + args = _args(**cell.spec, seed=_seed(cell)) + out = op(**args) + assert out.shape == args["hidden_states"].shape and out.dtype == torch.bfloat16 + ref = op.reference(**args) + elt, rms = _ulp(out, ref) + print(f" {cell.why[:44]}: {elt:.2f} ulp element-wise, {rms:.3f} ulp RMS") + _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" + + +def test_autotuned_tactic_matches_the_fallback_tactic(): + """A warm profiling cache changes the tactic, not the result.""" + from tensorrt_llm._torch.autotuner import AutoTuner, autotune + + a = _args(256, **CKPT, seed=42) + AutoTuner.get().clear_cache() + try: + cold = op(**a) + with autotune(): + op(**a) + entries = len(AutoTuner.get().profiling_cache) + warm = op(**a) + print(f" warming wrote {entries} cache entries") + assert entries > 0, "warming profiled nothing -- the comparison would be vacuous" + assert torch.equal(warm, cold), "a tuned tactic changed the result" + _compare(warm, op.reference(**a)) + finally: + 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 ──────────────────────────────────────────────────────────── + + +def _strided(t: torch.Tensor) -> torch.Tensor: + wide = torch.zeros(*t.shape[:-1], 2 * t.shape[-1], device=DEV, dtype=t.dtype) + wide[..., ::2] = t + view = wide[..., ::2] + assert not view.is_contiguous() + return view + + +def test_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(): + 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) 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 new file mode 100644 index 000000000000..5aaa3b5107d6 --- /dev/null +++ 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 @@ -0,0 +1,221 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the fused_sigmoid_gate_mul_add catalog entry. + +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. +""" + +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.fused_sigmoid_gate_mul_add import ( # noqa: E501 + fused_sigmoid_gate_mul_add as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "fused_sigmoid_gate_mul_add requires a CUDA device" + +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.""" + t, h, dt = spec["tokens"], spec["hidden"], spec["dtype"] + 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 + 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": + 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 + + +# ── the certified cells ─────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + final, gate, shared, output, buffers = _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) + + out = op(final, gate, shared, output) + + destination = final if output is None else output + assert out.data_ptr() == destination.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) + gate = torch.empty(0, 1, dtype=torch.bfloat16, device=DEV) + returned = op(final, gate, shared) + assert returned.shape == (0, HIDDEN) and returned.data_ptr() == final.data_ptr() + + +def test_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.""" + 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] + 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) 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 new file mode 100644 index 000000000000..fcf05d43e29e --- /dev/null +++ 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 @@ -0,0 +1,305 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the renorm_moe_routing_op catalog entry. + +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. +""" + +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.renorm_moe_routing_op import ( # noqa: E501 + renorm_moe_routing_op as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "renorm_moe_routing_op requires a CUDA device" + + +def _compare( + out: tuple[torch.Tensor, torch.Tensor], + ref: tuple[torch.Tensor, torch.Tensor], +) -> None: + """Ids equal and weights bit-exact. + + The reference is not an approximation of the kernel: over 2040 + configurations and 16384-token runs at 512 / top-10, not one id and not + one weight bit differed. A one-ulp `assert_close` runs first only so a + drift reports a readable diff. The nearest wrong answers sit 0.4476 + away (un-renormalized softmax, `1/topk`), and a `torch.topk` reference + differs on a third of the rows at 512 bf16 experts. + """ + ids, weights = out + ref_ids, ref_weights = ref + assert ids.dtype == ref_ids.dtype == torch.int32, (ids.dtype, ref_ids.dtype) + assert weights.dtype == ref_weights.dtype, (weights.dtype, ref_weights.dtype) + assert ids.shape == ref_ids.shape == weights.shape, (ids.shape, ref_ids.shape) + torch.testing.assert_close(ids, ref_ids, rtol=0, atol=0) + eps = torch.finfo(weights.dtype).eps + torch.testing.assert_close(weights, ref_weights, rtol=eps, atol=0.0) + mismatched = int((weights != ref_weights).sum()) + assert mismatched == 0, f"{mismatched} of {weights.numel()} weights are not bit-exact" + + +DEV = "cuda" + + +def _logits(tokens, num_experts, logits_dtype, dist="randn", magnitude=1.0, seed=0): + g = torch.Generator(device=DEV).manual_seed(seed) + if dist == "coarse": + x = torch.randint(0, 3, (tokens, num_experts), device=DEV, generator=g).float() + else: + x = torch.randn(tokens, num_experts, device=DEV, generator=g) * magnitude + return x.to(logits_dtype) + + +def _check(logits, topk, output_dtype=None): + """Drive the op and gate it through the entry's own reference and `_compare`.""" + ids, weights = op(logits, topk, output_dtype) + assert ids.is_contiguous() and weights.is_contiguous() + assert ids.device == weights.device == logits.device + _compare((ids, weights), op.reference(logits, topk, output_dtype)) + return ids, weights + + +# ── the certified cells ─────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + s = cell.spec + logits = _logits( + s["tokens"], + s["num_experts"], + s["logits_dtype"], + s["dist"], + s.get("magnitude", 1.0), + seed=zlib.crc32(cell.why.encode()), + ) + ids, weights = _check(logits, s["topk"], s["output_dtype"]) + assert weights.dtype == s["output_dtype"] + # renormalized: every row sums to 1 + torch.testing.assert_close( + weights.float().sum(-1), + torch.ones(s["tokens"], device=DEV), + rtol=8e-3 if s["output_dtype"] is torch.bfloat16 else 1e-6, + atol=0.0, + ) + + +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)) + + +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 + torch.testing.assert_close( + other_weights / row_sums.unsqueeze(-1), weights, rtol=1e-6, atol=1e-7 + ) + + +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()) + + +def test_guards_refuse(): + 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) 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 new file mode 100644 index 000000000000..a6491f715bd3 --- /dev/null +++ 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 @@ -0,0 +1,126 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the flashinfer_gemma_rmsnorm entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.norm.flashinfer_gemma_rmsnorm import ( # noqa: E501 + flashinfer_gemma_rmsnorm as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "flashinfer_gemma_rmsnorm requires a CUDA device" + + +def _seed(text: str) -> int: + return zlib.crc32(text.encode()) + + +def _build(spec, seed): + """`(buffer, x, weight, eps)`: `x` is a view of `buffer` laid out as the spec says.""" + g = torch.Generator(device="cuda").manual_seed(seed) + shape, dtype = spec["shape"], spec["dtype"] + width = shape[-1] + + def randn(*s): + return torch.randn(*s, generator=g, device="cuda").to(dtype) + + if "row_width" in spec: + 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 + ) + return buffer, x, weight, spec["eps"] + + +@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, x, weight, eps = _build(cell.spec, seed=_seed(cell.why)) + assert tuple(x.shape) == cell.spec["shape"] and x.stride(-1) == 1 + if "row_width" in cell.spec and x.shape[0] > 1: + assert x.stride(0) == cell.spec["row_width"] and not x.is_contiguous() + buffer_before = buffer.clone() + + out = op(x, weight, eps) + + assert out.shape == x.shape and out.dtype == x.dtype + assert out.data_ptr() != x.data_ptr(), "the op must return a new tensor" + assert torch.equal(buffer, buffer_before), "x (and the stream it slices) is read-only" + torch.testing.assert_close(out, op.reference(x, weight, eps)) + + +def test_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): + 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 new file mode 100644 index 000000000000..e297e7674808 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/norm/test_modeling_v2_layer_norm.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 qwen4_exp layer_norm entry. + +Weight and bias are drawn over this checkpoint's own 110 vision LayerNorm +tensors' ranges: the affine is wide enough (weights to 17.5, biases to -25.6) +that a wrong normalization is not masked by a near-unit scale. +""" + +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.norm.layer_norm import ( # noqa: E501 + layer_norm as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "layer_norm requires a CUDA device" + +DEVICE = torch.device("cuda") +WIDTH = 1152 +EPS = 1e-6 + +#: Over all of model.visual.*.norm{1,2} and model.visual.merger.norm (bf16). +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) + + +def _affine(gen, width=WIDTH, dtype=torch.bfloat16): + weight = torch.rand(width, generator=gen, device=DEVICE) * (WEIGHT_HI - WEIGHT_LO) + WEIGHT_LO + bias = torch.rand(width, generator=gen, device=DEVICE) * (BIAS_HI - BIAS_LO) + BIAS_LO + return weight.to(dtype), bias.to(dtype) + + +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 + + +@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] + 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] + 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) + x = torch.randn(8, WIDTH, generator=gen, device=DEVICE).to(torch.bfloat16) + weight, bias = _affine(gen) + with pytest.raises(AssertionError): + op(x, weight[:-1].contiguous(), bias, EPS) + 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 new file mode 100644 index 000000000000..381ef675e1fe --- /dev/null +++ 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 @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the rms_norm_gated_token_major entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.norm.rms_norm_gated_token_major import ( # noqa: E501 + rms_norm_gated_token_major as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "rms_norm_gated_token_major requires a CUDA device" + + +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 GDN in_proj row the old column-slice measurement used: [q k v | z] = 10240 | 6144. +_PROJ_WIDTH, _Z_OFFSET = 16384, 10240 + + +def _seed(text: str) -> int: + return zlib.crc32(text.encode()) + + +def _build(spec, seed): + """`__call__`'s positional arguments for one spec.""" + g = torch.Generator(device="cuda").manual_seed(seed) + t, heads, n, dtype = spec["tokens"], spec["heads"], spec["width"], spec["dtype"] + + 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"] + + 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] + else: + 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"] + + +@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() + 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) + 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") + ) + + +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) + 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 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) 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 new file mode 100644 index 000000000000..e08dc99a1dfe --- /dev/null +++ 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 @@ -0,0 +1,293 @@ +# 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. + +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 +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.gated_ngram_injection import ( # noqa: E501 + gate_ratio, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.gated_ngram_injection import ( # noqa: E501 + gated_ngram_injection as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "gated_ngram_injection requires a CUDA device" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """16-bit: `rtol = eps`, `atol` one subnormal spacing (`gate_ratio <= 1`) + -- half torch's bf16 `rtol`. Nothing cancels (each element is a sigmoid + times one `value` element), so the budget is the final rounding; with + the reference also rounded once, op and reference agree or sit one ulp + apart. One ulp is at most `eps * |ref|`, so this band always admits it + and never two. Measured at 2048 tokens: 0.992 (a one-ulp element just + above a binade boundary), eight wrong readings at 105x..1115x. The + `atol` is for Gaussian `value` elements that land subnormal (measured + on fp16). + + fp32: torch's default band. The 2560-term signed dot is accumulated at + the width it is reported in, so its conditioning, not the cast, sets + the error; a custom pair would be fitted to a result. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + assert out.is_contiguous() + if ref.dtype == torch.float32: + torch.testing.assert_close(out, ref) + return + ratio = gate_ratio(out, ref) + assert ratio <= 1.0, f"{ratio:.4g}x the band" + + +DEV = "cuda" +HC_COUNT = 4 +HIDDEN = 2560 +WIDTH = HC_COUNT * HIDDEN +RMS_NORM_EPS = 1e-6 +CLAMP_MIN = 1e-6 + +# This checkpoint's PLE tensors: norm_key / norm_query gammas (element RMS +# 0.1358, extremes 0.863 and -0.326), and the value's scale derived as +# sqrt(2560) * 0.015295 (value_proj RMS) * 0.00747 (dequantized table RMS). +# The entry is exactly linear in `value`, so its absolute scale decides nothing. +CKPT_NORM_GAMMA_RMS = 0.1358 +CKPT_NORM_GAMMA_MAX = 0.863281 +CKPT_NORM_GAMMA_MIN = -0.326172 +CKPT_VALUE_RMS = 0.0058 + + +def _gemma_grouped_rmsnorm(x, gamma, hc_count): + """Where key/query come from -- here only for their distribution.""" + hidden = x.shape[-1] // hc_count + xf = x.float().reshape(*x.shape[:-1], hc_count, hidden) + xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + RMS_NORM_EPS) + return (xf.flatten(-2) * (1.0 + gamma.float())).to(x.dtype) + + +def _make_operands(seed, dtype, leading, hc_count=HC_COUNT, value_dtype=None): + g = torch.Generator(device=DEV).manual_seed(seed) + width = hc_count * HIDDEN + + def normal(*dims, scale, dt=dtype): + return (torch.randn(*dims, generator=g, device=DEV) * scale).to(dt) + + gammas = [] + for _ in range(2): + gamma = normal(width, scale=CKPT_NORM_GAMMA_RMS) + gamma[0] = CKPT_NORM_GAMMA_MAX + gamma[width // 2] = CKPT_NORM_GAMMA_MIN + gammas.append(gamma) + key = _gemma_grouped_rmsnorm(normal(*leading, width, scale=1.0), gammas[0], hc_count) + query = _gemma_grouped_rmsnorm(normal(*leading, width, scale=1.0), gammas[1], hc_count) + value = normal(*leading, HIDDEN, scale=CKPT_VALUE_RMS, dt=value_dtype or dtype) + return key, query, value + + +def _single_channel_operands(dots, value_dtype): + """One nonzero channel per stream: each dot is exactly `dots[s]`.""" + hc_count = len(dots) + key = torch.zeros(8, hc_count * HIDDEN, dtype=torch.bfloat16, device=DEV) + query = torch.zeros_like(key) + for stream, dot in enumerate(dots): + key[:, stream * HIDDEN] = 1.0 + query[:, stream * HIDDEN] = dot + value = torch.full((8, HIDDEN), 0.5, dtype=value_dtype, device=DEV) + return key, query, value + + +def _build(spec, seed): + if spec.get("kind") == "single_channel": + dots = spec["dots"] + key, query, value = _single_channel_operands(dots, spec["value_dtype"]) + return key, query, value, len(dots), spec["clamp_min"] + 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") + ) + return key, query, value, hc_count, spec.get("clamp_min", CLAMP_MIN) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + before = [a.clone() for a in args[:3]] + out = op(*args) + key, _, value, hc_count, _ = args + assert out.shape == (*key.shape[:-1], hc_count * value.shape[-1]) + assert out.dtype == value.dtype + for a, b in zip(args[:3], before): + assert torch.equal(a, b), "the operands are read-only" + _compare(out, op.reference(*args)) + + +def test_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,)) + unguarded = op(key, query[:1], value, HC_COUNT) + assert unguarded.shape == (64, WIDTH), "the hazard this guard covers is gone" + assert not torch.equal(unguarded, op.reference(key, query, value, HC_COUNT)) + assert op(key, query, value[:1], HC_COUNT).shape == (64, WIDTH) + + with pytest.raises(AssertionError, match="broadcasts silently"): + with validating(op): + op(key, query[:1], value, HC_COUNT) + with pytest.raises(AssertionError, match="broadcasts silently"): + with validating(op): + op(key, query, value[:1], HC_COUNT) + + +def test_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() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(3): + op(key, query, value, HC_COUNT, CLAMP_MIN) + torch.cuda.current_stream().wait_stream(warmup) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(key, query, value, HC_COUNT, CLAMP_MIN) + + fresh = _make_operands(101, torch.bfloat16, (8,)) + for static, new in zip((key, query, value), fresh): + static.copy_(new) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(captured, op(*fresh, HC_COUNT, CLAMP_MIN)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/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 new file mode 100644 index 000000000000..b6069fe13f5b --- /dev/null +++ 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 @@ -0,0 +1,406 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the ngram_token_history entry. + +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. +""" + +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.ngram_token_history import ( # noqa: E501 + CKPT_EOS_TOKEN_ID, + CKPT_NGRAM_SIZE, + CKPT_VOCAB_SIZE, + ENGINE_SLOTS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ple.ngram_token_history import ( # noqa: E501 + ngram_token_history as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "ngram_token_history requires a CUDA device" + + +def _compare( + out: tuple[torch.Tensor, torch.Tensor], ref: tuple[torch.Tensor, torch.Tensor] +) -> None: + """`(window, context_pool)` -- the returned window with the pool as the + call left it -- against `reference`'s pair, both bitwise: integer + bookkeeping, no tolerance. The whole pool is compared, which is the + lock on every slot the batch did not name.""" + (window, pool), (ref_window, ref_pool) = out, ref + assert window.dtype == ref_window.dtype == torch.int64, (window.dtype, ref_window.dtype) + assert window.shape == ref_window.shape, (window.shape, ref_window.shape) + assert window.is_contiguous() + assert torch.equal(window, ref_window), "window disagrees with the reference" + assert torch.equal(pool, ref_pool), "published context pool disagrees with the reference" + + +DEV = "cuda" +EOS = CKPT_EOS_TOKEN_ID +NGRAM = CKPT_NGRAM_SIZE +CONTEXT_LEN = NGRAM - 1 + + +def _qsl(lengths, dtype, qsl_from="literal"): + if qsl_from == "seq_lens": + # The call site under MTP verify: an int32 prefix of the metadata's + # seq_lens buffer, scanned into an int64 vector with a leading zero. + buf = torch.full((len(lengths) + 9,), 99, dtype=torch.int32, device=DEV) + buf[: len(lengths)] = torch.tensor(lengths, dtype=torch.int32, device=DEV) + seq_lens = buf[: len(lengths)] + out = torch.zeros(len(lengths) + 1, dtype=torch.int64, device=DEV) + torch.cumsum(seq_lens, dim=0, dtype=torch.int64, out=out[1:]) + return out + starts = [0] + for length in lengths: + starts.append(starts[-1] + length) + return torch.tensor(starts, dtype=dtype, device=DEV) + + +def _build(spec, seed): + """`(args, backing_pool)`: the op's positional arguments, and the tensor + the pool argument views (itself, unless `layout == "pool_view"`).""" + 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 + ).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, + 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, + ) + return args, backing + + +def _expect(args): + """The reference, driven from a snapshot of the pool before the call.""" + return op.reference(args[0], args[1].clone(), *args[2:]) + + +@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())) + 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) + _compare((window, args[1]), expected) + for a, b in zip((args[0], args[2], args[3], args[4]), read_only): + assert torch.equal(a, b), "only context_pool may be written" + # 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) + ids = torch.randint(0, CKPT_VOCAB_SIZE, (96,), generator=g).to(DEV) + start = torch.randint(0, CKPT_VOCAB_SIZE, (3, CONTEXT_LEN), generator=g).to(DEV) + slot = torch.tensor([1], dtype=torch.int32, device=DEV) + + def call(pool, chunk, carry): + flags = torch.tensor([carry], device=DEV) + qsl = torch.tensor([0, chunk.shape[0]], dtype=torch.int32, device=DEV) + return op(chunk, pool, slot, flags, qsl, EOS, NGRAM) + + one_shot_pool = start.clone() + one_shot = call(one_shot_pool, ids, False) + pool = start.clone() + split = torch.cat([call(pool, ids[:41], False), call(pool, ids[41:], True)]) + assert torch.equal(split, one_shot) + assert torch.equal(pool, one_shot_pool) + + +# ----------------------------------------------------------- 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 + + +def _two_sequences(): + pool = torch.randint(0, CKPT_VOCAB_SIZE, (4, CONTEXT_LEN), device=DEV) + return ( + torch.tensor([11, 12, 13, 21], device=DEV), + pool, + torch.tensor([0, 1], dtype=torch.int32, device=DEV), + torch.tensor([True, False], device=DEV), + torch.tensor([0, 3, 4], dtype=torch.int32, device=DEV), + EOS, + NGRAM, + ) + + +def _misses(args, bad_index, bad_value): + """Run the op unguarded with one argument swapped; True if it returned + and its answer differs from the reference on the intended arguments.""" + expected = _expect(args) + bad = list(args) + bad[1] = args[1].clone() + bad[bad_index] = bad_value + window = op(*bad) + return not (torch.equal(window, expected[0]) and torch.equal(bad[1], expected[1])) + + +def test_guard_refuses_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() + short = torch.tensor([0, 3], dtype=torch.int32, device=DEV) + assert _misses(args, 4, short), "the hazard this guard covers is gone" + bad = list(args) + bad[4] = short + with pytest.raises(AssertionError, match=r"query_start_loc must be \[num_seqs"): + with validating(op): + op(*bad) + + +@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 + + +def test_cuda_graph_capture_and_replay() -> None: + """The captured decode shape, 8 sequences at the engine pool width.""" + spec = dict( + lengths=(1,) * 8, + slots=(0, ENGINE_SLOTS - 1, 4, 17, 30, 43, 56, 69), + carries=(True,) * 8, + pool_slots=ENGINE_SLOTS, + ids_dtype=torch.int32, + qsl_dtype=torch.int64, + ) + args, _ = _build(spec, seed=11001) + start = args[1].clone() + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(3): + op(*args) + torch.cuda.current_stream().wait_stream(side) + + args[1].copy_(start) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(*args) + + replay_ids = torch.randint(0, CKPT_VOCAB_SIZE, (8,), dtype=torch.int32, device=DEV) + args[0].copy_(replay_ids) + args[1].copy_(start) + expected = op.reference(replay_ids, start.clone(), *args[2:]) + graph.replay() + torch.cuda.synchronize() + _compare((captured, args[1]), expected) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/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 new file mode 100644 index 000000000000..a58ee2bdde8b --- /dev/null +++ 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 @@ -0,0 +1,251 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the dequantize_e4m3_per_tensor catalog entry. + +The entry owns the contract. This file turns a cell into operands and holds +the claims no cell carries: the direction, the round trip through the +quantizer, degenerate and overflowing scales, graph capture, and the loud +rejections. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.quantization.dequantize_e4m3_per_tensor import ( # noqa: E501 + PLE_SCALE, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.quantization.dequantize_e4m3_per_tensor import ( # noqa: E501 + dequantize_e4m3_per_tensor as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "dequantize_e4m3_per_tensor requires a CUDA device" + + +_BIT_VIEW = { + torch.bfloat16: torch.int16, + torch.float16: torch.int16, + torch.float32: torch.int32, +} + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Bit-exact, NaN for NaN. + + Both sides form the same exact fp32 product and round it once, so any + difference is a defect. Measured bit-exact on every case, including all + 256 codes x three output dtypes x five scales, with -0.0 and the + subnormals compared on bit patterns. NaN payloads are not compared. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + nan = torch.isnan(ref) + assert torch.equal(torch.isnan(out), nan), "NaN lanes differ" + view = _BIT_VIEW[ref.dtype] + assert torch.equal(out.contiguous().view(view)[~nan], ref.contiguous().view(view)[~nan]), ( + "dequantization is exact; a bit difference is a defect" + ) + + +DEV = "cuda" +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 = {} + + +def _fp8(x: torch.Tensor) -> torch.Tensor: + return x.to(torch.float8_e4m3fn) + + +def _scale(v: float, rank: int, dtype=torch.bfloat16) -> torch.Tensor: + """The only shape the op accepts: `input`'s rank, every dimension 1.""" + return torch.full((1,) * rank, v, dtype=dtype, device=DEV) + + +def _input(spec, seed: int) -> torch.Tensor: + g = torch.Generator(device=DEV).manual_seed(seed) + source, shape = spec["source"], spec["shape"] + if source == "gather": + if "table" not in _TABLE: + tg = torch.Generator(device=DEV).manual_seed(11) + _TABLE["table"] = _fp8(torch.randn(1 << 20, 160, device=DEV, generator=tg) * 40.0) + idx = torch.randint(0, 1 << 20, shape[:2], device=DEV, generator=g) + x = _TABLE["table"][idx] + 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 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]) + + +# ── the certified cells ─────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + x = _input(cell.spec, zlib.crc32(cell.why.encode())) + for value in cell.spec["scales"]: + scales = _scales(cell.spec, value, x.dim()) + out = op(x, scales) + assert out.shape == x.shape and out.dtype == scales.dtype and out.is_contiguous() + _compare(out, op.reference(x, scales)) + if cell.spec["source"] == "gather": + flat = out.reshape(x.shape[0], -1) + assert flat.shape[1] == 2560 and flat.data_ptr() == out.data_ptr(), "flatten is a view" + + +# ── what a cell cannot carry ────────────────────────────────────────────── + + +def test_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) + static_in = _fp8(torch.randn(256, 16, 160, device=DEV, generator=g) * 8.0) + static_scales = _scale(PLE_SCALE, 3) + eager = op(static_in, static_scales).clone() + warm = torch.cuda.Stream() + warm.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warm): + for _ in range(3): + op(static_in, static_scales) + torch.cuda.current_stream().wait_stream(warm) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_out = op(static_in, static_scales) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(_bits(static_out), _bits(eager)) + fresh = _fp8(torch.randn(256, 16, 160, device=DEV, generator=g) * 8.0) + static_in.copy_(fresh) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(_bits(static_out), _bits(op(fresh, static_scales))) + static_scales.copy_(_scale(0.5, 3)) + graph.replay() + torch.cuda.synchronize() + _compare(static_out, op.reference(fresh, _scale(0.5, 3))) + + +def test_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() 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 new file mode 100644 index 000000000000..2a65ca954ee7 --- /dev/null +++ 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 @@ -0,0 +1,398 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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. +""" + +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.fp8_quantize_1x128 import ( # noqa: E501 + AMAX_FLOOR, + E4M3_MAX, + SCALE_ULP_GATE, + block_scales, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.quantization.fp8_quantize_1x128 import ( # noqa: E501 + fp8_quantize_1x128 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "fp8_quantize_1x128 requires a CUDA device" + + +def _ordinal(codes: torch.Tensor) -> torch.Tensor: + """Sign-magnitude ordinal of an e4m3 byte: adjacent codes differ by 1.""" + b = codes.view(torch.uint8).int() + return torch.where(b >= 128, -(b & 0x7F), b & 0x7F) + + +def _compare( + out: tuple[torch.Tensor, torch.Tensor], + ref: tuple[torch.Tensor, torch.Tensor], +) -> None: + """Scales within 4 fp32 ulp; codes at most one e4m3 step apart, >=99.5% bit-exact. + + This build compiles the kernel's two fp32 divisions approximately + (`--use_fast_math`), each within CUDA's documented 2 ulp, so neither + output is reproducible to the last bit with `use_ue8m0=False`. Measured + over the certified column: scales within 2 ulp of both `amax / 448` and + `1 / (448 / amax)` (~80% bit-exact at T=1024); codes at worst + 99.6875% bit-exact at H 2560 (99.8566% elsewhere), every miss exactly + one code away with the exact product within 1.122e-07 relative of the + midpoint between the two -- a rounding tie the scale's last ulp decides. + `use_ue8m0=True` is bit-exact (the power-of-two scale makes the + multiply exact), and passes this gate trivially. + + Wrong recipes at T=256, H=2560 (bytes agreeing / worst relative scale + error): power-of-two 0.043 / 0.99, per-tensor 0.0002 / 2.0, per-row + 0.083 / 1.6, blocks down the token axis 0.19 / 1.0, blocks offset by 64 + 0.48 / 0.93, e4m3 max as 240 0 / 0.87, signed max 0.58 / 0.58; the + correct one 0.9993 / 1.4e-07. + """ + data, scales = out + ref_data, ref_scales = ref + assert data.dtype == ref_data.dtype == torch.float8_e4m3fn, (data.dtype, ref_data.dtype) + assert scales.dtype == ref_scales.dtype == torch.float32, (scales.dtype, ref_scales.dtype) + assert data.shape == ref_data.shape, (data.shape, ref_data.shape) + assert scales.shape == ref_scales.shape, (scales.shape, ref_scales.shape) + if scales.numel(): + ulp = ( + scales.contiguous().view(torch.int32).long() + - ref_scales.contiguous().view(torch.int32).long() + ) + worst = int(ulp.abs().max()) + assert worst <= SCALE_ULP_GATE, f"scale {worst} fp32 ulp off > {SCALE_ULP_GATE}" + if data.numel(): + step = (_ordinal(data) - _ordinal(ref_data)).abs() + assert int(step.max()) <= 1, f"a code is {int(step.max())} e4m3 steps off" + exact = 1.0 - (step != 0).float().mean().item() + assert exact >= 0.995, f"only {100 * exact:.4f}% of codes bit-exact" + + +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 + + +def _expand(v: torch.Tensor, n: int) -> torch.Tensor: + """`[T, nb]` block values -> `[T, n]` per-column values.""" + return v.repeat_interleave(128, dim=1)[:, :n].contiguous() + + +def _input(tokens: int, hidden: int, seed: int, magnitude: float = 3.0) -> torch.Tensor: + g = torch.Generator(device=DEV).manual_seed(seed) + 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.""" + qs, _ = block_scales(x, False) + prod = x.float() * _expand(qs, x.shape[1]) + diff = data.float() != ref_data.float() # -0 and +0 are the same code here + if not bool(diff.any()): + return 0.0 + mid = ((data.float() + ref_data.float()) / 2).double() + rel = ((prod.double()[diff] - mid[diff]).abs() / mid[diff].abs()).max().item() + assert rel <= TIE_REL_GATE, f"a miss {rel:.3e} from its midpoint is not a tie" + return rel + + +# ── the certified cells ─────────────────────────────────────────────────── + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + spec = cell.spec + x = _input(spec["tokens"], spec["hidden"], seed=zlib.crc32(cell.why.encode())) + data, scales = op(x, spec["use_ue8m0"]) + nb = (spec["hidden"] + 127) // 128 + assert data.shape == x.shape and data.is_contiguous() and data.device == x.device + assert scales.shape == (nb, spec["tokens"]) and scales.is_contiguous() + ref = op.reference(x, spec["use_ue8m0"]) + _compare((data, scales), ref) + rel = _assert_misses_are_ties(x, data, ref[0]) if not spec["use_ue8m0"] else 0.0 + exact = (data.view(torch.uint8) == ref[0].view(torch.uint8)).float().mean().item() + print(f" {cell.why[:44]}: {100 * exact:.4f}% bit-exact, worst tie {rel:.3e}") + + +# ── 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)]) + + +def test_cuda_graph_capture_and_replay(): + """A replay reproduces a fresh eager call bitwise; so does a side stream.""" + x = _input(64, 2560, seed=17, magnitude=1.0) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(3): + op(x) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + g_data, g_sf = op(x) + x.copy_(_input(64, 2560, seed=18, magnitude=1.0)) + e_data, e_sf = op(x) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(g_data.view(torch.uint8), e_data.view(torch.uint8)) + assert torch.equal(g_sf.view(torch.int32), e_sf.view(torch.int32)) + alt = torch.cuda.Stream() + with torch.cuda.stream(alt): + a_data, a_sf = op(x) + alt.synchronize() + assert torch.equal(a_data.view(torch.uint8), e_data.view(torch.uint8)) + 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() 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 new file mode 100644 index 000000000000..2e691b0f38fe --- /dev/null +++ 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 @@ -0,0 +1,263 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the causal_conv1d_fwd entry.""" + +import os +import zlib +from typing import Optional + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_fwd import ( # noqa: E501 + ENGINE_SLOTS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_fwd import ( # noqa: E501 + causal_conv1d_fwd as op, +) + +# Cells sized at the engine's pool get a non-vacuity check below. +assert ENGINE_SLOTS == 461 + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "causal_conv1d_fwd requires a CUDA device" + + +def _compare( + out: tuple[torch.Tensor, Optional[torch.Tensor]], + ref: tuple[torch.Tensor, Optional[torch.Tensor]], +) -> None: + """`(activation, conv_states)` pairs: the returned activation with the + pool as the call left it, against `reference`'s pair. + + Activation at torch's default band: the kernel accumulates the taps in + fp32 and rounds once, so it sits within one ulp of the fp32 reference + (checkpoint geometry: max abs 0.015625, bitwise fraction 0.999996; fp16 + one ulp; fp32 1.4e-6). The state is a copy of the input tail, so it is + gated bitwise over the whole pool -- which is also the lock on every + slot the batch did not name. + """ + (act, states), (ref_act, ref_states) = out, ref + assert act.dtype == ref_act.dtype and act.shape == ref_act.shape + torch.testing.assert_close(act, ref_act) + assert (states is None) == (ref_states is None) + if states is not None: + assert torch.equal(states, ref_states), "the published state must be bitwise exact" + + +PAD_SLOT_ID = -1 + + +def _seed(text: str) -> int: + return zlib.crc32(text.encode()) + + +def _build(spec, seed): + """`__call__`'s positional arguments for one spec, freshly allocated.""" + g = torch.Generator(device="cuda").manual_seed(seed) + dim, width = spec["dim"], spec.get("width", 4) + dtype, layout = spec.get("dtype", torch.bfloat16), spec.get("layout", "in_place") + + 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) + 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] + + +def _clone(args): + """The mutable operands cloned (x, conv_states, out); strides preserved.""" + got = list(args) + for i in (0, 3, 9): + if got[i] is not None: + got[i] = got[i].clone() + return got + + +def _run(args): + """Drive the op on clones; return (returned activation, pool after, clones).""" + got = _clone(args) + returned = op(*got) + return returned, got[3], got + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + args = _build(cell.spec, seed=_seed(cell.why)) + ref = op.reference(*args) + returned, pool, got = _run(args) + assert returned is (got[0] 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" + _compare((returned, pool), ref) + if cell.spec["pool_slots"] == ENGINE_SLOTS: + # The whole-pool bitwise gate above already locks the unnamed slots; + # this keeps it from being vacuous on the large pool. + for slot in cell.spec["slots"]: + assert not torch.equal(pool[slot], args[3][slot]), slot + + +def 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.""" + 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] + 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: + """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 + returned = op(*unguarded) + with pytest.raises(AssertionError): + _compare((returned, unguarded[3]), ref) + + with pytest.raises(AssertionError): + with validating(op): + op(*_guard_case(case)) 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 new file mode 100644 index 000000000000..2d465304f39b --- /dev/null +++ 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 @@ -0,0 +1,225 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the causal_conv1d_triton_update entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.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 + causal_conv1d_triton_update as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "causal_conv1d_triton_update requires a CUDA device" + + +def _compare(out: tuple, ref: tuple) -> None: + """`(out, conv_state, intermediate_conv_window)` triples: the returned + tensor with the two buffers as the call left them, against + `reference`'s triple. + + The result at torch's default band over the rows the op writes (NaN in + the reference marks the rest): against the operand-dtype-product + reference bf16/fp16 were bit-exact at dim 512 and at dim 10240 off by + one ulp in at most 9 of 2.6M elements (the fp32 `exp`); fp32 is off by + 1.4e-6. The pool and the snapshots are copies, gated bitwise over the + whole buffer -- the lock on every unnamed slot and snapshot row. + """ + (act, state, inter), (ref_act, ref_state, ref_inter) = out, ref + assert act.dtype == ref_act.dtype and act.shape == ref_act.shape + written = ~torch.isnan(ref_act) + torch.testing.assert_close(act[written], ref_act[written]) + assert torch.equal(state, ref_state), "the advanced state must be bitwise exact" + assert (inter is None) == (ref_inter is None) + if inter is not None: + assert torch.equal(inter, ref_inter), "the snapshots must be bitwise exact" + + +PAD_SLOT_ID = -1 + + +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") + + 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) + 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") + return [ + x, + conv_state, + weight, + bias, + spec.get("activation", "silu"), + indices, + spec.get("pad_slot_id", PAD_SLOT_ID), + accepted, + window, + rows_idx, + spec.get("launch_dependent_kernels", 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() + returned = op(*got) + return returned, got[1], got[8] + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + x, weight = args[0].clone(), args[2].clone() + ref = op.reference(*args) + got = _run(args) + assert got[0] is not args[0] and got[0].shape == args[0].shape + assert torch.equal(args[0], x) and torch.equal(args[2], weight), "x and weight are read-only" + _compare(got, ref) + if cell.spec["pool_slots"] == ENGINE_SLOTS: + for slot in cell.spec["slots"]: # non-vacuity of the whole-pool gate + assert not torch.equal(got[1][slot], args[1][slot]), slot + + +def test_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 {}) + 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): + 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 new file mode 100644 index 000000000000..4b61e62ec676 --- /dev/null +++ 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 @@ -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 causal_conv1d_update entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_update import ( # noqa: E501 + ENGINE_SLOTS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.causal_conv1d_update import ( # noqa: E501 + causal_conv1d_update as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "causal_conv1d_update requires a CUDA device" + + +def _compare( + out: tuple[torch.Tensor, torch.Tensor], ref: tuple[torch.Tensor, torch.Tensor] +) -> None: + """`(x, conv_state)` pairs: the returned `x` with the pool as the call + left it, against `reference`'s pair. + + Activation at torch's default band: bf16/fp16 were measured bitwise + equal to the fp32 reference in every cell; fp32 differs by tap order + (max 1.43e-6, inside the default fp32 atol). The pool is a copy of + input values, so it is gated bitwise -- over the whole pool, which is + the lock on every slot the batch did not name. + """ + (act, state), (ref_act, ref_state) = out, ref + assert act.dtype == ref_act.dtype and act.shape == ref_act.shape + torch.testing.assert_close(act, ref_act) + assert torch.equal(state, ref_state), "the advanced state must be bitwise exact" + + +PAD_SLOT_ID = -1 +#: Rows of prefill ahead of the decode rows in the mixed-batch cell. +_PREFILL_ROWS = 37 + + +def _build(spec, seed): + """`__call__`'s arguments for one spec, plus the buffers some layouts 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") + + 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 + 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 + + +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() + 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]) + return returned, got[1] + + +@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, 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() + 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 + + +def _small(**overrides): + 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"\)")): + 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 new file mode 100644 index 000000000000..92fff412c4f4 --- /dev/null +++ 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 @@ -0,0 +1,383 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the chunk_gated_delta_rule entry. + +The first indexed in-place call per autotune key in a process must be as correct as a +warm one: the state kernel's autotuner restores the indexed state its benchmark launches +write. "First call in the process" is a property of the process, so the cold-call tests +re-run this file as a child: `python `. +""" + +import os +import subprocess +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 + HEAD_K_DIM, + HEAD_V_DIM, + NUM_K_HEADS, + NUM_V_HEADS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.chunk_gated_delta_rule import ( # noqa: E501 + chunk_gated_delta_rule as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "chunk_gated_delta_rule requires a CUDA device" + + +#: `_compare`'s band; see its docstring. +RTOL, ATOL_OUTPUT, ATOL_STATE = 1.6e-2, 2e-3, 2e-2 + + +def _compare(out: tuple, ref: tuple) -> None: + """`(o, final_state, pool)`: the call's two returns plus the + `initial_state` it was handed, against `reference`'s triple; `None` + components must match. + + Wider than bf16's default because the op is a chunked reformulation + whose WY operands are rounded to bf16 -- the same recurrence, not the + same sum. rtol stays at the bf16 default; atol is 2e-3 on `o` and 2e-2 + on states, about 4x and 6x the largest deviations measured warm over + every cell against this reference (`o` 4.88e-4, one bf16 ulp, on + magnitudes up to 7.9e-2; states 3.57e-3 on up to 8.3e-1). + Discriminated: seven wrong references (no gating, beta = 1, q/k swapped, + initial state dropped, scale = 1, key heads reversed, `[K, V]` + transposed state) all fail the `o` gate, the closest 11x outside it. + """ + for got, want, atol in zip(out, ref, (ATOL_OUTPUT, ATOL_STATE, ATOL_STATE)): + assert (got is None) == (want is None), (got is None, want is None) + if got is None: + continue + assert got.dtype == want.dtype and got.shape == want.shape, ( + (got.dtype, got.shape), + (want.dtype, want.shape), + ) + torch.testing.assert_close(got.float(), want.float(), rtol=RTOL, atol=atol) + + +DEV = "cuda" +SCALE = HEAD_K_DIM**-0.5 + + +def _l2norm(x: torch.Tensor) -> torch.Tensor: + xf = x.float() + return xf / torch.sqrt(xf.pow(2).sum(-1, keepdim=True) + 1e-6) + + +def _operands( + seqlens: 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, +): + """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) + 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) + v = torch.randn( + batch, 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) + 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() + 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).""" + 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 + + +def _gate_output(o: torch.Tensor, ref: tuple) -> None: + """`_compare`, on `o` alone.""" + _compare((o, None, None), (ref[0], None, None)) + + +def _gate_pool(pool: torch.Tensor, ref: tuple) -> None: + """`_compare`, on the pool alone.""" + _compare((ref[0], None, pool), (ref[0], None, ref[2])) + + +def _build(spec, seed): + """`__call__`'s 13 arguments for one spec, 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) + if spec.get("output") == "buffer": + 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 + + +@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) + pool = args[6] + pool_before = None if pool is None else pool.clone() + ref = op.reference(*args) + + o, final_state = op(*args) + + _compare((o, final_state, pool), ref) + if args[12] is not None: + assert o.data_ptr() == args[12].data_ptr(), "output= must be written in place" + 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) + + +def test_guard_refuses_a_noncontiguous_output() -> None: + """Unguarded, a head-slice view of a wider buffer is written through + shape-derived strides and the result is wrong.""" + q, k, v, g, beta, cu = _operands((96,), seed=61) + ref = op.reference(q, k, v, g, beta, cu_seqlens=cu) + head_slice = torch.zeros(1, 96, 2 * NUM_V_HEADS, HEAD_V_DIM, device=DEV, dtype=torch.bfloat16)[ + :, :, :NUM_V_HEADS + ] + assert not head_slice.is_contiguous() + o, _ = op(q, k, v, g, beta, cu_seqlens=cu, output=head_slice) + with pytest.raises(AssertionError): + _gate_output(o, ref) + with pytest.raises(AssertionError, match="output must be contiguous"): + with validating(op): + op(q, k, v, g, beta, cu_seqlens=cu, output=head_slice) + + +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. + +CHILD_TIMEOUT_S = 420 +#: The GDN prefill call shape, packed varlen, state by slot index. +CHILD_SEQLENS = (256, 96) + + +def _state_kernel_autotuner(): + """The `triton.autotune` wrapper around `chunk_gated_delta_rule_fwd_kernel_h_blockdim64`; + a key in its `cache` is a key whose benchmark launches already happened.""" + from tensorrt_llm._torch.modules.fla.chunk_delta_h import ( + chunk_gated_delta_rule_fwd_kernel_h_blockdim64 as node, + ) + + while node is not None: + if isinstance(getattr(node, "cache", None), dict): + return node + node = getattr(node, "fn", None) + raise AssertionError("no triton autotuner found around the chunked state kernel") + + +def _checkpoint_case(seed: int = 101): + q, k, v, g, beta, cu = _operands(CHILD_SEQLENS, seed=seed) + pool = _pool(4, seed=seed + 1) + indices = torch.tensor([2, 0], dtype=torch.int32, device=DEV) + buf = torch.empty( + 1, sum(CHILD_SEQLENS), NUM_V_HEADS, HEAD_V_DIM, device=DEV, dtype=torch.bfloat16 + ) + return (q, k, v, g, beta, None, pool, indices, True, False, cu, False, buf) + + +def _child_cold_call() -> None: + autotuner = _state_kernel_autotuner() + assert not autotuner.cache, "the child must be the first indexed call in its process" + args = _checkpoint_case() + ref = op.reference(*args) + o, _ = op(*args) + assert len(autotuner.cache) == 1, "the call must have benchmarked exactly one key" + _gate_output(o, ref) + _gate_pool(args[6], ref) + + +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, +} + + +def _run_child(mode: str) -> None: + # `_validating` lives four directories up; the child has no pytest to + # resolve `__extra_import_path__` for it. + root = str(Path(__file__).resolve().parents[4]) + existing = os.environ.get("PYTHONPATH") + env = dict(os.environ, PYTHONPATH=root if not existing else os.pathsep.join([root, existing])) + proc = subprocess.run( + [sys.executable, __file__, mode], + capture_output=True, + text=True, + timeout=CHILD_TIMEOUT_S, + env=env, + ) + print(proc.stdout) + assert proc.returncode == 0, ( + f"child `{mode}` exited {proc.returncode}\n--- stdout ---\n{proc.stdout}\n" + f"--- stderr ---\n{proc.stderr}" + ) + + +def test_cold_indexed_call_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") + + +if __name__ == "__main__": + _CHILD_MODES[sys.argv[1]]() diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ssm/test_modeling_v2_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 new file mode 100644 index 000000000000..5e011da0ac63 --- /dev/null +++ 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 @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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. +""" + +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.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, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "cu_seqlens_from_seq_lens requires a CUDA device" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Bitwise: integer arithmetic, no tolerance to loosen.""" + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + assert out.is_contiguous() + assert torch.equal(out, ref), f"{out.tolist()} != {ref.tolist()}" + + +DEV = "cuda" + + +def _seq_lens(lengths, dtype=torch.int32): + return torch.tensor(list(lengths), dtype=dtype, device=DEV) + + +def _build(spec): + lengths, dtype = spec["lengths"], spec["in_dtype"] + layout = spec.get("layout", "contiguous") + if layout == "prefix": + # The metadata buffer is sized for max_num_requests; the call site + # passes its first `batch` entries. + 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) + + +@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 + seq_lens = _build(spec) + before = seq_lens.clone() + out = op(seq_lens, spec["out_dtype"]) + assert torch.equal(seq_lens, before), "seq_lens is read-only" + assert out.device.type == "cuda" + _compare(out, op.reference(seq_lens, spec["out_dtype"])) + + # The entry's purpose: slicing a packed buffer with the vector recovers + # exactly the sequences packed into it. + lengths = spec["lengths"] + g = torch.Generator(device=DEV).manual_seed(zlib.crc32(cell.why.encode())) + pieces = [torch.randint(0, 10_000, (n,), generator=g, device=DEV) for n in lengths] + packed = torch.cat(pieces) if pieces else torch.empty(0, dtype=torch.int64, device=DEV) + bounds = out.tolist() + assert bounds[-1] == packed.shape[0] + for s, piece in enumerate(pieces): + assert torch.equal(packed[bounds[s] : bounds[s + 1]], piece), f"sequence {s} mis-sliced" + + +def test_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) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(3): + op(static_lens, torch.int32) + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_out = op(static_lens, torch.int32) + for lengths in ([1] * 8, [2] * 8, [3, 1, 2, 1, 1, 1, 1, 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) 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 new file mode 100644 index 000000000000..a7d49b0d7aa5 --- /dev/null +++ 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 @@ -0,0 +1,393 @@ +# 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. +""" + +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, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.dilated_short_conv_paged import ( # noqa: E501 + dilated_short_conv_paged as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "dilated_short_conv_paged requires a CUDA device" + + +def _compare( + out: tuple[torch.Tensor, torch.Tensor], ref: tuple[torch.Tensor, torch.Tensor] +) -> None: + """`(out, conv_states)` pairs: the returned activation with the pool as + the call left it, against `reference`'s pair. + + Activation at torch's default band: bf16 and fp16 were measured bitwise + equal to the reference (at 10240 channels and 461 slots too); fp32 + differs only by tap order. The pool is a copy, gated bitwise over the + whole pool -- the lock on every slot the batch did not name. + """ + (act, pool), (ref_act, ref_pool) = out, ref + assert act.dtype == ref_act.dtype and act.shape == ref_act.shape + assert act.is_contiguous() + torch.testing.assert_close(act, ref_act) + assert torch.equal(pool, ref_pool), "the published state must be bitwise exact" + + +KERNEL_SIZE = 4 +DILATION = 3 +STATE_LEN = (KERNEL_SIZE - 1) * DILATION +PAD_SLOT_ID = -1 + + +def _build(spec, seed): + """`__call__`'s positional arguments for one spec, freshly allocated. + + 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 + + def randn(*shape, scale=1.0): + return (scale * torch.randn(*shape, generator=g, device="cuda")).to(dtype) + + 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") + 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] + + +def _run(args): + """Drive the op on a copy of the pool; return (out, pool after).""" + got = list(args) + got[3] = args[3].clone() + return op(*got), got[3] + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + x, weight = args[0].clone(), args[1].clone() + ref = op.reference(*args) + out, pool = _run(args) + assert out.data_ptr() != args[0].data_ptr() + assert torch.equal(args[0], x) and torch.equal(args[1], weight), "x and weight are read-only" + _compare((out, pool), ref) + if cell.spec.get("pool_slots") == ENGINE_SLOTS: + for slot in cell.spec["slots"]: # non-vacuity of the whole-pool gate + assert not torch.equal(pool[slot], args[3][slot]), slot + + +def _small(lengths, slots, seed, **spec): + return _build(dict(channels=256, lengths=lengths, slots=slots, **spec), seed=seed) + + +def test_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) + x, weight, _, pool, qsl, cidx, his = args[:7] + live = pool.clone() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + op(x, weight, None, live, qsl, cidx, his, DILATION, True, PAD_SLOT_ID) + torch.cuda.current_stream().wait_stream(stream) + + graph = torch.cuda.CUDAGraph() + live.copy_(pool) + with torch.cuda.graph(graph): + captured = op(x, weight, None, live, qsl, cidx, his, DILATION, True, PAD_SLOT_ID) + + g = torch.Generator(device="cuda").manual_seed(26) + 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")) + x.copy_(torch.randn(x.shape, generator=g, device="cuda").to(x.dtype)) + live.copy_(pool) + graph.replay() + torch.cuda.synchronize() + + ref_out, ref_pool = op.reference( + x, weight, None, pool, qsl, cidx, his, DILATION, True, PAD_SLOT_ID + ) + torch.testing.assert_close(captured, ref_out) + assert torch.equal(live, ref_pool) + + +@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: + """Refused loudly by the torch calls the entry is composed from -- why it + has no `is_valid`.""" + 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") + with pytest.raises(RuntimeError) as excinfo: + with validating(op): + op(*args) + 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 new file mode 100644 index 000000000000..233386dcf207 --- /dev/null +++ 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 @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the extract_transpose_prefill_slice entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.extract_transpose_prefill_slice import ( # noqa: E501 + extract_transpose_prefill_slice as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "extract_transpose_prefill_slice requires a CUDA device" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Bitwise: the kernel moves elements with no cast and no arithmetic, + so any tolerance would only hide a wrong index. Compared as raw bytes so + dtypes without an equality kernel (`float8_e4m3fn`) are gated the same.""" + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + assert out.is_contiguous() + if out.numel(): + assert torch.equal(out.reshape(-1).view(torch.uint8), ref.reshape(-1).view(torch.uint8)) + + +def _src(num_rows: int, num_cols: int, dtype: torch.dtype, seed: int) -> torch.Tensor: + g = torch.Generator(device="cuda").manual_seed(seed) + 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) + + +def _build(spec, seed): + src = _src(spec["num_rows"], spec["num_cols"], spec["dtype"], seed) + if "row_stop" in spec: + src = src[: spec["row_stop"]] + if "col_slice" in spec: + lo, hi = spec["col_slice"] + src = src[:, lo:hi] + assert not src.is_contiguous() and src.stride(1) == 1 + return src, spec["n"], spec["start_col"], spec["width"] + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + src, n, start_col, width = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + before = src.contiguous().view(torch.uint8).clone() + out = op(src, n, start_col, width) + assert out.data_ptr() != src.data_ptr() or out.numel() == 0 + assert torch.equal(src.contiguous().view(torch.uint8), before), "src is read-only" + _compare(out, op.reference(src, n, start_col, width)) + + +def test_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: + """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") + 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) 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 new file mode 100644 index 000000000000..38a6cf3f54d1 --- /dev/null +++ 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 @@ -0,0 +1,254 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the fused_gdn_gating entry. + +Geometry is a Qwen3.8-Flash-Next linear-attention layer: 48 value heads, `a` +bf16, `A_log` / `dt_bias` fp32 drawn across the checkpoint's own ranges. +""" + +import math +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_gdn_gating import ( # noqa: E501 + A_LOG_RANGE, + APPROX_ULP, + BA_WIDTH, + DT_BIAS_RANGE, + NUM_V_HEADS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_gdn_gating import ( # noqa: E501 + fused_gdn_gating as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "fused_gdn_gating requires a CUDA device" + + +#: The smallest softplus `beta` a cell drives; the tail error scales as 1 / beta. +_MIN_CERTIFIED_BETA = 0.5 + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Both fp32. `g` multiplies two approximate transcendentals, so the + relative term is `4 * 2^-22`. The absolute floor is the unstabilised + softplus's tail error -- about one approximate ulp of the log's + argument -- scaled by `1 / beta` and by the `exp(A_log)` it multiplies; + `_compare` cannot see either, so it takes the certified extremes + (exp(A_log) <= 139.42, beta >= 0.5): atol = 4 * 2^-22 * 139.42 / 0.5 = + 2.66e-04. + + Measured against this gate on 256 checkpoint-range rows: worst error / + allowance 0.215. Five wrong formulas (dt_bias dropped, sign dropped, + A_log without exp, exp(-A_log), sigmoid for softplus) land 8.7e+05 to + 1.5e+07 allowances away. + """ + assert out.dtype == ref.dtype == torch.float32, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + atol = 4.0 * APPROX_ULP * math.exp(A_LOG_RANGE[1]) / _MIN_CERTIFIED_BETA + torch.testing.assert_close(out, ref, rtol=4.0 * APPROX_ULP, atol=atol) + + +DEV = "cuda" +SOFTPLUS_THRESHOLD = 20.0 +#: 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): + lo, hi = A_LOG_RANGE + A_log = torch.rand(num_heads, generator=g, device=DEV) * (hi - lo) + lo + lo, hi = DT_BIAS_RANGE + dt_bias = torch.rand(num_heads, generator=g, device=DEV) * (hi - lo) + lo + return A_log.to(dtype), dt_bias.to(dtype) + + +def _build(spec, seed): + """`__call__`'s positional arguments for one spec.""" + g = torch.Generator(device=DEV).manual_seed(seed) + rows, heads = spec["rows"], spec.get("num_heads", NUM_V_HEADS) + a_dtype = spec.get("a_dtype", torch.bfloat16) + layout = spec.get("layout", "contiguous") + + def randn(*shape): + return (torch.randn(*shape, generator=g, device=DEV) * spec.get("a_scale", 1.0)).to(a_dtype) + + if "fill" in spec: + a = torch.full((rows, heads), spec["fill"], dtype=a_dtype, device=DEV) + elif layout == "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) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + args = _build(cell.spec, seed=zlib.crc32(cell.why.encode())) + g = op(*args) + _compare(g, op.reference(*args)) + assert g.is_contiguous() and torch.isfinite(g).all() + 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.""" + 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) + 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 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 new file mode 100644 index 000000000000..5cdbfe632874 --- /dev/null +++ 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 @@ -0,0 +1,221 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the fused_gdn_post_conv entry. + +Geometry is a Qwen3.8-Flash-Next linear-attention layer: 16 x 128 key heads, +48 x 128 value heads, qkv_dim 10240, bf16 activations. +""" + +import math +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_gdn_post_conv import ( # noqa: E501 + A_LOG_RANGE, + DT_BIAS_RANGE, + HEAD_K_DIM, + HEAD_V_DIM, + NUM_K_HEADS, + NUM_V_HEADS, + QKV_DIM, + SOFTPLUS_THRESHOLD, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_gdn_post_conv import ( # noqa: E501 + fused_gdn_post_conv as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "fused_gdn_post_conv requires a CUDA device" + + +#: Relative error of the SASS approximate transcendentals (`ex2.approx.f32`, +#: `lg2.approx.f32`) Triton lowers `tl.exp` / `tl.log` to. +_APPROX_ULP = 2.0**-22 + + +def _compare(out: tuple, ref: tuple) -> None: + """Per output, each gate derived rather than tuned. + + q, k -- torch's bf16 band: both sides round one fp32 value, so they + differ by at most one bf16 ulp. v -- bitwise: the kernel moves bits. + g -- a product of two SASS approximate transcendentals (relative error + ~2^-22 each) with the softplus magnitude bounded by the threshold and + every term scaled by `exp(A_log)`: + `atol = 4 * 2^-22 * softplus_threshold * exp(max A_log)`, `rtol = 8 * 2^-22`. + At this checkpoint's A_log ceiling that is 6.7e-04 against a measured + worst case of 7.1e-05; five wrong gate formulas land at 1.2e+02 and up. + beta -- fp32: torch's fp32 band (one well-conditioned transcendental); + narrower: bitwise, the fp32 sigmoid rounded exactly once. + """ + assert len(out) == len(ref) == 5 + for got, want in zip(out, ref): + assert got.dtype == want.dtype and got.shape == want.shape, ( + (got.dtype, got.shape), + (want.dtype, want.shape), + ) + assert got.is_contiguous() + q, k, v, g, beta = out + q_ref, k_ref, v_ref, g_ref, beta_ref = ref + torch.testing.assert_close(q, q_ref) + torch.testing.assert_close(k, k_ref) + assert torch.equal(v, v_ref), "v must be a bit-exact copy" + g_atol = 4 * _APPROX_ULP * SOFTPLUS_THRESHOLD * math.exp(A_LOG_RANGE[1]) + torch.testing.assert_close(g, g_ref, rtol=8 * _APPROX_ULP, atol=g_atol) + if beta.dtype == torch.float32: + torch.testing.assert_close(beta, beta_ref) + else: + assert torch.equal(beta, beta_ref), "beta must be the fp32 sigmoid rounded once" + + +DEV = "cuda" +DTYPE = torch.bfloat16 +HEADS = (NUM_K_HEADS, HEAD_K_DIM, NUM_V_HEADS, HEAD_V_DIM) +L2_NORM_EPS = 1e-6 +#: The stock mixer's `z` half of `in_proj_qkvz`, and its `[b|a]` projection width. +Z_WIDTH = NUM_V_HEADS * HEAD_V_DIM +BA_WIDTH = 2 * NUM_V_HEADS + + +def _build(spec, seed): + """`__call__`'s positional arguments, then its keyword `beta_dtype`.""" + 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) + + 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) + lo, hi = DT_BIAS_RANGE + dt_bias = (torch.rand(NUM_V_HEADS, generator=g, device=DEV) * (hi - lo) + lo).to(gate_dtype) + 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) + + +@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())) + 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) + 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) + + +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) 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 new file mode 100644 index 000000000000..980c4be80374 --- /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_update.py @@ -0,0 +1,580 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the fused_recurrent_gated_delta_rule_update entry. + +The axis this entry exists for is more than one token per sequence per step -- +the speculative verify call -- so the cells sweep `P`, and the two state sinks +(write-back into the pool, per-step states into the rollback buffer) are +driven separately and together. + +Two tests answer "is this harness blind" rather than a correctness question: +the neighbouring `chunk_gated_delta_rule` is 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. +""" + +from __future__ import annotations + +import json +import subprocess +import sys +import zlib +from pathlib import Path + +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, + NUM_K_HEADS, + NUM_V_HEADS, + SLOTS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_recurrent_gated_delta_rule_update import ( # noqa: E501 + fused_recurrent_gated_delta_rule_update as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "fused_recurrent_gated_delta_rule_update requires a CUDA device" + + +def _compare(out: tuple, ref: tuple) -> None: + """`(returned, initial_state_source, intermediate_states_buffer)` as the + call left them, against `reference`'s triple, each at torch's default + dtype-aware band -- nothing is loosened. + + The kernel and the reference both accumulate in fp32 and differ only in + reduction order. Measured over every cell shape: output <= 2.27e-4 at a + bf16 absmax of 3.9e-2 .. 8.2e-2 (under one bf16 ulp there), written-back + state <= 2.98e-7 at absmax ~1, per-step states <= 1.8e-7; a bf16 + per-step buffer costs one bf16 ulp (1.94e-3 at absmax 0.89). Eight wrong + recurrences (q/k heads tiled not interleaved, q/k swapped, beta 1, no + gating, reversed tokens, every token from the initial state, a + transposed or absent initial state) land 1.56e-2 .. 8.17e-2 away on the + output -- 128x the correct result's own deviation at the closest. + """ + (o, pool, cache), (o_ref, pool_ref, cache_ref) = out, ref + assert o.dtype == o_ref.dtype and o.shape == o_ref.shape, ( + (o.dtype, o.shape), + (o_ref.dtype, o_ref.shape), + ) + torch.testing.assert_close(o, o_ref) + assert (pool is None) == (pool_ref is None) + if pool is not None: + torch.testing.assert_close(pool, pool_ref) + assert (cache is None) == (cache_ref is None) + if cache is not None: + torch.testing.assert_close(cache, cache_ref) + + +DEV = "cuda" +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. + + 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 + + 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) + 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) + 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() + + 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() + + 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)) + + buffer = cache_steps = None + if mode in ("verify", "both"): + rows = spec.get("cache_rows", SLOTS if pool_slots is None else pool_slots) + cache_steps = tokens + buffer = torch.full( + (rows, cache_steps, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM), + CACHE_SENTINEL, + dtype=spec.get("cache_dtype", torch.float32), + device=DEV, + ) + + 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, + v=v, + g=g, + beta=beta, + initial_state_source=pool, + initial_state_indices=indices, + scale=None, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + 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] + ) + + +def _seed(why): + return zlib.crc32(why.encode()) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + kw, 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() + + 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" + + +_COLD_SPEC = dict(mode="write_back", n=2, p=3, slots=(5, 2)) + + +def cold_call_measurement() -> dict: + """The first indexed in-place call of the process, then the second. + + Public because the child process below calls it by name. + """ + kw, _ = _build(_COLD_SPEC, seed=7) + pool = kw["initial_state_source"] + pool_before = pool.clone() + ref = op.reference(**kw) + cold_out = op(**kw).clone() + cold_pool = pool.clone() + try: + _compare((cold_out, cold_pool, None), ref) + cold_ok = True + except AssertionError: + cold_ok = False + pool.copy_(pool_before) + warm_out = op(**kw) + return { + "cold_compare_ok": cold_ok, + "cold_out_err": (cold_out.float() - ref[0].float()).abs().max().item(), + "cold_pool_err": (cold_pool - ref[1]).abs().max().item(), + "out_bitwise_equal": bool(torch.equal(cold_out, warm_out)), + "pool_bitwise_equal": bool(torch.equal(cold_pool, pool)), + } + + +_COLD_CHILD = """ +import importlib.util, json, sys +sys.path.insert(0, {root!r}) +spec = importlib.util.spec_from_file_location("_cold_entry_test", {path!r}) +mod = importlib.util.module_from_spec(spec) +sys.modules["_cold_entry_test"] = mod +spec.loader.exec_module(mod) +print("COLD_RESULT " + json.dumps(mod.cold_call_measurement())) +""" + + +def test_cold_call_in_a_fresh_process_is_correct() -> None: + """Measured in a child that has launched nothing: a warm-up here would + hide exactly the defect the sibling prefill core has. Cold must pass the + entry's own gate *and* be bitwise the warm call.""" + root = str(Path(__file__).resolve().parents[4]) # where `_validating` lives + proc = subprocess.run( + [sys.executable, "-c", _COLD_CHILD.format(root=root, path=__file__)], + capture_output=True, + text=True, + timeout=480, + check=False, + ) + assert proc.returncode == 0, ( + f"cold-call child exited {proc.returncode}\n{proc.stdout[-4000:]}\n{proc.stderr[-4000:]}" + ) + line = [ln for ln in proc.stdout.splitlines() if ln.startswith("COLD_RESULT ")] + assert line, f"child produced no result\n{proc.stdout[-4000:]}" + result = json.loads(line[0][len("COLD_RESULT ") :]) + print(f"\ncold call: {result}") + assert result["cold_compare_ok"], result + assert result["out_bitwise_equal"] and result["pool_bitwise_equal"], result + + +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)) + + +def test_cache_steps_off_the_buffer_step_axis_misplaces_rows_and_is_refused() -> None: + """Unguarded, a 4-step buffer driven at cache_steps 2: row 1's states land in + row 0's steps 2..3 and row 1 is never written -- the rows the promotion + reads are not the ones the kernel filled.""" + kw = _build(dict(mode="verify", n=2, p=2), seed=75)[0] + deep = torch.full((SLOTS, 4, NUM_V_HEADS, HEAD_V_DIM, HEAD_K_DIM), CACHE_SENTINEL, device=DEV) + op(**dict(kw, intermediate_states_buffer=deep, cache_steps=2)) + assert (deep[1] == CACHE_SENTINEL).all() and not (deep[0, 2:] == CACHE_SENTINEL).any(), ( + "the hazard this guard covers is gone" + ) + with pytest.raises(AssertionError, match="must equal intermediate_states_buffer's step axis"): + with validating(op): + op(**dict(kw, intermediate_states_buffer=deep.clone(), cache_steps=2)) + + +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 new file mode 100644 index 000000000000..114df72f3257 --- /dev/null +++ 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 @@ -0,0 +1,261 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the 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. +""" + +import zlib +from typing import Optional + +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, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.fused_sigmoid_gating_delta_rule_update import ( # noqa: E501 + fused_sigmoid_gating_delta_rule_update as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "fused_sigmoid_gating_delta_rule_update requires a CUDA device" + + +def _compare( + out: tuple[torch.Tensor, Optional[torch.Tensor]], + ref: tuple[torch.Tensor, Optional[torch.Tensor]], +) -> None: + """`(o, pool)`: the returned output with the pool as the call left it. + + Both at torch's default dtype-aware band: the kernel and the reference + accumulate the same fp32 sequence and round once (bf16 `o`: measured + 3.05e-05 max abs at N=32, 461 slots; fp32 pool 2.4e-07). Five wrong + variants (a/b swapped, norm dropped, pool `(V, K)` transposed, initial + state ignored, q-head order reversed) land 1.7e-2 to 4.5 away, three + orders of magnitude outside. The pool is compared whole, so the + unnamed slots are gated too (the reference leaves them bit-identical). + """ + (o, pool), (ref_o, ref_pool) = out, ref + assert o.dtype == ref_o.dtype and o.shape == ref_o.shape, ( + (o.dtype, o.shape), + (ref_o.dtype, ref_o.shape), + ) + assert torch.isfinite(o).all() + torch.testing.assert_close(o, ref_o) + assert (pool is None) == (ref_pool is None) + if pool is not None: + assert pool.dtype == ref_pool.dtype and pool.shape == ref_pool.shape + torch.testing.assert_close(pool, ref_pool) + + +DEV = "cuda" +NUM_K_HEADS, NUM_V_HEADS, HEAD_K_DIM, HEAD_V_DIM = 16, 48, 128, 128 +MIXED_QKV_WIDTH = 2 * NUM_K_HEADS * HEAD_K_DIM + NUM_V_HEADS * HEAD_V_DIM # 10240 +SOFTPLUS_BETA, SOFTPLUS_THRESHOLD = 1.0, 20.0 +_LAYERS = 3 +_TAIL = 5 + + +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.""" + 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) + 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": + 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] + 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)) + 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, + dt_bias, + SOFTPLUS_BETA, + SOFTPLUS_THRESHOLD, + q, + k, + v, + b, + pool, + indices, + None, + spec.get("l2norm", True), + cu_seqlens, + output, + ] + return args, extras + + +@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())) + 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 + 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) + 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) + 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])) 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 new file mode 100644 index 000000000000..54fe1ac9e2d0 --- /dev/null +++ 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 @@ -0,0 +1,209 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the pack_gdn_decode_qkv entry. + +The op is a pure copy, so every comparison is bitwise; the discrimination test +is what shows that gate is not vacuous on these operands. +""" + +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, + HEAD_K_DIM, + HEAD_V_DIM, + NUM_K_HEADS, + NUM_V_HEADS, + QKV_DIM, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.pack_gdn_decode_qkv import ( # noqa: E501 + pack_gdn_decode_qkv as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "pack_gdn_decode_qkv requires a CUDA device" + + +def _compare(out: tuple, ref: tuple) -> None: + """Bitwise, all three: the op moves bits and changes nothing.""" + assert len(out) == len(ref) == 3 + for name, got, want in zip("qkv", out, ref): + assert got.dtype == want.dtype and got.shape == want.shape, ( + name, + (got.dtype, got.shape), + (want.dtype, want.shape), + ) + assert torch.equal(got, want), f"{name} must be a bit-exact copy" + + +DEV = "cuda" +Q_WIDTH = NUM_K_HEADS * HEAD_K_DIM # 2048 +V_WIDTH = NUM_V_HEADS * HEAD_V_DIM # 6144 +#: 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 + `[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) + + 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() + 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 + + +@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)) + + 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 + + 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), + ) + assert all(gt.data_ptr() == t.data_ptr() for gt, t in zip(grouped, out)) + per_token = conv_out.permute(0, 2, 1) # [B, P, qkv_dim], straight off the conv + want = ( + per_token.narrow(2, 0, Q_WIDTH).reshape(b, p, num_k_heads, head_k_dim), + per_token.narrow(2, Q_WIDTH, Q_WIDTH).reshape(b, p, num_k_heads, head_k_dim), + per_token.narrow(2, 2 * Q_WIDTH, V_WIDTH).reshape(b, p, num_v_heads, head_v_dim), + ) + assert all(torch.equal(gt, w) for gt, w in zip(grouped, want)) + + +def test_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. + + The input is carved from a longer flat buffer so the over-read lands in + known memory rather than in allocator garbage. + """ + rows, width = 32, QKV_DIM - 1024 + flat = torch.randn(rows * width + 1024, device=DEV).to(torch.bfloat16) + mixed = flat[: rows * width].view(rows, width) + _, _, v = op(mixed, *CKPT_HEADS) + torch.cuda.synchronize() + tail = v[0].reshape(rows, V_WIDTH)[:, V_WIDTH - 1024 :] + assert torch.equal(tail[:-1], mixed[1:, :1024]), "the hazard this guard covers is gone" + + with pytest.raises(AssertionError, match="must be at least qkv_dim=10240 wide"): + with validating(op): + op(mixed, *CKPT_HEADS) + + +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 new file mode 100644 index 000000000000..84c95ef21b64 --- /dev/null +++ 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 @@ -0,0 +1,238 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the reset_recurrent_state_rows entry.""" + +import zlib +from typing import Optional + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.reset_recurrent_state_rows import ( # noqa: E501 + CONV_SHAPE, + REC_SHAPE, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.ssm.reset_recurrent_state_rows import ( # noqa: E501 + reset_recurrent_state_rows as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "reset_recurrent_state_rows requires a CUDA device" + + +def _compare( + out: tuple[torch.Tensor, Optional[torch.Tensor]], + ref: tuple[torch.Tensor, Optional[torch.Tensor]], +) -> None: + """Bitwise, both pools whole: every element is its old value or +0.0, + so a tolerance could only hide a wrong row -- and comparing the whole + pool is the lock on every slot the batch did not name. A zeroed row + must be +0.0 by bit pattern; `torch.equal` alone would accept -0.0.""" + for got, want in zip(out, ref): + assert (got is None) == (want is None) + if got is None: + continue + assert got.dtype == want.dtype and got.shape == want.shape + assert torch.equal(got, want), "pool differs from the reference" + zero = want.reshape(want.shape[0], -1).eq(0).all(dim=1) + if zero.any(): + bits = {1: torch.uint8, 2: torch.int16, 4: torch.int32}[got.element_size()] + assert (got[zero].contiguous().view(bits) == 0).all(), "a reset row is not +0.0" + + +DEV = "cuda" +_LAYERS = 3 + + +def _build(spec, seed): + """`__call__`'s arguments, plus the backing allocations a layout 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") + + 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] + 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) + return [rec, indices, flags, conv], backing + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(op) + spec = cell.spec + args, backing = _build(spec, seed=zlib.crc32(cell.why.encode())) + rec, indices, flags, conv = args + before = {name: t.clone() for name, t in backing.items()} + indices_before, flags_before = indices.clone(), flags.clone() + ref = op.reference(*args) + + out = op(*args) + + assert out[0] is rec and out[1] is conv, "returns the pools it was handed" + _compare(out, ref) + assert torch.equal(indices, indices_before) and torch.equal(flags, flags_before) + + named = {s for s in indices.tolist() if 0 <= s < spec["slots"]} + reset = {s for s, f in zip(indices.tolist(), flags.tolist()) if not f and s in named} + if spec["slots"] > 64: + # Non-vacuity: the engine-width cells reset at least one row, so the + # whole-pool compare is also the lock on the ~430 unnamed slots. + assert reset + if spec.get("layout") == "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(): + torch.manual_seed(7) + rec = torch.randn(6, *REC_SHAPE[1:], dtype=torch.float32, device=DEV) + conv = torch.randn(6, *CONV_SHAPE, dtype=torch.bfloat16, device=DEV) + return rec, conv + + +def test_guard_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.""" + 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) + 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 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/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 new file mode 100644 index 000000000000..8724f9b5997d --- /dev/null +++ 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 @@ -0,0 +1,259 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp patch_embed_conv3d entry. + +The reference is a true-fp32 matmul of the flattened patch against the +flattened filter: a different op, accumulation order and memory layout from +the convolution under test, so it checks the contraction pairing rather than +restating it. +""" + +import zlib + +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, +) + +__extra_import_path__ = ["../../../.."] +from _validating import ULP, certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "patch_embed_conv3d requires a CUDA device" + + +def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Default rtol; atol raised to 4 ulps of the dtype at the result's RMS. + + Summing 1536 signed products leaves a few outputs per thousand near + zero, where the relative gate measures nothing and torch's 1e-5 atol is + three orders below one bf16 ulp. Measured at the checkpoint's geometry: + max |conv - fp32 reference| was 4.31 / 4.38 / 7.81 ulps at the RMS for + 64 / 256 / 4096 bf16 patches (1-2 ulps of the element's own magnitude, + which the untouched rtol covers); fp16 4.36 / 4.39 / 8.15. Wrong + variants land 325-392x outside this gate (bias dropped: 10.9x, bounded + by `max|bias|` ~ 0.02). For fp32 the 1e-5 floor binds, so nothing is + loosened there. + """ + assert out.dtype == ref.dtype, (out.dtype, ref.dtype) + assert out.shape == ref.shape, (out.shape, ref.shape) + rms = ref.float().pow(2).mean().sqrt().item() + rtol = {torch.bfloat16: 1.6e-2, torch.float16: 1e-3, torch.float32: 1.3e-6}[ref.dtype] + atol = max(1e-5, 4.0 * rms * ULP[ref.dtype]) + torch.testing.assert_close(out, ref, rtol=rtol, atol=atol) + + +IN_CHANNELS = 3 +KT, KH, KW = 2, 16, 16 +EMBED_DIM = 1152 +PATCH_NUMEL = IN_CHANNELS * KT * KH * KW # 1536 + + +def _operands( + num_patches, + dtype=torch.bfloat16, + seed=0, + in_channels=IN_CHANNELS, + kernel=(KT, KH, KW), + embed_dim=EMBED_DIM, + bias=True, +): + """Pixels in [-1, 1] (processor mean/std 0.5), weights at the tower's + initializer_range 0.02.""" + g = torch.Generator(device="cuda").manual_seed(seed) + kt, kh, kw = kernel + pixel_values = torch.empty(num_patches, in_channels * kt * kh * kw, device="cuda") + pixel_values = pixel_values.uniform_(-1.0, 1.0, generator=g).to(dtype) + weight = ( + 0.02 * torch.randn(embed_dim, in_channels, kt, kh, kw, generator=g, device="cuda") + ).to(dtype) + bias_t = (0.02 * torch.randn(embed_dim, generator=g, device="cuda")).to(dtype) if bias else None + return pixel_values, weight, bias_t + + +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())) + 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 + 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) + + +def test_cuda_graph_capture_and_replay() -> None: + pixel_values, weight, bias = _operands(64, seed=12) + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(3): + op(pixel_values, weight, bias) + torch.cuda.current_stream().wait_stream(warmup) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(pixel_values, weight, bias) + + pixel_values.copy_(_operands(64, seed=13)[0]) + graph.replay() + torch.cuda.synchronize() + _compare(captured, op.reference(pixel_values, weight, bias)) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_pos_embed_interpolate.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/vision/test_modeling_v2_pos_embed_interpolate.py new file mode 100644 index 000000000000..78611e4d4c96 --- /dev/null +++ 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 @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the qwen4_exp pos_embed_interpolate entry. + +`op.reference` is the explicit four-corner sum in fp64 -- exact, but sharing +the entry's sample positions, so it gates numerics. `_ref_resize` here runs +torch's own `F.interpolate` in fp64 and shares no arithmetic with the entry, so +it gates the structure: which axis is which, how the table is read, and the +spatial-merge row order. It sits up to 8.6e-6 from the exact value (its sample +coordinates are float), which is why fp32 cells gate on `op.reference`. +""" + +import zlib + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.vision.pos_embed_interpolate import ( # noqa: E501 + pos_embed_interpolate as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 + +assert torch.cuda.is_available(), "pos_embed_interpolate requires a CUDA device" + +NUM_GRID = 48 +NUM_POS = NUM_GRID * NUM_GRID +HIDDEN = 1152 +MERGE = 2 + + +def _table(rows=NUM_POS, hidden=HIDDEN, dtype=torch.bfloat16, seed=0): + g = torch.Generator(device="cuda").manual_seed(seed) + return torch.randn(rows, hidden, generator=g, device="cuda").to(dtype) + + +def _merge_order(x, height, width, merge): + hidden = x.shape[-1] + return ( + x.reshape(height // merge, merge, width // merge, merge, hidden) + .permute(0, 2, 1, 3, 4) + .reshape(height * width, hidden) + ) + + +def _ref_resize(table, grid_thw, merge, mode="bilinear", align_corners=True): + """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) + rows = _merge_order(resized[0].permute(1, 2, 0), height, width, merge) + outs.append(rows.repeat(frames, 1) if frames > 1 else rows) + return torch.cat(outs) if len(outs) > 1 else outs[0] + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying): # noqa: F811 + certifying(op) + spec = cell.spec + table = _table(spec["rows"], spec["hidden"], spec["dtype"], seed=zlib.crc32(cell.why.encode())) + grid = [list(g) for g in spec["grid"]] + before = table.clone() + out = op(table, grid, spec["merge"]) + assert out.shape == (sum(t * h * w for t, h, w in grid), spec["hidden"]) + assert out.dtype == table.dtype + assert torch.equal(table, before) and out.data_ptr() != table.data_ptr() + torch.testing.assert_close(out, op.reference(table, grid, spec["merge"])) + 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)) + + +def test_guard_refuses_a_non_square_table() -> None: + """The hazard, unguarded: 2300 rows -> isqrt 47, every index in range, a + correctly shaped result off the wrong grid.""" + table = _table(rows=2300, hidden=64, seed=10) + out = op(table, [(1, 4, 4)], MERGE) + assert out.shape == (16, 64) + assert torch.equal(out, op(table[: 47 * 47], [(1, 4, 4)], MERGE)), "read as a 47x47 grid" + + with pytest.raises(AssertionError, match="must hold a square grid"): + with validating(op): + op(table, [(1, 4, 4)], MERGE) + + +def test_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) + grid = [(1, 8, 12)] + + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(3): + op(table, grid, MERGE) + torch.cuda.current_stream().wait_stream(warmup) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(table, grid, MERGE) + + table.copy_(_table(seed=13)) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(captured, op.reference(table, grid, MERGE)) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/moe/test_modeling_v2_fp4_block_scale_moe_runner.py b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py similarity index 80% rename from tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/moe/test_modeling_v2_fp4_block_scale_moe_runner.py rename to tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py index c17d948bebbe..02d286fa949c 100644 --- a/tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/moe/test_modeling_v2_fp4_block_scale_moe_runner.py +++ b/tests/unittest/_torch/modeling_v2/moe/test_modeling_v2_fp4_block_scale_moe_runner.py @@ -7,13 +7,13 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v3.r1_0528_nvfp4__sm_103__dep4.ops.moe.fp4_block_scale_moe_runner import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( # noqa: E501 fp4_block_scale_moe_runner as moe, ) from tensorrt_llm._torch.autotuner import AutoTuner, autotune -__extra_import_path__ = ["../../../.."] -from _validating import validating # noqa: E402 — needs the path declared above +__extra_import_path__ = [".."] +from _validating import certifying, validating # noqa: E402,F401 — needs the path declared above assert torch.cuda.is_available(), "fp4_block_scale_moe_runner requires a CUDA device" @@ -69,7 +69,10 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: torch's default band cannot express that: its bf16 `atol` of 1e-5 sits three orders of magnitude below one output ulp of a two-GEMM chain, and per-element `rtol` is meaningless where cancellation drives `|ref|` to - zero. + zero. That requantization is the only lossy stage between the GEMMs, and + exactly one: read straight out of the kernel it is bit-exact at H = I = + 256 and again at H 2560 / I 640 over 512 experts, so the band absorbs no + second, unmodelled quantization. Both numbers are measured, not chosen. Worst over every configuration the entry's test covers: 9.97 ulp element-wise and 0.88 ulp RMS; the R1 @@ -77,13 +80,19 @@ def _compare(out: torch.Tensor, ref: torch.Tensor) -> None: the 256-expert stack. Only the RMS gate is scale-free -- the element figure is a max over `T*H` elements, so its tail grows with `T` at constant accuracy and reaches 16.33 at the top of the certified token - column. Size a caller's tolerance off the RMS number. + column. Size a caller's tolerance off the RMS number. The qwen4_exp + geometry (512 experts, top-10, H 2560 / I 640) is the least strained: + 3.82 / 0.70 over its column, element tails re-drawn over 12 seeds up to + 10.67 at T=4096 with the RMS figure at 0.67-0.69, and 5.39 / 0.67 with + uniform rather than renormalized weights. Both were sized against deliberately wrong computations: a gate/up swap (263-278 ulp), a scalar-role swap (177/48), a missing scale swizzle (408/357), a missing row shuffle (508/432), and a reference skipping the FC1 requantization (43/28 element, 23/23 RMS). The smallest RMS distance - any of them reaches is 39.8. + any of them reaches is 39.8 at the deepseek geometries; at qwen4_exp's the + same five land 223, 62, 265, 362 and 23 RMS ulp out, the last the closest + any wrong variant comes. """ assert out.dtype == ref.dtype == torch.bfloat16, (out.dtype, ref.dtype) assert out.shape == ref.shape, (out.shape, ref.shape) @@ -466,6 +475,16 @@ def _assert_moe_close(y: torch.Tensor, ref: torch.Tensor) -> None: 3.20 (H=2560) and 9.62 / 4.68 (H=7168), but only the RMS gate is scale-free — read it, not the element one, when sizing a caller's tolerance. + The many-small-experts geometry (E=512, H=2560, I=640, top-10) is the + least strained of the three: the seeded sweep in + `test_many_experts_top10_geometry` tops out at 3.82 elt / 0.70 rms over + the whole token column (3.82 / 0.67 at T=8192), and re-drawn over 12 fresh + weight+activation+routing seeds the element tail reaches 10.67 at T=4096 + (p90 4.36, mean 3.57) and 5.03 at T=8192 (p90 3.60, mean 3.20), with the + RMS figure between 0.670 and 0.694 at both counts. Ten renormalised slot + weights summing to 1 are not what holds it: the same geometry with uniform + [0, 1) weights lands at 5.39 elt / 0.67 rms over the two counts that test + repeats. Both gates were sized against deliberately wrong computations when they were set. At H=256/I=128 and at H=7168/I=2048 respectively: a gate/up swap (263 / 278 ulp), a scalar-role swap (177 / 48), a missing scale @@ -473,6 +492,10 @@ def _assert_moe_close(y: torch.Tensor, ref: torch.Tensor) -> None: missing row shuffle (508 / 432), and a reference that skips the FC1-output requantization (43 / 28 ulp element, 23 / 23 ulp RMS); the smallest RMS distance any of those reaches is 39.8, at the R1 shape's scalar-role swap. + `test_many_experts_top10_geometry` measures the same five at E=512 / + H=2560 / I=640 and asserts each is more than 10 ulp RMS out: 223, 62, 265, + 362 and 23 respectively, the last being the closest any wrong variant comes + anywhere in this file. `test_ep_window_18_of_72` and `test_ep_window_64_of_256` add a mis-set expert-parallel window (304-363 ulp RMS) and one window left out of the sum (>=124 ulp RMS overall, >=128 at T=8192). @@ -526,6 +549,47 @@ def _dsr1(): _R1_TOKENS = (1, 2, 3, 5, 7, 8, 16, 24, 32, 33, 40, 64, 128, 256, 1024, 4096, 8192) +# qwen4_exp's many-small-experts routed geometry, built once and shared: 512 +# experts, H = 2560, I = 640, top-10. Three axes at once that neither DeepSeek +# geometry reaches -- 512 experts, a top_k past 8, and an intermediate width of +# 640 -- with every alignment rule still satisfied (H % 256 == 0, 2*I = 1280 % +# 128 == 0, H/16 = 160 % 4 == 0, I/16 = 40 % 4 == 0), so nothing is padded. The +# kernel-ready stack is ~1.4 GiB (FC1 [512, 1280, 1280] + [512, 1280, 160], FC2 +# [512, 2560, 320] + [512, 2560, 40]) and the reference operands take it to +# ~5 GiB. +_MANY = None + + +def _many_experts(): + global _MANY + if _MANY is None: + _MANY = _build(512, 2560, 640, seed=512640) + return _MANY + + +# The token column swept at that geometry. The pre-routed path picks its +# permutation kernel by row count: `BlockKernelMaxNumTokens` takes 4 rows and +# below, the dynamic-block kernel is out of reach here because +# `DynBlockKernelMaxNumExperts` is 256, and `MaxNumTokensSingleCluster` carries +# everything from 5 rows up to 8192 (all three in +# `cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingCustomPolicy.cuh`). So +# the column brackets the 4/5 boundary as well as the decode and prefill ends. +_MANY_TOKENS = (1, 2, 3, 4, 5, 8, 16, 64, 256, 1024, 4096, 8192) + + +def _renorm_routing(num_tokens, num_experts, top_k, gen): + """Pre-routed ids/weights in the form a renormalising router emits. + + Softmax over every expert, top-k of that, and the k weights rescaled to sum + to 1 -- the `norm_topk_prob` convention. `torch.topk` returns distinct + indices, which is what this op requires of a row. + """ + logits = torch.randn(num_tokens, num_experts, device=DEV, dtype=torch.float32, generator=gen) + w, idx = torch.topk(torch.softmax(logits, dim=-1), top_k, dim=-1) + w = w / w.sum(-1, keepdim=True) + return idx.to(torch.int32).contiguous(), w.to(torch.bfloat16).contiguous() + + # ── tests ───────────────────────────────────────────────────────────────── @@ -1472,7 +1536,259 @@ def test_ep_window_64_of_256(): ) -def test_r1_autotuner_cache_is_inert(): +def test_many_experts_top10_geometry(): + """E=512, H=2560, I=640, top-10: qwen4_exp's many-small-experts routed layer. + + Three axes at once that neither DeepSeek geometry above reaches — 512 + experts, a top_k past 8, and an intermediate width of 640 — swept over the + whole token column, 1 row through 8192 (trtllm's default `max_num_tokens`, + the widest unchunked prefill a stock engine issues). Token counts are + appended in ascending order, so every count below the top draws exactly the + values it drew before. + + Slot weights are renormalised to sum to 1, which is what a + `norm_topk_prob` router emits, and two counts are repeated with plain + uniform weights so the gate is not resting on that convex combination. + Three further things are pinned here: `routing_method_type` is still inert + at the (512, top-10) pair where the routing tier tables would matter, two + identical calls are bitwise equal, and the wrong computations land far + outside the gate at this geometry too. + """ + args, ref = _many_experts() + num_experts, hidden, inter, top_k = 512, 2560, 640, 10 + gen = torch.Generator(device=DEV).manual_seed(51264) + g1 = 137.0 + worst = (0.0, 0.0) + keep = {} + for num_tokens in _MANY_TOKENS: + data, sf, x = _rand_activation(num_tokens, hidden, g1, gen) + ids, wts = _renorm_routing(num_tokens, num_experts, top_k, gen) + assert (ids.sort(dim=-1).values.diff(dim=-1) != 0).all(), ( + "the routing draw repeated an expert id inside one token's row, which " + "this op does not reliably deduplicate" + ) + g2 = _calibrate_g2(x, ids, ref) + s1, sg, s2 = _scalars(num_experts, g1, g2) + kw = dict( + topk_ids=ids, + topk_weights=wts, + output1_scale_scalar=s1, + output1_scale_gate_scalar=sg, + output2_scale_scalar=s2, + ) + out = _call(data, sf, args, num_experts, top_k, **kw) + assert len(out) == 1, len(out) + y = out[0] + assert y.shape == (num_tokens, hidden) and y.dtype == torch.bfloat16, y.shape + exp = moe.reference(x, ids, wts, ref["up"], ref["gate"], ref["down"], g2).to(torch.bfloat16) + _compare(y, exp) + top = _dev(y, exp) # after the loop: the largest token count's numbers + worst = tuple(max(a, b) for a, b in zip(worst, top)) + if num_tokens in (256, 8192): + keep[num_tokens] = (data, sf, x, ids, wts, g2, s1, sg, s2, y, exp) + + # Un-renormalised slot weights, so the sum-to-1 combine is not what is + # holding the gate: ten weights drawn uniformly from [0, 1) instead. + uniform = (0.0, 0.0) + for num_tokens in (256, 8192): + data, sf, x, ids, _, g2, s1, sg, s2, _, _ = keep[num_tokens] + wts = torch.rand(num_tokens, top_k, device=DEV, generator=gen).to(torch.bfloat16) + y = _call( + data, + sf, + args, + num_experts, + top_k, + topk_ids=ids, + topk_weights=wts, + output1_scale_scalar=s1, + output1_scale_gate_scalar=sg, + output2_scale_scalar=s2, + )[0] + exp = moe.reference(x, ids, wts, ref["up"], ref["gate"], ref["down"], g2).to(torch.bfloat16) + _compare(y, exp) + uniform = tuple(max(a, b) for a, b in zip(uniform, _dev(y, exp))) + + data, sf, x, ids, wts, g2, s1, sg, s2, y256, exp256 = keep[256] + kw = dict( + topk_ids=ids, + topk_weights=wts, + output1_scale_scalar=s1, + output1_scale_gate_scalar=sg, + output2_scale_scalar=s2, + ) + # Which (expert count, top_k) pairs the routing kernels are compiled for is + # a per-policy tier list (`PolicyTraits` / `TierList` in + # `cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingCustomPolicy.cuh`), + # and a pair with no tier is a logged error rather than a raise. That would + # bite at (512, top-10) if the caller's `routing_method_type` chose the + # policy — but on the pre-routed path it does not: `runPostTopKPipeline` + # (`RoutingFromTopKIds.cu`) substitutes its own preprocess/postprocess pair + # before dispatching. Checked at the pair where it would show. + for rmt in (0, 2, 4, 5, 6): + alt = _call(data, sf, args, num_experts, top_k, routing_method_type=rmt, **kw)[0] + assert torch.equal(alt, y256), ( + f"routing_method_type={rmt} changed the pre-routed result at " + f"num_experts=512, top_k=10: {_dev(alt, y256.float())} (elt / rms ulp)" + ) + assert torch.equal(_call(data, sf, args, num_experts, top_k, **kw)[0], y256), ( + "two identical calls at this geometry were not bitwise equal" + ) + + # How far the wrong computations land, at this geometry, from the result + # the gate accepts. The RMS gate is 2 ulp; the closest of these is an order + # of magnitude outside it. + dist = {} + for name, wrong in ( + ("gate/up swapped", _ref_moe(x, ids, wts, ref, g2, swap_gate_up=True)), + ("output1_* roles swapped", _ref_moe(x, ids, wts, ref, g2, swap_scalars=True)), + ( + "no FC1-output requantization", + _ref_moe(x, ids, wts, ref, g2, quantize_intermediate=False), + ), + ): + dist[name] = _dev(y256, wrong.to(torch.bfloat16)) + up_p, gt_p, up_s, gt_s, dn_p, dn_s = ref["raw"] + for name, prep in (("no scale swizzle", "swizzle"), ("no row shuffle", "shuffle")): + w1, sw1 = _prep_fc1(up_p, gt_p, up_s, gt_s, **{prep: False}) + w2, sw2 = _prep_fc2(dn_p, dn_s, **{prep: False}) + bad = _call( + data, + sf, + dict( + gemm1_weights=w1, + gemm1_weights_scale=sw1, + gemm2_weights=w2, + gemm2_weights_scale=sw2, + intermediate_size=inter, + ), + num_experts, + top_k, + **kw, + )[0] + dist[name] = _dev(bad, exp256.float()) + del w1, sw1, w2, sw2, bad + torch.cuda.empty_cache() + for name, (elt, rms) in dist.items(): + assert rms > 10.0, ( + f"{name} is only {rms:.1f} rms ulp from the correct result, against a " + f"2 ulp gate — this test would not catch it" + ) + print( + f" test_many_experts_top10_geometry OK (worst {worst[0]:.2f} elt / " + f"{worst[1]:.2f} rms ulp; at T={_MANY_TOKENS[-1]} {top[0]:.2f} elt / " + f"{top[1]:.2f} rms ulp; uniform weights {uniform[0]:.2f} / " + f"{uniform[1]:.2f}; wrong variants " + + ", ".join(f"{n} {v[1]:.0f}" for n, v in dist.items()) + + " rms ulp)" + ) + + +def test_narrow_intermediate_is_nvfp4_requantized(): + """The FC1 epilogue quantizes once, and still exactly, at I = 640. + + `test_intermediate_is_nvfp4_requantized` pins the formula with an identity + down projection, which forces I = H; 640 is not a legal hidden size here + (`H % 256 != 0`), so this repeats the measurement at H = 2560 / I = 640 and + 512 experts with a down projection that is the identity on its first 640 + output rows and zero on the rest. The kernel's own output then carries the + requantized intermediate in columns [0, 640) and must be exactly zero past + it. Bit-exact, not toleranced: e2m1 x e4m3 products are exact in bf16. + + This is what the entry's tolerance rests on. Other runners in this family + requantize more than once per slot, and a reference modelling a single + stage cannot match those; here one modelled stage reproduces every element, + so the tolerance is absorbing accumulation order and nothing else. + """ + num_experts, hidden, inter, num_tokens = 512, 2560, 640, 64 + gen = torch.Generator(device=DEV).manual_seed(640640) + up_p, up_s, up_c = _rand_nvfp4(num_experts, inter, hidden, gen) + gt_p, gt_s, gt_c = _rand_nvfp4(num_experts, inter, hidden, gen) + w1, sw1 = _prep_fc1(up_p, gt_p, up_s, gt_s) + sel = torch.zeros(num_experts, hidden, inter, dtype=torch.uint8, device=DEV) + d = torch.arange(inter, device=DEV) + sel[:, d, d] = 2 # e2m1 code 2 == +1.0 + dn_p = (sel[..., 0::2] | (sel[..., 1::2] << 4)).contiguous() + dn_s = torch.ones(num_experts, hidden, inter // SV, device=DEV).to(torch.float8_e4m3fn) + w2, sw2 = _prep_fc2(dn_p, dn_s) + del sel, up_p, gt_p, dn_p + torch.cuda.empty_cache() + + g1 = 137.0 + data, sf, x = _rand_activation(num_tokens, hidden, g1, gen) + ids = torch.randint( + 0, num_experts, (num_tokens, 1), dtype=torch.int32, device=DEV, generator=gen + ) + wts = torch.ones(num_tokens, 1, dtype=torch.bfloat16, device=DEV) + + act = torch.zeros(num_tokens, inter, dtype=torch.float32, device=DEV) + for e in ids[:, 0].unique().tolist(): + tok = (ids[:, 0] == e).nonzero(as_tuple=True)[0] + xe = x[tok] + up = xe @ _dequant(up_c[e], up_s[e]).t() + gate = xe @ _dequant(gt_c[e], gt_s[e]).t() + act[tok] = up * gate * torch.sigmoid(gate) + g2 = E4M3_MAX * E2M1_MAX / act.abs().max().item() + a1 = torch.full((num_experts,), 1.0 / g1, dtype=torch.float32, device=DEV) + y = _call( + data, + sf, + dict( + gemm1_weights=w1, + gemm1_weights_scale=sw1, + gemm2_weights=w2, + gemm2_weights_scale=sw2, + intermediate_size=inter, + ), + num_experts, + 1, + topk_ids=ids, + topk_weights=wts, + output1_scale_scalar=a1 * g2, + output1_scale_gate_scalar=a1, + output2_scale_scalar=torch.ones(num_experts, dtype=torch.float32, device=DEV), + )[0].float() + + head, tail = y[:, :inter], y[:, inter:] + exact = _q_nvfp4(act, g2) + assert torch.equal(head, exact), ( + f"at I={inter} the FC1 epilogue is not sf=e4m3(g*amax/6), data=e2m1(g*x/sf): " + f"{(head != exact).sum().item()} of {head.numel()} elements differ" + ) + assert torch.equal(tail, torch.zeros_like(tail)), ( + "the zero rows of the selection down projection did not produce zeros — " + "the columns past I do not isolate the intermediate" + ) + # A second quantization anywhere in the chain, an unrounded block scale, or + # no quantization at all would each break the equality above; these two + # record how far the near misses are from matching. + b = (act * g2).reshape(num_tokens, inter // SV, SV) + sfx = b.abs().amax(-1, keepdim=True) / E2M1_MAX + unrounded = (_e2m1_rne(b / sfx.clamp_min(1e-30)) * sfx).reshape(num_tokens, inter) + fracs = {} + for name, cand in ( + ("no requantization", act * g2), + ("exact (unrounded) scale", unrounded), + ): + fracs[name] = (head == cand).float().mean().item() + assert fracs[name] < 0.9, f"{name} also matches {fracs[name]:.2%} — no discrimination" + print( + f" test_narrow_intermediate_is_nvfp4_requantized OK (bit-exact on all " + f"{head.numel()} elements at I={inter}; unrounded scale " + f"{fracs['exact (unrounded) scale']:.1%}, no requantization " + f"{fracs['no requantization']:.1%})" + ) + + +@pytest.fixture +def clear_autotuner_cache(): + """Leave the profiling cache cold even when the test fails part-way, so a + failure here does not turn into a failure of the test that runs next.""" + yield + AutoTuner.get().clear_cache() + + +def test_r1_autotuner_cache_is_inert(clear_autotuner_cache): """A warm autotuner profiling cache does not change the R1-geometry result. Every other call in this file runs with a **cold** cache, where the tuner @@ -1571,7 +1887,6 @@ def entries(): f"local_expert_offset={off}, local_num_experts={n}, " f"T={num_tokens}: {_dev(warm, ref_cold.float())} (elt / rms ulp)" ) - tuner.clear_cache() print( f" test_r1_autotuner_cache_is_inert OK ({len(keys)} tuned entries over " f"unique_ids {unique_ids}, {len(tactics)} distinct tactics, " @@ -1581,6 +1896,88 @@ def entries(): ) +def test_many_experts_autotuner_cache_is_inert(clear_autotuner_cache): + """A warm autotuner profiling cache does not change the 512-expert result. + + `local_num_experts` separates the tuner's cache entries — the R1 test + above measures that by counting profiling sweeps — so a cache warmed at 256 + or 64 experts says nothing about 512, and a serving engine's cache is warm + while every other call in this file runs cold. Same shape as the R1 test: + take the cold results, warm through trtllm's own `autotune()` context at + this layout, re-run, demand **bitwise** equality. It clears the cache at + the end, as the R1 test does, so the cold-cache receipt the rest of this + file records survives either ordering of the two. + """ + args, ref = _many_experts() + num_experts, hidden, top_k = 512, 2560, 10 + gen = torch.Generator(device=DEV).manual_seed(5125) + g1 = 137.0 + token_counts = (1, 8, 256, 1024, 4096, 8192) + cases = {} + for num_tokens in token_counts: + data, sf, x = _rand_activation(num_tokens, hidden, g1, gen) + ids, wts = _renorm_routing(num_tokens, num_experts, top_k, gen) + cases[num_tokens] = ( + data, + sf, + ids, + wts, + _scalars(num_experts, g1, _calibrate_g2(x, ids, ref)), + ) + + def run(num_tokens): + data, sf, ids, wts, (s1, sg, s2) = cases[num_tokens] + return _call( + data, + sf, + args, + num_experts, + top_k, + topk_ids=ids, + topk_weights=wts, + output1_scale_scalar=s1, + output1_scale_gate_scalar=sg, + output2_scale_scalar=s2, + )[0] + + tuner = AutoTuner.get() + op = "trtllm::fp4_block_scale_moe_runner" + + def entries(): + return [k for k in tuner.profiling_cache.cache if k[0] == op] + + assert not entries(), ( + "the profiling cache already holds entries for this op — the cold-cache " + "receipt every other test in this file records would be void" + ) + cold = {t: run(t) for t in token_counts} + assert not entries(), ( + "an ordinary call wrote a profiling-cache entry — calls outside " + "autotune() are not cold after all" + ) + with autotune(): + run(token_counts[-1]) + keys = entries() + assert keys, "autotune() recorded nothing for this op — nothing was warmed" + tactics = {str(tuner.profiling_cache.cache[k][1]) for k in keys} + assert "-1" not in tactics, ( + f"the tuner recorded the fallback tactic, so warm and cold are the same " + f"execution and this comparison proves nothing: {sorted(tactics)}" + ) + unique_ids = sorted({str(k[2]) for k in keys}) + for num_tokens in token_counts: + warm = run(num_tokens) + assert torch.equal(warm, cold[num_tokens]), ( + f"a warm autotuner cache changed the result at T={num_tokens}: " + f"{_dev(warm, cold[num_tokens].float())} (elt / rms ulp)" + ) + print( + f" test_many_experts_autotuner_cache_is_inert OK ({len(keys)} tuned " + f"entries over unique_ids {unique_ids}, {len(tactics)} distinct tactics; " + f"all {len(token_counts)} token counts bitwise unchanged)" + ) + + def test_do_finalize_and_output_buffer(): """Both output modes, and what `do_finalize=False` actually returns.""" num_experts, hidden, inter, top_k, num_tokens = 8, 256, 128, 2, 16 @@ -1897,8 +2294,12 @@ def test_strided_inputs_are_silently_wrong(): print(" test_strided_inputs_are_silently_wrong OK") +# The shared layer each cell's (num_experts, hidden, intermediate) names. +_GEOMETRIES = {(256, 7168, 2048): _dsr1, (512, 2560, 640): _many_experts} + + @pytest.mark.parametrize("cell", moe.CELLS, ids=[c.why[:44] for c in moe.CELLS]) -def test_certified_cells(cell): +def test_certified_cells(cell, certifying): # noqa: F811 """Drive every configuration the entry claims, through the entry's own reference and its own gate. @@ -1910,15 +2311,17 @@ def test_certified_cells(cell): different failure from the kernel moving, and collapsing them into one assert loses which happened. """ + certifying(moe) spec = cell.spec - args, ref = _dsr1() + args, ref = _GEOMETRIES[(spec["num_experts"], spec["hidden"], spec["intermediate"])]() off, n = spec["local_expert_offset"], spec["local_num_experts"] num_experts, hidden, top_k = spec["num_experts"], spec["hidden"], spec["top_k"] gen = torch.Generator(device=DEV).manual_seed(zlib.crc32(cell.why.encode()) % 2**31) g1 = 137.0 data, sf, x = _rand_activation(spec["tokens"], hidden, g1, gen) - ids, wts = _routing(spec["tokens"], num_experts, top_k, gen) + routing = _renorm_routing if spec.get("routing") == "renormalized" else _routing + ids, wts = routing(spec["tokens"], num_experts, top_k, gen) g2 = _calibrate_g2(x, ids, ref) s1, sg, s2 = _scalars(num_experts, g1, g2) diff --git a/tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/quantization/test_modeling_v2_fp4_quantize.py b/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py similarity index 86% rename from tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/quantization/test_modeling_v2_fp4_quantize.py rename to tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py index 6717edd6d35a..83f3b95fb046 100644 --- a/tests/unittest/_torch/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4/quantization/test_modeling_v2_fp4_quantize.py +++ b/tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py @@ -8,11 +8,11 @@ import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.models.deepseek_v3.r1_0528_nvfp4__sm_103__dep4.ops.quantization.fp4_quantize import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( # noqa: E501 fp4_quantize as op, ) -__extra_import_path__ = ["../../../.."] +__extra_import_path__ = [".."] from _validating import validating # noqa: E402 — needs the path declared above assert torch.cuda.is_available(), "fp4_quantize requires a CUDA device" @@ -93,6 +93,18 @@ def test_the_quantization_round_trips() -> None: assert torch.all(err <= block_amax * 0.5 + 1e-6) +def test_the_two_layouts_differ_at_equal_length() -> None: + """At a multiple of 128 rows the linear and swizzled buffers are the same + length, so a consumer handed the wrong one accepts it. The linear cells' + exact scale gate is what tells them apart -- this shows it can.""" + x, gs, _, _ = _build(dict(shape=(128, 2560), dtype=torch.bfloat16, vec=16, swizzled=False), 0) + _, linear = op(x, gs, 16, False, False) + _, swizzled = op(x, gs, 16, False, True) + assert linear.numel() == swizzled.numel() == 128 * 2560 // 16 + assert not torch.equal(linear, swizzled) + assert torch.equal(linear, op.reference(x, gs, 16, False, False)[1]) + + def test_guard_refuses_a_per_token_global_scale() -> None: """Element 0 is applied to every row and the rest are ignored.""" x = torch.randn(8, 128, dtype=torch.bfloat16, device="cuda") 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 new file mode 100644 index 000000000000..d89a528d1417 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py @@ -0,0 +1,141 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""qwen4_exp's KV cache pool contract rejects the mismatches the kernels would not. + +The checks read only `.shape` / `.dtype`, so these run on meta tensors with no GPU and +no engine. Each case is one wrong pool that the kernel consuming it would accept and +mis-address. +""" + +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, +) + +G = M._PoolGeometry( + tokens_per_block=32, + heads_kv=2, + head_dim=256, + index_kv_heads=1, + index_head_dim=128, + conv_dim=8192, + conv_kernel=4, + num_v_heads=48, + head_v_dim=128, + head_k_dim=128, + stream=10240, + ple_state_len=9, + ngram_context_len=3, + dtype=torch.bfloat16, +) +SLOTS = 6 + + +def _t(*shape, dtype=torch.bfloat16): + return torch.empty(*shape, dtype=dtype, device="meta") + + +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 test_valid_pools_pass() -> None: + M.check_position_plane(G, _t(9, 32, 3, dtype=torch.int32)) + 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_ple_pools( + G, + _t(SLOTS, G.stream, G.ple_state_len), + _t(SLOTS, G.ngram_context_len, dtype=torch.int64), + _conv(), + ) + M.check_step_buffers( + 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), + ) + + +@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_attention_pools_reject_wrong_index_plane_and_kv_layout() -> None: + kv = _t(9, 2, G.heads_kv, G.tokens_per_block, G.head_dim) + with pytest.raises(AssertionError, match="QSA index plane"): + M.check_attention_pools(G, 3, _t(9, 32, 1, 64), kv) + with pytest.raises(AssertionError, match="QSA index plane"): + M.check_attention_pools(G, 3, _t(9, 32, 1, 128, dtype=torch.float16), kv) + with pytest.raises(AssertionError, match="paged K/V"): + M.check_attention_pools(G, 3, _t(9, 32, 1, 128), _t(9, 2, 4, 32, 256)) + # NHD instead of HND: heads and block swapped. + with pytest.raises(AssertionError, match="paged K/V"): + M.check_attention_pools(G, 3, _t(9, 32, 1, 128), _t(9, 2, 32, 2, 256)) + # K and V not adjacent role planes. + with pytest.raises(AssertionError, match="role planes"): + M.check_attention_pools(G, 3, _t(9, 32, 1, 128), _t(9, 3, 2, 32, 256)) + + +def test_recurrent_pools_reject_wrong_key_dim_order_and_dtype() -> None: + with pytest.raises(AssertionError, match="recurrent state pool"): + M.check_recurrent_pools( + G, _conv(), _t(SLOTS, G.num_v_heads, G.head_k_dim, G.head_v_dim + 1) + ) + with pytest.raises(AssertionError, match="conv state pool"): + M.check_recurrent_pools(G, _conv(dtype=torch.float32), _ssm()) + + +def test_ple_pools_reject_slot_mismatch_and_int32_history() -> None: + ple_conv = _t(SLOTS, G.stream, G.ple_state_len) + ngram = _t(SLOTS, G.ngram_context_len, dtype=torch.int64) + with pytest.raises(AssertionError, match="one slot space"): + M.check_ple_pools( + G, ple_conv, _t(SLOTS + 1, G.ngram_context_len, dtype=torch.int64), _conv() + ) + with pytest.raises(AssertionError, match="n-gram context pool"): + M.check_ple_pools(G, ple_conv, _t(SLOTS, G.ngram_context_len, dtype=torch.int32), _conv()) + with pytest.raises(AssertionError, match="did not allocate"): + M.check_ple_pools(G, None, ngram, _conv()) + + +def test_step_buffers_reject_missing_dtype_drift_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) + # 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) + with pytest.raises(AssertionError, match="intermediate_state_indices"): + M.check_step_buffers(G, int_conv, int_ssm, _conv(), _ssm(), idx.to(torch.int64)) + with pytest.raises(AssertionError, match="intermediate_state_indices"): + M.check_step_buffers(G, int_conv, int_ssm, _conv(), _ssm(), None) 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 new file mode 100644 index 000000000000..bdb99113b301 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_weight_table.py @@ -0,0 +1,250 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""qwen4_exp's weight table declares what the hand-written code did. + +The declaration reference below is a transcription of the `__init__` block the table +replaced, keyed the way the table keys it: vision block `b`'s weights are `l{b}_vis_*` +and the MTP module's are `l{num_hidden_layers}_mtp_*`, its layer index in the checkpoint +and the KV pool. The n-gram table is not in either: its height is a checkpoint fact, +declared at load. + +No GPU: nothing here allocates, and `dims` reads the configs alone. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +__extra_import_path__ = [".."] # noqa: F841 -- repo's file-scoped import hook + +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 + weights as W, +) + +# Qwen3.8-Flash-Next's own numbers, so the comparison is against the shapes this target +# actually ships rather than a toy configuration. +_TEXT = dict( + num_hidden_layers=48, + layer_types=["full_attention" if i % 4 == 3 else "linear_attention" for i in range(48)], + hidden_size=2560, + hc_count=4, + hc_lowrank=320, + head_dim=256, + num_attention_heads=24, + num_key_value_heads=2, + indexer_n_heads=4, + indexer_kv_heads=1, + indexer_head_dim=128, + linear_num_key_heads=16, + linear_num_value_heads=48, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_conv_kernel_dim=4, + num_experts=512, + moe_intermediate_size=640, + shared_expert_intermediate_size=640, + vocab_size=248320, + ple_layer_ids=[2], + ple_embed_dim=2560, + ple_conv_kernel_size=4, + ngram_size=3, + heads_per_ngram=8, +) +_VISION = dict( + depth=27, + hidden_size=1152, + intermediate_size=4304, + in_channels=3, + temporal_patch_size=2, + patch_size=16, + num_position_embeddings=2304, + spatial_merge_size=2, +) + + +def _core(mtp: bool): + """The configuration surface the table reads.""" + return SimpleNamespace( + model_config=SimpleNamespace( + pretrained_config=SimpleNamespace(**_TEXT), + torch_dtype=torch.bfloat16, + spec_config=object() if mtp else None, + ), + vision_config=SimpleNamespace(**_VISION), + mtp_enabled=mtp, + ) + + +@pytest.mark.parametrize("mtp", [False, True], ids=["no_mtp", "mtp"]) +def test_every_declared_weight_has_a_source(mtp): + """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) + manifest = W.MODEL_WEIGHTS.manifest(core) + for key, sources in manifest.items(): + assert sources, f"{key} is declared with no checkpoint source" + with torch.device("meta"): + declared = W.MODEL_WEIGHTS.declare(W.MODEL_WEIGHTS.dims(core)) + assert set(manifest) == set(declared) + + +def test_mtp_rows_appear_only_when_mtp_is_live(): + off = W.MODEL_WEIGHTS.manifest(_core(False)) + on = W.MODEL_WEIGHTS.manifest(_core(True)) + added = set(on) - set(off) + assert set(off) <= set(on) + 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]) + + +def _reference_declaration(mtp: bool) -> dict: + """The declaration block from modeling.py, transcribed.""" + t, v = SimpleNamespace(**_TEXT), SimpleNamespace(**_VISION) + bf16, u8, f8, f32, i64 = ( + torch.bfloat16, + torch.uint8, + torch.float8_e4m3fn, + torch.float32, + torch.int64, + ) + hidden, stream, rank = t.hidden_size, t.hc_count * t.hidden_size, t.hc_lowrank + q_width = t.num_attention_heads * t.head_dim + kv_width = t.num_key_value_heads * t.head_dim + index_width = (t.indexer_n_heads + t.indexer_kv_heads) * t.indexer_head_dim + key_dim = t.linear_num_key_heads * t.linear_key_head_dim + value_dim = t.linear_num_value_heads * t.linear_value_head_dim + conv_dim = 2 * key_dim + value_dim + experts, inter, shared = ( + t.num_experts, + t.moe_intermediate_size, + t.shared_expert_intermediate_size, + ) + ngram_heads = t.heads_per_ngram * (t.ngram_size - 1) + vh, vi = v.hidden_size, v.intermediate_size + vm = vh * v.spatial_merge_size**2 + nl = t.num_hidden_layers + w = {} + + def P(name, *shape, dtype=bf16): + w[name] = (tuple(shape), dtype) + + def hyper_connection(prefix): + P(f"{prefix}_norm", stream) + P(f"{prefix}_down", rank, stream) + P(f"{prefix}_up", stream, rank) + P(f"{prefix}_inject", t.hc_count, stream) + + def gated_attention(prefix): + 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) + 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) + 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) + P(f"l{i}_A_log", t.linear_num_value_heads) + P(f"l{i}_dt_bias", t.linear_num_value_heads) + P(f"l{i}_conv_w", conv_dim, 1, t.linear_conv_kernel_dim) + P(f"l{i}_gnorm", t.linear_value_head_dim) + P(f"l{i}_out_w", hidden, value_dim) + P(f"l{i}_router", experts, hidden) + P(f"l{i}_fc1_w", experts, 2 * inter, hidden // 2, dtype=u8) + P(f"l{i}_fc1_s", experts, 2 * inter, hidden // 16, dtype=u8) + P(f"l{i}_fc2_w", experts, hidden, inter // 2, dtype=u8) + P(f"l{i}_fc2_s", experts, hidden, inter // 16, dtype=u8) + for name in ("e_isc1", "e_isc1_up", "e_ws2_1", "e_ws2_1_up", "e_isc2", "e_ws2_2"): + P(f"l{i}_{name}", experts, dtype=f32) + shared_expert(f"l{i}_") + P("fhc_norm", stream) + P("fhc_down", rank, stream) + P("fhc_up", stream, rank) + P("embed", t.vocab_size, hidden) + P("ple_key_w", stream, hidden) + P("ple_val_w", t.ple_embed_dim, hidden) + P("ple_nk", stream) + P("ple_nq", stream) + P("ple_nc", stream) + P("ple_conv_w", stream, 1, t.ple_conv_kernel_size) + P("ple_mult", t.ngram_size, dtype=i64) + P("ple_vocab", ngram_heads, dtype=i64) + P("ple_offs", ngram_heads, dtype=i64) + P("ple_tscale", 1) + P("vis_patch_w", vh, v.in_channels, v.temporal_patch_size, v.patch_size, v.patch_size) + P("vis_patch_b", vh) + P("vis_pos", v.num_position_embeddings, vh) + for b in range(v.depth): + P(f"l{b}_vis_n1_w", vh) + P(f"l{b}_vis_n1_b", vh) + P(f"l{b}_vis_qkv_w", 3 * vh, vh) + P(f"l{b}_vis_qkv_b", 3 * vh) + P(f"l{b}_vis_proj_w", vh, vh) + P(f"l{b}_vis_proj_b", vh) + P(f"l{b}_vis_n2_w", vh) + P(f"l{b}_vis_n2_b", vh) + P(f"l{b}_vis_fc1_w", vi, vh) + P(f"l{b}_vis_fc1_b", vi) + P(f"l{b}_vis_fc2_w", vh, vi) + P(f"l{b}_vis_fc2_b", vh) + P("vis_m_n_w", vh) + P("vis_m_n_b", vh) + P("vis_m_fc1_w", vm, vm) + P("vis_m_fc1_b", vm) + P("vis_m_fc2_w", hidden, vm) + P("vis_m_fc2_b", hidden) + if mtp: + m = f"l{nl}_mtp_" + P(f"{m}pre_e_norm", hidden) + P(f"{m}pre_h_norm", stream) + P(f"{m}fc_e", hidden, hidden) + P(f"{m}fc_h", hidden, hidden) + hyper_connection(f"{m}ahc") + hyper_connection(f"{m}mhc") + gated_attention(m) + P(f"{m}router", experts, hidden) + P(f"{m}fc1_w", experts, 2 * inter, hidden, dtype=f8) + P(f"{m}fc1_s", experts, (2 * inter) // 128, hidden // 128, dtype=f32) + P(f"{m}fc2_w", experts, hidden, inter, dtype=f8) + P(f"{m}fc2_s", experts, hidden // 128, inter // 128, dtype=f32) + shared_expert(m) + P(f"{m}fhc_norm", stream) + P(f"{m}fhc_down", rank, stream) + P(f"{m}fhc_up", stream, rank) + return w + + +@pytest.mark.parametrize("mtp", [False, True], ids=["no_mtp", "mtp"]) +def test_the_table_declares_what_the_init_block_did(mtp): + core = _core(mtp) + 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) + 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] 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 90532e2b552d..81256a1c2c97 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -107,6 +107,34 @@ def _model_config(pretrained_config, spec_config=None, **mapping_kwargs): _NGRAM = NGramDecodingConfig(max_draft_len=3, max_matching_ngram_size=2) +def _qwen4_exp_config(**vision_overrides): + """The shape Qwen3.8-Flash-Next-NVFP4's normalized composite config declares.""" + text = PretrainedConfig( + num_hidden_layers=48, + hidden_size=2560, + num_experts=512, + num_attention_heads=24, + linear_num_value_heads=48, + vocab_size=248320, + ) + vision = dict( + depth=27, + hidden_size=1152, + num_heads=16, + spatial_merge_size=2, + out_hidden_size=2560, + num_position_embeddings=2304, + deepstack_visual_indexes=[], + hidden_act="gelu_pytorch_tanh", + ) + vision.update(vision_overrides) + return PretrainedConfig( + architectures=["Qwen4ExpForConditionalGeneration"], + text_config=text, + vision_config=PretrainedConfig(**vision), + ) + + _DEP4 = dict(world_size=4, tp_size=4, moe_ep_size=4, moe_tp_size=1, enable_attention_dp=True) @@ -250,3 +278,63 @@ def test_r1_dep4_with_vanilla_mtp_does_not_match(): modeling_v2_resolve(config, _llm_args("require")) with pytest.raises(ValueError, match="MTP"): modeling_v2_resolve(config, _llm_args("require")) + + +# --- qwen4_exp: one-model MTP-Eagle over a linear chain, one rank --------------- + +_QWEN4_EXP_TARGET = "models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.modeling." + + +@pytest.mark.parametrize( + "spec_config", + [None, MTPDecodingConfig(max_draft_len=1), MTPDecodingConfig(max_draft_len=3)], + ids=["no-spec", "mtp1", "mtp3"], +) +def test_qwen4_exp_tp1_matches(spec_config): + 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) + + +def _mtp(max_draft_len=3, **attributes): + """An MTP config with fields the engine fills in after load set directly: the + constructor remaps or recomputes them, so passing them would test nothing.""" + config = MTPDecodingConfig(max_draft_len=max_draft_len) + for name, value in attributes.items(): + setattr(config, name, value) + return config + + +@pytest.mark.parametrize( + "spec_config", + [ + _NGRAM, + MTPDecodingConfig(max_draft_len=4), + MTPDecodingConfig(max_draft_len=3, use_mtp_vanilla=True), + _mtp(max_total_draft_tokens=5), + _mtp(num_nextn_predict_layers=2), + ], + ids=[ + "ngram", + "draft4", + "vanilla-mtp", + "max-total-draft", + "two-predictors", + ], +) +def test_qwen4_exp_unsupported_speculation_falls_back(spec_config): + """The target asserts each of these at construction; route must decline first.""" + config = _model_config(_qwen4_exp_config(), spec_config=spec_config) + assert modeling_v2_resolve(config, _llm_args("auto")) is None + with pytest.raises(ValueError): + modeling_v2_resolve(config, _llm_args("require")) + + +def test_qwen4_exp_deepstack_vision_tower_does_not_match(): + config = _model_config(_qwen4_exp_config(deepstack_visual_indexes=[5, 11, 17])) + assert modeling_v2_resolve(config, _llm_args("auto")) is None + + +def test_qwen4_exp_on_two_ranks_does_not_match(): + config = _model_config(_qwen4_exp_config(), world_size=2, tp_size=2) + assert modeling_v2_resolve(config, _llm_args("auto")) is None diff --git a/tests/unittest/scripts/test_cbts.py b/tests/unittest/scripts/test_cbts.py index 7b859d18797f..1caf89a0e9b2 100644 --- a/tests/unittest/scripts/test_cbts.py +++ b/tests/unittest/scripts/test_cbts.py @@ -64,7 +64,7 @@ ) from rules.base import PRInputs # noqa: E402 from rules.docs_rule import CPU_TEST_YAML_STEM, DOCS_STAGE, DocsRule, is_docs_path # noqa: E402 -from rules.modeling_v2_rule import _is_mv2_claim # noqa: E402 +from rules.modeling_v2_rule import _entry_is_mv2, _is_mv2_claim # noqa: E402 from rules.openengine_rule import _is_openengine_claim # noqa: E402 from rules.out_of_scope_rule import is_out_of_scope # noqa: E402 from rules.spec_dec_rule import _is_spec_claim # noqa: E402 @@ -239,6 +239,16 @@ def test_documentation_does_not_trigger_backend_rules() -> None: assert not _is_vg_claim("examples/visual_gen/guide.rst") +def test_modeling_v2_entries_cover_unit_and_accuracy_tests() -> None: + assert _entry_is_mv2("unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm") + assert _entry_is_mv2("unittest/_torch/modeling_v2/comm") + assert _entry_is_mv2( + "accuracy/modeling_v2/test_qwen4_exp.py::TestModelingV2Qwen38FlashNextNvfp4Sm103Tp1::test_gsm8k" + ) + assert _entry_is_mv2("accuracy/modeling_v2/test_gpt_oss.py::TestModelingV2GptOss120bSm103Tp1") + assert not _entry_is_mv2("accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_bf16") + + def test_documentation_under_tests_is_left_to_docs_rule() -> None: rule = CbtsTestsDefRule(YAMLIndex(), {}, repo_root=REPO_ROOT)